@@ -79,7 +79,7 @@ def get_conditions_provider_for_video(
     video = Path(video)
 
     # Load pickle for multi-animal projects
-    cond_file = video.parent / f"{video.stem}{cond_provider.scorer}_full.pickle"
+    cond_file = video.parent / f"{video.stem}{cond_provider.scorer}_assemblies.pickle"
     if not cond_file.exists():
 
         # Load h5 for single-animal projects
 
@@ -225,6 +225,7 @@ def load_conditions_h5(
                 Each array has shape (num_conditions, num_bodyparts, 3).
             If "images" is None: a list containing the conditions for each frame.
         """
+
         def _parse_row(df_row) -> np.ndarray:
             # Row to numpy and reshape
             pose = df_row.to_numpy().reshape((num_conditions, num_bodyparts, 3))
@@ -418,7 +419,7 @@ def load_conditions_pickle(filepath: str | Path) -> list[np.ndarray]:
             data = pickle.load(f)
 
         frames = [f for f in data.keys() if isinstance(f, int)]
-        n_frames = max(*frames)
+        n_frames = max(*frames) + 1
 
         parsed = []
         for i in range(n_frames):
 
@@ -108,6 +108,10 @@ def build_bottom_up_postprocessor(
                 "identity_scores": max_individuals,
             },
             pad_value=-1,
+            expected_shapes={
+                "bodyparts": (num_bodyparts, 3),
+                "identity_scores": (num_bodyparts, max_individuals),
+            },
         ),
     ]
 
@@ -172,6 +176,11 @@ def build_top_down_postprocessor(
                     "bbox_scores": max_individuals,
                 },
                 pad_value=-1,
+                expected_shapes={
+                    "bodyparts": (num_bodyparts, 3),
+                    "bboxes": (4,),
+                    "bbox_scores": (),  # scalar
+                },
             ),
         ]
     )
@@ -271,26 +280,42 @@ def __init__(
         self,
         max_individuals: dict[str, int],
         pad_value: int,
+        expected_shapes: dict[str, tuple[int, ...]],
     ):
         self.max_individuals = max_individuals
         self.pad_value = pad_value
+        self.expected_shapes = expected_shapes
 
     def __call__(
         self, predictions: dict[str, np.ndarray], context: Context
     ) -> tuple[dict[str, np.ndarray], Context]:
         for name in predictions:
             output = predictions[name]
-            if isinstance(output, list):
-                output = np.array(output)
+            output = np.array(output)  # Normalize all inputs to np.ndarray
+
+            expected_shape = self.expected_shapes.get(name, ())
+            expected_ndim = 1 + len(
+                expected_shape
+            )  # individuals_dimension + expected shape for single individual
+
+            # Special handling for empty arrays
+            if len(output) == 0:
+                output = np.empty((0, *expected_shape), dtype=float)
+            elif output.ndim < expected_ndim:
+                output = np.reshape(output, (len(output), *expected_shape))
 
             if (
                 name in self.max_individuals
                 and len(output) < self.max_individuals[name]
             ):
                 pad_size = self.max_individuals[name] - len(output)
                 tail_shape = output.shape[1:]
-                padding = self.pad_value * np.ones((pad_size, *tail_shape))
-                predictions[name] = np.concatenate([output, padding])
+                padding = self.pad_value * np.ones(
+                    (pad_size, *tail_shape), dtype=output.dtype
+                )
+                output = np.concatenate([output, padding], axis=0)
+
+            predictions[name] = output
 
         return predictions, context
 
 
@@ -370,7 +370,7 @@ class FilterLowConfidencePoses(Preprocessor):
     def __init__(
         self,
         confidence_threshold: float = 0.05,
-        aggregate_func: Callable[[np.ndarray], float] = lambda arr: np.max(arr, axis=1),
+        aggregate_func: Callable[[np.ndarray], float] = lambda arr: np.nanmax(arr, axis=1),
     ) -> None:
         self.confidence_threshold = confidence_threshold
         self.aggregate_func = aggregate_func
@@ -382,7 +382,14 @@ def __call__(
             raise ValueError(f"Must include cond_kpts, found {context}")
 
         keypoints = context["cond_kpts"]
-        mask = self.aggregate_func(keypoints[:, :, 2]) >= self.confidence_threshold
+
+        if 0 in keypoints.shape:
+            # No poses to filter; return early
+            return image, context
+
+        confidences = keypoints[:, :, 2]
+        aggregated_confidence = self.aggregate_func(confidences)
+        mask = aggregated_confidence >= self.confidence_threshold
         context["cond_kpts"] = keypoints[mask]
 
         return image, context
@@ -462,7 +469,9 @@ def __call__(
 
         # can have no bounding boxes if detector made no detections
         if len(images) == 0:
-            images = np.zeros((0, *image.shape))
+            h, w = self.output_size[1], self.output_size[0]  # output_size = (w, h)
+            c = image.shape[2] if image.ndim == 3 else 1
+            images = np.zeros((0, h, w, c), dtype=image.dtype)
         else:
             images = np.stack(images, axis=0)
 
@@ -508,12 +517,11 @@ def __call__(
         self, image: np.ndarray, context: Context
     ) -> tuple[np.ndarray, Context]:
         cond_keypoints = context[self.cond_kpt_key]
-        if len(cond_keypoints) == 0:
-            return image, context
 
         rescaled = cond_keypoints.copy()
-        rescaled[..., :2] = (
-            rescaled[..., :2] - np.array(context["offsets"])[:, None]
-        ) / np.array(context["scales"])[:, None]
+        if rescaled.size > 0:  # only rescale if non-empty
+            rescaled[..., :2] = (
+                rescaled[..., :2] - np.array(context["offsets"])[:, None]
+            ) / np.array(context["scales"])[:, None]
         context["model_kwargs"] = {"cond_kpts": np.expand_dims(rescaled, axis=1)}
         return image, context
@@ -581,6 +581,12 @@ def predict(
                 }
             ]
         """
+        cond_kpts = kwargs.get("cond_kpts", None)
+        if cond_kpts is not None and cond_kpts.shape[0] == 0:
+            # No conditions, so just return an empty prediction list
+            return []
+
+        # Normal prediction path
         if self.device and "cuda" in str(self.device):
             with torch.autocast(device_type=str(self.device)):
                 outputs = self.model(inputs.to(self.device), **kwargs)
 
@@ -1053,8 +1053,8 @@ def create_video(
         fname=video,
         sname=output_path,
         codec=codec,
-        sw=bbox[1]-bbox[0] if display_cropped else "",
-        sh=bbox[3]-bbox[2] if display_cropped else "",
+        sw=bbox[1] - bbox[0] if display_cropped else "",
+        sh=bbox[3] - bbox[2] if display_cropped else "",
         fps=fps,
     )
 
@@ -1330,7 +1330,11 @@ def create_video_with_all_detections(
                     ind = frames.index(n)
 
                     # Draw bounding boxes of required and present
-                    if plot_bboxes and "bboxes" in data[frame_names[ind]]:
+                    if (
+                        plot_bboxes
+                        and "bboxes" in data[frame_names[ind]]
+                        and "bbox_scores" in data[frame_names[ind]]
+                    ):
                         bboxes = data[frame_names[ind]]["bboxes"]
                         bbox_scores = data[frame_names[ind]]["bbox_scores"]
                         n_bboxes = bboxes.shape[0]
 
@@ -15,7 +15,10 @@
 
 import deeplabcut.utils.auxiliaryfunctions as af
 from deeplabcut.compat import Engine
-from deeplabcut.pose_estimation_pytorch.config.utils import is_model_top_down
+from deeplabcut.pose_estimation_pytorch.config.utils import (
+    is_model_top_down,
+    is_model_cond_top_down,
+)
 
 from utils import (
     cleanup,
@@ -38,6 +41,7 @@ def main(
     max_snapshots_to_keep: int = 5,
     device: str = "cpu",
     logger: dict | None = None,
+    conditions_shuffle: int = 0,
     create_labeled_videos: bool = False,
     delete_after_test_run: bool = False,
 ) -> None:
@@ -55,6 +59,23 @@ def main(
             if is_model_top_down(net_type):
                 epochs_ = top_down_epochs
             try:
+                pytorch_cfg_updates = {
+                    "data.conditions.snapshot_index": -1,
+                    "train_settings.display_iters": 50,
+                    "train_settings.epochs": epochs_,
+                    "train_settings.batch_size": batch_size,
+                    "runner.device": device,
+                    "runner.snapshots.save_epochs": save_epochs,
+                    "runner.snapshots.max_snapshots": max_snapshots_to_keep,
+                    "detector.train_settings.display_iters": 1,
+                    "detector.train_settings.epochs": detector_epochs,
+                    "detector.train_settings.batch_size": detector_batch_size,
+                    "detector.runner.snapshots.save_epochs": save_epochs,
+                    "detector.runner.snapshots.max_snapshots": max_snapshots_to_keep,
+                    "logger": logger,
+                }
+                if is_model_cond_top_down(net_type):
+                    pytorch_cfg_updates["data.conditions.shuffle"] = conditions_shuffle
                 run(
                     config_path=config_path,
                     train_fraction=train_frac,
@@ -63,20 +84,7 @@ def main(
                     videos=[str(project_path / "videos" / "video.mp4")],
                     device=device,
                     engine=engine,
-                    pytorch_cfg_updates={
-                        "train_settings.display_iters": 50,
-                        "train_settings.epochs": epochs_,
-                        "train_settings.batch_size": batch_size,
-                        "runner.device": device,
-                        "runner.snapshots.save_epochs": save_epochs,
-                        "runner.snapshots.max_snapshots": max_snapshots_to_keep,
-                        "detector.train_settings.display_iters": 1,
-                        "detector.train_settings.epochs": detector_epochs,
-                        "detector.train_settings.batch_size": detector_batch_size,
-                        "detector.runner.snapshots.save_epochs": save_epochs,
-                        "detector.runner.snapshots.max_snapshots": max_snapshots_to_keep,
-                        "logger": logger,
-                    },
+                    pytorch_cfg_updates=pytorch_cfg_updates,
                     create_labeled_videos=create_labeled_videos,
                 )
             except Exception as err:
@@ -96,8 +104,15 @@ def main(
         "project_name": "testscript-dev",
         "run_name": "test-logging",
     }
+    net_types = [
+        "top_down_resnet_50",
+        "resnet_50",
+        "dekr_w32",
+        "rtmpose_m",
+        "ctd_coam_w32",
+    ]
     main(
-        net_types=["top_down_resnet_50", "resnet_50", "dekr_w32", "rtmpose_m"],
+        net_types=net_types,
         params=SyntheticProjectParameters(
             multianimal=True,
             num_bodyparts=4,
@@ -115,6 +130,8 @@ def main(
         max_snapshots_to_keep=2,
         device="cpu",  # "cpu", "cuda:0", "mps"
         logger=None,
+        conditions_shuffle=net_types.index("resnet_50")
+        + 1,  # shuffles start at index 1
         create_labeled_videos=True,
         delete_after_test_run=True,
     )
Original file line number	Diff line number	Diff line change
`@@ -581,6 +581,12 @@ def predict(`
`581`	`581`	`}`
`582`	`582`	`]`
`583`	`583`	`"""`
	`584`	`+ cond_kpts = kwargs.get("cond_kpts", None)`
	`585`	`+ if cond_kpts is not None and cond_kpts.shape[0] == 0:`
	`586`	`+ # No conditions, so just return an empty prediction list`
	`587`	`+ return []`
	`588`	`+`
	`589`	`+ # Normal prediction path`
`584`	`590`	`if self.device and "cuda" in str(self.device):`
`585`	`591`	`with torch.autocast(device_type=str(self.device)):`
`586`	`592`	`outputs = self.model(inputs.to(self.device), **kwargs)`