Add ground truth tracking video

crabs/tracker/evaluate_tracker.py

@@ -1,17 +1,15 @@
-import csv
 import logging
 from typing import Any, Dict, Optional, Tuple
 
 import numpy as np
 
-from crabs.tracker.utils.tracking import extract_bounding_box_info
+from crabs.tracker.utils.tracking import get_predicted_data
 
 
 class TrackerEvaluate:
     def __init__(
         self,
         gt_dir: str,
-        predicted_boxes_id: list[np.ndarray],
         iou_threshold: float,
     ):
         """
@@ -21,93 +19,12 @@ def __init__(
         ----------
         gt_dir : str
             Directory path of the ground truth CSV file.
-        tracked_list : List[np.ndarray]
-            A list where each element is a numpy array representing tracked objects in a frame.
-            Each numpy array has shape (N, 5), where N is the number of objects.
-            The columns are [x1, y1, x2, y2, id], where (x1, y1) and (x2, y2)
-            define the bounding box and id is the object ID.
         iou_threshold : float
             Intersection over Union (IoU) threshold for evaluating tracking performance.
         """
         self.gt_dir = gt_dir
-        self.predicted_boxes_id = predicted_boxes_id
         self.iou_threshold = iou_threshold
 
-    def get_predicted_data(self) -> Dict[int, Dict[str, Any]]:
-        """
-        Convert predicted bounding box and ID into a dictionary organized by frame number.
-
-        Returns
-        -------
-        Dict[int, Dict[str, Any]]:
-            A dictionary where the key is the frame number and the value is another dictionary containing:
-            - 'bbox': A numpy array with shape (N, 4) containing coordinates of the bounding boxes
-            [x, y, x + width, y + height] for every object in the frame.
-            - 'id': A numpy array containing the IDs of the tracked objects.
-        """
-        predicted_dict: Dict[int, Dict[str, Any]] = {}
-
-        for frame_number, frame_data in enumerate(self.predicted_boxes_id):
-            if frame_data.size == 0:
-                continue
-
-            bboxes = frame_data[:, :4]
-            ids = frame_data[:, 4]
-
-            predicted_dict[frame_number] = {"bbox": bboxes, "id": ids}
-
-        return predicted_dict
-
-    def get_ground_truth_data(self) -> Dict[int, Dict[str, Any]]:
-        """
-        Extract ground truth bounding box data from a CSV file and organize it by frame number.
-
-        Returns
-        -------
-        Dict[int, Dict[str, Any]]:
-            A dictionary where the key is the frame number and the value is another dictionary containing:
-            - 'bbox': A numpy arrays with shape of (N, 4) containing coordinates of the bounding box
-                [x, y, x + width, y + height] for every crabs in the frame.
-            - 'id': The ground truth ID
-        """
-        with open(self.gt_dir, "r") as csvfile:
-            csvreader = csv.reader(csvfile)
-            next(csvreader)  # Skip the header row
-            ground_truth_data = [
-                extract_bounding_box_info(row) for row in csvreader
-            ]
-
-        # Format as a dictionary with key = frame number
-        ground_truth_dict: dict = {}
-        for data in ground_truth_data:
-            frame_number = data["frame_number"]
-            bbox = np.array(
-                [
-                    data["x"],
-                    data["y"],
-                    data["x"] + data["width"],
-                    data["y"] + data["height"],
-                ],
-                dtype=np.float32,
-            )
-            track_id = int(float(data["id"]))
-
-            if frame_number not in ground_truth_dict:
-                ground_truth_dict[frame_number] = {"bbox": [], "id": []}
-
-            ground_truth_dict[frame_number]["bbox"].append(bbox)
-            ground_truth_dict[frame_number]["id"].append(track_id)
-
-        # format as numpy arrays
-        for frame_number in ground_truth_dict:
-            ground_truth_dict[frame_number]["bbox"] = np.array(
-                ground_truth_dict[frame_number]["bbox"], dtype=np.float32
-            )
-            ground_truth_dict[frame_number]["id"] = np.array(
-                ground_truth_dict[frame_number]["id"], dtype=np.float32
-            )
-        return ground_truth_dict
-
     def calculate_iou(self, box1: np.ndarray, box2: np.ndarray) -> float:
         """
         Calculate IoU (Intersection over Union) of two bounding boxes.
@@ -339,12 +256,11 @@ def evaluate_tracking(
 
         return mota_values
 
-    def run_evaluation(self) -> None:
+    def run_evaluation(self, predicted_boxes_id, ground_truth_dict) -> None:
         """
         Run evaluation of tracking based on tracking ground truth.
         """
-        predicted_dict = self.get_predicted_data()
-        ground_truth_dict = self.get_ground_truth_data()
+        predicted_dict = get_predicted_data(predicted_boxes_id)
         mota_values = self.evaluate_tracking(ground_truth_dict, predicted_dict)
         overall_mota = np.mean(mota_values)
         logging.info("Overall MOTA: %f" % overall_mota)

crabs/tracker/track_video.py

@@ -20,7 +20,7 @@
     release_video,
     save_required_output,
 )
-from crabs.tracker.utils.tracking import prep_sort
+from crabs.tracker.utils.tracking import get_ground_truth_data, prep_sort
 
 
 class Tracking:
@@ -89,7 +89,7 @@ def prep_outputs(self):
             self.tracking_output_dir,
         ) = prep_csv_writer(self.args.output_dir, self.video_file_root)
 
-        if self.args.save_video:
+        if self.args.gt_path or self.args.save_video:
             frame_width = int(self.video.get(cv2.CAP_PROP_FRAME_WIDTH))
             frame_height = int(self.video.get(cv2.CAP_PROP_FRAME_HEIGHT))
             cap_fps = self.video.get(cv2.CAP_PROP_FPS)
@@ -153,11 +153,19 @@ def run_tracking(self):
         Run object detection + tracking on the video frames.
         """
         # If we pass ground truth: check the path exist
