Lightning Pose

repository·main·Indexed 43 days ago

https://github.com/paninski-lab/lightning-pose

An end-to-end package for robust multi-view and single-view animal pose estimation using advanced transformer architectures and PyTorch Lightning. Version 2.4.2 features Multi-View Transformers to handle occlusions and temporal context for single-view datasets. It includes a GUI for labeling, model management, and evaluation, as well as CLI workflows and tools for automated hyperparameter sweeps.

Tokens
49.7K
Snippets
159
Records
243
Agent score
55%

What's inside Lightning Pose

  1. Overview of Lightning Pose capabilities

    main

    Lightning Pose supports two primary modes of operation:

    Multi-View Capabilities

    Designed for 3D tracking using multiple camera views:

    • Multi-View Transformer (MVT): Unified architecture for simultaneous processing across views via early-feature fusion.
    • Patch Masking: Training scheme to force the model to learn robust cross-view correspondences.
    • Geometric Consistency: Uses 3D triangulation losses and geometrically-aware 3D data augmentation for calibrated setups.
    • Variance Inflation: Outlier detection for identifying geometrically inconsistent predictions.

    Single-View Capabilities

    Designed for robust 2D estimation from individual video streams:

    • Temporal Context Networks: Uses surrounding frames to resolve anatomical ambiguities and handle brief occlusions.
    • Unsupervised Learning: Training objectives that penalize predictions violating physical laws.
    • Pretrained Backbone Support: Optimized for generic Vision Transformer (ViT) backbones.
  2. Overview of Inference Timing Benchmark components

    main

    The benchmark suite consists of several specialized scripts designed to run a single-GPU inference speed comparison:

    • benchmark_single_config.py: The worker script. It runs exactly one combination (model panel, precision/runtime variant, and decoder) and appends timing results to a CSV. It can be called standalone for debugging.
    • run_benchmark.sh: The orchestrator. It reads the YAML config, loops through all combinations, calls the worker script as a fresh process for each, and finally triggers the plotting script.
    • load_config.py: A utility that parses timing_config.yaml and outputs shell variable assignments for the orchestrator.
    • timing_config.yaml: The configuration template for defining the sweep parameters.
    • plot_results.py: The plotting script. It reads all *.csv files in a directory and generates grouped-bar subplots.
  3. Overview of Multiview Lightning Pose

    main

    Lightning Pose supports multi-camera setups where a single scene is captured from multiple angles using different cameras. It utilizes a multi-view transformer solution that processes all views simultaneously. This approach allows the model to learn cross-view correlations, which improves pose estimation performance compared to single-view processing.

    Limitations (as of October 2025):

    • Multi-view mode does not support context frames.
    • Multi-view mode does not support unsupervised losses.
  4. Cloud and Headless Workflow

    main

    The Lightning Pose application is designed for both local and cloud deployment. It features a browser-based interface that is compatible with headless GPU servers, allowing for:

    • Cloud-based labeling: Specialized multi-view annotation tools that utilize camera calibration.
    • Unified Multi-view Viewer: Integrated visualization to inspect and compare predictions across all camera views simultaneously.
  5. Configure batch sizes for Multi-GPU training

    main

    Lightning Pose maintains the same effective batch size across multiple GPUs as it would on a single GPU.

    Important: If you previously reduced your batch size to make a model fit on a single GPU, you should increase it back to your desired effective batch size when moving to multi-GPU training.

    The following configuration parameters are affected by the num_gpus setting:

    • training.train_batch_size (labeled frames)
    • training.val_batch_size (labeled frames)
    • dali.train.base.sequence_length (unlabeled video frames)
    • dali.train.context.batch_size (unlabeled video frames)
  6. Transfer a subset of a project directory

    main

    To reduce transfer size for large projects, you can selectively include or omit specific directories. Common strategies include:

    • Omit the models directory: If you do not need to perform model operations.
    • Selective models: Include only the specific models required for your current task.
    • Omit videos: If you do not need to view videos or extract frames.

    Refer to the Directory Structure Reference to understand the purpose of each file before deciding what to omit.

  7. Use Labeled-frames mode for sparse bounding box files

    main

    When cropping labeled frames, Lightning Pose uses Labeled-frames mode. In this mode, rows are matched to images using the path provided in the first column.

    Requirements:

    • File Naming: Must be named bbox.csv and located in the bbox directory.
    • Sparsity: Unlike video mode, this mode can be sparse. You only need to provide entries for the frames that have labels.
    • Path Matching: The path in the first column must match the image paths used in your labeled data CSV exactly.
    ,x,y,h,w
    labeled-data/session0_view0/img00000005.png,1230,117,391,391
    labeled-data/session0_view0/img00000010.png,482,138,425,425
    labeled-data/session0_view0/img00000230.png,1230,117,391,391
  8. Choosing between the Lightning Pose App and CLI

    main

    Lightning Pose provides two primary interfaces built on top of the lightning_pose Python package. They are interoperable because they both operate on the same Project structure.

    • The App (litpose run_app): Best for interactive tasks such as labeling data, training models, or viewing predictions to evaluate model performance.
    • The CLI (litpose): Best for building custom automated pipelines or accessing advanced features not exposed in the graphical App interface.
  9. Understand the lightning-pose CLI help and error formatting

    main

    The lightning-pose CLI uses customized argparse classes to provide a more readable user experience. When you run the CLI with --help or encounter an error, the output is formatted with specific improvements:

    • Welcome Message: The top-level parser displays a welcome banner: Welcome to the lightning-pose CLI!.
    • Improved Readability: Help text is wrapped to fit your terminal width, preserving paragraphs and long words (like hyphens) to prevent awkward breaks.
    • Argument Formatting: Each command-line argument is followed by a newline to prevent dense, hard-to-read blocks of text.
    • Error Reporting: Errors are printed in red to stderr, followed by a separator line and the relevant help documentation to assist in immediate troubleshooting.
    • Documentation Link: The CLI help output includes a direct link to the official user guide: https://lightning-pose.readthedocs.io/en/latest/source/user_guide/index.html.
  10. Use mirrored frames for multiview data

    main

    In a mirrored setup, multiple views (e.g., side view and bottom view) are captured and stored as separate frames. This approach simplifies labeling and processing because each view is treated as an independent image.

    When labeling mirrored data, treat each keypoint as a unique entity for each view. For example, if tracking a nose from two views, you should use two distinct keypoint names: nose_top (for the top view) and nose_bot (for the bottom view).

    Lightning Pose utilizes this structure via the Multiview consistency loss when processing unlabeled video data, which leverages the fact that these separate keypoints represent the same physical body part.

  11. Understand the LightningPose model directory structure

    main

    A LightningPose model directory contains configuration, weights, exported runtimes, and prediction results. The structure varies slightly depending on whether you are using Singleview or Multiview modes.

    Singleview Structure

    /models/model_dir/
    ├── config.yaml
    ├── predictions.csv
    ├── tb_logs/ ...
    ├── exports_onnx/
    ├── exports_trt/
    ├── video_preds/
    └── image_preds/

    Multiview Structure

    /models/model_dir/
    ├── config.yaml
    ├── predictions_view0.csv
    ├── predictions_view1.csv
    ├── tb_logs/ ...
    ├── exports_onnx/
    ├── exports_trt/
    ├── video_preds/
    └── image_preds/
    /models/model_dir/
    ├── config.yaml
    |
    # Singleview
    ├── predictions.csv
    # Multiview
    ├── predictions_view0.csv
    ├── predictions_view1.csv
    |
    ├── tb_logs/ ...
    ├── exports_onnx/
    │   └── epoch=214-step=12685-best_fp16.onnx
    ├── exports_trt/
    │   └── epoch=214-step=12685-best_fp16/
    │       └── trt_metadata.json
    ├── video_preds/
    │   ├── session0_view0.mp4/
    │   │   └── predictions.csv
    │   └── session0_view1.mp4/
    │   └── predictions.csv
    └── image_preds/
        └── CollectedData.csv/
             └── predictions.csv