inference api: batch-wise image analysis for superanimal prepare_2d

Author: deruyter92

fmpose3d/inference_api/fmpose3d.py

@@ -187,28 +187,28 @@ def predict(
                 cv2.imwrite(p, frames[idx])
                 paths.append(p)
 
-            # Run DeepLabCut on each frame individually (the API
-            # expects a single image path).
+            # Run DeepLabCut once for all frames.
+            predictions = superanimal_analyze_images(
+                superanimal_name=cfg.superanimal_name,
+                model_name=cfg.sa_model_name,
+                detector_name=cfg.detector_name,
+                images=paths,
+                max_individuals=cfg.max_individuals,
+                out_folder=tmpdir,
+            )
+            # predictions: {image_path: {"bodyparts": (N_ind, K, 3), ...}}
+            # Iterate in input order to keep frame alignment stable.
             for img_path in paths:
-                predictions = superanimal_analyze_images(
-                    superanimal_name=cfg.superanimal_name,
-                    model_name=cfg.sa_model_name,
-                    detector_name=cfg.detector_name,
-                    images=img_path,
-                    max_individuals=cfg.max_individuals,
-                    out_folder=tmpdir,
-                )
-                # predictions: {image_path: {"bodyparts": (N_ind, K, 3), ...}}
-                for _path, payload in predictions.items():
-                    bodyparts = payload.get("bodyparts")
-                    if bodyparts is None:
-                        # No detection -- fill with zeros.
-                        all_mapped.append(np.zeros((1, 26, 2), dtype="float32"))
-                        continue
-                    xy = bodyparts[..., :2]  # (N_ind, K, 2)
-                    mapped = self._map_keypoints(xy)
-                    # Take only the first individual.
-                    all_mapped.append(mapped[:1])
+                payload = predictions.get(img_path)
+                bodyparts = None if payload is None else payload.get("bodyparts")
+                if bodyparts is None:
+                    # No detection -- fill with zeros.
+                    all_mapped.append(np.zeros((1, 26, 2), dtype="float32"))
+                    continue
+                xy = bodyparts[..., :2]  # (N_ind, K, 2)
+                mapped = self._map_keypoints(xy)
+                # Take only the first individual.
+                all_mapped.append(mapped[:1])
 
         # Stack along frame axis → (1, N, 26, 2)
         kpts = np.stack(all_mapped, axis=1)  # (1, N, 26, 2)