Post proc code cleaning

social-distancing/config-skeleton.ini

@@ -10,3 +10,7 @@ ImageSize: 300,300,3
 ModelPath: 
 ClassID: 0
 MinScore: 0.25
+
+[PostProcessor]
+MaxTrackFrame: 5
+NMSThreshold: 0.98

social-distancing/libs/Core.py

@@ -1,10 +1,9 @@
-import os, re
 import time
-
-import numpy as np
 import cv2 as cv
-from scipy.spatial import distance as dist
+import numpy as np
 from libs.centroid_object_tracker import CentroidTracker
+from scipy.spatial import distance as dist
+
 
 class Distancing:
 
@@ -14,7 +13,8 @@ def __init__(self, config):
         self.detector = None
         self.device = self.config.get_section_dict('Detector')['Device']
         self.running_video = False
-        self.tracker = CentroidTracker(maxDisappeared=5)
+        self.tracker = CentroidTracker(
+            maxDisappeared=int(self.config.get_section_dict("PostProcessor")["MaxTrackFrame"]))
         if self.device == 'Jetson':
             from libs.detectors.jetson.Detector import Detector
             self.detector = Detector(self.config)
@@ -25,7 +25,7 @@ def __init__(self, config):
             self.detector = None
 
         self.image_size = [int(i) for i in self.config.get_section_dict('Detector')['ImageSize'].split(',')]
-        
+
         if self.device != 'Dummy':
             print('Device is: ', self.device)
             print('Detector is: ', self.detector.name)
