lightning-pose 2.4.2__tar.gz → 2.4.3__tar.gz

{lightning_pose-2.4.2 → lightning_pose-2.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lightning-pose
-Version: 2.4.2
+Version: 2.4.3
 Summary: Semi-supervised pose estimation using pytorch lightning
 License-Expression: MIT
 License-File: LICENSE

{lightning_pose-2.4.2 → lightning_pose-2.4.3}/lightning_pose/api/model.py

@@ -728,7 +728,8 @@ class Model(_RuntimeMixin):  # pyright: ignore[reportGeneralTypeIssues]
             bbox_file (str | Path, optional): Path to a per-frame bbox CSV (columns x, y, h, w;
                 one row per frame). When provided, each frame is cropped to its bounding box
                 before being passed to the model, and predictions are returned in the original
-                coordinate space. Single-view only. Defaults to None.
+                coordinate space. Single-view only; for multiview models use
+                ``predict_on_video_file_multiview(bbox_file_per_view=...)``. Defaults to None.
 
         Returns:
             PredictionResult: A PredictionResult object containing the predictions and metrics.
@@ -798,6 +799,7 @@ class Model(_RuntimeMixin):  # pyright: ignore[reportGeneralTypeIssues]
         generate_labeled_video: bool = False,
         progress_file: Path | None = None,
         reader: _Reader | None = None,
+        bbox_file_per_view: list[str] | list[Path] | None = None,
     ) -> MultiviewPredictionResult:
         """Version of ``predict_on_video_file`` that accesses multiple camera views of each frame.
 
@@ -813,10 +815,20 @@ class Model(_RuntimeMixin):  # pyright: ignore[reportGeneralTypeIssues]
             reader: which video-reading backend to use, "dali", "pynvvc", or "opencv". None
                 (default) auto-selects pynvvc if it's usable on this machine for this video,
                 else dali if it's installed, else opencv (the portable fallback, always
-                available).
+                available). Multiview bbox cropping currently requires (and defaults to)
+                "dali".
+            bbox_file_per_view: optional list of per-frame bbox CSVs (columns x, y, h, w; one
+                row per frame), in the same order as ``video_file_per_view``. When provided,
+                each view's frames are cropped to that view's bboxes before being passed to
+                the model, and predictions are returned in the original coordinate space.
 
         Returns:
             object containing the predictions and metrics for each view.
+
+        Raises:
+            ValueError: if the model is not multiview, the number of video or bbox files does
+                not match the number of views, or a bbox CSV's row count differs from its
+                video's frame count.
         """
         if not self.config.is_multi_view():
             raise ValueError('predict_on_video_file_multiview requires a multi-view model')
@@ -830,6 +842,16 @@ class Model(_RuntimeMixin):  # pyright: ignore[reportGeneralTypeIssues]
             )
 
         video_file_per_view = [Path(f) for f in video_file_per_view]
+        if bbox_file_per_view is not None and len(bbox_file_per_view) != len(view_names):
+            raise ValueError(
+                f'expected {len(view_names)} bbox files (one per view), '
+                f'got {len(bbox_file_per_view)}'
+            )
+        # pair each bbox file with its video before the videos are put in view order
+        video_to_bbox_file = (
+            dict(zip(video_file_per_view, map(Path, bbox_file_per_view), strict=True))
+            if bbox_file_per_view is not None else None
+        )
 
         if output_dir == self.__class__.UNSPECIFIED:
             output_dir = self.video_preds_dir()
@@ -859,6 +881,10 @@ class Model(_RuntimeMixin):  # pyright: ignore[reportGeneralTypeIssues]
             output_pred_file=prediction_csv_file_list,
             progress_file=progress_file,
             reader=reader,
+            bbox_file=(
+                [str(video_to_bbox_file[v]) for v in video_file_per_view]
+                if video_to_bbox_file is not None else None
+            ),
         )
         if generate_labeled_video:
             for video_file, preds_df in zip(video_file_per_view, df_list, strict=True):

