add bbox color parameter where possible

n-poulsen · Dec 13, 2024 · Nov 13, 2024 · Nov 15, 2024 · Nov 18, 2024 · Nov 19, 2024
commit 895cc4a7a3abfa0e789484f7b9fefcfddf25c2fc
diff --git a/deeplabcut/pose_estimation_pytorch/apis/evaluate.py b/deeplabcut/pose_estimation_pytorch/apis/evaluate.py
@@ -323,8 +323,8 @@ def plot_gt_and_predictions(
     alpha_value: float = 0.7,
     p_cutoff: float = 0.6,
     bounding_boxes: tuple[np.ndarray, np.ndarray] | None = None,
-    bounding_boxes_color="k",
     bboxes_pcutoff: float = 0.6,
+    bounding_boxes_color: str = "auto",
 ):
     """Plot ground truth and predictions on an image.
 
@@ -341,8 +341,11 @@ def plot_gt_and_predictions(
         alpha_value: Transparency of the points
         p_cutoff: Confidence threshold for showing predictions
         bounding_boxes:  bounding boxes (top-left corner, size) and their respective confidence levels,
-        bounding_boxes_color: If bounding_boxes is not None, this is the color that will be used for plotting them
         bboxes_cutoff: bounding boxes confidence cutoff threshold.
+        bounding_boxes_color: If plotting bounding boxes, this is the color that will be used for bounding boxes.
+            If set to "auto" (default value):
+                - if mode is "bodypart", the bbox color will be a default color
+                - if mode is "individual", each individual's color will be used for its bounding box
     """
     # Ensure output directory exists
     output_dir = Path(output_dir)
