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