{lightning_pose-2.4.2 → lightning_pose-2.4.3}/lightning_pose/cli/commands/predict.py

@@ -120,7 +120,8 @@ def register_parser(subparsers: Any) -> argparse.ArgumentParser:
             "directory containing bbox CSV files produced by ``litpose create_bbox`` or a "
             "compatible external source. For CSV inputs, looks for ``bbox.csv`` inside this "
             "directory. For video inputs, looks for ``<video_stem>_bbox.csv`` inside this "
-            "directory. When provided, each frame is cropped to its bounding box before "
+            "directory; for multiview models, one ``<video_stem>_bbox.csv`` per view video. "
+            "When provided, each frame is cropped to its bounding box before "
             "being passed to the model, and predictions are saved in the original coordinate "
             "space."
         ),
@@ -235,6 +236,7 @@ def handle(args: argparse.Namespace) -> None:
         _predict_multi_type_multi_view(
             model, input_paths, args.skip_viz, not args.overwrite,
             progress_file=args.progress_file, reader=args.reader,
+            bbox_dir=args.bbox_dir,
         )
     else:
         for p in input_paths:
@@ -320,6 +322,7 @@ def _predict_multi_type_multi_view(
     skip_existing: bool,
     progress_file: Path | None = None,
     reader: _Reader | None = None,
+    bbox_dir: Path | None = None,
 ) -> None:
     """Run multi-view prediction on a list of paths (videos or directories).
 
@@ -329,6 +332,9 @@ def _predict_multi_type_multi_view(
         skip_viz: if True, skip generating labeled visualization outputs.
         skip_existing: if True, skip sessions for which output CSVs already exist.
         progress_file: optional path to write prediction progress as JSON.
+        reader: optional video-reading backend; None auto-selects.
+        bbox_dir: optional directory containing a ``<video_stem>_bbox.csv`` for every view's
+            video. When provided, each view is cropped to its bboxes before inference.
     """
     # delay this import because it's slow
     from lightning_pose.utils.io import (
@@ -362,11 +368,21 @@ def _predict_multi_type_multi_view(
                 logger.info(f'skipping {session_name} (prediction file already exists)')
                 continue
 
+            bbox_file_per_view = None
+            if bbox_dir is not None:
+                bbox_file_per_view = [
+                    bbox_dir / f'{Path(video).stem}_bbox.csv' for video in video_file_per_view
+                ]
+                missing = [f for f in bbox_file_per_view if not f.exists()]
+                if missing:
+                    raise FileNotFoundError(f'bbox file(s) not found: {missing}')
+
             model.predict_on_video_file_multiview(
                 video_file_per_view,
                 generate_labeled_video=not skip_viz,
                 progress_file=progress_file,
                 reader=reader,
+                bbox_file_per_view=bbox_file_per_view,
             )
     # if we have a list of directories, we process the videos in each separately
     elif all(path.is_dir() for path in paths):
@@ -379,7 +395,7 @@ def _predict_multi_type_multi_view(
 
                 _predict_multi_type_multi_view(
                     model, video_files, skip_viz, skip_existing,
-                    progress_file=progress_file, reader=reader,
+                    progress_file=progress_file, reader=reader, bbox_dir=bbox_dir,
                 )
             else:
                 logger.info(f'skipping {path}: no videos found')

{lightning_pose-2.4.2 → lightning_pose-2.4.3}/lightning_pose/data/bboxes.py

@@ -24,6 +24,15 @@ dimensions; handle both single-view and multi-view):
 
 - :func:`frame_to_model_batch` — frame → model  (called before loss computation)
 - :func:`model_to_frame_batch` — model → frame  (called after inference)
+
+Video-reader helpers (used by all three readers in ``lightning_pose.data.video`` to crop
+per-frame bboxes out of full-resolution frames before inference):
+
+- :func:`as_bbox_df_list`     — normalize a reader's ``bbox_df`` to one DataFrame per view
+- :func:`slice_bbox_rows`     — select (and pad) the bbox rows for one batch
+- :func:`crop_and_resize_frames` — crop + resize one view's frames to its bboxes
+- :func:`crop_views`          — apply the above to every view of a batch independently
+- :func:`full_frame_bbox`     — default bbox covering the whole decoded frame
 """
 
 import pandas as pd
@@ -288,6 +297,24 @@ def model_to_frame_batch(
     return model_keypoints_.reshape((-1, num_targets))
 
 
+def slice_bbox_rows(bbox_df: pd.DataFrame, idx_start: int, seq_len: int) -> pd.DataFrame:
+    """Select the bbox rows for one batch, padding a short final batch with the last row.
+
+    Args:
+        bbox_df: DataFrame with columns ``["x", "y", "h", "w"]``, one row per video frame.
+        idx_start: index of the first frame in the batch.
+        seq_len: number of frames in the batch.
+
+    Returns:
+        DataFrame with exactly ``seq_len`` rows, ready for :func:`crop_and_resize_frames`.
+    """
+    rows = bbox_df.iloc[idx_start:idx_start + seq_len]
+    if len(rows) < seq_len:
+        last_row = bbox_df.iloc[[-1]]
+        rows = pd.concat([rows] + [last_row] * (seq_len - len(rows)), ignore_index=True)
+    return rows
+
+
 def crop_and_resize_frames(
     frames: Float[torch.Tensor, "seq_len rgb h w"],
     bbox_rows: pd.DataFrame,
@@ -299,9 +326,10 @@ def crop_and_resize_frames(
     """Crop each frame to its per-frame bounding box and resize to a common size.
 
     Shared by all video-decoding backends: each hands this function a full-resolution ``frames``
-    tensor plus one already-sliced-and-padded bbox row per frame (cursor/step bookkeeping to select
-    those rows lives in each backend's own wrapper, not here —
-    see ``LitDaliWrapper._apply_bbox_crop``).
+    tensor plus one already-sliced-and-padded bbox row per frame. Row selection is done by
+    :func:`slice_bbox_rows` (via :func:`crop_views`); the cursor/step bookkeeping that decides
+    where a batch starts lives in each reader's wrapper, not here (see
+    ``LitDaliWrapper._dali_output_to_tensors`` and ``LitPredictWrapperBase.__next__``).
 
     Args:
         frames: full-resolution frames, shape (seq_len, 3, H, W).
@@ -341,3 +369,83 @@ def crop_and_resize_frames(
             )
         )
     return torch.stack(cropped_frames), torch.stack(bboxes)
