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
  19 class RGB :
  20     def __init__(self) :
  21         self.context = Context()
  22         self.context.init()
  23         self.imgGene = ImageGenerator()
  24         self.imgGene.create(self.context)
  25         self.imgGene.set_resolution_preset(RES_VGA)
  26         self.imgGene.fps = FPS
  27         self.context.start_generating_all()
  28
  29     def capture(self) :
  30         rgb_frame = numpy.fromstring(self.imgGene.get_raw_image_map_bgr(), dtype=numpy.uint8).reshape(480, 640, 3)
  31         image = cv.fromarray(rgb_frame)
  32         cv.CvtColor(cv.fromarray(rgb_frame), image, cv.CV_BGR2RGB)
  33         pyimage = pygame.image.frombuffer(image.tostring(), cv.GetSize(image), 'RGB')
  34
  35         return pyimage
  36
  37     def update(self) :
  38         return self.context.wait_one_update_all(self.imgGene)
  39
  40
  41 class RGBSprite(pygame.sprite.DirtySprite, RGB) :
  42
  43     def __init__(self, alpha=255, size=SCREEN_SIZE) :
  44         pygame.sprite.DirtySprite.__init__(self)
  45         RGB.__init__(self)
  46         self.dirty = 2 # toujours dirty !
  47         self.size = size
  48         self.image = pygame.Surface(size)
  49         self.workSur = pygame.Surface(SCREEN_SIZE)
  50         self.image.set_alpha(alpha)
  51         self.rect = pygame.Rect((0, 0), (0, 0))
  52
  53     def update(self) :
  54         RGB.update(self)
  55         img = self.capture()
  56         self.workSur.blit(img, (0, 0))
  57         self.workSur = pygame.transform.flip(self.workSur, True, False) # miroir
  58         if self.size != SCREEN_SIZE :
  59             pygame.transform.scale(self.workSur, self.size, self.image) # étirement, blit implicite
  60         else :
  61             self.image.blit(self.workSur, (0, 0))
  62
  63
  64 def main() :
  65     pygame.init()
  66     screen = pygame.display.set_mode(SCREEN_SIZE)
  67     pygame.display.set_caption(SCREEN_TITLE)
  68
  69     rgb = RGB()
  70
  71     sur = pygame.Surface((640, 480))
  72     sur.fill((255, 255, 255))
  73
  74     while True :
  75         for event in pygame.event.get():
  76             pass
  77
  78         rgb.update()
  79
  80         rgbImg = rgb.capture()
  81         sur.blit(rgbImg, (0, 0))
  82         screen.blit(pygame.transform.flip(sur, True, False), (0, 0))
  83         pygame.display.flip()
  84
  85
  86 if __name__ == "__main__" :
  87     main()