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

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

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

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

@@ -720,10 +720,11 @@ class Model(_RuntimeMixin):  # pyright: ignore[reportGeneralTypeIssues]
                 Defaults to False.
             progress_file (Path, optional): Path to a file to save progress information for the
                 App. Defaults to None.
-            reader (optional): which video-reading backend to use, "dali" or "pynvvc".
-                None (default) auto-selects pynvvc if it's usable on this machine for this
-                video, else falls back to dali. Independent of the model's runtime
-                (eager/onnx) and torch.compile -- this only controls video ingestion.
+            reader (optional): 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). Independent of the model's runtime (eager/onnx)
+                and torch.compile -- this only controls video ingestion.
             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
@@ -809,9 +810,10 @@ class Model(_RuntimeMixin):  # pyright: ignore[reportGeneralTypeIssues]
             compute_metrics: whether to compute pixel error and loss metrics on predictions.
             generate_labeled_video: whether to save a labeled video.
             progress_file: path to a file to save progress information for the App.
-            reader: which video-reading backend to use, "dali" or "pynvvc". None (default)
-                auto-selects pynvvc if it's usable on this machine for this video, else falls
-                back to dali.
+            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).
 
         Returns:
             object containing the predictions and metrics for each view.
@@ -905,7 +907,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.0 → lightning_pose-2.4.2}/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.0 → lightning_pose-2.4.2}/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.0 → lightning_pose-2.4.2}/lightning_pose/cli/commands/__init__.py

@@ -19,14 +19,18 @@ command module must define three top-level functions:
 
 **Adding a new command**: create ``<name>.py`` here with all three functions (copy an existing
 module, e.g. ``train.py``, as a template), add it to :data:`COMMANDS` below, and add a matching
-``docs/source/cli_reference/<name>.rst`` page.
+``docs/source/cli_reference/<name>.rst`` page (with a toctree entry in that directory's
+``index.rst``). ``tests/cli/test_main.py::TestCliReferenceDocs`` fails if either is missing, so
+a forgotten doc page is caught by CI rather than relying on this docstring being read.
 """
 
 from lightning_pose.cli.commands import (
+    convert,
     create_bbox,
     crop,
     export,
     predict,
+    recommend,
     remap,
     run_app,
     smooth_bbox,
@@ -43,4 +47,6 @@ COMMANDS = {
     'crop': crop,
     'remap': remap,
     'run_app': run_app,
+    'recommend': recommend,
+    'convert': convert,
 }

lightning_pose-2.4.2/lightning_pose/cli/commands/convert.py

@@ -0,0 +1,72 @@
+"""Convert command for the lightning-pose CLI."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Any
+
+
+def register_parser(subparsers: Any) -> argparse.ArgumentParser:
+    """Register the convert command parser."""
+    from textwrap import dedent
+
+    description_text = dedent(
+        """\
+            Converts a labeled dataset from another pose estimation tool into a Lightning
+            Pose project directory.
+
+            ``dataset_path`` is dispatched by type:
+
+            * a directory -- treated as a DeepLabCut project (``labeled-data/<video>/
+              CollectedData*.{csv,h5}``).
+            * a ``.slp`` file -- treated as a SLEAP ``.pkg.slp`` package. Export it from the
+              SLEAP GUI via Predict -> Export Labels Package. Only single-view, single-animal
+              SLEAP projects are supported.
+            """
+    )
+
+    convert_parser = subparsers.add_parser(
+        'convert',
+        description=description_text,
+        usage='litpose convert <dataset_path> --lp_dir LP_DIR',
+    )
+    convert_parser.add_argument(
+        'dataset_path',
+        type=Path,
+        help='a DeepLabCut project directory, or a SLEAP .pkg.slp file',
+    )
+    convert_parser.add_argument(
+        '--lp_dir',
+        type=Path,
+        required=True,
+        help='destination Lightning Pose project directory; created if it does not exist',
+    )
+    return convert_parser
+
+
+def get_parser() -> argparse.ArgumentParser:
+    """Return an ArgumentParser for the `litpose convert` subcommand (for docs)."""
+    parser = argparse.ArgumentParser(prog='litpose')
+    subparsers = parser.add_subparsers(dest='command')
+    return register_parser(subparsers)
+
+
+def handle(args: argparse.Namespace) -> None:
+    """Handle the convert command."""
+    dataset_path: Path = args.dataset_path
+    lp_dir: Path = args.lp_dir
+
+    if dataset_path.is_dir():
+        from lightning_pose.converters import dlc
+
+        dlc.convert(dataset_path, lp_dir)
+    elif dataset_path.suffix == '.slp':
+        from lightning_pose.converters import sleap
+
+        sleap.convert(dataset_path, lp_dir)
+    else:
+        raise ValueError(
+            f'could not determine dataset type for {dataset_path}: expected a DeepLabCut '
+            f'project directory or a SLEAP .slp file'
+        )

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

