Add support for cropzoom pipeline (#250)

* Cropzoom PR

* try re-enabling cz test

* Revert "try re-enabling cz test"

This reverts commit 8bb3bbf.

* fixups

* further docs

* bump version

Author: ksikka

cli/main.py

@@ -5,8 +5,15 @@
 import os
 import sys
 from pathlib import Path
+from textwrap import dedent
 from typing import TYPE_CHECKING
 
+from omegaconf import OmegaConf
+
+# Don't import anything from torch or lightning_pose until needed.
+# These imports are slow and delay CLI help text outputs.
+# if TYPE_CHECKING allows use of imports for type annotations, without
+# actually invoking the import at runtime.
 if TYPE_CHECKING:
     from lightning_pose.model import Model
 
@@ -45,6 +52,11 @@ def _build_parser():
         "If not specified, defaults to "
         "./outputs/{YYYY-MM-DD}/{HH:MM:SS}/",
     )
+    train_parser.add_argument(
+        "--detector_model",
+        type=types.existing_model_dir,
+        help="If specified, uses cropped training data in the detector model's directory.",
+    )
     train_parser.add_argument(
         "--overrides",
         nargs="*",
@@ -94,13 +106,86 @@ def _build_parser():
         "              uses the labels to compute pixel error.\n"
         "              saves outputs to `image_preds/<csv_file_name>`\n",
     )
+    predict_parser.add_argument(
+        "--overrides",
+        nargs="*",
+        metavar="KEY=VALUE",
+        help="overrides attributes of the config file. Uses hydra syntax:\n"
+        "https://hydra.cc/docs/advanced/override_grammar/basic/",
+    )
 
     post_prediction_args = predict_parser.add_argument_group("post-prediction")
     post_prediction_args.add_argument(
         "--skip_viz",
         action="store_true",
         help="skip generating prediction-annotated images/videos",
     )
+
+    # Crop command
+    crop_parser = subparsers.add_parser(
+        "crop",
+        description=dedent(
+            """\
+            Crops a video or labeled frames based on model predictions.
+            Requires model predictions to already have been generated using `litpose predict`.
+    
+            Cropped videos are saved to:
+                <model_dir>/
+                └── video_preds/
+                    ├── <video_filename>.csv              (predictions)
+                    ├── <video_filename>_bbox.csv         (bbox)
+                    └── remapped_<video_filename>.csv     (TODO move to remap command)
+                └── cropped_videos/
+                    └── cropped_<video_filename>.mp4      (cropped video)
+
+            Cropped images are saved to:
+                <model_dir>/
+                └── image_preds/
+                    └── <csv_file_name>/
+                        ├── predictions.csv
+                        ├── bbox.csv                      (bbox)
+                        └── cropped_<csv_file_name>.csv   (cropped labels)
+                └── cropped_images/
+                        └── a/b/c/<image_name>.png        (cropped images)\
+            """
+        ),
+        usage="litpose crop <model_dir> <input_path:video|csv>... --crop_ratio=CROP_RATIO --anchor_keypoints=x,y,z",
+    )
+    crop_parser.add_argument(
+        "model_dir", type=types.existing_model_dir, help="path to a model directory"
+    )
+
+    crop_parser.add_argument(
+        "input_path", type=Path, nargs="+", help="one or more files"
+    )
+    crop_parser.add_argument(
+        "--crop_ratio",
+        type=float,
+        default=2.0,
+        help="Crop a bounding box this much larger than the animal. Default is 2.",
+    )
+    crop_parser.add_argument(
+        "--anchor_keypoints",
+        type=str,
+        default="",  # Or a reasonable default like "0,0,0" if appropriate
+        help="Comma-separated list of anchor keypoint names, defaults to all keypoints",
+    )
+
+    remap_parser = subparsers.add_parser(
+        "remap",
+        description=dedent(
+            """\
+            Remaps predictions from cropped to original coordinate space.
+            Requires model predictions to already have been generated using `litpose predict`.
+
+            Remapped predictions are saved as "remapped_{preds_file}" in the same folder as preds_file.
+            """
+        ),
+        usage="litpose remap <preds_file> <bbox_file>",
+    )
+    remap_parser.add_argument("preds_file", type=Path, help="path to a prediction file")
+    remap_parser.add_argument("bbox_file", type=Path, help="path to a bbox file")
+
     return parser
 
 
@@ -120,6 +205,84 @@ def main():
     elif args.command == "predict":
         _predict(args)
 
