Particle Filter

#!/usr/bin/python3
#
#Demo of a particle filter for object tracking:
#The goal to track is the white square moving about in the image
#Several 'distractors' noise dots are also there to make it harder.
#A 'particle' is a set of coordinates (x,y) with a value attached to it.
#The value is the measure of a probability distribution function (pdf), here of the location of the object.
#The main idea is that the density of the particles represents the pdf itself.
#So where the pdf is high, we want to have more particles.
#To get to this situation, we do the following:
#   -start from a uniform distribution (prior) with some particles at random positions. Each particle is just as likely to be at the objects location: p(x,y)=1/n
#   -measure the actual value of the pdf at the position of the particles.This is done by measuring the number of white pixels around a particle's location.
#   -calculate the importance weight w of each particle
#   -now resample the particles with replacement, according to their weights. Particles with higher weights get sampled more often.
#   -make the position estimate as the centre-of-mass of the particles or (smarter) of the best (highest weighted) particles.
#   -perturb the particles a bit (depending on the problem I guess)
#   -repeat
#
#Stuff to play with:
# -the size and number of noise dots
# -number of particles
# -particle perturbation stepsize and object velocity (must be +- moving-dot-velocity)
# -Idea: make a 2nd order particle filter --> estimate the velocity of the white square too, and make the particles follow

#Based loosely on: https://scipy-cookbook.readthedocs.io/items/ParticleFilter.html
#
#larsupilami73


from scipy import *
import matplotlib.pyplot as plt
import os

random.seed(42)


# parameters
N = 100     #image size
d = 8       #square object size
K = 250     #number of images in the animation
Kn = 15     #number of noise dots
dn = 2      #noise dot size, increase to make it harder to follow the object

#this generates a moving square that is to be followed by the particle filter
def imagegen():
    x, y = 80, 20   #position of topleft object corner
    px = +2.9       #initial velocity
    py = -2.7
    counter = K
    image = zeros((N,N))
    yield image,x,y

    while(counter):
        counter-=1
        x = x + px
        y = y + py
        if x>=N-d:
            x = N-d
            px = -px
        elif x<=0:
            x = 0
            px = - px

        if y>=N-d:
            y = N-d
            py = -py
        elif y<=0:
            y = 0
            py = -py


        image = zeros((N,N))
        image[int(x):int(x+d),int(y):int(y+d)] = 1.0 #make pixels white
        #px = px + random.uniform(-0.5,0.5) #perturb the velocity too
        #py = py + random.uniform(-0.5,0.5)
        yield image,x+d//2,y+d//2


#this is to add random-walking noise to the image,
#to make it more difficult for the particle filter
def noisegen():
    noise = random.randint(low=0,high=N-1,size=(Kn,2))
    pnoise = zeros((Kn,2)) #noise dots velocity
    for k_ in range(Kn): #make sure none have zero velocity
        while(not pnoise[k_].any()):
            pnoise[k_] = random.randint(low=-2,high=3,size=(2))
    noisefield = zeros((N,N))
    for nx,ny in noise:
        noisefield[nx,ny] = 1.0
    yield noisefield
    while(True):
        noise = noise + pnoise
        noise = noise % (ones((Kn,2))*N) #loop the noisedots around the corners of the image
        noise = noise.astype(int)
        noisefield = zeros((N,N))
        for nx,ny in noise:
            noisefield[nx:nx+dn,ny:ny+dn] = 1.0 #make white (1.0) or gray (0< .. <1.0)
        yield noisefield



