ellipse_tracking.py

#!/usr/bin/env python2
# -*- coding: utf-8 -*-
#%%
"""
Created on Sat May 20 17:42:52 2017

Fits shapes to ellipses in image
Uses a kalman filter to track shape centers
"""
import load_pedestrian as ld
from MOGbkgsubtract import mog_bkg_subtractor
from frame_diff_bkgsubtract import frame_diff_bkgsubtract
from scipy.signal import savgol_filter

import cv2
import matplotlib.pyplot as plt
import numpy as np
import math as math
import matplotlib.cm as cm
import pandas as pd
from filterpy.kalman import KalmanFilter
from heatmappy import Heatmapper
from PIL import Image

class Track:
    """
    A person and a list of historical points
    Points are generated by feeding measurements into
    a kalman filter
    """
    def __init__(self, cr, process_noise=0.0005):
        self.kf = self.intialize_kalman_filter(cr, process_noise)
        self.estimated_points = [cr]
        self.measured_points = [cr]
        
    def intialize_kalman_filter(self, cr, process_noise):
        f1 = KalmanFilter(dim_x=4, dim_z=2)
        dt = 1.# time step
        f1.F = np.array ([[1, dt, 0,  0],
                        [0,  1, 0,  0],
                        [0,  0, 1, dt],
                        [0,  0, 0,  1]])
        f1.u = 0 # no info about rotors turning
        f1.H = np.array ([[1, 0, 0, 0], #measurement to state.
                          [0, 0, 1, 0]])
        f1.R = np.array([[5,0], #  variance in measurments
                         [0, 5]])
        f1.Q = np.eye(4) * process_noise # process noise
        f1.x = np.array([[cr[0],0,cr[1],0]]).T # initial position
        f1.P = np.eye(4) * 500 #covariance matrix
        return(f1)

    def update(self, cr, frame_idx):
        self.kf.predict()
        z = np.array([[cr[0]],[cr[1]]])
        self.kf.update(z)
        est_loc = self.kf.x
        self.estimated_points.append( (est_loc[0,0], est_loc[2,0], frame_idx ) )
        self.measured_points.append( (cr[0], cr[1], frame_idx) )

        
class points_tracker:
    def __init__(self, min_points):
        self.min_points = min_points
        self.min_idle_time = 30
        self.min_distance_from_last_point = 100
        self.smooth_window_len = 9
        self.archive_tracks = []
        self.active_tracks = []
        self.frame_idx = 0
        
    def update_tracks(self, shape_centers):
        self.add_centers_to_tracks(shape_centers)
        self.archive_old_tracks()
            
    def add_centers_to_tracks(self, shape_centers):
        if len(self.active_tracks) == 0:
            for cr in shape_centers:
                track = Track(cr)
                self.active_tracks.append(track)
        else:
            track_ends = [tr.estimated_points[-1] for tr in self.active_tracks]
            for cr in shape_centers:
                min_index, min_len = self.find_closest_end(cr, track_ends)
                if min_len < self.min_distance_from_last_point: 
                    self.active_tracks[min_index].update(cr, self.frame_idx)
                else:
                    track = Track(cr)
                    self.active_tracks.append(track)

    def find_closest_end(self, cr, track_ends):        
        min_len = 999999
        min_index = 0
        for i, end in enumerate(track_ends):
            dist = self.dist_between_points(cr, end)
            if dist < min_len:
                min_len = dist
                min_index = i
        return(min_index, min_len)

    def dist_between_points(self, p0, p1):
        return math.sqrt((p0[0] - p1[0])**2 + (p0[1] - p1[1])**2)
    
    def archive_old_tracks(self):
        for tr in self.active_tracks:
            if self.track_is_idling(tr):
                self.archive_tracks.append(tr)
                self.active_tracks.remove(tr)    

    def track_is_idling(self,tr):
        time_since_last_update = self.frame_idx - tr.estimated_points[-1][2]
        return time_since_last_update > self.min_idle_time

    def draw_tracks(self, frame_clr):
        drawable_tracks = self.get_drawable_tracks()
        for index, tr in enumerate(drawable_tracks):
            track_color = tuple(256 * x for x in cm.Spectral(index * 100 % 255)[0:3])
            cv2.polylines(frame_clr, [np.int32([tup[0:2] for tup in tr])], False, track_color)
    
    def get_drawable_tracks(self, smooth=True):
        """
        return valid archive and active tracks
        """
        drawable_tracks = []
        for track in self.active_tracks:
            if len(track.estimated_points) >= self.min_points:
                drawable_tracks.append(track.estimated_points)
        for track in self.archive_tracks:
            if len(track.estimated_points) >= self.min_points:
                drawable_tracks.append(track.estimated_points)
        if smooth:
            drawable_tracks = self.smooth_savgol(drawable_tracks)
            
        return(drawable_tracks)    

    def smooth_savgol(self, rough_tracks):
        #similar to a box window, using polynomials
        smooth_tracks= []
        for tr in rough_tracks:
            tr = np.array(tr)
            
            if(len(tr) < 2*self.smooth_window_len):
                smooth_tracks.append(tr)
                continue
            
            xnew = savgol_filter(tr[:,0], self.smooth_window_len, 1)
            ynew = savgol_filter(tr[:,1], self.smooth_window_len, 1)
            tr = list( zip(xnew, ynew, tr[:,2] ))
            smooth_tracks.append(tr)
        return(smooth_tracks)
        
    def get_drawable_tracks_npstack(self):
        tracks = self.get_drawable_tracks()
        for index, track in enumerate(tracks):
            track_num = index * np.ones((len(track),1), dtype = 'int32')
            track = np.append(track, track_num, axis=1)

            if index == 0:
                all_tracks = track
            else: 
                all_tracks = np.append(all_tracks, track, axis=0)

        return(all_tracks)

