chore: improve Norfair parameters for examples

app.py

@@ -1,5 +1,4 @@
 import os
-import tempfile
 
 import gradio as gr
 

demo_utils/configuration.py

@@ -1,7 +1,12 @@
-DISTANCE_THRESHOLD_BBOX: float = 3.33
-DISTANCE_THRESHOLD_CENTROID: int = 30
+DISTANCE_THRESHOLD_CENTROID: int = 0.08
 MAX_DISTANCE: int = 10000
 
 models_path = {"YOLOv7": "custom_models/yolov7.pt", "YOLOv7 Tiny": "custom_models/yolov7-tiny.pt"}
 
 style = {"Bounding box": "bbox", "Centroid": "centroid"}
+
+examples = {
+    "soccer": {"distance_threshold": 0.1, "absolute_path": True, "classes": [0]},
+    "oxford_town_center": {"distance_threshold": 0.05, "absolute_path": False, "classes": [0]},
+    "traffic_1_A": {"distance_threshold": 0.1, "absolute_path": False, "classes": [2, 3, 5, 7]},
+}

demo_utils/distance_function.py

@@ -1,54 +0,0 @@
-import numpy as np
-import torch
-import torchvision.ops.boxes as bops
-
-from demo_utils.configuration import MAX_DISTANCE
-
-
-def euclidean_distance(detection, tracked_object):
-    return np.linalg.norm(detection.points - tracked_object.estimate)
-
-
-def iou_pytorch(detection, tracked_object):
-    # Slower but simplier version of iou
-
-    detection_points = np.concatenate([detection.points[0], detection.points[1]])
-    tracked_object_points = np.concatenate([tracked_object.estimate[0], tracked_object.estimate[1]])
-
-    box_a = torch.tensor([detection_points], dtype=torch.float)
-    box_b = torch.tensor([tracked_object_points], dtype=torch.float)
-    iou = bops.box_iou(box_a, box_b)
-
-    # Since 0 <= IoU <= 1, we define 1/IoU as a distance.
-    # Distance values will be in [1, inf)
-    return np.float(1 / iou if iou else MAX_DISTANCE)
-
-
-def iou(detection, tracked_object):
-    # Detection points will be box A
-    # Tracked objects point will be box B.
-
-    box_a = np.concatenate([detection.points[0], detection.points[1]])
-    box_b = np.concatenate([tracked_object.estimate[0], tracked_object.estimate[1]])
-
-    x_a = max(box_a[0], box_b[0])
-    y_a = max(box_a[1], box_b[1])
-    x_b = min(box_a[2], box_b[2])
-    y_b = min(box_a[3], box_b[3])
-
-    # Compute the area of intersection rectangle
-    inter_area = max(0, x_b - x_a + 1) * max(0, y_b - y_a + 1)
-
-    # Compute the area of both the prediction and tracker
-    # rectangles
-    box_a_area = (box_a[2] - box_a[0] + 1) * (box_a[3] - box_a[1] + 1)
-    box_b_area = (box_b[2] - box_b[0] + 1) * (box_b[3] - box_b[1] + 1)
-
-    # Compute the intersection over union by taking the intersection
-    # area and dividing it by the sum of prediction + tracker
-    # areas - the interesection area
-    iou = inter_area / float(box_a_area + box_b_area - inter_area)
-
-    # Since 0 <= IoU <= 1, we define 1/IoU as a distance.
-    # Distance values will be in [1, inf)
-    return 1 / iou if iou else (MAX_DISTANCE)

demo_utils/files.py

@@ -1,8 +0,0 @@
-import os
-from typing import List
-
-
-def get_files(folder: str) -> List:
-    files_list = [file.split(".")[0] for file in os.listdir(folder)]
-
-    return files_list

examples/{traffic_1_C.mp4 → traffic_1_A.mp4}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cdd8096f71a7c650410de01507655745eef1cb3dbbc2bc1536664a675cec6ebb
-size 25907576
+oid sha256:fec3593a2a30bbcf0c614597dbd64ce5065d27d4ca0a13f3b444a44b5cc289ef
+size 32455820

inference.py

