diff --git a/crabs/tracker/track_video.py b/crabs/tracker/track_video.py index 1e19d48e..99d1c211 100644 --- 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 index fc37c400..36d0e544 100644 --- 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 index 172f6fc5..3954ab14 100644 --- 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/guides/ManualLabellingSteps.md b/guides/ManualLabellingSteps.md index 6e968356..8d496861 100755 --- a/guides/ManualLabellingSteps.md +++ b/guides/ManualLabellingSteps.md @@ -26,9 +26,10 @@ Below we outline the steps on how to add annotations to these extracted frames. ### 1. Download and launch the VIA annotation tool -- Access the online version of the tool [here](https://www.robots.ox.ac.uk/~vgg/software/via/via.html) or download it locally [via this link](https://www.robots.ox.ac.uk/~vgg/software/via/) (click on `Downloads` > `Image annotator` section and download the zip file). +- We use the VIA image annotator tool version 2.0.12. +- We recommend downloading it locally from [this link](https://www.robots.ox.ac.uk/~vgg/software/via/downloads/via-2.0.12.zip) (or from [here](https://www.robots.ox.ac.uk/~vgg/software/via/), under `Downloads` > `Version 2` > `via-2.0.12.zip`). - It is very lightweight, browser-based and can be run offline. -- If you download the tool locally, launch it by selecting the `via.html` file from the expanded zip archive +- If you download the tool locally, launch it by selecting the `via.html` file from the expanded zip archive. ### 2. Define the project settings diff --git a/tests/test_unit/test_tracking_utils.py b/tests/test_unit/test_tracking_utils.py index 10582e76..37f581b0 100644 --- 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