321578d35389144803bd9c3d0c477fe62ab4e052
[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 MOUSE_SCREEN_COORDINATES = 0
18 MOUSE_REAL_COORDINATES = 1
19 DEFAULT_HAND_DISTANCE_RANGE = 150 # distance de débatement, autour du centre, en mm
20
21 class RGB :
22 def __init__(self, mouseMode=MOUSE_REAL_COORDINATES) :
23 self.mouseMode = mouseMode
24 self.context = Context()
25 self.context.init()
26 self.imgGene = ImageGenerator()
27 self.imgGene.create(self.context)
28 self.imgGene.set_resolution_preset(RES_VGA)
29 self.imgGene.fps = FPS
30
31 self.depthGene = DepthGenerator()
32 self.depthGene.create(self.context)
33 self.depthGene.set_resolution_preset(RES_VGA)
34 self.depthGene.fps = FPS
35
36 self.handsGene = HandsGenerator()
37 self.handsGene.create(self.context)
38 self.handsGene.register_hand_cb(self.handCreateCB, self.handUpdateCB, self.handDestroyCB)
39
40 self.gestGene = GestureGenerator()
41 self.gestGene.create(self.context)
42 self.gestGene.add_gesture('Click')
43 self.gestGene.register_gesture_cb(self.gestureRecognizedCB, self.gestureProgressCB)
44
45
46 self.userGene = UserGenerator()
47 self.userGene.create(self.context)
48 self.userGene.register_user_cb(self.newUserCB, self.lostUserCB)
49
50 self.context.start_generating_all()
51
52 screen = pygame.display.get_surface()
53 self.xratio = float(screen.get_size()[0]) / SCREEN_SIZE[0]
54 self.yratio = float(screen.get_size()[1]) / SCREEN_SIZE[1]
55 self.handCenteredPosition = None
56 self.pygameScreenWidth, self.pygameScreenHeight = screen.get_size()
57
58 def getProjPos(self, realPos) :
59 return self.depthGene.to_projective([realPos])[0]
60
61 def setMousePosition(self, realPos) :
62 if self.mouseMode == MOUSE_SCREEN_COORDINATES :
63 x, y, z = self.getProjPos(realPos)
64 # mirror
65 x = SCREEN_SIZE[0] - x
66 # scale
67 x, y = x * self.xratio, y * self.yratio
68 elif self.mouseMode == MOUSE_REAL_COORDINATES :
69 x, y = realPos[:2]
70 x, y = x - self.handCenteredPosition[0], y - self.handCenteredPosition[1]
71 x = -x # miroir
72 x, y = x + DEFAULT_HAND_DISTANCE_RANGE, DEFAULT_HAND_DISTANCE_RANGE - y
73 x, y = map(lambda i : numpy.clip(i, 0, DEFAULT_HAND_DISTANCE_RANGE * 2), [x, y])
74 x = x * (self.pygameScreenWidth / (2. * DEFAULT_HAND_DISTANCE_RANGE))
75 y = y * (self.pygameScreenHeight / (2. * DEFAULT_HAND_DISTANCE_RANGE))
76 print x, y
77 pygame.mouse.set_pos(int(x), int(y))
78
79 def handCreateCB(self, src, id, pos, time):
80 print 'Create ', id, pos
81
82 def handUpdateCB(self, src, id, pos, time):
83 self.setMousePosition(pos)
84
85 def handDestroyCB(self, src, id, time):
86 print 'Destroy ', id
87
88 def gestureRecognizedCB(self, src, gesture, id, end_point) :
89 print 'gestureDetected', src, gesture, id, end_point
90
91 def gestureProgressCB(self, src, gesture, point, progress) :
92 print 'gestureProgress', src, gesture, point, progress
93 self.handsGene.start_generating()
94 self.handsGene.start_tracking(point)
95 self.handCenteredPosition = point
96
97 def newUserCB(self, *args, **kw) :
98 print 'newUserCB', args, kw
99
100 def lostUserCB(self, *args, **kw) :
101 print 'lostUserCB', args, kw
102
103
104 def capture(self) :
105 rgb_frame = numpy.fromstring(self.imgGene.get_raw_image_map_bgr(), dtype=numpy.uint8).reshape(480, 640, 3)
106 image = cv.fromarray(rgb_frame)
107 cv.CvtColor(cv.fromarray(rgb_frame), image, cv.CV_BGR2RGB)
108 pyimage = pygame.image.frombuffer(image.tostring(), cv.GetSize(image), 'RGB')
109 return pyimage
110
111 def update(self) :
112 self.context.wait_any_update_all()
113 #self.context.wait_and_update_all()
114 #return self.context.wait_one_update_all(self.imgGene)
115
116
117 class RGBSprite(pygame.sprite.DirtySprite, RGB) :
118
119 def __init__(self, alpha=255, size=SCREEN_SIZE) :
120 pygame.sprite.DirtySprite.__init__(self)
121 RGB.__init__(self)
122 self.dirty = 2 # toujours dirty !
123 self.size = size
124 self.image = pygame.Surface(size)
125 self.workSur = pygame.Surface(SCREEN_SIZE)
126 self.image.set_alpha(alpha)
127 self.rect = pygame.Rect((0, 0), (0, 0))
128
129 def update(self) :
130 RGB.update(self)
131 img = self.capture()
132 self.workSur.blit(img, (0, 0))
133 self.workSur = pygame.transform.flip(self.workSur, True, False) # miroir
134 if self.size != SCREEN_SIZE :
135 pygame.transform.scale(self.workSur, self.size, self.image) # étirement, blit implicite
136 else :
137 self.image.blit(self.workSur, (0, 0))
138
139
140 def main() :
141 pygame.init()
142 screen = pygame.display.set_mode(SCREEN_SIZE)
143 pygame.display.set_caption(SCREEN_TITLE)
144
145 rgb = RGB()
146
147 sur = pygame.Surface((640, 480))
148 sur.fill((255, 255, 255))
149
150 while True :
151 for event in pygame.event.get():
152 pass
153
154 rgb.update()
155
156 rgbImg = rgb.capture()
157 sur.blit(rgbImg, (0, 0))
158 screen.blit(pygame.transform.flip(sur, True, False), (0, 0))
159 pygame.display.flip()
160
161
162 if __name__ == "__main__" :
163 main()