+
+
+def as_bbox_df_list(
+    bbox_df: pd.DataFrame | list[pd.DataFrame] | None,
+    num_views: int,
+) -> list[pd.DataFrame] | None:
+    """Normalize a video reader's ``bbox_df`` to a list with one DataFrame per view.
+
+    Args:
+        bbox_df: a single DataFrame (single-view), a list with one DataFrame per view
+            (multiview), or None when no bbox cropping is configured.
+        num_views: number of views in the batch being built (1 for single-view).
+
+    Returns:
+        list of per-view bbox DataFrames, or None if ``bbox_df`` is None.
+
+    Raises:
+        ValueError: if the number of DataFrames does not match ``num_views``.
+    """
+    if bbox_df is None:
+        return None
+    bbox_dfs = [bbox_df] if isinstance(bbox_df, pd.DataFrame) else list(bbox_df)
+    if len(bbox_dfs) != num_views:
+        raise ValueError(f'expected {num_views} bbox DataFrame(s), got {len(bbox_dfs)}')
+    return bbox_dfs
+
+
+def full_frame_bbox(
+    height: int,
+    width: int,
+    seq_len: int,
+    device: torch.device | str,
+) -> Float[torch.Tensor, "seq_len xyhw"]:
+    """Bbox covering an entire decoded frame (x=0, y=0, h=height, w=width), per frame.
+
+    Args:
+        height: height of the decoded (pre-resize) frame in pixels.
+        width: width of the decoded (pre-resize) frame in pixels.
+        seq_len: number of frames in the batch.
+        device: device to create the tensor on.
+
+    Returns:
+        float32 tensor of shape (seq_len, 4) in [x, y, h, w] order.
+    """
+    return torch.tensor([0, 0, height, width], dtype=torch.float32, device=device).repeat(
+        seq_len, 1,
+    )
+
+
+def crop_views(
+    frames_list: list[torch.Tensor],
+    bbox_dfs: list[pd.DataFrame],
+    idx_start: int,
+    resize_dims: list[int],
+) -> tuple[list[torch.Tensor], list[torch.Tensor]]:
+    """Crop and resize each view of a batch to its own per-frame bboxes.
+
+    Views are processed independently, so they may have different native resolutions; after
+    the crop+resize every view has shape ``resize_dims`` and the results can be stacked.
+
+    Args:
+        frames_list: full-resolution frames, one tensor of shape (seq_len, 3, H, W) per view.
+        bbox_dfs: one bbox DataFrame per view (columns ``["x", "y", "h", "w"]``, one row per
+            video frame), in the same order as ``frames_list``.
+        idx_start: video frame index of the first frame in this batch.
+        resize_dims: target ``[height, width]`` to resize each crop to.
+
+    Returns:
+        Tuple of (cropped+resized frames per view, actual bbox coordinates used per view); each
+        frames tensor has shape (seq_len, 3, *resize_dims) and each bbox tensor (seq_len, 4).
+    """
+    seq_len = frames_list[0].shape[0]
+    cropped_list = []
+    bbox_list = []
+    for frames, bbox_df in zip(frames_list, bbox_dfs, strict=True):
+        rows = slice_bbox_rows(bbox_df, idx_start, seq_len)
+        cropped, bboxes = crop_and_resize_frames(frames, rows, resize_dims)
+        cropped_list.append(cropped)
+        bbox_list.append(bboxes)
+    return cropped_list, bbox_list

