adding scores (#207)

* adding scores * small changes how we save the score * revert config * fix test * Update tests/test_unit/test_tracking_utils.py Co-authored-by: sfmig <[email protected]> Signed-off-by: nikk-nikaznan <[email protected]> * changes based on review --------- Signed-off-by: nikk-nikaznan <[email protected]> Co-authored-by: sfmig <[email protected]>
SainsburyWellcomeCentre · Jul 4, 2024 · 492d72a · 492d72a
1 parent e12b1ec
commit 492d72a
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 17 deletions.
diff --git a/crabs/tracker/track_video.py b/crabs/tracker/track_video.py
@@ -193,6 +193,7 @@ def run_tracking(self):
 
             # predict bounding boxes
             prediction = self.get_prediction(frame)
+            pred_scores = prediction[0]["scores"].detach().cpu().numpy()
 
             # run tracking
             tracked_boxes = self.update_tracking(prediction)
@@ -206,6 +207,7 @@ def run_tracking(self):
                 tracked_boxes,
                 frame,
                 frame_number,
+                pred_scores,
             )
 
             # update frame number

diff --git a/crabs/tracker/utils/io.py b/crabs/tracker/utils/io.py
@@ -7,7 +7,7 @@
 
 from crabs.detector.utils.visualization import draw_bbox
 from crabs.tracker.utils.tracking import (
-    save_frame_and_csv,
+    save_output_frames,
     write_tracked_bbox_to_csv,
 )
 
@@ -107,6 +107,7 @@ def save_required_output(
     tracked_boxes: list[list[float]],
     frame: np.ndarray,
     frame_number: int,
+    pred_scores: np.ndarray,
 ) -> None:
     """
     Handle the output based on argument options.
@@ -131,20 +132,23 @@ def save_required_output(
         The current frame.
     frame_number : int
         The frame number.
+    pred_scores : np.ndarray
+        The prediction score from detector
     """
     frame_name = f"{video_file_root}_frame_{frame_number:08d}.png"
+
+    for bbox, pred_score in zip(tracked_boxes, pred_scores):
+        write_tracked_bbox_to_csv(
+            bbox, frame, frame_name, csv_writer, pred_score
+        )
+
     if save_frames:
-        save_frame_and_csv(
+        save_output_frames(
             frame_name,
             tracking_output_dir,
-            tracked_boxes,
             frame,
             frame_number,
-            csv_writer,
         )
-    else:
-        for bbox in tracked_boxes:
-            write_tracked_bbox_to_csv(bbox, frame, frame_name, csv_writer)
 
     if save_video:
         frame_copy = frame.copy()

diff --git a/crabs/tracker/utils/tracking.py b/crabs/tracker/utils/tracking.py
@@ -47,6 +47,7 @@ def write_tracked_bbox_to_csv(
     frame: np.ndarray,
     frame_name: str,
     csv_writer: Any,
+    pred_score: np.ndarray,
 ) -> None:
     """
     Write bounding box annotation to a CSV file.
@@ -62,6 +63,8 @@ def write_tracked_bbox_to_csv(
         The name of the frame.
     csv_writer : Any
         The CSV writer object to write the annotation.
+    pred_score : np.ndarray
+        The prediction score from detector.
     """
     # Bounding box geometry
     xmin, ymin, xmax, ymax, id = bbox
@@ -79,21 +82,19 @@ def write_tracked_bbox_to_csv(
             '{{"name":"rect","x":{},"y":{},"width":{},"height":{}}}'.format(
                 xmin, ymin, width_box, height_box
             ),
-            '{{"track":"{}"}}'.format(int(id)),
+            '{{"track":"{}", "confidence":"{}"}}'.format(int(id), pred_score),
         )
     )
 
 
-def save_frame_and_csv(
+def save_output_frames(
     frame_name: str,
     tracking_output_dir: Path,
-    tracked_boxes: list[list[float]],
     frame: np.ndarray,
     frame_number: int,
-    csv_writer: Any,
 ) -> None:
     """
-    Save tracked bounding boxes as frames and write to a CSV file.
+    Save tracked bounding boxes as frames.
 
     Parameters
     ----------
@@ -109,14 +110,13 @@ def save_frame_and_csv(
         The frame number.
     csv_writer : Any
         CSV writer object for writing bounding box data.
+    pred_scores : np.ndarray
+        The prediction score from detector
 
     Returns
     -------
     None
     """
-    for bbox in tracked_boxes:
-        # Add bbox to csv
-        write_tracked_bbox_to_csv(bbox, frame, frame_name, csv_writer)
 
     # Save frame as PNG - once as per frame
     frame_path = tracking_output_dir / frame_name

diff --git a/tests/test_unit/test_tracking_utils.py b/tests/test_unit/test_tracking_utils.py
@@ -49,8 +49,9 @@ def test_write_tracked_bbox_to_csv(csv_writer, csv_output):
     bbox = np.array([10, 20, 50, 80, 1])
     frame = np.zeros((100, 100, 3), dtype=np.uint8)
     frame_name = "frame_0001.png"
+    pred_score = 0.900
 
-    write_tracked_bbox_to_csv(bbox, frame, frame_name, csv_writer)
+    write_tracked_bbox_to_csv(bbox, frame, frame_name, csv_writer, pred_score)
 
     expected_row = (
         "frame_0001.png",
@@ -59,7 +60,7 @@ def test_write_tracked_bbox_to_csv(csv_writer, csv_output):
         1,
         0,
         '"{""name"":""rect"",""x"":10,""y"":20,""width"":40,""height"":60}"',
-        '"{""track"":""1""}"',
+        '"{""track"":""1"", ""confidence"":""0.9""}"',
     )
     expected_row_str = ",".join(map(str, expected_row))
     assert csv_output.getvalue().strip() == expected_row_str