Computer vision smapler.

Main

from Viewer import Viewer
from VideoAnalyzer import VideoAnalyzer
from Constants import *
import time, pygame
import os

os.environ['SDL_VIDEO_CENTERED'] = '1'      # Set the window in the center

vw = Viewer()
va = VideoAnalyzer(vw)

vw.start()
va.start()

Media

"""
All of the graphics, playing of sounds, and input are handled here
"""
from Constants import *
import threading
import pygame
import time
from pygame import mixer
import random

def get_tri_pts(center, size):
    p1 = (center[0]-size, center[1]+size)
    p2 = (center[0], center[1]-size)
    p3 = (center[0]+size, center[1]+size)
    return [p1,p2,p3]

def closest_angle_tick(angle):
    c = 2*PI/32
    i = int(angle/c)
    return i*(c)

def closest_radial_tick(radius):
    i = int(radius/RADIAL_QUANT)
    return i*RADIAL_QUANT

def duration2angle(duration):
    return duration*ANGLE_UPDATE/UPDATE_PERIOD

def polar2rect(r,t):
    center_x = VIEWER_SIZE[0]/2
    center_y = VIEWER_SIZE[1]/2
    x = center_x + r*np.cos(t)
    y = center_y + r*np.sin(t)
    return (int(x), int(y))

def polar2rect2(r,t):
    center_x = CENTER_DRAW[0]
    center_y = CENTER_DRAW[1]
    x = center_x + r*np.cos(t)
    y = center_y + r*np.sin(t)
    return (int(x), int(y))

def rect2polar(x,y):
    center_x = VIEWER_SIZE[0]/2
    center_y = VIEWER_SIZE[1]/2
    r = np.hypot(x-center_x, y-center_y)
    t = np.arctan2(y-center_y, x-center_x)
    return (int(r), t%(2*PI))


def draw_sector(screen, color, r1, r2, start, stop, width):
    if r1 <= 5:
        return

    center_x = CENTER_DRAW[0]
    center_y = CENTER_DRAW[1]
    bounding_rect1 = pygame.Rect(center_x - r1,
                                  center_y - r1,
                                  2*r1,
                                  2*r1)
    bounding_rect2 = pygame.Rect(center_x - r2,
                                  center_y - r2,
                                  2*r2,
                                  2*r2)
    p1 = polar2rect2(r1, start)
    p2 = polar2rect2(r2, start)
    p3 = polar2rect2(r1, stop)
    p4 = polar2rect2(r2, stop)

    pygame.draw.arc(screen, color, bounding_rect1, -stop, -start, width)
    pygame.draw.arc(screen, color, bounding_rect2, -stop, -start, width)
    pygame.draw.line(screen, color, p1, p2, width) 
    pygame.draw.line(screen, color, p3, p4, width) 

class SampleObject(object):
    def __init__(self, angle, num, duration, radius, color):
        super(SampleObject, self).__init__()
        self.angle = angle
        self.num = num
        self.duration = duration
        self.radius = radius
        self.color = color

    def __hash__(self):
        return hash((self.angle, self.num, self.duration, self.radius, self.color))
    def __eq__(self, other):
        return (self.angle == other.angle 
               and self.num == other.num 
               and self.duration == other.duration
               and self.radius == other.radius
               and self.color == other.color)
        