@@ -1,20 +1,16 @@
-import os
-import tempfile
-
 import numpy as np
 from norfair import AbsolutePaths, Paths, Tracker, Video
 from norfair.camera_motion import HomographyTransformationGetter, MotionEstimator
+from norfair.distances import create_normalized_mean_euclidean_distance
 
 from custom_models import YOLO, yolo_detections_to_norfair_detections
 from demo_utils.configuration import (
-    DISTANCE_THRESHOLD_BBOX,
     DISTANCE_THRESHOLD_CENTROID,
+    examples,
     models_path,
     style,
 )
-from demo_utils.distance_function import euclidean_distance, iou
 from demo_utils.draw import center, draw
-from demo_utils.files import get_files
 
 
 def inference(
@@ -24,12 +20,10 @@ def inference(
     track_points: str = "Bounding box",
     model_threshold: float = 0.25,
 ):
-    # temp_dir = tempfile.TemporaryDirectory()
-    # output_path = temp_dir.name
-        
     coord_transformations = None
     paths_drawer = None
     fix_paths = False
+    classes = None
     track_points = style[track_points]
     model = YOLO(models_path[model])
     video = Video(input_path=input_video)
@@ -42,9 +36,6 @@ def inference(
         features[0] == 1 or (len(features) > 1 and features[1] == 1)
     )
 
-    if motion_estimation and drawing_paths:
-        fix_paths = True
-
     if motion_estimation:
         transformations_getter = HomographyTransformationGetter()
 
@@ -52,10 +43,24 @@ def inference(
             max_points=500, min_distance=7, transformations_getter=transformations_getter
         )
 
-    distance_function = iou if track_points == "bbox" else euclidean_distance
-    distance_threshold = (
-        DISTANCE_THRESHOLD_BBOX if track_points == "bbox" else DISTANCE_THRESHOLD_CENTROID
+    distance_function = create_normalized_mean_euclidean_distance(
+        video.input_height, video.input_width
     )
+    distance_threshold = DISTANCE_THRESHOLD_CENTROID
+
+    if motion_estimation and drawing_paths:
+        fix_paths = True
+
+    # Examples configuration
+    for example in examples:
+        if example not in input_video:
+            continue
+        fix_paths = examples[example]["absolute_path"]
+        distance_threshold = examples[example]["distance_threshold"]
+        classes = examples[example]["classes"]
+
+        print(f"Set config to {example}: {fix_paths} {distance_threshold} {classes}")
+        break
 
     tracker = Tracker(
         distance_function=distance_function,
@@ -66,18 +71,17 @@ def inference(
         paths_drawer = Paths(center, attenuation=0.01)
 
     if fix_paths:
-        paths_drawer = AbsolutePaths(max_history=15, thickness=2)
+        paths_drawer = AbsolutePaths(max_history=50, thickness=2)
 
     for frame in video:
         yolo_detections = model(
-            frame, conf_threshold=model_threshold, iou_threshold=0.45, image_size=720
+            frame,
+            conf_threshold=model_threshold,
+            iou_threshold=0.45,
+            image_size=720,
+            classes=classes,
         )
 
-        mask = np.ones(frame.shape[:2], frame.dtype)
-
-        if motion_estimation:
-            coord_transformations = motion_estimator.update(frame, mask)
-
         detections = yolo_detections_to_norfair_detections(
             yolo_detections, track_points=track_points
         )
@@ -86,6 +90,15 @@ def inference(
             detections=detections, coord_transformations=coord_transformations
         )
 
+        mask = np.ones(frame.shape[:2], frame.dtype)
+        if track_points == "bbox":
+            for det in detections:
+                i = det.points.astype(int)
+                mask[i[0, 1] : i[1, 1], i[0, 0] : i[1, 0]] = 0
+
+        if motion_estimation:
+            coord_transformations = motion_estimator.update(frame, mask)
+
         frame = draw(
             paths_drawer,
             track_points,
@@ -99,7 +112,5 @@ def inference(
 
     base_file_name = input_video.split("/")[-1].split(".")[0]
     file_name = base_file_name + "_out.mp4"
-    
-    # return os.path.join(output_path, file_name)
 
     return file_name