improved animal_name handling when analyzing videos (#2884)

n-poulsen · web-flow · commit fa2ebee4bf9b · 2025-02-17T17:30:06.000+01:00
diff --git a/deeplabcut/pose_estimation_tensorflow/predict_videos.py b/deeplabcut/pose_estimation_tensorflow/predict_videos.py
@@ -271,6 +271,7 @@ def analyze_videos(
     use_shelve=False,
     auto_track=True,
     n_tracks=None,
+    animal_names=None,
     calibrate=False,
     identity_only=False,
     use_openvino="CPU" if is_openvino_available else None,
@@ -397,6 +398,13 @@ def analyze_videos(
         animals in the video is different from the number of animals the model was
         trained on.
 
+    animal_names: list[str], optional
+        If you want the names given to individuals in the labeled data file, you can
+        specify those names as a list here. If given and `n_tracks` is None, `n_tracks`
+        will be set to `len(animal_names)`. If `n_tracks` is not None, then it must be
+        equal to `len(animal_names)`. If it is not given, then `animal_names` will
+        be loaded from the `individuals` in the project config.yaml file.
+
     use_openvino: str, optional
         Use "CPU" for inference if OpenVINO is available in the Python environment.
 
@@ -630,6 +638,7 @@ def analyze_videos(
                         trainingsetindex,
                         destfolder=destfolder,
                         n_tracks=n_tracks,
+                        animal_names=animal_names,
                         modelprefix=modelprefix,
                         save_as_csv=save_as_csv,
                     )
diff --git a/deeplabcut/refine_training_dataset/stitch.py b/deeplabcut/refine_training_dataset/stitch.py
@@ -8,6 +8,8 @@
 #
 # Licensed under GNU Lesser General Public License v3.0
 #
+from typing import List, Optional
+
 import matplotlib.pyplot as plt
 import networkx as nx
 import numpy as np
@@ -891,8 +893,11 @@ def concatenate_data(self):
 
     def format_df(self, animal_names=None):
         data = self.concatenate_data()
-        if not animal_names or len(animal_names) != self.n_tracks:
+        if not animal_names or len(animal_names) < self.n_tracks:
             animal_names = [f"ind{i}" for i in range(1, self.n_tracks + 1)]
+        elif len(animal_names) > self.n_tracks:
+            animal_names = animal_names[:self.n_tracks]
+
         coords = ["x", "y", "likelihood"]
         n_multi_bpts = data.shape[1] // (len(animal_names) * len(coords))
         n_unique_bpts = 0 if self.single is None else self.single.data.shape[1]
@@ -1031,6 +1036,7 @@ def stitch_tracklets(
     shuffle=1,
     trainingsetindex=0,
     n_tracks=None,
+    animal_names: Optional[List[str]] = None,
     min_length=10,
     split_tracklets=True,
     prestitch_residuals=True,
@@ -1071,6 +1077,13 @@ def stitch_tracklets(
         passed if the number of animals in the video is different from
         the number of animals the model was trained on.
 
+    animal_names: list, optional
+        If you want the names given to individuals in the labeled data file, you can
+        specify those names as a list here. If given and `n_tracks` is None, `n_tracks`
+        will be set to `len(animal_names)`. If `n_tracks` is not None, then it must be
+        equal to `len(animal_names)`. If it is not given, then `animal_names` will
+        be loaded from the `individuals` in the project config.yaml file.
+
     min_length : int, optional
         Tracklets less than `min_length` frames of length
         are considered to be residuals; i.e., they do not participate
@@ -1107,8 +1120,8 @@ def stitch_tracklets(
         tracklets should be stitched together, the lower the returned value.
 
     destfolder: string, optional
-        Specifies the destination folder for analysis data (default is the path of the video). Note that for subsequent analysis this
-        folder also needs to be passed.
+        Specifies the destination folder for analysis data (default is the path of the
+        video). Note that for subsequent analysis this folder also needs to be passed.
 
     track_method: string, optional
          Specifies the tracker used to generate the pose estimation data.
@@ -1135,7 +1148,14 @@ def stitch_tracklets(
     cfg = auxiliaryfunctions.read_config(config_path)
     track_method = auxfun_multianimal.get_track_method(cfg, track_method=track_method)
 
-    animal_names = cfg["individuals"]
+    if animal_names is None:
+        animal_names = cfg["individuals"]
+    elif n_tracks is not None and n_tracks != len(animal_names):
+        raise ValueError(
+            "When setting both `n_tracks` and `animal_names`, `n_tracks` must be equal "
+            f"to len(animal_names)`. Found `n_tracks`={n_tracks} and `animal_names`="
+            f"{animal_names} of length {len(animal_names)}.`")
+
     if n_tracks is None:
         n_tracks = len(animal_names)
 
diff --git a/deeplabcut/utils/make_labeled_video.py b/deeplabcut/utils/make_labeled_video.py
@@ -455,7 +455,7 @@ def create_labeled_video(
 
     displayedindividuals: list[str] or str, optional, default="all"
         Individuals plotted in the video.
-        By default, all individuals present in the config will be showed.
+        By default, all individuals present in the config will be shown.
 
     codec: str, optional, default="mp4v"
         Codec for labeled video. For available options, see
@@ -613,9 +613,6 @@ def create_labeled_video(
             )
         )
 
-    individuals = auxfun_multianimal.IntersectionofIndividualsandOnesGivenbyUser(
-        cfg, displayedindividuals
-    )
     if draw_skeleton:
         bodyparts2connect = cfg["skeleton"]
         if displayedbodyparts != "all":
@@ -644,7 +641,7 @@ def create_labeled_video(
         DLCscorerlegacy,
         track_method,
         cfg,
-        individuals,
+        displayedindividuals,
         color_by,
         bodyparts,
         codec,
@@ -757,8 +754,14 @@ def proc_video(
                 print("Labeled video already created. Skipping...")
                 return
 
-            if all(individuals):
-                df = df.loc(axis=1)[:, individuals]
+            if individuals != "all":
+                if isinstance(individuals, str):
+                    individuals = [individuals]
+
+                if all(individuals) and "individuals" in df.columns.names:
+                    mask = df.columns.get_level_values("individuals").isin(individuals)
+                    df = df.loc[:, mask]
+
             cropping = metadata["data"]["cropping"]
             [x1, x2, y1, y2] = metadata["data"]["cropping_parameters"]
             labeled_bpts = [