@@ -377,6 +380,16 @@ def plot_gt_and_predictions(
     else:
         raise ValueError(f"Invalid mode: {mode}")
 
+    if bounding_boxes_color == "auto":
+        if mode == "bodypart":
+            bboxes_color = None
+        elif mode == "individual":
+            bboxes_color = get_cmap(num_pred + 1, name=colormap)
+        else:
+            raise ValueError(f"Invalid mode: {mode}")
+    else:
+        bboxes_color = bounding_boxes_color
+
     # Plot regular bodyparts
     ax = make_multianimal_labeled_image(
         frame,
@@ -389,8 +402,8 @@ def plot_gt_and_predictions(
         p_cutoff,
         ax=ax,
         bounding_boxes=bounding_boxes,
-        bounding_boxes_color=bounding_boxes_color,
         bboxes_cutoff=bboxes_pcutoff,
+        bboxes_color=bboxes_color,
     )
 
     # Plot unique bodyparts if present

diff --git a/deeplabcut/utils/make_labeled_video.py b/deeplabcut/utils/make_labeled_video.py
@@ -90,6 +90,7 @@ def CreateVideo(
     plot_bboxes=True,
     bboxes_list=None,
     bboxes_pcutoff=0.6,
+    bboxes_color: tuple | None = None,
 ):
     """Creating individual frames with labeled body parts and making a video"""
     bpts = Dataframe.columns.get_level_values("bodyparts")
@@ -156,7 +157,8 @@ def CreateVideo(
         C = colorclass.to_rgba(np.linspace(0, 1, nindividuals))
     colors = (C[:, :3] * 255).astype(np.uint8)
 
-    bboxes_color = (0, 0, 0)
+    if bboxes_color is None:
+        bboxes_color = (255, 0, 0)
 
     with np.errstate(invalid="ignore"):
         for index in trange(min(nframes, len(Dataframe))):
@@ -256,6 +258,7 @@ def CreateVideoSlow(
     plot_bboxes=True,
     bboxes_list=None,
     bboxes_pcutoff=0.6,
+    bboxes_color: str | None = None,
 ):
     """Creating individual frames with labeled body parts and making a video"""
 
@@ -314,7 +317,8 @@ def CreateVideoSlow(
     else:
         colors = visualization.get_cmap(nbodyparts, name=colormap)
 
-    bounding_boxes_color = "k"
+    if bboxes_color is None:
+        bboxes_color = "red"
 
     nframes_digits = int(np.ceil(np.log10(nframes)))
     if nframes_digits > 9:
@@ -361,7 +365,7 @@ def CreateVideoSlow(
                             bbox_width,
                             bbox_height,
                             linewidth=1,
-                            edgecolor=bounding_boxes_color,
+                            edgecolor=bboxes_color,
                             facecolor="none",
                         )
                         ax.add_patch(rectangle)
@@ -1029,6 +1033,7 @@ def create_video(
     plot_bboxes=True,
     bboxes_list=None,
     bboxes_pcutoff=0.6,
+    bboxes_color: tuple | None = None,
 ):
     if color_by not in ("bodypart", "individual"):
         raise ValueError("`color_by` should be either 'bodypart' or 'individual'.")
@@ -1086,6 +1091,7 @@ def create_video(
         plot_bboxes=plot_bboxes,
         bboxes_list=bboxes_list,
         bboxes_pcutoff=bboxes_pcutoff,
+        bboxes_color=bboxes_color,
     )
 
 
@@ -1306,7 +1312,7 @@ def create_video_with_all_detections(
                 .get("model", {})
                 .get("box_score_thresh", 0.6)
             )
-            bboxes_color = (0, 0, 0)
+            bboxes_color = (255, 0, 0)
 
             for n in trange(clip.nframes):
                 frame = clip.load_frame()

diff --git a/deeplabcut/utils/visualization.py b/deeplabcut/utils/visualization.py
@@ -125,8 +125,8 @@ def make_multianimal_labeled_image(
     labels: list = ["+", ".", "x"],
     ax: plt.Axes | None = None,
     bounding_boxes: tuple[np.ndarray, np.ndarray] | None = None,
-    bounding_boxes_color="k",
     bboxes_cutoff: float = 0.6,
+    bboxes_color: Colormap | str | None = None,
 ) -> plt.Axes:
     """
     Plots groundtruth labels and predictions onto the matplotlib's axes, with the specified graphical parameters.
@@ -142,9 +142,12 @@ def make_multianimal_labeled_image(
         pcutoff: cut-off confidence value
         labels: labels to use for ground truth, reliable predictions, and not reliable predictions (confidence below cut-off value)
         ax: matplotlib plot's axes object
-        bounding_boxes:  bounding boxes (top-left corner, size) and their respective confidence levels,
-        bounding_boxes_color: If bounding_boxes is not None, this is the color that will be used for plotting them
+        bounding_boxes: bounding boxes (top-left corner, size) and their respective confidence levels,
         bboxes_cutoff: bounding boxes confidence cutoff threshold.
+        bboxes_color: color(s) for the bounding boxes.
+            If Colormap is passed -> each bounding box will be colored into its own color from the colormap.
+            If string is passed -> all bboxes will be of string's defined color.
+            If None -> all bboxes will be colored into a default color.
 
     Returns:
         matplotlib Axes object with plotted labels and predictions.
@@ -156,15 +159,21 @@ def make_multianimal_labeled_image(
     ax.imshow(frame, "gray")
 
     if bounding_boxes is not None:
-        for bbox, bbox_score in zip(bounding_boxes[0], bounding_boxes[1]):
+        for i, (bbox, bbox_score) in enumerate(zip(bounding_boxes[0], bounding_boxes[1])):
             bbox_origin = (bbox[0], bbox[1])
             (bbox_width, bbox_height) = (bbox[2], bbox[3])
+            if isinstance(bboxes_color, Colormap):
+                bbox_color = bboxes_color(i)
+            elif bboxes_color is None:
+                bbox_color = "red"
+            else:
+                bbox_color = bboxes_color
             rectangle = patches.Rectangle(
                 bbox_origin,
                 bbox_width,
                 bbox_height,
                 linewidth=1,
-                edgecolor=bounding_boxes_color,
+                edgecolor=bbox_color,
                 facecolor="none",
                 linestyle="--" if bbox_score < bboxes_cutoff else "-",
             )
@@ -446,6 +455,7 @@ def plot_evaluation_results(
     p_cutoff: float = 0.6,
     bounding_boxes: dict | None = None,
     bboxes_cutoff: float = 0.6,
+    bounding_boxes_color: str = "auto",
 ) -> None:
     """
     Creates labeled images using the results of inference, and saves them to an output
@@ -471,6 +481,11 @@ def plot_evaluation_results(
             (np array for coordinates and np array for confidence).
             None corresponds to no bounding boxes.
         bboxes_cutoff: bounding boxes confidence cutoff threshold.
+        bounding_boxes_color: If plotting bounding boxes, this is the color that will be used for bounding boxes.
+            If set to "auto" (default value):
+                - if mode is "bodypart", the bbox color will be a default color
+                - if mode is "individual", each individual's color will be used for its bounding box
+
     """
     if bounding_boxes is None:
         bounding_boxes = {}
@@ -540,6 +555,16 @@ def plot_evaluation_results(
         else:
             colors = []
 
+        if bounding_boxes_color == "auto":
+            if mode == "bodypart":
+                bboxes_color = None
+            elif mode == "individual":
+                bboxes_color = get_cmap(individuals + 1, name=colormap)
+            else:
+                raise ValueError(f"Invalid mode: {mode}")
+        else:
+            bboxes_color = bounding_boxes_color
+
         ax = make_multianimal_labeled_image(
             frame=frame,
             coords_truth=ground_truth,
@@ -552,6 +577,7 @@ def plot_evaluation_results(
             ax=ax,
             bounding_boxes=bboxes,
             bboxes_cutoff=bboxes_cutoff,
+            bboxes_color=bboxes_color,
         )
         if plot_unique_bodyparts:
             unique_predictions = unique_predictions.swapaxes(0, 1)