-        if self.args.gt_path and not os.path.exists(self.args.gt_path):
+        if not os.path.exists(str(self.args.gt_path)):
             logging.info(
                 f"Ground truth file {self.args.gt_path} does not exist. Exiting..."
             )
             return
+        else:
+            evaluation = TrackerEvaluate(
+                self.args.gt_path,
+                self.config["iou_threshold"],
+            )
+            ground_truth_dict = get_ground_truth_data(
+                self.args.gt_path,
+            )
 
         # initialisation
         frame_idx = 0
@@ -203,18 +211,14 @@ def run_tracking(self):
                 frame,
                 frame_idx + 1,
                 pred_scores,
+                ground_truth_dict,
             )
 
             # update frame number
             frame_idx += 1
 
         if self.args.gt_path:
-            evaluation = TrackerEvaluate(
-                self.args.gt_path,
-                self.tracked_bbox_id,
-                self.config["iou_threshold"],
-            )
-            evaluation.run_evaluation()
+            evaluation.run_evaluation(self.tracked_bbox_id, ground_truth_dict)
 
         # Close input video
         self.video.release()

crabs/tracker/utils/io.py

@@ -2,6 +2,7 @@
 import os
 from datetime import datetime
 from pathlib import Path
+from typing import Any, Dict, Optional
 
 import cv2
 import numpy as np
@@ -108,6 +109,7 @@ def save_required_output(
     frame: np.ndarray,
     frame_number: int,
     pred_scores: np.ndarray,
+    ground_truth_dict: Optional[Dict[int, Dict[str, Any]]] = None,
 ) -> None:
     """
     Handle the output based on argument options.
@@ -134,6 +136,8 @@ def save_required_output(
         The frame number.
     pred_scores : np.ndarray
         The prediction score from detector
+    ground_truth_dict : dict
+        Dictionary containing ground truth bounding boxes and IDs for each frame, organized by frame number.
     """
     frame_name = f"{video_file_root}_frame_{frame_number:08d}.png"
 
@@ -150,7 +154,34 @@ def save_required_output(
             frame_number,
         )
 