@@ -39,22 +39,22 @@ def __process(self, cv_image):
         return object_list list of  dict for each obj, 
         obj["bbox"] is normalized coordinations for [x0, y0, x1, y1] of box
         """
-        if self.device == 'Dummy': 
+        if self.device == 'Dummy':
             return cv_image, [], None
 
         resized_image = cv.resize(cv_image, tuple(self.image_size[:2]))
         rgb_resized_image = cv.cvtColor(resized_image, cv.COLOR_BGR2RGB)
         tmp_objects_list = self.detector.inference(rgb_resized_image)
-        hscale = cv_image.shape[0]/resized_image.shape[0]
-        wscale = cv_image.shape[1]/resized_image.shape[1]
+        hscale = cv_image.shape[0] / resized_image.shape[0]
+        wscale = cv_image.shape[1] / resized_image.shape[1]
 
         for obj in tmp_objects_list:
             box = obj["bbox"]
             x0 = box[1]
             y0 = box[0]
             x1 = box[3]
             y1 = box[2]
-            obj["centroid"] = [(x0+x1)/2, (y0+y1)/2, x1 - x0, y1 - y0]
+            obj["centroid"] = [(x0 + x1) / 2, (y0 + y1) / 2, x1 - x0, y1 - y0]
             obj["bbox"] = [x0, y0, x1, y1]
 
         objects_list, distancings = self.calculate_distancing(tmp_objects_list)
@@ -68,93 +68,121 @@ def process_video(self, video_uri):
             print('opened video ', video_uri)
         else:
             print('failed to load video ', video_uri)
-            return 
+            return
 
         while input_cap.isOpened() and self.running_video:
             _, cv_image = input_cap.read()
             _, objects, distancings = self.__process(cv_image)
             self.ui.update(cv_image, objects, distancings)
-            time.sleep(0.030) 
 
         input_cap.release()
         self.running_video = False
 
     def process_image(self, image_path):
         cv_image = cv.imread(image_path)
         _, objects, distancings = self.__process(cv_image)
-        self.ui.update(cv_image, objects, distancings) 
+        self.ui.update(cv_image, objects, distancings)
 
     def calculate_distancing(self, objects_list):
+        """
+        this function post-process the raw boxes of object detector and calculate a distance matrix
+        for detected bounding boxes.
+        post processing is consist of:
+        1. omitting large boxes by filtering boxes which are biger than the 1/4 of the size the image.
+        2. omitting duplicated boxes by applying an auxilary non-maximum-suppression.
+        3. apply a simple object tracker to make the detection more robust.
+
+        params:
+        object_list: a list of dictionaries. each dictionary has attributes of a detected object such as
+        "id", "centroid" (a tuple of the normalized centroid coordinates (cx,cy,w,h) of the box) and "bbox" (a tuple
+        of the normalized (xmin,ymin,xmax,ymax) coordinate of the box)
+
+        returns:
+        object_list: the post processed version of the input
+        distances: a NxN ndarray which i,j element is distance between i-th and l-th bounding box
+
+        """
         new_objects_list = self.ignore_large_boxes(objects_list)
-        new_objects_list = self.non_max_suppression_fast(new_objects_list, 0.98)
+        new_objects_list = self.non_max_suppression_fast(new_objects_list,
+                                                         float(self.config.get_section_dict("PostProcessor")[
+                                                                   "NMSThreshold"]))
         tracked_boxes = self.tracker.update(new_objects_list)
         new_objects_list = [tracked_boxes[i] for i in tracked_boxes.keys()]
         for i, item in enumerate(new_objects_list):
             item["id"] = item["id"].split("-")[0] + "-" + str(i)
 
-        centroids = np.array( [obj["centroid"] for obj in new_objects_list] ) 
+        centroids = np.array([obj["centroid"] for obj in new_objects_list])
         distances = dist.cdist(centroids, centroids)
         return new_objects_list, distances
 
     @staticmethod
     def ignore_large_boxes(object_list):
+
+        """
+        filtering boxes which are biger than the 1/4 of the size the image
+        params:
+            object_list: a list of dictionaries. each dictionary has attributes of a detected object such as
+            "id", "centroid" (a tuple of the normalized centroid coordinates (cx,cy,w,h) of the box) and "bbox" (a tuple
+            of the normalized (xmin,ymin,xmax,ymax) coordinate of the box)
+        returns:
+        object_list: input object list without large boxes
+        """
         large_boxes = []
         for i in range(len(object_list)):
             if (object_list[i]["centroid"][2] * object_list[i]["centroid"][3]) > 0.25:
                 large_boxes.append(i)
-        updated_object_list = [j for i,j in enumerate(object_list) if i not in large_boxes]
+        updated_object_list = [j for i, j in enumerate(object_list) if i not in large_boxes]
         return updated_object_list
 
     @staticmethod
     def non_max_suppression_fast(object_list, overlapThresh):
+
+        """
+        omitting duplicated boxes by applying an auxilary non-maximum-suppression.
+        params:
+        object_list: a list of dictionaries. each dictionary has attributes of a detected object such
+        "id", "centroid" (a tuple of the normalized centroid coordinates (cx,cy,w,h) of the box) and "bbox" (a tuple
+        of the normalized (xmin,ymin,xmax,ymax) coordinate of the box)
+
+        overlapThresh: threshold of minimum IoU of to detect two box as duplicated.
+
+        returns:
+        object_list: input object list without duplicated boxes
+        """
         # if there are no boxes, return an empty list
         boxes = np.array([item["centroid"] for item in object_list])
+        corners = np.array([item["bbox"] for item in object_list])
         if len(boxes) == 0:
             return []
-        # if the bounding boxes integers, convert them to floats --
-        # this is important since we'll be doing a bunch of divisions
         if boxes.dtype.kind == "i":
             boxes = boxes.astype("float")
-        # initialize the list of picked indexes 
+        # initialize the list of picked indexes
         pick = []
-        # grab the coordinates of the bounding boxes
-        cy = boxes[:,1]
-        cx = boxes[:,0]
-        h = boxes[:,3]
-        w = boxes[:,2]
-        x1 = cx - (w/2)
-        x2 = cx + (w/2)
-        y1 = cy - (h/2)
-        y2 = cy + (h/2)
-        # compute the area of the bounding boxes and sort the bounding
-        # boxes by the bottom-right y-coordinate of the bounding box
+        cy = boxes[:, 1]
+        cx = boxes[:, 0]
+        h = boxes[:, 3]
+        w = boxes[:, 2]
+        x1 = corners[:, 0]
+        x2 = corners[:, 2]
+        y1 = corners[:, 1]
+        y2 = corners[:, 3]
         area = (h + 1) * (w + 1)
-        idxs = np.argsort(cy + (h/2))
-        # keep looping while some indexes still remain in the indexes
-        # list
+        idxs = np.argsort(cy + (h / 2))
         while len(idxs) > 0:
-            # grab the last index in the indexes list and add the
-            # index value to the list of picked indexes
             last = len(idxs) - 1
             i = idxs[last]
             pick.append(i)
-            # find the largest (x, y) coordinates for the start of
-            # the bounding box and the smallest (x, y) coordinates
-            # for the end of the bounding box
             xx1 = np.maximum(x1[i], x1[idxs[:last]])
             yy1 = np.maximum(y1[i], y1[idxs[:last]])
             xx2 = np.minimum(x2[i], x2[idxs[:last]])
             yy2 = np.minimum(y2[i], y2[idxs[:last]])
-            # compute the width and height of the bounding box
+
             w = np.maximum(0, xx2 - xx1 + 1)
             h = np.maximum(0, yy2 - yy1 + 1)
             # compute the ratio of overlap
             overlap = (w * h) / area[idxs[:last]]
             # delete all indexes from the index list that have
             idxs = np.delete(idxs, np.concatenate(([last],
-                np.where(overlap > overlapThresh)[0])))
-        # return only the bounding boxes that were picked using the
-        #integer data 
-        updated_object_list = [j for i,j in enumerate(object_list) if i in pick]
+                                                   np.where(overlap > overlapThresh)[0])))
+        updated_object_list = [j for i, j in enumerate(object_list) if i in pick]
         return updated_object_list
-

social-distancing/libs/centroid_object_tracker.py

@@ -1,71 +1,83 @@
 # import the necessary packages
-from scipy.spatial import distance as dist
 from collections import OrderedDict
+
 import numpy as np
+from scipy.spatial import distance as dist
+
+
+class CentroidTracker:
+    """
+    a simple object tracker based on Euclidian distance of bounding boxes centroid of two consecutive frames.
+    if a box is losted betweeb two frames the tracker keep the box for next maxDisappeared frames.
+    """
+
+    def __init__(self, maxDisappeared=50):
 
-class CentroidTracker():
-	def __init__(self, maxDisappeared=50):
+        """
+        maxDisappeared:if a box is losted betweeb two frames the tracker keep the box for next maxDisappeared frames.
 