#a simple particlefilter
class ParticleFilter:
    def __init__(self,numparticles=100,low=0.0,high=1.0):
        self.numparticles = numparticles
        self.particles = random.uniform(low=low,high=high,size=(numparticles,2)) #prior distribution is uniform, each particle is a position (x,y), density of particles needs to reflect the pdf
        self.samples = zeros(numparticles)
        self.weights = ones(numparticles)/numparticles

    def update(self, image):
        """
        Update and resample all particles.
        The goal is to have the density of the particles reflect the pdf.
        """

        #wiggle the particles a bit, otherwise, we end up with just the same ones after every update
        #the perturbation should be related to the velocity of the object we want to track
        #if the perturbation is too low, compared to the velocity the particles drag behind the object
        #if the perturbation is too high, the cloud of particles is to wide (density does not represent the pdf very well)
        self.particles = self.particles + random.uniform(-4,4,size=self.particles.shape)
        clip(self.particles,0,N)

        #measure at the position of the particles
        particles_ = self.particles.astype(int) #must be ints
        for k in range(self.numparticles):
            self.samples[k] = image[particles_[k,1]-d//2:particles_[k,1]+d//2, particles_[k,0]-d//2:particles_[k,0]+d//2].sum() #measure how much of the pixels are 'on' in a square of d x d

        #update the importance weights
        self.weights = 1./(1. + (self.samples-d*d)**2)

        #normalize importance weights
        self.weights /=self.weights.sum()

        #choose particles with replacement, according to their importance weight
        the_choice = random.choice(self.numparticles, size=self.numparticles,replace=True,p=self.weights)
        self.particles = self.particles[the_choice]
        self.weights = self.weights[the_choice]

    def estimate(self):
        """
        The estimate of the optimum could be:
        1. the position of the particle with the highest weight
        2. the weighted mean position of all particles
        3. the weighted mean position of the n-most weighty particles
        We go with 3 here.
        """
        numparticles = self.numparticles//4 #25% heaviest particles
        positions_heaviest_particles = self.weights.argsort()[-numparticles:]
        return average(self.particles[positions_heaviest_particles],axis=0,weights=self.weights[positions_heaviest_particles])



#image and noise are generators
noise_g = noisegen()
image_g = imagegen()
#particlefilter is an object that gets fed the images
particlefilter = ParticleFilter(numparticles=100,low=0,high=N)


#temp directory for the frames
if not os.path.exists('./frames'):
    os.mkdir('./frames')
try:
    os.system('rm ./frames/frame*.png')
except:
    pass


#make the figure
fig = plt.figure()
ax = fig.add_subplot(111)
ax.set_title('Particle Filter Object Tracking Demo')
image_ax = ax.imshow(random.normal(size=(N,N)), cmap='gray', vmin=0,vmax=1)  #add imshow artist
particle_ax = ax.scatter(particlefilter.particles[:,0],particlefilter.particles[:,1],c='blue',marker='+',alpha=0.7) #add scatter artist
estimate_ax = ax.scatter(0,0,c='red',s=50, marker='o',alpha=0.9)
text_ax = ax.text(1,5,'real pos.: %i,%i' %(0,0),color='lightgreen',zorder=1,fontsize=14) #add text artist
text2_ax = ax.text(1,1,'estimated pos.: %i,%i' %(0,0), color='lightgreen',zorder=1,fontsize=14) #add text artist
ax.set_xlim([0,N-0.5])
ax.set_ylim([0,N])
plt.ion()

#here the figure update loop starts
framecounter=0
for (image_,real_x,real_y),noise_ in zip(image_g,noise_g):

    #image + noise
    X_ = image_ + noise_

    #show the image
    image_ax.set_array(X_)

    #update the particlefilter
    particlefilter.update(X_)

    #overlay the particles
    particle_ax.set_offsets(particlefilter.particles)

    #overlay the estimate of the position of the square
    estimate_y, estimate_x = particlefilter.estimate()

    #draw the estimate
    estimate_ax.set_offsets([estimate_y,estimate_x])

    #update text
    text_ax.set_text('real pos.: %i,%i' %(real_x,real_y))
    text2_ax.set_text('estimated pos.: %i,%i' %(estimate_x,estimate_y))

    plt.draw()
    plt.pause(0.2) #update image every .. seconds
    plt.savefig('./frames/frame%05d.png'%framecounter)
    print('Frame number: %d/%d' %(framecounter,K))
    framecounter +=1

plt.close()

print('converting to mp4!')
os.system("ffmpeg -y -r 5 -i ./frames/frame%05d.png -c:v libx264 -vf fps=5 particlefilter.mp4")