Class that used in car counting during CV traffic analysis

vehicle_counter.py

import logging
import math
import numpy as np
from scipy import spatial
import cv2

# ============================================================================

CAR_COLOURS = [ (0,0,255), (0,106,255), (0,216,255), (0,255,182), (0,255,76)
    , (144,255,0), (255,255,0), (255,148,0), (255,0,178), (220,0,255) ]

class VehicleCounter(object):

    def __init__(self, shape, divider):
        self.log = logging.getLogger("vehicle_counter")

        self.height, self.width = shape
        self.divider = divider

        self.vehicle_count = 0
        self.path_size = 10
        self.pathes = []
        self.max_dst = 30


    def distance(self, x,y,type='euclidian'):
        if type == 'euclidian':
            return math.sqrt((x[0]-y[0])**2+(x[1]-y[1])**2)

    def update_count(self, matches, output_image = None):
        self.log.debug("Updating count using %d matches...", len(matches))
        self.update_pathes(matches)
        self.draw_pathes(output_image)
        
        cv2.putText(output_image, ("%02d" % self.vehicle_count), (150, 10)
                , cv2.FONT_HERSHEY_PLAIN, 0.7, (127, 255, 255), 1)
                
        self.log.debug("VEHICLES FOUND: %s" % self.vehicle_count)
        
    def update_pathes(self, matches):
        matches = np.array(matches)[:,-1]
        matches = matches.tolist()
        if not matches: 
            return
            
        # add new points if pathes is empty
        if not self.pathes:
            for match in matches:
                self.pathes.append([match])
            
        else:
            # link new points with old pathes based on minimum distance between points
            count = len(matches)
            new_pathes = []
            for i,path in enumerate(self.pathes):
                _min = 999999
                _match = None
                for match in matches:
                    d = self.distance(match,path[-1])
                    if d < _min:
                        _min = d
                        _match = match

                if _match and _min <= self.max_dst:
                    matches.remove(_match)
                    path.append(_match)
                    
                # drop pathes that not linked
                if _min <= self.max_dst:
                    new_pathes.append(path)  
               
            self.pathes = new_pathes
                              
            # add new pathes
            if len(matches):
                for match in matches:
                    # do not add matches that already should be counted
                    if match[1] > self.divider:
                        continue
                    self.pathes.append([match])
               
                    
        # save only last N points in path
        for i,_ in enumerate(self.pathes):
            self.pathes[i] = self.pathes[i][self.path_size*-1:]
            
        # count vechiles and drop counted pathes:
        new_pathes = []
        for i,path in enumerate(self.pathes):
            d = path[-2:]
            if len(d) < 2:
                new_pathes.append(path)
                continue
            
            if d[0][1] <= self.divider and d[1][1] > self.divider:
                self.vehicle_count += 1
            else:
                new_pathes.append(path)  
            
        self.pathes = new_pathes      
        
    def draw_pathes(self, output_image=None):
        if not output_image.any():
            return
            
        for i, path in enumerate(self.pathes):
            for point in path:
                cv2.circle(output_image, point, 2, CAR_COLOURS[i], -1)
                cv2.polylines(output_image, [np.int32(path)]
                , False, CAR_COLOURS[i], 1)
                
        
# ============================================================================