-		self.nextObjectID = 0
-		self.objects = OrderedDict()
-		self.disappeared = OrderedDict()
-		self.maxDisappeared = maxDisappeared
+        """
+        self.nextObjectID = 0
+        self.objects = OrderedDict()
+        self.disappeared = OrderedDict()
+        self.maxDisappeared = maxDisappeared
 
-	def register(self, object_item):
-		self.objects[self.nextObjectID] = object_item
-		self.disappeared[self.nextObjectID] = 0
-		self.nextObjectID += 1
+    def register(self, object_item):
+        self.objects[self.nextObjectID] = object_item
+        self.disappeared[self.nextObjectID] = 0
+        self.nextObjectID += 1
 
-	def deregister(self, objectID):
-		del self.objects[objectID]
-		del self.disappeared[objectID]
+    def deregister(self, objectID):
+        del self.objects[objectID]
+        del self.disappeared[objectID]
 
-	def update(self, object_list):
-		if len(object_list) == 0:
-			for objectID in list(self.disappeared.keys()):
-				self.disappeared[objectID] += 1
-				if self.disappeared[objectID] > self.maxDisappeared:
-					self.deregister(objectID)
-			return self.objects
-		inputCentroids = np.zeros((len(object_list), 2))
-		for i, object_item in enumerate(object_list):
-			inputCentroids[i] = (object_item["centroid"][0], object_item["centroid"][1])
-		if len(self.objects) == 0:
-			for i in range(0, len(inputCentroids)):
-				self.register(object_list[i])
-		else:
-			objectIDs = list(self.objects.keys())
-			objectCentroids = [object_item["centroid"][0:2] for object_item in self.objects.values()]
-			D = dist.cdist(np.array(objectCentroids), inputCentroids)
-			rows = D.min(axis=1).argsort()
-			cols = D.argmin(axis=1)[rows]
-			usedRows = set()
-			usedCols = set()
-			for (row, col) in zip(rows, cols):
-				if row in usedRows or col in usedCols:
-					continue
-				objectID = objectIDs[row]
-				self.objects[objectID] = object_list[col]
-				self.disappeared[objectID] = 0
-				usedRows.add(row)
-				usedCols.add(col)
+    def update(self, object_list):
+        if len(object_list) == 0:
+            for objectID in list(self.disappeared.keys()):
+                self.disappeared[objectID] += 1
+                if self.disappeared[objectID] > self.maxDisappeared:
+                    self.deregister(objectID)
+            return self.objects
+        inputCentroids = np.zeros((len(object_list), 2))
+        for i, object_item in enumerate(object_list):
+            inputCentroids[i] = (object_item["centroid"][0], object_item["centroid"][1])
+        if len(self.objects) == 0:
+            for i in range(0, len(inputCentroids)):
+                self.register(object_list[i])
+        else:
+            objectIDs = list(self.objects.keys())
+            objectCentroids = [object_item["centroid"][0:2] for object_item in self.objects.values()]
+            D = dist.cdist(np.array(objectCentroids), inputCentroids)
+            rows = D.min(axis=1).argsort()
+            cols = D.argmin(axis=1)[rows]
+            usedRows = set()
+            usedCols = set()
+            for (row, col) in zip(rows, cols):
+                if row in usedRows or col in usedCols:
+                    continue
+                objectID = objectIDs[row]
+                self.objects[objectID] = object_list[col]
+                self.disappeared[objectID] = 0
+                usedRows.add(row)
+                usedCols.add(col)
 
-			unusedRows = set(range(0, D.shape[0])).difference(usedRows)
-			unusedCols = set(range(0, D.shape[1])).difference(usedCols)
+            unusedRows = set(range(0, D.shape[0])).difference(usedRows)
+            unusedCols = set(range(0, D.shape[1])).difference(usedCols)
 
-			if D.shape[0] >= D.shape[1]:
-				biggest_existing_id = int(object_list[-1]["id"].split("-")[-1])
+            if D.shape[0] >= D.shape[1]:
+                biggest_existing_id = int(object_list[-1]["id"].split("-")[-1])
 
-				for row in unusedRows:             
-					objectID = objectIDs[row]
-					self.objects[objectID]["id"] = f"0-{biggest_existing_id+1}"
-					self.disappeared[objectID] += 1
-					biggest_existing_id += 1
-					if self.disappeared[objectID] > self.maxDisappeared:
-						self.deregister(objectID)
+                for row in unusedRows:
+                    objectID = objectIDs[row]
+                    self.objects[objectID]["id"] = self.objects[objectID]["id"].split("-")[0] + "-" + str(
+                        biggest_existing_id + 1)
+                    self.disappeared[objectID] += 1
+                    biggest_existing_id += 1
+                    if self.disappeared[objectID] > self.maxDisappeared:
+                        self.deregister(objectID)
 
-			else:
-				for col in unusedCols: 
-					self.register(object_list[col])
+            else:
+                for col in unusedCols:
+                    self.register(object_list[col])
 
-		return self.objects
+        return self.objects