+    elif args.command == "crop":
+        _crop(args)
+
+    elif args.command == "remap":
+        _remap_preds(args)
+
+
+def _crop(args: argparse.Namespace):
+    import lightning_pose.utils.cropzoom as cz
+    from lightning_pose.model import Model
+
+    model_dir = args.model_dir
+    model = Model.from_dir(model_dir)
+
+    # Make both cropped_images and cropped_videos dirs. Reason: After this, the user
+    # will train a pose model, and current code in io utils checks that both
+    # data_dir and videos_dir are present. if we just create one or the other,
+    # the check will fail.
+    model.cropped_data_dir().mkdir(parents=True, exist_ok=True)
+    model.cropped_videos_dir().mkdir(parents=True, exist_ok=True)
+
+    input_paths = [Path(p) for p in args.input_path]
+
+    detector_cfg = OmegaConf.create(
+        {
+            "crop_ratio": args.crop_ratio,
+            "anchor_keypoints": args.anchor_keypoints.split(",") if args.anchor_keypoints else [],
+        }
+    )
+    assert detector_cfg.crop_ratio > 1
+
+    for input_path in input_paths:
+        if input_path.suffix == ".mp4":
+            input_preds_file = model.video_preds_dir() / (input_path.stem + ".csv")
+            output_bbox_file = model.video_preds_dir() / (
+                input_path.stem + "_bbox.csv"
+            )
+            output_file = model.cropped_videos_dir() / ("cropped_" + input_path.name)
+
+            cz.generate_cropped_video(
+                input_video_file=input_path,
+                input_preds_file=input_preds_file,
+                detector_cfg=detector_cfg,
+                output_bbox_file=output_bbox_file,
+                output_file=output_file,
+            )
+        elif input_path.suffix == ".csv":
+            preds_dir = model.image_preds_dir() / input_path.name
+            input_data_dir = Path(model.config.cfg.data.data_dir)
+            cropped_data_dir = model.cropped_data_dir()
+
+            output_bbox_file = preds_dir / "bbox.csv"
+            output_csv_file_path = preds_dir / ("cropped_" + input_path.name)
+            input_preds_file = preds_dir / "predictions.csv"
+            cz.generate_cropped_labeled_frames(
+                input_data_dir=input_data_dir,
+                input_csv_file=input_path,
+                input_preds_file=input_preds_file,
+                detector_cfg=detector_cfg,
+                output_data_dir=cropped_data_dir,
+                output_bbox_file=output_bbox_file,
+                output_csv_file=output_csv_file_path,
+            )
+        else:
+            raise NotImplementedError("Only mp4 and csv files are supported.")
+
+
+def _remap_preds(args: argparse.Namespace):
+    import lightning_pose.utils.cropzoom as cz
+
+    output_file = args.preds_file.with_name("remapped_" + args.preds_file.name)
+
+    cz.generate_cropped_csv_file(
+        input_csv_file=args.preds_file,
+        input_bbox_file=args.bbox_file,
+        output_csv_file=output_file,
+    )
+
 
 def _train(args: argparse.Namespace):
     import hydra
@@ -142,11 +305,32 @@ def _train(args: argparse.Namespace):
         cfg = hydra.compose(config_name=args.config_file.stem, overrides=args.overrides)
 
         # Delay this import because it's slow.
+        from lightning_pose.model import Model
         from lightning_pose.train import train
 
         # TODO: Move some aspects of directory mgmt to the train function.
         output_dir.mkdir(parents=True, exist_ok=True)
         # Maintain legacy hydra chdir until downstream no longer depends on it.
+
+        if args.detector_model:
+            # create detector model object before chdir so that relative path is resolved correctly
+            detector_model = Model.from_dir(args.detector_model)
+            import copy
+
+            cfg = copy.deepcopy(cfg)
+            cfg.data.data_dir = str(detector_model.cropped_data_dir())
+            cfg.data.video_dir = str(detector_model.cropped_videos_dir())
+            if isinstance(cfg.data.csv_file, str):
+                cfg.data.csv_file = str(
+                    detector_model.cropped_csv_file_path(cfg.data.csv_file)
+                )
+            else:
+                cfg.data.csv_file = [
+                    str(detector_model.cropped_csv_file_path(f))
+                    for f in cfg.data.csv_file
+                ]
+            cfg.eval.test_videos_directory = cfg.data.video_dir
+
         os.chdir(output_dir)
         train(cfg)
 
@@ -155,7 +339,7 @@ def _predict(args: argparse.Namespace):
     # Delay this import because it's slow.
     from lightning_pose.model import Model
 
-    model = Model.from_dir(args.model_dir)
+    model = Model.from_dir2(args.model_dir, hydra_overrides=args.overrides)
     input_paths = [Path(p) for p in args.input_path]
 
     for p in input_paths:

docs/api/lightning_pose.utils.cropzoom.generate_cropped_csv_file.rst

@@ -0,0 +1,6 @@
+generate_cropped_csv_file
+=========================
+
+.. currentmodule:: lightning_pose.utils.cropzoom
+
+.. autofunction:: generate_cropped_csv_file

docs/api/lightning_pose.utils.scripts.calculate_steps_per_epoch.rst

@@ -0,0 +1,6 @@
+calculate_steps_per_epoch
+=========================
+
+.. currentmodule:: lightning_pose.utils.scripts
+
+.. autofunction:: calculate_steps_per_epoch