{lightning_pose-2.4.2 → lightning_pose-2.4.3}/lightning_pose/data/video/__init__.py

@@ -26,6 +26,12 @@ Backends
   ``opencv-python-headless``, an unconditional cross-platform dependency -- works on any
   platform/hardware, including macOS, Windows, and GPU-less machines. Slower than DALI/pynvvc on
   a supported NVIDIA box; the guaranteed fallback everywhere else. Prediction only.
+- ``utils.py`` -- :class:`~lightning_pose.data.video.utils.PreparePredictBase` /
+  :class:`~lightning_pose.data.video.utils.LitPredictWrapperBase`, the validation, windowing,
+  normalization, bbox-crop and batch-assembly logic shared by the two prediction-only backends
+  (``pynvvc.py`` and ``opencv.py`` implement only decoding); also the batch-assembly and
+  normalization helper functions that ``dali.py`` reuses. Imports neither DALI nor
+  PyNvVideoCodec.
 - ``factory.py`` -- :func:`~lightning_pose.data.video.factory.build_video_reader`, the reader-
   selection dispatch used by ``predict_video``. See its docstring for the fallback chain and
   the recipe for adding a new backend.

{lightning_pose-2.4.2 → lightning_pose-2.4.3}/lightning_pose/data/video/dali.py

@@ -29,10 +29,14 @@ Standard mode (no ``bbox_df``):
 
 Bbox-crop mode (``bbox_df`` supplied to ``PrepareDALI``):
     DALI delivers full-resolution frames (``resize_dims=None`` in the predict pipe) so that
-    ``LitDaliWrapper._apply_bbox_crop`` can crop each frame to its per-frame bounding box
-    and resize to the original ``resize_dims`` using ``torch.nn.functional.interpolate``.
+    ``LitDaliWrapper._dali_output_to_tensors`` can crop each frame to its per-frame bounding
+    box and resize to the original ``resize_dims`` using ``torch.nn.functional.interpolate``.
     The ``bbox`` field of each batch contains the actual crop coordinates, so downstream