class Viewer(threading.Thread):
    def __init__(self):
        super(Viewer, self).__init__()
        self.draw = True
        self.objects = set() #contains objs -> [angle, id, duration, radius]
        self.votes = dict()
        self.objects_mutex = threading.Lock()
        self.cur_angle = 0
        self.cur_angle_display = 0
        self.update_time = time.time() + UPDATE_PERIOD
        self.draw_time = time.time() + RENDER_PERIOD
        self.done = False
        self.mouse_pos = [0,0]
        self.cur_sample = 0
        self.alpha = 255
        self.alpha_update = 20

    def load_sounds(self):
        BASE_FILENAME = "pclips\\%s.wav"#"C:\\Users\\marcd_000\\Downloads\\thtf\\original\\%s.flac"
        filenames = ["DA3W", "DB3W", "DD3W",
                     "BA3W", "BB3W", "GA3W",
                     "VA3W", "VB3W", "VC2W"]
        volumes = [ 1.0, 1.0, 1.0,
                    1.0, 1.0, 1.0,
                    1.0, 1.0, 1.0 ]
        self.sounds = []
        try:
            for i in xrange(len(filenames)):
                filename = filenames[i]
                self.sounds.append(mixer.Sound(BASE_FILENAME%(filename)))
                self.sounds[-1].set_volume(volumes[i])
        except Exception, e:
            print "Failure loading"
            print e
            exit()

        if len(self.sounds) != 9:
            print "Not all sounds loaded"
            exit()

    def run(self):
        print "Starting Media"
        mixer.pre_init(frequency=44100)
        pygame.init()
        mixer.init()
        print mixer.get_init()
        self.load_sounds()
        self.screen = pygame.display.set_mode([ 700, 500],pygame.FULLSCREEN | pygame.RESIZABLE | pygame.DOUBLEBUF | pygame.HWSURFACE)
        self.clock = pygame.time.Clock()

        while not self.done:
            self.update(time.time())
            self.render(time.time())
            self.clock.tick(VIEWER_FPS)

        pygame.quit()

    def update(self, time):
        if time >= self.update_time:
            self.update_time += UPDATE_PERIOD            
        else:
            return

        # Check for events
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.quit()
            elif event.type == pygame.MOUSEBUTTONDOWN:
                self.handle_mouse(event.pos)
            elif event.type == pygame.MOUSEMOTION:
                self.mouse_pos = event.pos
            elif event.type == pygame.KEYDOWN:
                global ANGULAR_OFFSET
                global CENTER_DRAW
                print CDIFF,ANGULAR_OFFSET
                if event.key == pygame.K_q:
                    self.quit()
                elif event.key == pygame.K_f:
                    self.cur_sample = (self.cur_sample + 1)%12
                elif event.key == pygame.K_r:
                    self.cur_angle = 0
                elif event.key == pygame.K_c:
                    self.objects = set()
                elif event.key == pygame.K_t:
                    LOWER_BLUE[1] -= 10
                elif event.key == pygame.K_y:
                    LOWER_BLUE[1] += 10
                elif event.key == pygame.K_g:
                    LOWER_BLUE[2] -= 10
                elif event.key == pygame.K_h:
                    LOWER_BLUE[2] += 10
                elif event.key == pygame.K_i:
                    CDIFF[1] -= 10
                elif event.key == pygame.K_l:
                    CDIFF[0] += 10
                elif event.key == pygame.K_j:
                    CDIFF[0] -= 10
                elif event.key == pygame.K_k:
                    CDIFF[1] += 10
                elif event.key == pygame.K_z:
                    global ANGULAR_OFFSET
                    ANGULAR_OFFSET -= 1*PI/180.0
                elif event.key == pygame.K_x:
                    global ANGULAR_OFFSET 
                    ANGULAR_OFFSET += 1*PI/180.0

                CENTER_DRAW = [VIEWER_SIZE[0]/2 + CDIFF[0], VIEWER_SIZE[1]/2 + CDIFF[1]]

            elif event.type == pygame.VIDEORESIZE:
                VIEWER_SIZE[0] = event.dict['size'][0]
                VIEWER_SIZE[1] = event.dict['size'][1]


        # Update our board ??
        self.play_inview()
        

        self.cur_angle = (self.cur_angle + ANGLE_UPDATE) % (2*PI)
        self.cur_angle_display = self.cur_angle

        self.alpha += self.alpha_update
        self.alpha_update += 1
        if self.alpha > 255:
            self.alpha = 50
            self.alpha_update = 50

    def render(self, time):
        if time >= self.draw_time:
            self.draw_time += RENDER_PERIOD            
        else:
            return

        # Draw Graphics
        self.draw_screen()
        pygame.display.flip()
        self.cur_angle_display += ANGLE_UPDATE/(UPDATE_PERIOD/RENDER_PERIOD)

    def quit(self):
        self.draw = False
        pygame.quit()
        exit()

    # Sound stuff
    """
    Returns a copy of the objects as a, iterable list
    """
    def get_objects(self):
        objects = []
        self.objects_mutex.acquire()
        for obj in self.objects:
            if self.votes[obj] >= REQ_VOTES:
                objects.append(obj)
        self.objects_mutex.release()
        return objects

    def add_objects(self, new_objects):
        new_objects_set = set(new_objects)
        self.objects_mutex.acquire()
        missing = self.objects.difference(new_objects)

        # Decrement those no longer in same spot
        for obj in missing:
            self.votes[obj]-=1
        
        for newobj in new_objects:
            self.votes[newobj] = 5*REQ_VOTES
            if newobj not in self.objects:
                self.objects.add(newobj)

        self.objects_mutex.release()

    def clear_objects(self):
        self.objects_mutex.acquire()
        self.objects.clear()
        self.objects_mutex.release()

    def shape2board(self, shapes):
        #self.clear_objects()
        objects = []
        for shape in shapes:
            style = shape[0]
            color = shape[1]
            m = np.mean(shape[2], 0)
            polar = rect2polar(m[0]*VIEWER_SIZE[0], m[1]*VIEWER_SIZE[1])
            num = SHAPE_TO_ID[style][color]
            duration = min(1.0, 0.5+(polar[0]/(VIEWER_SIZE[1]/2.0)))
            duration = duration*self.sounds[num].get_length()
            obj = SampleObject(closest_angle_tick(polar[1]), 
                               num, 
                               duration, 
                               closest_radial_tick(polar[0]),
                               color)
            objects.append(obj)
        self.add_objects(objects)

    def play_inview(self):
        min_angle = (self.cur_angle) % (2*PI)
        max_angle = (self.cur_angle + ANGLE_UPDATE) % (2*PI)
        if max_angle < min_angle:
            min_angle -= 2*PI

        objects = self.get_objects()
        for obj in objects:
            if min_angle <= obj.angle and obj.angle <= max_angle:
                self.sounds[obj.num].play(maxtime=int(1000*obj.duration))

    # I/O Stuff
    def handle_mouse(self, pos):
        polar = rect2polar(pos[0], pos[1]) 
        duration = min(1.0, polar[1]/(1.0*VIEWER_SIZE[1]/3.0))
        duration = duration*self.sounds[self.cur_sample].get_length()
        obj = SampleObject(closest_angle_tick(polar[1]), 
                           self.cur_sample, 
                           duration, 
                           closest_radial_tick(polar[0]),
                           GREEN)
        self.votes[obj] = 1000*REQ_VOTES
        self.objects.add(obj)
        #self.add_objects( [obj] )


    # Drawing Stuff
    def draw_screen(self):
        # Clear Screen
        self.screen.fill(BLACK_RGB)
        
        # Draw inner circle
        pygame.draw.circle(self.screen, WHITE_RGB, CENTER_DRAW, 10, 3)
        pygame.draw.circle(self.screen, WHITE_RGB, CENTER_DRAW, VIEWER_SIZE[1]/2, 3)
        c = 2*PI/32
        for i in xrange(53*4):
            start_angle = ANGULAR_OFFSET+closest_angle_tick(i*c)
            stop_angle = ANGULAR_OFFSET+closest_angle_tick(start_angle)
            radius1 = 10
            radius2 = VIEWER_SIZE[1]/2
            draw_sector(self.screen, (150,150,150), radius1, radius2, start_angle, stop_angle, 1)
     

        objects = self.get_objects()
        line_width = 2
        size = 3
        for obj in objects:
            xy = polar2rect(obj.radius, ANGULAR_OFFSET + obj.angle)
            cxy = [0,0]
            cxy[0] = xy[0] + CDIFF[0]
            cxy[1] = xy[1] + CDIFF[1]
            color = COLORS[obj.color]
            pygame.draw.circle(self.screen, color, cxy, 6, line_width)
         
            start_angle = ANGULAR_OFFSET + obj.angle
            stop_angle = start_angle + duration2angle(obj.duration)
            diff_angle = self.alpha/255.0
            color = tuple( (np.array(color)*diff_angle).astype(int) )
            radius = obj.radius
            draw_sector(self.screen, color, radius, radius, start_angle, stop_angle, line_width)


        # Draw inner circle
        # Draw current spot
        line_width = 1
        radius1 = 10
        radius2 = VIEWER_SIZE[1]/2
        #start_angle = self.cur_angle - ANGLE_UPDATE/2
        #stop_angle = self.cur_angle + ANGLE_UPDATE/2
        #draw_sector(self.screen, BLACK_RGB, radius1, radius2, start_angle, stop_angle, line_width)
        start_angle = ANGULAR_OFFSET + self.cur_angle_display 
        stop_angle = ANGULAR_OFFSET + self.cur_angle_display 
        draw_sector(self.screen, WHITE_RGB, radius1, radius2, start_angle, stop_angle, line_width)

        
        # Draw Mouse
        """
        color = COLORS[self.cur_sample/4]
        rect = [self.mouse_pos[0]-5, self.mouse_pos[1]-5, 10, 10]
        pygame.draw.rect(self.screen, color, rect, line_width) 

        # Draw Mouse
        polar_mouse = rect2polar(self.mouse_pos[0], self.mouse_pos[1])
        radius1 = 10
        radius2 = closest_radial_tick(VIEWER_SIZE[1]/2)
        start_angle = ANGULAR_OFFSET + closest_angle_tick(polar_mouse[1])
        stop_angle = start_angle
        line_width = 1
        draw_sector(self.screen, WHITE_RGB, radius1, radius2, start_angle, stop_angle, line_width)
        """

