chore: improve code

app.py

@@ -10,6 +10,13 @@ dd_model = gr.Dropdown(choices=["YoloV7", "YoloV7 Tiny"], value="YoloV7", label=
 
 cb_motion_estimation = gr.Checkbox(value=True, label="Track camera movement")
 
+features = gr.CheckboxGroup(
+    choices=["Track camera movement", "Draw objects paths"],
+    value=["Track camera movement", "Draw objects paths"],
+    label="Features",
+    type="index"
+)
+
 cb_path_draw = gr.Checkbox(value=True, label="Draw objects paths")
 
 dd_track_points = gr.Dropdown(
@@ -33,8 +40,7 @@ iface = gr.Interface(
     inputs=[
         input_video,
         dd_model,
-        cb_motion_estimation,
-        cb_path_draw,
+        features,
         dd_track_points,
         slide_threshold,
     ],

inference_utils.py → custom_models/Yolo.py

@@ -1,23 +1,9 @@
-import argparse
-import glob
 import os
-from enum import Enum
 from typing import List, Optional, Union
 
-import norfair
 import numpy as np
 import torch
-import torchvision.ops.boxes as bops
-from norfair import Detection, draw_absolute_grid
-
-DISTANCE_THRESHOLD_BBOX: float = 3.33
-DISTANCE_THRESHOLD_CENTROID: int = 30
-MAX_DISTANCE: int = 10000
-
-
-models_path = {"YoloV7": "models/yolov7.pt", "YoloV7 Tiny": "models/yolov7-tiny.pt"}
-
-style = {"Bounding box": "bbox", "Centroid": "centroid"}
+from norfair import Detection
 
 
 class YOLO:
@@ -56,59 +42,6 @@ class YOLO:
         return detections
 
 
-def euclidean_distance(detection, tracked_object):
-    return np.linalg.norm(detection.points - tracked_object.estimate)
-
-
-def center(points):
-    return [np.mean(np.array(points), axis=0)]
-
-
-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)
-
-
 def yolo_detections_to_norfair_detections(
     yolo_detections: torch.tensor, track_points: str = "centroid"  # bbox or centroid
 ) -> List[Detection]:
@@ -134,35 +67,3 @@ def yolo_detections_to_norfair_detections(
             norfair_detections.append(Detection(points=bbox, scores=scores))
 
     return norfair_detections
-
-
-def clean_videos(path: str):
-    # Remove past videos
-    files = glob.glob(f"{path}/*")
-    for file in files:
-        if file.endswith(".mp4"):
-            os.remove(file)
-
-
-def draw(
-    paths_drawer,
-    track_points,
-    frame,
-    detections,
-    tracked_objects,
-    coord_transformations,
-    fix_paths,
-):
-    if track_points == "centroid":
-        norfair.draw_points(frame, detections)
-        norfair.draw_tracked_objects(frame, tracked_objects)
-    elif track_points == "bbox":
-        norfair.draw_boxes(frame, detections)
-        norfair.draw_tracked_boxes(frame, tracked_objects)
-
-    if fix_paths:
-        frame = paths_drawer.draw(frame, tracked_objects, coord_transformations)
-    elif paths_drawer is not None:
-        frame = paths_drawer.draw(frame, tracked_objects)
-
-    return frame

custom_models/__init__.py

@@ -0,0 +1 @@
+from .Yolo import YOLO, yolo_detections_to_norfair_detections

demo_utils/configuration.py

@@ -0,0 +1,7 @@
+DISTANCE_THRESHOLD_BBOX: float = 3.33
+DISTANCE_THRESHOLD_CENTROID: int = 30
+MAX_DISTANCE: int = 10000
+
+models_path = {"YoloV7": "custom_models/yolov7.pt", "YoloV7 Tiny": "custom_models/yolov7-tiny.pt"}
+
+style = {"Bounding box": "bbox", "Centroid": "centroid"}

demo_utils/distance_function.py

@@ -0,0 +1,54 @@
+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/draw.py

@@ -0,0 +1,30 @@
+import norfair
+import numpy as np
+
+
+def draw(
+    paths_drawer,
+    track_points,
+    frame,
+    detections,
+    tracked_objects,
+    coord_transformations,
+    fix_paths,
+):
+    if track_points == "centroid":
+        norfair.draw_points(frame, detections)
+        norfair.draw_tracked_objects(frame, tracked_objects)
+    elif track_points == "bbox":
+        norfair.draw_boxes(frame, detections)
+        norfair.draw_tracked_boxes(frame, tracked_objects)
+
+    if fix_paths:
+        frame = paths_drawer.draw(frame, tracked_objects, coord_transformations)
+    elif paths_drawer is not None:
+        frame = paths_drawer.draw(frame, tracked_objects)
+
+    return frame
+
+
+def center(points):
+    return [np.mean(np.array(points), axis=0)]

demo_utils/files.py

@@ -0,0 +1,10 @@
+import glob
+import os
+
+
+def clean_videos(path: str):
+    # Remove past videos
+    files = glob.glob(f"{path}/*")
+    for file in files:
+        if file.endswith(".mp4"):
+            os.remove(file)

inference.py

@@ -3,31 +3,25 @@ import glob
 import os
 
 import numpy as np
-from norfair import AbsolutePaths, FixedCamera, Paths, Tracker, Video
+from norfair import AbsolutePaths, Paths, Tracker, Video
 from norfair.camera_motion import HomographyTransformationGetter, MotionEstimator
 
-from inference_utils import (
-    YOLO,
-    center,
-    clean_videos,
-    draw,
-    euclidean_distance,
-    iou,
+from custom_models import YOLO, yolo_detections_to_norfair_detections
+from demo_utils.configuration import (
+    DISTANCE_THRESHOLD_BBOX,
+    DISTANCE_THRESHOLD_CENTROID,
     models_path,
     style,
-    yolo_detections_to_norfair_detections,
 )
-
-DISTANCE_THRESHOLD_BBOX: float = 3.33
-DISTANCE_THRESHOLD_CENTROID: int = 30
-MAX_DISTANCE: int = 10000
+from demo_utils.distance_function import euclidean_distance, iou
+from demo_utils.draw import center, draw
+from demo_utils.files import clean_videos
 
 
 def inference(
     input_video: str,
     model: str,
-    motion_estimation: bool,
-    drawing_paths: bool,
+    features: str,
     track_points: str,
     model_threshold: str,
 ):
@@ -42,6 +36,14 @@ def inference(
     model = YOLO(models_path[model])
     video = Video(input_path=input_video, output_path=output_path)
 
+    motion_estimation = len(features) > 0 and (
+        features[0] == 0 or (len(features) > 1 and features[1] == 0)
+    )
+
+    drawing_paths = len(features) > 0 and (
+        features[0] == 1 or (len(features) > 1 and features[1] == 1)
+    )
+
     if motion_estimation and drawing_paths:
         fix_paths = True
 
@@ -56,6 +58,7 @@ def inference(
     distance_threshold = (
         DISTANCE_THRESHOLD_BBOX if track_points == "bbox" else DISTANCE_THRESHOLD_CENTROID
     )
+
     tracker = Tracker(
         distance_function=distance_function,
         distance_threshold=distance_threshold,