docs/api/lightning_pose.utils.scripts.calculate_train_batches.rst

@@ -67,7 +67,7 @@ API Reference:
 
 .. autoclass:: lightning_pose.model.Model
     :members:
-    :exclude-members: __init__, PredictionResult, predict_on_label_csv_internal
+    :exclude-members: __init__, PredictionResult, from_dir2
 
 
 Lightning Pose Internal API

docs/source/api.rst

@@ -0,0 +1,123 @@
+########################
+Cropzoom pipeline
+########################
+
+For setups where an animal is freely moving in a large arena,
+it's advantageous to crop around the animal before running pose estimation.
+Lightning-pose calls this technique "cropzoom". This document describes how
+to set up a such a pipeline.
+
+Conceptual overview
+===================
+
+A cropzoom pipeline consists of two lightning-pose models:
+a "detector model" and a "pose model".
+
+* The detector model operates on the full image of the arena.
+* The pose model operates on the cropped animal.
+
+These two models are trained and predicted like any other
+lightning pose model. We provide additional tools that help you compose these models:
+
+* ``litpose crop``: Given the detector model's predictions, crops around the animal.
+* ``litpose remap``: Given the pose model's predictions and the crop bounding boxes,
+  remaps the predictions to the original coordinate space.
+
+Training
+--------
+
+Training involves:
+
+1. Train a "detector model"
+2. Crop training data for the pose model.
+3. Train a "pose model".
+
+Inference
+---------
+
+Inference involves:
+
+1. Predict using the "detector model"
+2. Crop the data using the above predictions
+3. Predict on the cropped data using the "pose model".
+4. Remap the pose model's predictions to the original coordinate space.
+
+
+Example
+=======
+
+This is a basic example of how you can setup a cropzoom pipeline.
+Paths to CSV and MP4 files below should be replaced with your files.
+The example is illustrative only. In reality you might be interested in
+making modifications to this such as:
+
+1. Using different model type, backbone, image_resize_dims for
+   your detector model and pose model. This can be accomplished using
+   different config files for the detector and pose model.
+2. Limiting ``train_frames`` and ``max_epochs`` for testing purposes.
+
+We'll use some bash variables to avoid repeating paths below:
+
+.. code-block:: bash
+
+    MODEL_DIR=outputs/chickadee/cropzoom
+    DETECTOR_MODEL=detector_0
+    POSE_MODEL=pose_supervised_0
+
+Training script
+---------------
+
+.. code-block:: bash
+
+    #!/bin/bash
+
+    # Train the detector model.
+    litpose train config.yaml --output_dir $MODEL_DIR/$DETECTOR_MODEL
+
+    # Crop data for pose model training.
+    litpose crop $MODEL_DIR/$DETECTOR_MODEL data/CollectedData.csv
+
+    # Train the pose model.
+    litpose train config.yaml --output_dir $MODEL_DIR/$POSE_MODEL \
+        --detector_model=$MODEL_DIR/$DETECTOR_MODEL
+
+Prediction on videos script
+---------------------------
+
+.. code-block:: bash
+
+    #!/bin/bash
+
+    litpose predict $MODEL_DIR/$DETECTOR_MODEL data/videos/test_vid.short.mp4
+
+    litpose crop $MODEL_DIR/$DETECTOR_MODEL data/videos/test_vid.short.mp4
+
+    litpose predict $MODEL_DIR/$POSE_MODEL $MODEL_DIR/$DETECTOR_MODEL/cropped_videos/cropped_test_vid.short.mp4
+
+    litpose remap $MODEL_DIR/$POSE_MODEL/video_preds/cropped_TRQ177_200624_112234_lBack.short.csv \
+        $MODEL_DIR/$DETECTOR_MODEL/video_preds/test_vid.short_bbox.csv
+
+Prediction on OOD Labeled Data
+------------------------------
+
+Say you have new labeled data for OoD animals, at `data/CollectedData_new.csv`,
+and you want to predict on these frames as well as compute pixel error.
+
+.. code-block:: bash
+
+    #!/bin/bash
+
+    litpose predict $MODEL_DIR/$DETECTOR_MODEL data/CollectedData_new.csv
+
+    litpose crop $MODEL_DIR/$DETECTOR_MODEL data/CollectedData_new.csv
+
+    litpose predict $MODEL_DIR/$POSE_MODEL \
+      $MODEL_DIR/$DETECTOR_MODEL/image_preds/CollectedData_new.csv/cropped_CollectedData_new.csv
+
+    litpose remap $MODEL_DIR/$POSE_MODEL/image_preds/cropped_CollectedData_new.csv/predictions.csv \
+      $MODEL_DIR/$DETECTOR_MODEL/image_preds/CollectedData_new.csv/bbox.csv
+
+Limitations
+===========
+
+* Pose models do not yet support PCA Multiview loss.