VideoAnalyzer

"""
All of the computer vision is done here
"""
import threading
import cv2
import time
from Constants import *
def area_ok(cnt, mina, maxa):
    a = cv2.contourArea(cnt) 
    return mina <= a and a <= maxa

def two_point_distance(p0, p1):
    return np.hypot(p0[0]-p1[0], p0[1]-p1[1])

def three_point_angle(p0, p1, p2):
    d1, d2 = (p0-p1).astype('float'), (p2-p1).astype('float')
    return np.arccos( np.dot(d1, d2) / np.sqrt( np.dot(d1, d1)*np.dot(d2, d2) ) )

def dists(cnt):
    n = len(cnt)
    return np.array([two_point_distance(cnt[i], cnt[(i+1)%n]) for i in xrange(n)])

def int_angles(cnt):
    n = len(cnt)
    return np.array([three_point_angle(cnt[i], cnt[(i+1)%n], cnt[(i+2)%n]) for i in xrange(n)])

def extreme_points(cnt):
    leftmost = tuple(cnt[cnt[:,:,0].argmin()][0])
    rightmost = tuple(cnt[cnt[:,:,0].argmax()][0])
    topmost = tuple(cnt[cnt[:,:,1].argmin()][0])
    bottommost = tuple(cnt[cnt[:,:,1].argmax()][0])
    return np.array([leftmost, topmost, rightmost, bottommost])
    