def process_frame(frame_bw, frame_clr, Pt, area=0):
    ellipses = get_ellipses(frame_bw, Pt, area)
    shape_centers = draw_ellipses_and_centers(frame_clr, ellipses)
    
    Pt.update_tracks(shape_centers)  
    Pt.draw_tracks(frame_clr)
    
    
def get_ellipses(frame_bw, Pt, area=0):
    """
    find contours and match ellipses. 
    Uses an area lower bound for what to remove.
    """

    _, contours, hierarchy = cv2.findContours(frame_bw, 1, 2)

    ellipses = []
    for contour in contours:
        if cv2.contourArea(contour) <= area:
            continue
        try:
            ellipse = cv2.fitEllipse(contour)
            ellipses.append(ellipse)
        except Exception:
            continue
    return ellipses

def draw_ellipses_and_centers(frame_clr, ellipses):
    shape_centers = []
    for ellipse in ellipses:
        cv2.ellipse(frame_clr, ellipse, (255, 255, 255), 2)
        
        center = (int(ellipse[0][0]), int(ellipse[0][1]) )
        cv2.circle(frame_clr, center, 3, (0, 0, 255), thickness=3)
        shape_centers.append( (center[0], center[1], Pt.frame_idx) )

    return shape_centers

#%%
run_name = ""
frame_idx = 0

Dataset = ld.get_dataset('hallway','large')

noise_bkg = np.load(base_dir + 'noisebkg.npy')
Bkg = frame_diff_bkgsubtract(0.0001, 30, bkg=noise_bkg)

Pt = points_tracker(5)
while True:
    frame_idx += 1
    Pt.frame_idx = frame_idx
    print(frame_idx)
    
    ret, frame_clr = Dataset.get_next_frame(frame_idx)
    if not ret:
        print("end of reel")
        break
    
    #temp skip frames to improve speed
    if not frame_idx % 3 == 0:
        continue
    
    frame_bw = Bkg.process_frame(frame_clr)
    
    

    plt.figure()
    plt.title(run_name)
    plt.imshow(frame_bw, cmap="gray")
    plt.show()


    process_frame(frame_bw, frame_clr, Pt, area=500)


    plt.figure()
    plt.title(run_name)
    plt.imshow(frame_clr)
    plt.show()

#%% Save final view and heatmap

npstack = Pt.get_drawable_tracks_npstack()
df = pd.DataFrame(npstack, columns=['x','y','time','track'])
paths = df[['x','y']]


write_name = run_name +  str(frame_idx) + "_"
cv2.imwrite(base_dir + 'outputFiles/' + write_name + 'final_img.png', frame_clr)


ret, frame_clr = Dataset.get_next_frame(frame_idx)

pil_img = Image.fromarray(frame_clr)
heatmapper = Heatmapper()
heatmap = heatmapper.heatmap_on_img(paths.values.tolist(), pil_img)
heatmap.save(base_dir + 'outputFiles/' + write_name + 'heatmap.png')