Posted by SecretStoat on Fri 3rd Sep 00:04 (modification of post by view diff)
download | new post

import sys
import pygame
import pygame.camera
from pygame.locals import *

class Concept():
        def __init__(self, pSize=(640, 480)):
                # Set up a pygame window display surface
                self.window = pygame.display.set_mode(pSize, 0)
                pygame.display.set_caption('Webcam Colour Tracking: ...Waiting for webcam...')

                # Set up clock for FPS counting
                self.clock = pygame.time.Clock()

                # Initialise camera
                pygame.camera.init()
                self.clist = pygame.camera.list_cameras()
                if not self.clist:
                        raise ValueError("Sorry, no cameras detected.")
                self.webcam = pygame.camera.Camera(self.clist[0], pSize) # Currently just choose the first camera
                self.webcam.start()

                # Create a surface to capture to, same bit depth as window display surface
                self.snapshot = pygame.surface.Surface(pSize, 0, self.window)

                # Load identifiers
                self.icons = []
                self.icons.append(["red", pygame.image.load ("red.png").convert_alpha(), (200, 50, 20), (80, 20, 15), (None, None)])
                self.icons.append(["green", pygame.image.load ("green.png").convert_alpha(), (50, 150, 50), (20, 80, 20), (None, None)])
                self.icons.append(["blue", pygame.image.load ("blue.png").convert_alpha(), (50, 50, 180), (20, 20, 80), (None, None)])
                self.icons.append(["laser", pygame.image.load ("laser.png").convert_alpha(), (180, 200, 180), (70, 50, 70), (None, None)])

        def get_colour_location(self, pColour=(200, 50, 50), pColourThreshold=(80, 20, 20)):
                self.snapshot = self.webcam.get_image(self.snapshot)
                if self.flipX or self.flipY:
                        self.snapshot = pygame.transform.flip(self.snapshot, flipX, flipY)

                # Threshold against the colour we got before
                mask = pygame.mask.from_threshold(self.snapshot, pColour, pColourThreshold)

                overlay = self.snapshot.copy()
                overlay.set_alpha(40)
                overlay.lock()
                self.snapshot.lock()
                for m in mask.connected_components():
                        if m.count() > 50:
                                for r in m.get_bounding_rects():
                                        pygame.draw.rect(self.snapshot, pColour, r, 1)
                                pygame.draw.polygon(overlay, (pColour[0], pColour[1], pColour[2], 160), m.outline(3), 0)
                self.snapshot.unlock()
                overlay.unlock()
                self.snapshot.blit(overlay, (0, 0))


                # Keep only the largest blob of that colour
                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() > 50:
                        return mask.centroid()
                return (None,None)

        def main(self, pFlipX=False, pFlipY=False):
                running = True
                self.flipX = pFlipX
                self.flipY = pFlipY

                while running:
                        self.clock.tick()
                        pygame.display.set_caption('Webcam Colour Tracking: %d fps' % self.clock.get_fps())

                        # Check for events
                        for e in pygame.event.get():
                                if e.type == QUIT or (e.type == KEYDOWN and e.key == K_ESCAPE):
                                        # Exit cleanly
                                        self.webcam.stop()
                                        running = False

                        # Find icon positions
                        for icon in self.icons:
                                tmp_coord = self.get_colour_location(pColour=icon[2], pColourThreshold=icon[3])
                                if tmp_coord != (None,None):
                                        if icon[4] != (None,None):
                                                delta = sum( [(x-y)**2 for (x,y) in zip(tmp_coord, icon[4])]) # Kalman might be better
                                                if delta > 20:
                                                        icon[4] = tmp_coord
                                        else:
                                                icon[4] = tmp_coord
                                        self.snapshot.blit(icon[1], icon[4])

                        # Finally blit the snapshot to the window
                        self.window.blit(self.snapshot,(0,0))
                        #pygame.display.flip()
                        pygame.display.update()


# Run Proof of Concept
if __name__=='__main__':
        con = Concept()
        con.main()