-    code can remap predictions back to the original coordinate space.
+    code can remap predictions back to the original coordinate space. For multiview, one
+    bbox DataFrame per view is supplied and each view is cropped independently, so views may
+    have different native resolutions. The crop and batch assembly use the helpers shared with
+    the other readers (``lightning_pose.data.bboxes.crop_views``,
+    ``lightning_pose.data.video.utils.build_batch_dict``).
 """
 
 import os
@@ -55,12 +59,13 @@ except ImportError as e:
     ) from e
 
 from lightning_pose.data import _IMAGENET_MEAN, _IMAGENET_STD
-from lightning_pose.data.bboxes import crop_and_resize_frames
+from lightning_pose.data.bboxes import as_bbox_df_list, crop_views, full_frame_bbox
 from lightning_pose.data.datatypes import (
     MultiviewUnlabeledBatchDict,
     UnlabeledBatchDict,
 )
 from lightning_pose.data.utils import count_frames
+from lightning_pose.data.video.utils import build_batch_dict
 
 # to ignore imports for sphinx-autoapidoc
 __all__: list[str] = []
@@ -206,7 +211,7 @@ class LitDaliWrapper(DALIGenericIterator):
     to the model's input dimensions before being returned.
 
     ``_frame_idx`` tracks the iterator's position in the video (by frame number) so that
-    ``_apply_bbox_crop`` reads the correct rows from ``bbox_df``.  It advances by
+    ``_dali_output_to_tensors`` reads the correct rows from ``bbox_df``.  It advances by
     ``seq_len`` for base models (non-overlapping windows) and by ``seq_len - 4`` for
     context models, because the DALI reader for context prediction uses a step of
     ``seq_len - 4`` so that consecutive 5-frame windows overlap by 4 frames.
@@ -214,7 +219,7 @@ class LitDaliWrapper(DALIGenericIterator):
     Testing without a GPU
     +++++++++++++++++++++
     This class can be instantiated without a real DALI pipeline or GPU for unit tests
-    that only exercise the PyTorch post-processing logic (e.g. ``_apply_bbox_crop``)::
+    that only exercise the PyTorch post-processing logic (e.g. ``_dali_output_to_tensors``)::
 
         wrapper = object.__new__(LitDaliWrapper)
         wrapper.do_context = False
@@ -229,7 +234,7 @@ class LitDaliWrapper(DALIGenericIterator):
         eval_mode: Literal["train", "predict"],
         num_iters: int = 1,
         do_context: bool = False,
-        bbox_df: pd.DataFrame | None = None,
+        bbox_df: pd.DataFrame | list[pd.DataFrame] | None = None,
         resize_dims: list[int] | None = None,
         **kwargs: Any,
     ) -> None:
@@ -241,7 +246,8 @@ class LitDaliWrapper(DALIGenericIterator):
                 should be fixed by lightning/dali teams)
             do_context: whether model/loader use 5-frame context or not
             bbox_df: optional DataFrame with columns ``["x", "y", "h", "w"]``, one row per
-                frame. When provided, each batch's frames are cropped per-frame and resized
+                frame (or, for multiview, a list with one such DataFrame per view, in view
+                order). When provided, each batch's frames are cropped per-frame and resized
                 to ``resize_dims`` before being returned, and the ``bbox`` field of the
                 batch dict is populated with the actual bbox coordinates.
             resize_dims: target ``[height, width]`` for post-crop resize; required when
@@ -264,11 +270,30 @@ class LitDaliWrapper(DALIGenericIterator):
         return self.num_iters
 
     @staticmethod
+    def _frame_size_hw(frame_size: torch.Tensor) -> tuple[int, int]:
+        """Extract (height, width) of the decoded frames from DALI's ``frame_size`` output.
+
+        Args:
+            frame_size: DALI ``frame_size`` output for one view, shape (1, 4) with either
+                ``seq_len,H,W,C`` or ``seq_len,C,H,W`` ordering.
+
+        Returns:
+            Tuple of (height, width) in pixels.
+        """
+        if frame_size[0, -1] == 3:  # order is seq_len,H,W,C
+            return int(frame_size[0, 1]), int(frame_size[0, 2])
+        return int(frame_size[0, 2]), int(frame_size[0, 3])  # order is seq_len,C,H,W
+
     def _dali_output_to_tensors(
-        batch: list
+        self, batch: list
     ) -> UnlabeledBatchDict | MultiviewUnlabeledBatchDict:
         """Convert a raw DALI batch output to a typed batch dictionary.
 