class VideoAnalyzer(threading.Thread):
    def __init__(self, viewer):
        super(VideoAnalyzer, self).__init__()
        # Attach connection to viewer
        self.viewer = viewer
        # Open Webcam 
        self.camera = cv2.VideoCapture(CAMERA_INDEX)
        self.mask = None

    def run(self):
        # Create mask
        self.mask = cv2.circle(np.zeros((480,640,3)), (344,227), 480/2, (1,1,1), -1)
        self.mask = cv2.circle(self.mask, (344,227), 75, (0,0,0), -1)
        self.mask = self.mask.astype(int)
        """
        starimg = cv2.imread('star.png',0)
        squareimg = cv2.imread('square.png',0)
        pentagonimg = cv2.imread('pentagon.png',0)
        triangleimg = cv2.imread('triangle.png',0)
        starimg = cv2.cvtColor(starimg, cv2.COLOR_BGR2GRAY)
        squareimg = cv2.cvtColor(squareimg, cv2.COLOR_BGR2GRAY)
        triangleimg = cv2.cvtColor(triangleimg, cv2.COLOR_BGR2GRAY)
        pentagonimg = cv2.cvtColor(pentagonimg, cv2.COLOR_BGR2GRAY)
        ret, self.threshstar = cv2.threshold(starimg, 127, 255,0)
        ret, self.threshsquare = cv2.threshold(squareimg, 127, 255,0)
        ret, self.threshpentagon = cv2.threshold(pentagonimg, 127, 255,0)
        ret, self.threshtriangle = cv2.threshold(triangleimg, 127, 255,0)
        """
        # Check Camera
        if not self.camera.isOpened():
            print "No camera"
        else: # Begin main loop
            print "Starting Camera"
            while True:
                # Get Frames for each color
                frames = self.get_frames()

                # Temp: Threshold
                #gray = cv2.cvtColor(frames[0], cv2.COLOR_BGR2GRAY)
                #r,gray = cv2.threshold(gray, 100, 255, cv2.THRESH_BINARY)
                #frames[5] = gray

                # Show each frame in a window
                cv2.imshow('Framer', frames[1])
                #cv2.imshow('Frameg', frames[2])
                cv2.imshow('Frameb', frames[3])
                cv2.imshow('Framey', frames[4])
                cv2.imshow('Framew', frames[5])
                
                # Get list of shapes found
                shapes = self.extract_shapes(frames)
                # Update Viewers sound list
                self.viewer.shape2board(shapes)
                
                # Draw and Show shapes in image for debugging
                frames[0] = self.draw_shapes(frames[0], shapes)
                cv2.imshow('Frame', frames[0])                

                # Check if key 'q' is pressed, if so exit
                if(cv2.waitKey(20)&0xFF == ord('q')):
                    break
            
            # Done with loop, release camera and close windows
            self.camera.release()
            cv2.destroyAllWindows()


    """
    get_frames:
    Returns a list of binary images masked by each color
    """
    def get_frames(self):
        # Read frame from camera
        ret, frame1 = self.camera.read()
        t = frame1.dtype
        frame1 = frame1*self.mask
        frame1 = frame1.astype(t)

        # Blur image to remove noise
        frame1 = cv2.blur(frame1, (3,3))
        # Convert to HSV color space for color detection
        frame2 = cv2.cvtColor(frame1, cv2.COLOR_BGR2HSV)
        #max_val = np.max(frame2[:,:,1])
        #min_val = np.min(frame2[:,:,1])
        #frame2[:,:,1] = (155.0*((frame2[:,:,1]-min_val)/(max_val-min_val)))+100
        #frame1 = cv2.cvtColor(frame2, cv2.COLOR_HSV2BGR)
 
        # Change to binary image based on color
        frame_white = cv2.inRange(frame2, LOWER_WHITE, HIGHER_WHITE)
        frame_red = cv2.inRange(frame2, LOWER_RED, HIGHER_RED)
        frame_red = frame_red | cv2.inRange(frame2, LOWER_RED2, HIGHER_RED2)
        frame_blue = cv2.inRange(frame2, LOWER_BLUE, HIGHER_BLUE)
        frame_green = cv2.inRange(frame2, LOWER_GREEN, HIGHER_GREEN)
        frame_yellow = cv2.inRange(frame2, LOWER_YELLOW, HIGHER_YELLOW)
        return [frame1, frame_red, frame_green, frame_blue, frame_yellow, frame_white]

    """
    draw_shapes:
    Draws the shapes found on an image

    params:
    frame: the image to draw the shapes on
    shapes: the list of shapes to draw

    return:
    a new image with the shapes drawn
    """
    def draw_shapes(self, frame, shapes):
        for s in shapes:
            if s[0] == SQUARE:
                color = [0,0,255]
            if s[0] == TRIANGLE:
                color = [255,0,0]
            if s[0] == STAR:
                color = [0,255,0]
            if s[0] == PENTAGON:
                color = [255,255,0]
            
            cnt = s[2]
            cnt[:, 0] = cnt[:, 0]*WORKSPACE_SIZE[0]
            cnt[:, 1] = cnt[:, 1]*WORKSPACE_SIZE[1]
            cnt = cnt.astype(int)
            thickness = 2
            cv2.drawContours(frame, [cnt], 0, color, thickness)

        return frame


    
    def find_shapes(self, frame, color):
        min_area = 100
        max_area = 2000
        shapes = []
        frame, contours, hierarchy = cv2.findContours(frame, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
        for cnt in contours:
            cnt_len = cv2.arcLength(cnt, True)
            cnt = cv2.approxPolyDP(cnt, 0.05*cnt_len, True)
            
            if not area_ok(cnt, min_area, max_area):
                continue

            ecnt = extreme_points(cnt)
            cnt = cnt.reshape(-1, 2)
            npts = len(cnt)

            if npts == 3:
                shape = TRIANGLE
            elif npts == 4:
                max_int = np.max(int_angles(cnt))
                std = np.std(dists(cnt))
                if max_int > 100.0*PI/180.0:
                    shape = TRIANGLE
                elif std >= 1.6:
                    shape = PENTAGON
                else:
                    shape = SQUARE
            elif npts == 5:
                shape = STAR         
            else:
                shape = PENTAGON

            # Normalize
            cnt = cnt.astype(float)
            cnt[:, 0] = cnt[:, 0]/WORKSPACE_SIZE[0]
            cnt[:, 1] = cnt[:, 1]/WORKSPACE_SIZE[1]
            shapes.append([shape, color, cnt])
        return shapes

    def extract_shapes(self, frames):
        # list of shapes that will be returns
        shapes = []
        for color,i in [(RED,1), (YELLOW,4), (BLUE,3)]:
            #shapes.extend(self.find_circles(frames[i], color))
            #shapes.extend(self.find_squares_triangles(frames[i], color))
            #shapes.extend(self.find_stars(frames[i], color))
            #shapes.extend(self.find_pentagons(frames[i], color))
            #shapes.extend(self.find_squares(frames[i], color))
            shapes.extend(self.find_shapes(frames[i], color))
        return shapes