From bd88753a4a061d5657501ef40f1b6118c5a76de9 Mon Sep 17 00:00:00 2001
From: sfmig <33267254+sfmig@users.noreply.github.com>
Date: Tue, 29 Oct 2024 15:56:41 +0000
Subject: [PATCH] Plot name of individual next to bbox (rather than the
 individual index in the position array)

---
 scripts/output_video.py | 27 +++++++++++++++++++--------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/scripts/output_video.py b/scripts/output_video.py
index 55b4fc8c..37fdf6fb 100644
--- a/scripts/output_video.py
+++ b/scripts/output_video.py
@@ -1,14 +1,15 @@
 from pathlib import Path
 
 import cv2
+import numpy as np
 
 
-def create_opencv_video(
+def create_opencv_video(  # noqa: C901
     ds,
     input_video,
     output_video_path,
-    list_individuals_idcs=None,
-    list_frame_idcs=None,
+    list_individuals=None,  # IDs of individuals to plot
+    list_frame_idcs=None,  # list of frame **indices** to plot
 ):
     # Open the video file
     cap = cv2.VideoCapture(input_video)
@@ -24,6 +25,16 @@ def create_opencv_video(
     fourcc = cv2.VideoWriter_fourcc(*"mp4v")  # Codec for mp4
     out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
 
+    # Prepare list of indices of individuals to plot
+    if list_individuals:
+        list_individuals_idcs = [
+            np.argwhere(ds_pred.individuals.data == f"id_{str(ind)}").item()
+            for ind in list_individuals
+        ]
+
+    else:
+        list_individuals_idcs = None
+
     # bboxes format
     rectangle_color = (0, 255, 0)  # Green color in BGR format
 
@@ -71,10 +82,10 @@ def create_opencv_video(
                         3,  # rectangle_thickness
                     )
 
-                    # add ID
+                    # add bbox ID
                     cv2.putText(
                         frame,
-                        str(ind_idx),
+                        str(ind_idx),  # index in position array
                         tuple(
                             int(x) for x in bottom_right
                         ),  # location of text bottom left
@@ -169,14 +180,14 @@ def create_opencv_video(
     #     )
 
     # Create prediction video
-    list_individuals_idcs = [22]
+    list_individuals = [27, 110, 129, 150, 162, 165]
     output_video_path = str(
         Path(__file__).parent
-        / f"pred_id_{'_'.join([str(el) for el in list_individuals_idcs])}.mp4"
+        / f"pred_id_{'_'.join([str(el) for el in list_individuals])}_VIA.mp4"
     )
     create_opencv_video(
         ds=ds_pred,
         input_video=input_video,
         output_video_path=output_video_path,
-        list_individuals_idcs=list_individuals_idcs,
+        list_individuals=list_individuals,
     )