+        When ``self.bbox_df`` is set, each view is cropped to its per-frame bbox and resized
+        to ``self.resize_dims`` independently, so views may have different native resolutions.
+        Otherwise DALI has already resized the frames, and the ``bbox`` field spans the full
+        decoded frame (x=0, y=0, h=H, w=W) of each view.
+
         Args:
             batch: raw output list from DALI's ``DALIGenericIterator.__next__``.
 
@@ -277,107 +302,50 @@ class LitDaliWrapper(DALIGenericIterator):
             ``MultiviewUnlabeledBatchDict`` for multi-view pipelines.
         """
         # always batch_size=1
-
-        if len(batch[0].keys()) == 3:  # single view pipeline
-
-            # shape (sequence_length, 3, H, W)
-            frames = batch[0]["frames"][0, :, :, :, :]
-            # shape (1,) or (2, 3)
-            transforms = batch[0]["transforms"][0]
-            # get frame size
-            if batch[0]["frame_size"][0, -1] == 3:  # order is seq_len,H,W,C
-                height = batch[0]["frame_size"][0, 1]
-                width = batch[0]["frame_size"][0, 2]
-            else:  # order is seq_len,C,H,W
-                height = batch[0]["frame_size"][0, 2]
-                width = batch[0]["frame_size"][0, 3]
-            bbox = torch.tensor([0, 0, height, width], device=frames.device).repeat(
-                (frames.shape[0], 1))
-
-            return UnlabeledBatchDict(
-                frames=frames, transforms=transforms, bbox=bbox, is_multiview=False,
-            )
-
-        else:  # multiview pipeline
-
-            # final shape: ("seq_len", "num_views", "RGB":3, "image_height", "image_width")
-            frames = torch.stack(
-                [batch[0][key][0, :, :, :, :] for key in batch[0].keys() if
-                 "transforms" not in key and "frame_size" not in key],
-                dim=1,
-            )
-
+        output = batch[0]
+        is_multiview = len(output.keys()) != 3
+
+        # per-view lists of frames (each of shape (seq_len, 3, H, W)) and decoded frame sizes
+        if is_multiview:
+            frames_list = [
+                output[key][0] for key in output.keys()
+                if 'transforms' not in key and 'frame_size' not in key
+            ]
+            frame_sizes = [
+                self._frame_size_hw(output[key]) for key in output.keys() if 'frame_size' in key
+            ]
             # final shape: ("num_views", "h":2, "w":3)
             transforms = torch.stack(
-                [batch[0][key] for key in batch[0].keys() if "transforms" in key],
-                dim=0,
+                [output[key] for key in output.keys() if 'transforms' in key], dim=0,
             )
-
-            # final shape: ("seq_len", "num_views * xyhw")
-            bbox = torch.cat([
-                torch.tensor([
-                    0,
-                    0,
-                    batch[0][key][0, 1],
-                    batch[0][key][0, 2],
-                ],
-                    device=frames.device
-                ) for key in batch[0].keys() if "frame_size" in key
-            ], dim=0).repeat(frames.shape[0], 1)
-
-            return MultiviewUnlabeledBatchDict(
-                frames=frames, transforms=transforms, bbox=bbox, is_multiview=True,
-            )
-
-    def _apply_bbox_crop(self, batch_dict: UnlabeledBatchDict) -> UnlabeledBatchDict:
-        """Crop frames to per-frame bboxes and resize to the model's input dimensions.
-
-        Called only in bbox-crop mode (when ``bbox_df`` is not None). DALI has already
-        delivered full-resolution frames; this method crops each frame individually,
-        resizes to ``self.resize_dims``, and advances ``_frame_idx`` by the same step
-        that the DALI reader used, so the two cursors stay in sync.
-
-        Args:
-            batch_dict: single-view unlabeled batch with full-resolution frames from DALI.
-
-        Returns:
-            new ``UnlabeledBatchDict`` with cropped+resized frames and the actual bbox
-            coordinates in the ``bbox`` field (shape ``(seq_len, 4)``, order x y h w).
-        """
-        assert self.bbox_df is not None  # guarded by caller; assertion narrows type
-        frames = batch_dict['frames']  # (seq_len, 3, H, W)
-        seq_len = frames.shape[0]
-        # context-model DALI reader uses step=seq_len-4 so consecutive windows overlap by 4
-        # frames; _frame_idx must advance by the same amount to stay in sync with the reader
-        step = seq_len - 4 if self.do_context else seq_len
-
-        # slice bbox rows for this batch; pad last partial batch with the final row
-        rows = self.bbox_df.iloc[self._frame_idx:self._frame_idx + seq_len]
-        if len(rows) < seq_len:
-            last_row = self.bbox_df.iloc[[-1]]
-            rows = pd.concat(
-                [rows] + [last_row] * (seq_len - len(rows)),
-                ignore_index=True,
+        else:
+            frames_list = [output['frames'][0]]
+            frame_sizes = [self._frame_size_hw(output['frame_size'])]
+            # shape (1,) or (2, 3)
+            transforms = output['transforms'][0]
+
+        seq_len = frames_list[0].shape[0]
+        bbox_dfs = as_bbox_df_list(self.bbox_df, len(frames_list))
+        if bbox_dfs is None:
+            bbox_list = [
+                full_frame_bbox(height, width, seq_len, frames_list[0].device)
+                for height, width in frame_sizes
+            ]
+        else:
+            assert self.resize_dims is not None  # required when bbox_df set, see __init__
+            frames_list, bbox_list = crop_views(
+                frames_list, bbox_dfs, self._frame_idx, self.resize_dims,
             )
+            # context-model DALI reader uses step=seq_len-4 so consecutive windows overlap by 4
+            # frames; _frame_idx must advance by the same amount to stay in sync with the reader
+            self._frame_idx += seq_len - 4 if self.do_context else seq_len
 
-        assert self.resize_dims is not None  # required when bbox_df set, see docstring
-        cropped_frames, bboxes = crop_and_resize_frames(frames, rows, self.resize_dims)
-        self._frame_idx += step
-
-        return UnlabeledBatchDict(
-            frames=cropped_frames,
-            transforms=batch_dict['transforms'],
-            bbox=bboxes,
-            is_multiview=False,
-        )
+        return build_batch_dict(frames_list, bbox_list, transforms, is_multiview)
 
     def __next__(self) -> UnlabeledBatchDict | MultiviewUnlabeledBatchDict:
         """Fetch the next batch, applying per-frame bbox crop+resize when configured."""
         batch = super().__next__()
-        batch_dict = self._dali_output_to_tensors(batch=batch)
-        if self.bbox_df is not None and not batch_dict['is_multiview']:
-            return self._apply_bbox_crop(batch_dict)  # type: ignore[arg-type]
-        return batch_dict
+        return self._dali_output_to_tensors(batch=batch)
 
 
 class PrepareDALI:
@@ -405,7 +373,7 @@ class PrepareDALI:
         dali_config: dict | DictConfig | ListConfig | None = None,
         imgaug: str | None = "default",
         num_threads: int = 1,
-        bbox_df: pd.DataFrame | None = None,
+        bbox_df: pd.DataFrame | list[pd.DataFrame] | None = None,
     ) -> None:
         """Initialize DALI pipelines and dataloaders for training or prediction.
 
