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
  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()