@@ -156,9 +156,10 @@ def register_parser(subparsers: Any) -> argparse.ArgumentParser:
         choices=sorted(_READER_CHOICES),
         default=None,
         help=(
-            "video-reading backend for video inputs. 'dali' or 'pynvvc'. If omitted "
-            "(default), auto-selects 'pynvvc' when it's usable on this machine for the "
-            "given video, else falls back to 'dali'. Independent of --runtime/--compile "
+            "video-reading backend for video inputs. 'dali', 'pynvvc', or 'opencv'. If "
+            "omitted (default), auto-selects 'pynvvc' when it's usable on this machine "
+            "for the given video, else 'dali' when it's installed, else 'opencv' (the "
+            "portable fallback, always available). Independent of --runtime/--compile "
             "-- this only controls video ingestion, not model execution. Ignored for "
             "CSV/image inputs."
         ),

lightning_pose-2.4.2/lightning_pose/cli/commands/recommend.py

@@ -0,0 +1,120 @@
+"""Recommend command for the lightning-pose CLI."""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def register_parser(subparsers: Any) -> argparse.ArgumentParser:
+    """Register the recommend command parser."""
+    from textwrap import dedent
+
+    description_text = dedent(
+        """\
+            Analyzes a lightning-pose dataset and prints config recommendations.
+
+            Unlabeled videos are assumed to live in ``<dataset_path>/videos``.  A complete,
+            ready-to-use config YAML is always written; use ``--output`` to control where, or
+            ``--print_only`` to skip writing a file entirely.
+            """
+    )
+
+    recommend_parser = subparsers.add_parser(
+        'recommend',
+        description=description_text,
+        usage='litpose recommend <dataset_path> [--output OUTPUT | --print_only]',
+    )
+    recommend_parser.add_argument(
+        'dataset_path',
+        type=Path,
+        help='directory containing one or more label CSV files (auto-discovered; multiple '
+        'CSVs are treated as a multi-view dataset), or a single label CSV file',
+    )
+    output_group = recommend_parser.add_mutually_exclusive_group()
+    output_group.add_argument(
+        '--output',
+        type=Path,
+        default=None,
+        help='path to write the recommended config YAML. a directory writes config.yaml '
+        'inside it. defaults to <dataset_path>/config.yaml, or a timestamped filename if '
+        'that already exists.',
+    )
+    output_group.add_argument(
+        '--print_only',
+        action='store_true',
+        help='print the recommendation report only; do not write a config YAML.',
+    )
+    return recommend_parser
+
+
+def get_parser() -> argparse.ArgumentParser:
+    """Return an ArgumentParser for the `litpose recommend` subcommand (for docs)."""
+    parser = argparse.ArgumentParser(prog='litpose')
+    subparsers = parser.add_subparsers(dest='command')
+    return register_parser(subparsers)
+
+
+def _ask_top_down_freely_moving() -> bool:
+    """Interactively ask whether the dataset is a top-down view of a freely moving animal.
+
+    Returns False without prompting when stdin is not an interactive terminal.
+    """
+    if not sys.stdin.isatty():
+        return False
+    answer = input(
+        'Does your dataset have a top-down view with a freely moving animal? [y/N]: '
+    )
+    return answer.strip().lower() in ('y', 'yes')
+
+
+def _resolve_output_path(dataset_path: Path, output: Path | None) -> Path:
+    """determine where to write the recommended config yaml"""
+    base_dir = dataset_path if dataset_path.is_dir() else dataset_path.parent
+    if output is None:
+        candidate = base_dir / 'config.yaml'
+        if candidate.exists():
+            candidate = base_dir / f'config_{datetime.now():%Y%m%d_%H%M%S}.yaml'
+        return candidate
+    if output.is_dir():
+        return output / 'config.yaml'
+    return output
+
+
+def handle(args: argparse.Namespace) -> None:
+    """Handle the recommend command."""
+    from omegaconf import OmegaConf
+
+    from lightning_pose.utils import recommender
+
+    analysis = recommender.analyze_dataset(args.dataset_path)
+    gpu = recommender.get_gpu_info()
+    top_down_freely_moving = (
+        _ask_top_down_freely_moving() if analysis.view_names is None else False
+    )
+    rec = recommender.recommend(analysis, gpu, top_down_freely_moving=top_down_freely_moving)
+
+    logger.info(recommender.format_report(rec, analysis, gpu))
+
+    if args.print_only:
+        return
+
+    output_path = _resolve_output_path(args.dataset_path, args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    header = '\n'.join(
+        ['# generated by `litpose recommend`; rationale for each recommended field:']
+        + [f'# {name}: {reason}' for name, reason in rec.rationale.items()]
+    )
+    cfg = recommender.build_config(rec, analysis)
+    with open(output_path, 'w') as f:
+        f.write(header + '\n\n')
+        f.write(OmegaConf.to_yaml(cfg))
+
+    logger.info(f'\nwrote config to {output_path}')

lightning_pose-2.4.2/lightning_pose/converters/__init__.py

@@ -0,0 +1,18 @@
+"""Convert labeled datasets from other pose estimation tools into Lightning Pose projects.
+
+Two converters, one per source format:
+
+- :func:`lightning_pose.converters.dlc.convert` -- reads a DeepLabCut project directory
+  (``labeled-data/<video>/CollectedData*.{csv,h5}``) and writes a Lightning Pose project
+  directory (``CollectedData.csv`` plus copied ``labeled-data/`` and ``videos/``).
+- :func:`lightning_pose.converters.sleap.convert` -- reads a SLEAP ``.pkg.slp`` package
+  (frames and labels bundled in one HDF5 file) and writes a Lightning Pose project directory.
+
+Both converters are dispatched from ``litpose convert <dataset_path> --lp_dir <lp_dir>``
+(:mod:`lightning_pose.cli.commands.convert`), which decides which one to call based on
+whether ``dataset_path`` is a directory (DLC) or a ``.slp`` file (SLEAP).
+
+**Adding a new converter**: create ``<name>.py`` here with a ``convert(<source>: Path, lp_dir:
+Path) -> None`` function, then add a dispatch branch in
+``lightning_pose.cli.commands.convert.handle``.
+"""

lightning_pose-2.4.2/lightning_pose/converters/dlc.py

@@ -0,0 +1,123 @@
+"""Convert a DeepLabCut project directory into a Lightning Pose project directory."""
+
+from __future__ import annotations
+
+import logging
+import shutil
+from pathlib import Path
+from typing import cast
+
+import pandas as pd
+
+from lightning_pose.utils.io import fix_empty_first_row
+
+logger = logging.getLogger(__name__)
+
+
+def _load_labels_for_video_dir(dlc_dir: Path, video_dir_name: str) -> pd.DataFrame | None:
+    """Load the ``CollectedData`` labels for a single ``labeled-data/<video_dir_name>`` dir.
+
+    Tries the CSV export first, falling back to the HDF5 export. Both DLC labeling schemes
+    are supported: the older scheme indexes rows by ``labeled-data/<video>/<image>.png``
+    directly, and the newer scheme splits the video and image name into separate columns
+    (or separate MultiIndex levels for the HDF5 export), which are recombined into that same
+    index format here.
+
+    Args:
+        dlc_dir: root of the DLC project.
+        video_dir_name: name of the ``labeled-data`` subdirectory to load labels for.
+
+    Returns:
+        the labels as a DataFrame with a 3-level ``(scorer, bodyparts, coords)`` column
+        MultiIndex, or None if no ``CollectedData*`` file was found for this directory.
+    """
+    video_dir = dlc_dir / 'labeled-data' / video_dir_name
+    csv_files = sorted(video_dir.glob('CollectedData*.csv'))
+    if csv_files:
+        df = pd.read_csv(csv_files[0], header=[0, 1, 2], index_col=0)
+        df = fix_empty_first_row(df)
+        if len(df.index.unique()) != df.shape[0]:
+            # new DLC labeling scheme that splits video/image in different cells
+            vids = df.loc[:, ('Unnamed: 1_level_0', 'Unnamed: 1_level_1', 'Unnamed: 1_level_2')]
+            imgs = df.loc[:, ('Unnamed: 2_level_0', 'Unnamed: 2_level_1', 'Unnamed: 2_level_2')]
+            new_index = [f'labeled-data/{v}/{i}' for v, i in zip(vids, imgs, strict=True)]
+            df = df.drop(
+                columns=[
+                    ('Unnamed: 1_level_0', 'Unnamed: 1_level_1', 'Unnamed: 1_level_2'),
+                    ('Unnamed: 2_level_0', 'Unnamed: 2_level_1', 'Unnamed: 2_level_2'),
+                ],
+            )
+            df.index = new_index
+        return df
+
+    h5_files = sorted(video_dir.glob('CollectedData*.h5'))
+    if h5_files:
+        df = cast(pd.DataFrame, pd.read_hdf(h5_files[0]))
+        if isinstance(df.index, pd.MultiIndex):
+            # new DLC labeling scheme that splits video/image in different cells
+            imgs = [i[2] for i in df.index]
+            vids = [df.index[0][1] for _ in imgs]
+            new_index = [f'labeled-data/{v}/{i}' for v, i in zip(vids, imgs, strict=True)]
+            df = df.reset_index(drop=True)
+            df.index = new_index
+        return df
+
+    logger.warning(f'could not find labels for {video_dir_name}; skipping')
+    return None
+
+
+def convert(dlc_dir: Path, lp_dir: Path) -> None:
+    """Convert a DeepLabCut project directory into a Lightning Pose project directory.
+
+    Reads every ``labeled-data/<video>/CollectedData*.{csv,h5}`` file under ``dlc_dir``,
+    concatenates them into a single ``CollectedData.csv``, and copies the ``labeled-data/``
+    and ``videos/`` directories over to ``lp_dir``.
+
+    Args:
+        dlc_dir: root of the DLC project (must contain a ``labeled-data`` subdirectory).
+        lp_dir: destination Lightning Pose project directory; created if it does not exist.
+            Must not already contain a ``labeled-data`` or ``videos`` subdirectory.
+
+    Raises:
+        NotADirectoryError: if ``dlc_dir`` does not exist.
+        ValueError: if ``dlc_dir`` and ``lp_dir`` are the same path.
+    """
+    logger.info(f'converting DLC project at {dlc_dir} to LP project at {lp_dir}')
+
+    if not dlc_dir.is_dir():
+        raise NotADirectoryError(f'did not find the directory {dlc_dir}')
+    if dlc_dir.resolve() == lp_dir.resolve():
+        raise ValueError('dlc_dir and lp_dir cannot be the same')
+
+    video_dir_names = sorted(
+        p.name for p in (dlc_dir / 'labeled-data').iterdir()
+        if p.is_dir() and not p.name.startswith('.') and not p.name.endswith('_labeled')
+    )
+
+    dfs = []
+    for video_dir_name in video_dir_names:
+        logger.info(video_dir_name)
+        df = _load_labels_for_video_dir(dlc_dir, video_dir_name)
+        if df is not None:
+            dfs.append(df)
+    df_all = pd.concat(dfs)
+
+    lp_dir.mkdir(parents=True, exist_ok=True)
+    df_all.to_csv(lp_dir / 'CollectedData.csv')
+
+    shutil.copytree(dlc_dir / 'labeled-data', lp_dir / 'labeled-data')
+
+    src_videos = dlc_dir / 'videos'
+    dst_videos = lp_dir / 'videos'
+    if src_videos.exists():
+        logger.info('copying video files')
+        shutil.copytree(src_videos, dst_videos)
+    else:
+        logger.info('DLC video directory does not exist; creating empty video directory')
+        dst_videos.mkdir(parents=True, exist_ok=True)
+
+    for image_path in df_all.index:
+        if not (lp_dir / image_path).exists():
+            raise FileNotFoundError(
+                f'expected image file not found after conversion: {image_path}'
+            )