@@ -422,15 +390,17 @@ class PrepareDALI:
                 ``"dlc"``); pass ``"default"`` for resize-only or ``None`` to disable.
             num_threads: number of CPU threads used by DALI pipelines.
             bbox_df: optional DataFrame with columns ``["x", "y", "h", "w"]``, one row per
-                frame. When provided, the predict pipeline loads full-resolution frames
-                (DALI resize is disabled) and ``LitDaliWrapper`` crops each frame to its
-                bbox before resizing to ``resize_dims``.
+                frame; for multiview, a list with one such DataFrame per view, in the same
+                order as ``filenames``. When provided, the predict pipeline loads
+                full-resolution frames (DALI resize is disabled) and ``LitDaliWrapper``
+                crops each frame to its bbox before resizing to ``resize_dims``.
 
         Raises:
             FileNotFoundError: if any path in ``filenames`` does not exist or is not a file.
             ValueError: for multiview inputs, if views have differing numbers of sessions or if a
                 session has differing frame counts across views (which would desynchronize the
-                per-view readers).
+                per-view readers); if the number of ``bbox_df`` DataFrames does not match the
+                number of views.
         """
         # determine if we have a multiview pipeline
         if isinstance(filenames, list) and isinstance(filenames[0], list):
@@ -449,6 +419,14 @@ class PrepareDALI:
                 if not os.path.exists(vid) or not os.path.isfile(vid):
                     raise FileNotFoundError(f"{vid} is not a video file!")
 
+        if bbox_df is not None:
+            num_bbox_dfs = 1 if isinstance(bbox_df, pd.DataFrame) else len(bbox_df)
+            if num_bbox_dfs != len(filenames_2d):
+                raise ValueError(
+                    f'got {num_bbox_dfs} bbox DataFrame(s) for {len(filenames_2d)} view(s); '
+                    'multiview requires a list with one bbox DataFrame per view'
+                )
+
         # frame counts of view 0, per session; reused below for `self.frame_count`
         view0_frame_counts = list(map(count_frames, filenames_2d[0]))
 
@@ -549,7 +527,7 @@ class PrepareDALI:
 
         When ``bbox_df`` is set, predict entries have ``resize_dims`` forced to ``None``
         so DALI delivers full-resolution frames.  The per-frame resize happens in PyTorch
-        inside ``LitDaliWrapper._apply_bbox_crop`` after cropping, using the original
+        inside ``LitDaliWrapper._dali_output_to_tensors`` after cropping, using the original
         ``resize_dims`` stored on the wrapper.
         """
         assert self.dali_config is not None

