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

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

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

{lightning_pose-2.4.1 → 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):
@@ -907,7 +933,7 @@ def _build_datamodule_pred(cfg: DictConfig | ListConfig) -> BaseDataModule | Unl
     with open_dict(cfg_pred.training):
         cfg_pred.training.imgaug = "default"
         cfg_pred.training.imgaug_hflip = False
-    imgaug_transform_pred = get_imgaug_transform(cfg=cfg_pred)
+    imgaug_transform_pred = get_imgaug_transform(cfg=cfg_pred, data_dir=cfg_pred.data.data_dir)
     dataset_pred = get_dataset(
         cfg=cfg_pred,
         data_dir=cfg_pred.data.data_dir,

{lightning_pose-2.4.1 → lightning_pose-2.4.3}/lightning_pose/api/model_config.py

@@ -8,6 +8,7 @@ from typing import get_args
 
 from omegaconf import DictConfig, ListConfig, OmegaConf, open_dict
 
+from lightning_pose.data.datasets import has_calibration_files
 from lightning_pose.models import ALLOWED_MODEL_TYPES
 from lightning_pose.utils.io import (
     check_video_paths,
@@ -220,14 +221,60 @@ class ModelConfig:
 
         self._validate_steps_vs_epochs()
 
+    def _validate_camera_dependent_loss(self, loss_name: str) -> None:
+        """Validate a single camera-dependent multi-view loss, if active.
+
+        Both ``supervised_reprojection_heatmap_mse`` and ``supervised_pairwise_projections``
+        require the same setup: ``model.model_type == "heatmap_multiview_transformer"`` (only
+        that class's ``get_loss_inputs_labeled`` populates the 3D/reprojected keypoints these
+        losses need -- any other model type raises a confusing ``TypeError`` from deep inside
+        :class:`~lightning_pose.losses.factory.LossFactory` instead), a 3D-safe augmentation
+        pipeline (``training.imgaug == "dlc"``, which gets silently promoted to ``dlc-mv`` —
+        see :func:`lightning_pose.data.factory.get_imgaug_transform`), ``training.imgaug_3d``
+        set to ``true``, and resolvable camera parameters. Without all of these, the loss
+        either never gets added (see :func:`lightning_pose.losses.factory.get_loss_factories`)
+        or is trained against geometrically-inconsistent per-view augmentations.
+
+        Args:
+            loss_name: config key under ``losses``, e.g. ``"supervised_reprojection_heatmap_mse"``.
+
+        Raises:
+            AssertionError: if the loss is active and any requirement is not met.
+        """
+        loss_cfg = self.cfg.losses.get(loss_name)
+        if loss_cfg is None or loss_cfg.get('log_weight') is None:
+            return
+
+        assert self.cfg.model.model_type == 'heatmap_multiview_transformer', (
+            f"model.model_type must be 'heatmap_multiview_transformer' when losses.{loss_name} "
+            "is active -- only that model class populates the 3D keypoints this loss requires"
+        )
+        assert self.cfg.training.imgaug == 'dlc', (
+            f"training.imgaug must be 'dlc' when losses.{loss_name} is active"
+        )
+        assert self.cfg.training.get('imgaug_3d') is True, (
+            f"training.imgaug_3d must be true when losses.{loss_name} is active"
+        )
+        has_cam_params = bool(
+            self.cfg.data.get('camera_params_file')
+            or has_calibration_files(self.cfg.data.data_dir)
+        )
+        assert has_cam_params, (
+            f"losses.{loss_name} is active but no camera parameters were found: set "
+            "data.camera_params_file, or provide a calibration.toml (or "
+            "calibrations/<session>.toml) at data.data_dir for auto-discovery. Without "
+            "camera parameters this loss is silently never added at train time."
+        )
+
     def _validate_model(self) -> None:
         """Validate the ``model`` config section.
 
         Checks:
         - ``model_type`` is a recognised value.
         - Multi-view models use ``heatmap_multiview_transformer``.
-        - When ``losses.supervised_reprojection_heatmap_mse`` is active, ``training.imgaug``
-          must be ``"dlc"`` and ``training.imgaug_3d`` must be ``true``.
+        - When ``losses.supervised_reprojection_heatmap_mse`` or
+          ``losses.supervised_pairwise_projections`` is active, see
+          :meth:`_validate_camera_dependent_loss`.
 
         Raises:
             AssertionError: if any check fails.
@@ -251,16 +298,8 @@ class ModelConfig:
                     stacklevel=2,
                 )
 
-            reprojection = self.cfg.losses.get('supervised_reprojection_heatmap_mse')
-            if reprojection is not None and reprojection.get('log_weight') is not None:
-                assert self.cfg.training.imgaug == 'dlc', (
-                    "training.imgaug must be 'dlc' when "
-                    "losses.supervised_reprojection_heatmap_mse is active"
-                )
-                assert self.cfg.training.get('imgaug_3d') is True, (
-                    "training.imgaug_3d must be true when "
-                    "losses.supervised_reprojection_heatmap_mse is active"
-                )
+            self._validate_camera_dependent_loss('supervised_reprojection_heatmap_mse')
+            self._validate_camera_dependent_loss('supervised_pairwise_projections')
 
     def _validate_losses(self) -> None:
         """Validate the ``losses`` config section.

{lightning_pose-2.4.1 → lightning_pose-2.4.3}/lightning_pose/api/model_runtime.py

@@ -118,7 +118,7 @@ def load_model_from_checkpoint(
     if not data_module and not skip_data_module:
         delete_extras = True
         data_dir, video_dir = return_absolute_data_paths(data_cfg=cfg.data)
-        imgaug_transform = get_imgaug_transform(cfg=cfg)
+        imgaug_transform = get_imgaug_transform(cfg=cfg, data_dir=data_dir)
         dataset = get_dataset(cfg=cfg, data_dir=data_dir, imgaug_transform=imgaug_transform)
         data_module = get_data_module(cfg=cfg, dataset=dataset, video_dir=video_dir)
     if not data_module:

{lightning_pose-2.4.1 → 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.1 → lightning_pose-2.4.3}/lightning_pose/converters/sleap.py

@@ -19,6 +19,19 @@ from PIL import Image
 logger = logging.getLogger(__name__)
 
 
+def _decode_json_attr(raw: object) -> dict:
+    """Decode an h5py ``json`` attribute, which may come back as ``str`` or ``bytes``.
+
+    h5py returns ``numpy.bytes_`` for fixed-length HDF5 string attributes and native ``str``
+    for variable-length UTF-8 ones, depending on how the writer encoded them. Calling ``str()``
+    on the former yields its Python repr (e.g. ``"b'{...}'"``), not the decoded text, so bytes
+    must be decoded explicitly rather than passed through ``str()``.
+    """
+    if isinstance(raw, bytes):
+        return json.loads(raw.decode('utf-8'))
+    return json.loads(str(raw))
+
+
 def _extract_video_names(pkg_slp_file: Path) -> dict[str, str]:
     """Map each ``video<N>`` group in a ``.pkg.slp`` file to its source video filename."""
     video_names = {}
@@ -28,7 +41,7 @@ def _extract_video_names(pkg_slp_file: Path) -> dict[str, str]:
                 source_video_path = f'{video_group_name}/source_video'
                 if source_video_path in hdf_file:
                     source_video_ds = cast(h5py.Dataset, hdf_file[source_video_path])
-                    source_video_dict = json.loads(str(source_video_ds.attrs['json']))
+                    source_video_dict = _decode_json_attr(source_video_ds.attrs['json'])
                     video_names[video_group_name] = source_video_dict['backend']['filename']
     return video_names
 
@@ -130,7 +143,7 @@ def _extract_labels(pkg_slp_file: Path) -> pd.DataFrame | None:
 
             if data:
                 metadata_ds = cast(h5py.Dataset, hdf_file['metadata'])
-                metadata_dict = json.loads(str(metadata_ds.attrs['json']))
+                metadata_dict = _decode_json_attr(metadata_ds.attrs['json'])
                 keypoints = [node['name'] for node in metadata_dict['nodes']]
                 columns = pd.MultiIndex.from_product(
                     [['lightning_tracker'], keypoints, ['x', 'y']],

{lightning_pose-2.4.1 → 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.1 → lightning_pose-2.4.3}/lightning_pose/data/datasets.py

@@ -550,6 +550,27 @@ class HeatmapDataset(BaseTrackingDataset):
         return example_dict  # type: ignore[return-value]
 
 
+def has_calibration_files(root_directory: str | Path) -> bool:
+    """Check whether a project directory contains any camera calibration files.
+
+    Used both by :meth:`MultiviewHeatmapDataset._discover_cam_params_from_image_paths` (to
+    skip labeled-data folder-name parsing when there is nothing to discover) and by
+    :func:`lightning_pose.data.factory.get_dataset` (to decide, before the dataset is built,
+    whether camera params will end up available via auto-discovery).
+
+    Args:
+        root_directory: project root directory (``cfg.data.data_dir``).
+
+    Returns:
+        True if ``root_directory/calibrations/`` (a directory) or
+        ``root_directory/calibration.toml`` exists.
+    """
+    return (
+        (Path(root_directory) / 'calibrations').is_dir()
+        or (Path(root_directory) / 'calibration.toml').exists()
+    )
+
+
 class MultiviewHeatmapDataset(torch.utils.data.Dataset):
     """Heatmap dataset that aggregates one :class:`HeatmapDataset` per camera view.
 
@@ -735,6 +756,12 @@ class MultiviewHeatmapDataset(torch.utils.data.Dataset):
             tuple of (cam_params_df, cam_params_file_to_camgroup); both None if no calibration
             files are found
         """
+        # nothing to discover if the project has no calibration files at all; skip parsing
+        # labeled-data folder names so <session>_<view> naming is only enforced when it's
+        # actually needed
+        if not has_calibration_files(self.root_directory):
+            return None, None
+
         image_names = self.dataset[self.view_names[0]].image_names
         cam_params_file_to_camgroup = {}
         calib_files = []
@@ -760,7 +787,10 @@ class MultiviewHeatmapDataset(torch.utils.data.Dataset):
             calib_by_session = Path(self.root_directory) / 'calibrations' / f'{session_id}.toml'
             calib_fallback = Path(self.root_directory) / 'calibration.toml'
             if calib_by_session.exists():
-                calib_file = str(Path('calibrations') / f'{session_id}.toml')
+                # .as_posix() (not str()) -- calib_file is a portable identifier stored in
+                # cam_params_df and used as a dict key, not a filesystem call, so it must
+                # stay forward-slash on every OS like image_names elsewhere in this class.
+                calib_file = (Path('calibrations') / f'{session_id}.toml').as_posix()
             elif calib_fallback.exists():
                 calib_file = 'calibration.toml'
             else:
@@ -780,8 +810,13 @@ class MultiviewHeatmapDataset(torch.utils.data.Dataset):
             return cam_params_df, cam_params_file_to_camgroup
 
         if cam_params_file_to_camgroup and not all_found:
-            logger.warning(
-                'calibration file not found for some frames; disabling 3D for entire dataset'
+            raise ValueError(
+                'calibration file not found for some, but not all, sessions in this dataset. '
+                'Partial calibration coverage is always treated as an error, since some '
+                'calibration files existing usually means calibration was intended for the '
+                'whole dataset. Ensure every session referenced by the dataset has a '
+                'corresponding calibrations/<session>.toml (or a calibration.toml fallback at '
+                'the project root), or remove calibration entirely if it is not needed.'
             )
         return None, None
 

{lightning_pose-2.4.1 → lightning_pose-2.4.3}/lightning_pose/data/factory.py

@@ -23,6 +23,7 @@ model-side steps):
 """
 
 import warnings
+from pathlib import Path
 
 import imgaug.augmenters as iaa
 import numpy as np
@@ -38,13 +39,14 @@ from lightning_pose.data.datasets import (
     BaseTrackingDataset,
     HeatmapDataset,
     MultiviewHeatmapDataset,
+    has_calibration_files,
 )
 
 # to ignore imports for sphinx-autoapidoc
 __all__: list[str] = []
 
 
-def get_imgaug_transform(cfg: DictConfig | ListConfig) -> iaa.Sequential:
+def get_imgaug_transform(cfg: DictConfig | ListConfig, data_dir: str | Path) -> iaa.Sequential:
     """Create simple and flexible data transform pipeline that augments images and keypoints.
 
     Args:
@@ -63,7 +65,9 @@ def get_imgaug_transform(cfg: DictConfig | ListConfig) -> iaa.Sequential:
             - "cfg.training.imgaug_3d":
                 boolean flag to control 3D-compatible augmentations for multiview models;
                 set to False to disable automatic "dlc-mv" enforcement;
-                set to True to enable 3D augmentations for when camera params file exist.
+                set to True to enable 3D augmentations for when camera params exist.
+        data_dir: root directory that camera calibration files are resolved relative to (used
+            to check for auto-discoverable calibration files; see ``has_calibration_files``).
 
     Returns:
         imgaug pipeline
@@ -75,12 +79,18 @@ def get_imgaug_transform(cfg: DictConfig | ListConfig) -> iaa.Sequential:
         # Check if user explicitly wants to use 3D augmentations for multiview models
         imagug_3d = cfg.training.get('imgaug_3d', None)
 
+        # camera params may come from an explicit camera_params_file, or be auto-discovered
+        # from a calibration.toml/calibrations/ dir at data_dir (see has_calibration_files)
+        has_cam_params = bool(
+            cfg.data.get('camera_params_file') or has_calibration_files(data_dir)
+        )
+
         # enforce "dlc-mv" imgaug pipeline for multiview models (no 2D geometric transforms)
         # only if explicitly requested or if no preference is set and camera params exist
         if (
             params not in ['default', 'none']
             and cfg.model.model_type.find('multiview') > -1
-            and cfg.data.get('camera_params_file')
+            and has_cam_params
             and (imagug_3d is True or imagug_3d is None)
         ):
             params = 'dlc-mv'
@@ -158,10 +168,12 @@ def get_dataset(
                 'No precautions regarding the size of the images were considered here, '
                 'images will be resized accordingly to configs!'
             )
-            if (
-                cfg.training.imgaug in ['default', 'none']
-                or not cfg.data.get('camera_params_file')
-            ):
+            # camera params may come from an explicit camera_params_file, or be auto-discovered
+            # from a calibration.toml/calibrations/ dir at data_dir (see has_calibration_files)
+            has_cam_params = bool(
+                cfg.data.get('camera_params_file') or has_calibration_files(data_dir)
+            )
+            if cfg.training.imgaug in ['default', 'none'] or not has_cam_params:
                 # we are either
                 # 1. running inference on un-augmented data, and need to make sure to resize
                 # 2. using a multiview model w/o camera params, and need to take care of resizing
@@ -236,6 +248,8 @@ def get_data_module(
     Raises:
         ValidationError: if a context model is requested but
             ``dali.context.train.batch_size < 5 * num_gpus``.
+        RuntimeError: if semi-supervised losses are requested but no CUDA device is
+            available, or the ``nvidia-dali`` package isn't installed.
     """
 
     # Old configs may have num_gpus: 0. We will remove support in a future release.
@@ -255,7 +269,10 @@ def get_data_module(
     val_batch_size = int(np.ceil(cfg.training.val_batch_size / cfg.training.num_gpus))
 
     from lightning_pose.models import check_if_semi_supervised
+    from lightning_pose.utils.device import require_cuda_for_semi_supervised
+
     semi_supervised = check_if_semi_supervised(cfg.model.losses_to_use)
+    require_cuda_for_semi_supervised(cfg.model.losses_to_use)
     if not semi_supervised:
         data_module = BaseDataModule(
             dataset=dataset,

{lightning_pose-2.4.1 → 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.