Expérimentations (concluantes) de hand tracking.
[minwii.git] / src / kinect / pygamedisplay.py
1 # -*- coding: utf-8 -*-
2 """
3 Affichage vidéo (et autres) de la kinect pour expérimentations / debug.
4
5 $Id$
6 $URL$
7 """
8
9 from openni import *
10 import numpy
11 import cv
12 import pygame
13
14 SCREEN_SIZE = 640, 480
15 SCREEN_TITLE = "Kinect debug"
16 FPS = 30
17
18 class RGB :
19 def __init__(self) :
20 self.context = Context()
21 self.context.init()
22 self.imgGene = ImageGenerator()
23 self.imgGene.create(self.context)
24 self.imgGene.set_resolution_preset(RES_VGA)
25 self.imgGene.fps = FPS
26
27 self.depthGene = DepthGenerator()
28 self.depthGene.create(self.context)
29 self.depthGene.set_resolution_preset(RES_VGA)
30 self.depthGene.fps = FPS
31
32 self.handsGene = HandsGenerator()
33 self.handsGene.create(self.context)
34 self.handsGene.register_hand_cb(self.handCreateCB, self.handUpdateCB, self.handDestroyCB)
35
36 self.gestGene = GestureGenerator()
37 self.gestGene.create(self.context)
38 self.gestGene.add_gesture('Click')
39 self.gestGene.register_gesture_cb(self.gestureRecognizedCB, self.gestureProgressCB)
40
41
42 self.userGene = UserGenerator()
43 self.userGene.create(self.context)
44 self.userGene.register_user_cb(self.newUserCB, self.lostUserCB)
45
46 self.context.start_generating_all()
47
48 screen = pygame.display.get_surface()
49 self.xratio = float(screen.get_size()[0]) / SCREEN_SIZE[0]
50 self.yratio = float(screen.get_size()[1]) / SCREEN_SIZE[1]
51
52 def getProjPos(self, realPos) :
53 return self.depthGene.to_projective([realPos])[0]
54
55 def setMousePosition(self, realPos) :
56 x, y, z = self.getProjPos(realPos)
57 # mirror
58 x = SCREEN_SIZE[0] - x
59 # scale
60 x, y = x * self.xratio, y * self.yratio
61 pygame.mouse.set_pos(int(x), int(y))
62
63 def handCreateCB(self, src, id, pos, time):
64 print 'Create ', id, pos
65
66 def handUpdateCB(self, src, id, pos, time):
67 self.setMousePosition(pos)
68
69 def handDestroyCB(self, src, id, time):
70 print 'Destroy ', id
71
72 def gestureRecognizedCB(self, src, gesture, id, end_point) :
73 print 'gestureDetected', src, gesture, id, end_point
74
75 def gestureProgressCB(self, src, gesture, point, progress) :
76 print 'gestureProgress', src, gesture, point, progress
77 self.handsGene.start_generating()
78 self.handsGene.start_tracking(point)
79
80 def newUserCB(self, *args, **kw) :
81 print 'newUserCB', args, kw
82
83 def lostUserCB(self, *args, **kw) :
84 print 'lostUserCB', args, kw
85
86
87 def capture(self) :
88 rgb_frame = numpy.fromstring(self.imgGene.get_raw_image_map_bgr(), dtype=numpy.uint8).reshape(480, 640, 3)
89 image = cv.fromarray(rgb_frame)
90 cv.CvtColor(cv.fromarray(rgb_frame), image, cv.CV_BGR2RGB)
91 pyimage = pygame.image.frombuffer(image.tostring(), cv.GetSize(image), 'RGB')
92 return pyimage
93
94 def update(self) :
95 self.context.wait_any_update_all()
96 #self.context.wait_and_update_all()
97 #return self.context.wait_one_update_all(self.imgGene)
98
99
100 class RGBSprite(pygame.sprite.DirtySprite, RGB) :
101
102 def __init__(self, alpha=255, size=SCREEN_SIZE) :
103 pygame.sprite.DirtySprite.__init__(self)
104 RGB.__init__(self)
105 self.dirty = 2 # toujours dirty !
106 self.size = size
107 self.image = pygame.Surface(size)
108 self.workSur = pygame.Surface(SCREEN_SIZE)
109 self.image.set_alpha(alpha)
110 self.rect = pygame.Rect((0, 0), (0, 0))
111
112 def update(self) :
113 RGB.update(self)
114 img = self.capture()
115 self.workSur.blit(img, (0, 0))
116 self.workSur = pygame.transform.flip(self.workSur, True, False) # miroir
117 if self.size != SCREEN_SIZE :
118 pygame.transform.scale(self.workSur, self.size, self.image) # étirement, blit implicite
119 else :
120 self.image.blit(self.workSur, (0, 0))
121
122
123 def main() :
124 pygame.init()
125 screen = pygame.display.set_mode(SCREEN_SIZE)
126 pygame.display.set_caption(SCREEN_TITLE)
127
128 rgb = RGB()
129
130 sur = pygame.Surface((640, 480))
131 sur.fill((255, 255, 255))
132
133 while True :
134 for event in pygame.event.get():
135 pass
136
137 rgb.update()
138
139 rgbImg = rgb.capture()
140 sur.blit(rgbImg, (0, 0))
141 screen.blit(pygame.transform.flip(sur, True, False), (0, 0))
142 pygame.display.flip()
143
144
145 if __name__ == "__main__" :
146 main()