{lightning_pose-2.4.2 → lightning_pose-2.4.3}/lightning_pose/data/video/factory.py

@@ -9,7 +9,9 @@ Adding a new video reader
 -------------------------
 
 Every backend module (``dali.py``, ``pynvvc.py``, ``opencv.py``, ...) implements the
-same contract. To add one:
+same contract. A prediction-only backend can subclass ``PreparePredictBase`` and
+``LitPredictWrapperBase`` from ``utils.py`` and implement only decoding (see ``opencv.py``).
+To add one:
 
 1. **Two-phase prepare class**: ``Prepare<Name>.__init__(model_type, filenames,
    resize_dims, dali_config, bbox_df=None, ...)`` validates inputs and precomputes
@@ -83,7 +85,7 @@ def build_video_reader(
     dali_config: dict | DictConfig | ListConfig,
     filenames: list[str] | list[list[str]],
     resize_dims: list[int],
-    bbox_df: pd.DataFrame | None,
+    bbox_df: pd.DataFrame | list[pd.DataFrame] | None,
 ) -> LitDaliWrapper | LitPynvvcWrapper | LitOpenCVWrapper:
     """Resolve which video-reading backend to use and build its dataloader.
 
@@ -109,7 +111,8 @@ def build_video_reader(
             multiview, one single-element list per view (see each backend's
             ``Prepare<Name>`` class for the exact shape).
         resize_dims: ``[height, width]`` the model expects as input.
-        bbox_df: optional per-frame bbox DataFrame; single-view only.
+        bbox_df: optional per-frame bbox DataFrame; for multiview, a list with one
+            DataFrame per view (same order as ``filenames``).
 
     Returns:
         A ready-to-iterate loader (``LitDaliWrapper``, ``LitPynvvcWrapper``, or