Particle Filter Algorithm in Stereo Vision

import pygame
import pygame.locals
import pygame.image
import pygame.surfarray
import random
import math
import sys, os
import time
import numpy

GRAY = pygame.Color(90, 90, 90)
RED = pygame.Color(255, 0, 0)
WHITE = pygame.Color(255, 255, 255, 255)
GREEN = pygame.Color(0, 255, 0)
BLUE = pygame.Color(0, 0, 255)
YELLOW = pygame.Color(255, 255, 0)

W = 100
H = 100
SAMPLESIZE = 10
SIGMA = W/6.
N_PARTICLES = 2000

def svertka(matrix_a, matrix_b):
    squares = numpy.square(matrix_a - matrix_b)
    return numpy.sum(squares)

class Bitmap():

    def display(self):
        screen = pygame.display.get_surface()
        screen.blit(self.surface, self.position)

class ImageHalf(Bitmap):
    def __init__(self, surface, x, y):
        self.surface = surface
        self.position = [x, y]

class FromSurface(Bitmap):
    def __init__(self, source):
        self.array3d = pygame.surfarray.array3d(source.surface)
        self.position = [0, 0]

    def to_grayscale(self):
        array3d = numpy.zeros(self.array3d.shape)
        array3d[:, :, 0] = self.array3d[:, :, 0] * 0.21 + self.array3d[:, :, 1] * 0.71 + self.array3d[:, :, 2] * 0.07
        array3d[:, :, 1] = array3d[:, :, 0]
        array3d[:, :, 2] = array3d[:, :, 0]
        self.array3d = array3d
        self.surface = pygame.surfarray.make_surface(self.array3d)

    def get_probe_area(self, x, y):
        return SquareArea(x, y, GREEN, SAMPLESIZE, self.array3d)

    def test_area(self, x, y, probe_square):
        x1 = min(W - SAMPLESIZE, max(0, random.gauss(x, SIGMA)))
        y1 = min(H - SAMPLESIZE, max(0, random.gauss(y, SIGMA)))
        lookup_square = SquareArea(x1, y1, RED, SAMPLESIZE, self.array3d)
        lookup_square.position[0] = lookup_square.start_index[0] + W
        matrix_a = probe_square.array3d
        matrix_b = lookup_square.array3d
        error = svertka(matrix_a, matrix_b)
        return lookup_square, error

class SquareArea():
    def __init__(self, x, y, color, size, array3d):
        self.position = [x, y]
        self.start_index = [x, y]
        self.color = color
        self.size = size
        self.array3d = array3d[x:x+size, y:y+size, 0]

    def display(self, shift = (0, 0), kolor = None):
        screen = pygame.display.get_surface()
        if kolor is None:
            clr = self.color
        else:
            clr = kolor
        pygame.draw.rect(screen, clr,
                            (self.position[0] + shift[0],
                            self.position[1] + shift[1],
                            self.size,
                            self.size),
                        1)

def resample(particles):
    #first most important Particle Filter Algorithm part
    #resample particles with substitute, particles with most weight
    #are more likely to be resampled
    #code taken from Sebastian Thrun's course 'Programming a Robotic Car'
    resampled_particles = []
    N = len(particles)
    index = int(random.random() * N)
    beta = 0.0
    mw = max((p[1] for p in particles))
    for i in range(N):
        beta += random.random() * 2.0 * mw
        while beta > particles[index][1]:
            beta -= particles[index][1]
            index = (index + 1) % N
        resampled_particles.append(particles[index])
    return resampled_particles

def move(particles, dx, dy, probe_square, rite_panel):
    #second most important Particle Filter Algorithm part
    #move particles randomly by approx. (dx, dy)
    #and calculate importance weight of moved particles
    p2 = []
    for particle in particles:
        x1 = int(round(particle[0].start_index[0] + random.gauss(dx, dx/10.)))
        y1 = int(round(particle[0].start_index[1] + random.gauss(dy, dy/10.)))
        x1 = min(W - SAMPLESIZE, max(0, x1))
        y1 = min(H - SAMPLESIZE, max(0, y1))
        lookup_square = SquareArea(x1, y1, RED, SAMPLESIZE, rite_panel.array3d)
        lookup_square.position[0] = lookup_square.start_index[0] + W
        matrix_a = probe_square.array3d
        matrix_b = lookup_square.array3d
        try:
            error = svertka(matrix_a, matrix_b)
            p2.append((lookup_square, 1. / error))
        except:
            p2.append((lookup_square, 0.))
    return p2

