Fun with PyGame's camera module

OMG T2 starts tomorrow, I'm probably gonna fail in Random Processes this time.
But before I start preparing, I had to try out the tempting Camera Module of PyGame.
PyGame has pretty cool stuff under pygame.transform and pygame.mask, which makes the task of thresholding very easy.
I came up with an interactive Tux based on the "Capturing a Live Stream" code in the Camera Module Introduction.


My program tries to detect red colored objects, finds the centroid of such points and draws a ghost (find in your /usr/share/icons/oxygen/32x32/apps/) and makes Tux (/usr/share/icons/oxygen/128x128/apps/tux.png) run away from it.

Here's how the code looks like -

# -*- coding: utf-8 -*-

# An interactive Tux, interact with it using any red colored object
#
# based on PyGame camera module intro by Nirav Patel
#   http://www.pygame.org/docs/tut/camera/CameraIntro.html
#                         -- Abhishek Mishra (ideamonk # gmail.com)

import os
import pygame
import pygame.camera
from pygame.locals import *

class Capture(object):
    ''' A Capture class to get location of a desired blob '''

  def __init__(self, ccolor=(248, 111, 115), threshold=(60, 10, 10)):
      self.size = (640,480)
      # create a display surface. standard pygame stuff
      self.display = pygame.display.set_mode(self.size, 0)
      # initialize camera module
      pygame.camera.init()
      # this is the same as what we saw before
      self.clist = pygame.camera.list_cameras()
      if not self.clist:
          raise ValueError("Sorry, no cameras detected.")
      self.cam = pygame.camera.Camera(self.clist[0], self.size)
      self.cam.start()

      # create a surface to capture to.  for performance purposes
      # bit depth is the same as that of the display surface.
      self.snapshot = pygame.surface.Surface(self.size, 0, self.display)
      # target color to detect -- default is red
      self.ccolor = ccolor
      # by default we give more priority to shades of red
      self.threshold = threshold

  def get_blob_location(self):
      self.snapshot = self.cam.get_image(self.snapshot)
      # threshold against the color we got before
      mask = pygame.mask.from_threshold(self.snapshot, self.ccolor, self.threshold)
      # keep only the largest blob of that color
      connected = mask.connected_component()
      # these numbers are purely experimental and specific to your room and object
      # print mask.count() # use this to estimate
      # make sure the blob is big enough that it isn't just noise
      if mask.count() > 7:
          # find the center of the blob
          return mask.centroid()
      return (None,None)

class Ghost():
    ''' Ghost class, to have a rect for collision detection '''
  def __init__(self):
      self.image = pygame.image.load (os.path.join ("./","gv.png"))
      self.rect = self.image.get_rect()

  def set_rect(self,position):
      self.left,self.top=position
      self.rect = pygame.Rect(self.left,self.top,32,32)
    
class Tux(Capture):
    ''' Tux class extends Capture and does stuff using get_blob_location '''

  def __init__(self):
      Capture.__init__(self)
      self.location = [ x/2 for x in self.size ]
      self.set_rect (self.location)
      self.image = pygame.image.load (os.path.join ("./","tux.png"))
      self.ghost = Ghost()
      self.backbuffer = pygame.Surface (self.size)
      self.force = 5

  def set_rect(self,position):
      left,top=position
      self.rect = pygame.Rect(left,top,128,128)
    
  def interact_tux(self):
      if (pygame.sprite.collide_rect(self,self.ghost)):
          # ghost collides with tux
          if self.ghost.left<self.location[0]+64:
              self.location[0] += self.force
          if self.ghost.left>self.location[0]+64:
              self.location[0] -= self.force
          if self.ghost.top<self.location[1]+64:
              self.location[1] += self.force
          if self.ghost.top>self.location[1]+64:
              self.location[1] -= self.force
            
        
  def main(self):
      going = True
      old_coord = (0,0)
    
      while going:
          events = pygame.event.get()
          for e in events:
              if e.type == QUIT or (e.type == KEYDOWN and e.key == K_ESCAPE):
                  # close the camera safely
                  self.cam.stop()
                  going = False

          new_coord = self.get_blob_location()
          if new_coord != (None,None):
              # delta = sum( [(x-y)**2 for (x,y) in zip(new_coord,old_coord)])    # for less fuzziness
              # if delta>200:
              old_coord = new_coord
              self.ghost.set_rect(old_coord)
            
          self.set_rect (self.location)
          self.interact_tux()
          self.backbuffer.blit(self.snapshot,(0,0))
          self.backbuffer.blit(self.image,self.location)
          self.backbuffer.blit(self.ghost.image, old_coord)
          self.backbuffer = pygame.transform.flip(self.backbuffer,True,False)
          self.display.blit(self.backbuffer,(0,0))
          pygame.display.flip()
        
if __name__=='__main__':
  t = Tux()   # all default params
  t.main()
  

And that's me doing weird things with tux :D -




Kalman Filter would be more interesting.