-    if save_video:
+    if ground_truth_dict and frame_number in ground_truth_dict:
+        frame_copy = frame.copy()
+        for bbox, obj_id in zip(
+            ground_truth_dict[frame_number]["bbox"],
+            ground_truth_dict[frame_number]["id"],
+        ):
+            x1, y1, x2, y2 = map(int, bbox)
+            draw_bbox(
+                frame_copy,
+                (x1, y1),
+                (x2, y2),
+                (0, 255, 0),  # Green for ground truth
+                f"GT ID: {int(obj_id)}",
+            )
+
+        # Draw tracked bounding boxes and IDs
+        for bbox in tracked_boxes:
+            xmin, ymin, xmax, ymax, obj_id = bbox
+            draw_bbox(
+                frame_copy,
+                (int(xmin), int(ymin)),
+                (int(xmax), int(ymax)),
+                (0, 0, 255),  # Red for predictions
+                f"Pred ID: {int(obj_id)}",
+            )
+        video_output.write(frame_copy)
+
+    elif save_video:
         frame_copy = frame.copy()
         for bbox in tracked_boxes:
             xmin, ymin, xmax, ymax, id = bbox

crabs/tracker/utils/tracking.py

@@ -1,3 +1,4 @@
+import csv
 import json
 import logging
 from pathlib import Path
@@ -152,3 +153,80 @@ def prep_sort(prediction: dict, score_threshold: float) -> np.ndarray:
             pred_sort.append(bbox)
 
     return np.asarray(pred_sort)
+
+
+def get_predicted_data(predicted_boxes_id) -> Dict[int, Dict[str, Any]]:
+    """
+    Convert predicted bounding box and ID into a dictionary organized by frame number.
+
+    Returns
+    -------
+    Dict[int, Dict[str, Any]]:
+        A dictionary where the key is the frame number and the value is another dictionary containing:
+        - 'bbox': A numpy array with shape (N, 4) containing coordinates of the bounding boxes
+        [x, y, x + width, y + height] for every object in the frame.
+        - 'id': A numpy array containing the IDs of the tracked objects.
+    """
+    predicted_dict: Dict[int, Dict[str, Any]] = {}
+
+    for frame_idx, frame_data in enumerate(predicted_boxes_id):
+        if frame_data.size == 0:
+            continue
+
+        bboxes = frame_data[:, :4]
+        ids = frame_data[:, 4]
+
+        predicted_dict[frame_idx] = {"bbox": bboxes, "id": ids}
+
+    return predicted_dict
+
+
+def get_ground_truth_data(gt_dir) -> Dict[int, Dict[str, Any]]:
+    """
+    Extract ground truth bounding box data from a CSV file and organize it by frame number.
+
+    Returns
+    -------
+    Dict[int, Dict[str, Any]]:
+        A dictionary where the key is the frame number and the value is another dictionary containing:
+        - 'bbox': A numpy arrays with shape of (N, 4) containing coordinates of the bounding box
+            [x, y, x + width, y + height] for every crabs in the frame.
+        - 'id': The ground truth ID
+    """
+    with open(gt_dir, "r") as csvfile:
+        csvreader = csv.reader(csvfile)
+        next(csvreader)  # Skip the header row
+        ground_truth_data = [
+            extract_bounding_box_info(row) for row in csvreader
+        ]
+
+    # Format as a dictionary with key = frame number
+    ground_truth_dict: dict = {}
+    for data in ground_truth_data:
+        frame_idx = data["frame_number"]
+        bbox = np.array(
+            [
+                data["x"],
+                data["y"],
+                data["x"] + data["width"],
+                data["y"] + data["height"],
+            ],
+            dtype=np.float32,
+        )
+        track_id = int(float(data["id"]))
+
+        if frame_idx not in ground_truth_dict:
+            ground_truth_dict[frame_idx] = {"bbox": [], "id": []}
+
+        ground_truth_dict[frame_idx]["bbox"].append(bbox)
+        ground_truth_dict[frame_idx]["id"].append(track_id)
+
+        # format as numpy arrays
+    for frame_idx in ground_truth_dict:
+        ground_truth_dict[frame_idx]["bbox"] = np.array(
+            ground_truth_dict[frame_idx]["bbox"], dtype=np.float32
+        )
+        ground_truth_dict[frame_idx]["id"] = np.array(
+            ground_truth_dict[frame_idx]["id"], dtype=np.float32
+        )
+    return ground_truth_dict