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) :
  44         pygame.sprite.DirtySprite.__init__(self)
  45         self.dirty = 2 # toujours dirty !
  46         RGB.__init__(self)
  47
  48         self.image = pygame.Surface((640, 480))
  49         self._regular = pygame.Surface((640, 480))
  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._regular.blit(img, (0, 0))
  57         self.image.blit(pygame.transform.flip(self._regular, True, False), (0, 0))
  58
  59
  60 def main() :
  61     pygame.init()
  62     screen = pygame.display.set_mode(SCREEN_SIZE)
  63     pygame.display.set_caption(SCREEN_TITLE)
  64
  65     rgb = RGB()
  66
  67     sur = pygame.Surface((640, 480))
  68     sur.fill((255, 255, 255))
  69
  70     while True :
  71         for event in pygame.event.get():
  72             pass
  73
  74         rgb.update()
  75
  76         rgbImg = rgb.capture()
  77         sur.blit(rgbImg, (0, 0))
  78         screen.blit(pygame.transform.flip(sur, True, False), (0, 0))
  79         pygame.display.flip()
  80
  81
  82 if __name__ == "__main__" :
  83     main()