class Plan():
    def __init__(self):
        self.points = []
        self.kolor = WHITE
        self.zoom = 10.

    def display(self):
        screen = pygame.display.get_surface()
        pygame.draw.rect(screen, pygame.Color(0, 0, 0), (0, H, W, H), 0)
        for x, y in self.points:
            try:
                pygame.draw.circle(screen, self.kolor, (x, 2 * H - self.zoom * y), 1)
            except:
                pass

    def zoom_up(self):
        self.zoom *= 1.2

    def zoom_down(self):
        self.zoom /= 1.2

    def push(self, point):
        self.points.append(point)

def run_game():
    pygame.init()
    random.seed()
    global W, H
    window = pygame.display.set_mode((W, H))
    left_image = pygame.image.load(sys.argv[1]).convert()
    rite_image = pygame.image.load(sys.argv[2]).convert()
    W, H = left_image.get_size()
    window = pygame.display.set_mode((W*2, H*2 + SAMPLESIZE))
    pygame.display.set_caption("Vision")

    the_left = ImageHalf(left_image, 0, 0)
    the_rite = ImageHalf(rite_image, W, 0)
    bw_left = FromSurface(the_left)
    bw_left.position = [0, H]
    bw_left.to_grayscale()
    bw_rite = FromSurface(the_rite)
    bw_rite.position = [W, H]
    bw_rite.to_grayscale()

    skip_dislay = False
    x = random.randint(0, W - SAMPLESIZE)
    y = random.randint(0, H - SAMPLESIZE)
    x0, y0 = x, y
    squares = []
    plan = Plan()

    while 1:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                sys.exit()
            if event.type == pygame.MOUSEBUTTONUP:
                x0, y0 = x, y
                x = event.pos[0]
                y = event.pos[1]
                if x >= 0 and x <= W and y >= 0 and y <= H:
                    skip_dislay = False
            if event.type == pygame.KEYUP:
                if event.key == pygame.K_UP:
                    plan.zoom_up()
                elif event.key == pygame.K_DOWN:
                    plan.zoom_down()
                skip_dislay = False
                x, y = x0, y0
        if skip_dislay:
            continue
        skip_dislay = True
        the_left.display()
        the_rite.display()
        bw_rite.display()
        probe_square = bw_left.get_probe_area(x, y)
        probe_square.display()
        if len(squares):
            squares = move(squares, x-x0, y-y0, probe_square, bw_rite)
            squares = resample(squares)
            for square in squares:
                square[0].display(shift=(0, H))
        else:
            for j in xrange(N_PARTICLES):
                lookup_square, error = bw_rite.test_area(x, y, probe_square)
                weight = 1. / error
                squares.append((lookup_square, weight))
        #the following is not Particle Filtes algorithm, just presentation
        squares.sort(key=lambda o: o[1])
        M = 10
        for square in squares[-M-10:-M]:
            square[0].display(kolor=WHITE)
        for square in squares[-M:]:
            square[0].display()
        #simple calculation of fake distance to object
        x_arr = numpy.array([square[0].position[0] - W for square in squares])
        #average x coordinate of top weighted M squares
        x_average = numpy.average(x_arr[-M:])
        x_std = numpy.std(x_arr[-M:])
        #average x coordinate of all squares
        x_average1 = numpy.average(x_arr)
        x_std1 = numpy.std(x_arr)
        focL = 30.
        base = 1.
        try:
            distance = focL * base / abs(x_average - probe_square.position[0])
            distance1 = focL * base / abs(x_average1 - probe_square.position[0])
            print "d=%f d1=%f std=%f std1=%f" % (distance, distance1, x_std, x_std1)
            if x_std <= x_std1:
                plan.push((x, distance))
            else:
                plan.push((x, distance1))
        except:
            pass
        plan.display()
        pygame.display.flip()




if __name__ == "__main__":
    run_game()