From 7ab54215b25c99779ad7027a22b7eb821ca30f2b Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Thu, 19 Dec 2024 13:33:06 +0100
Subject: [PATCH 01/14] improve api __init__ imports

---
 .../pose_estimation_pytorch/apis/__init__.py   | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/deeplabcut/pose_estimation_pytorch/apis/__init__.py b/deeplabcut/pose_estimation_pytorch/apis/__init__.py
index ba7f6f9e7c..39456ccc4c 100644
--- a/deeplabcut/pose_estimation_pytorch/apis/__init__.py
+++ b/deeplabcut/pose_estimation_pytorch/apis/__init__.py
@@ -11,16 +11,30 @@
 
 from deeplabcut.pose_estimation_pytorch.apis.analyze_images import (
     analyze_images,
+    analyze_image_folder,
     superanimal_analyze_images,
 )
-from deeplabcut.pose_estimation_pytorch.apis.analyze_videos import analyze_videos
+from deeplabcut.pose_estimation_pytorch.apis.analyze_videos import (
+    analyze_videos,
+    video_inference,
+    VideoIterator,
+)
 from deeplabcut.pose_estimation_pytorch.apis.convert_detections_to_tracklets import (
     convert_detections2tracklets,
 )
-from deeplabcut.pose_estimation_pytorch.apis.evaluate import evaluate_network
+from deeplabcut.pose_estimation_pytorch.apis.evaluate import (
+    predict,
+    evaluate,
+    evaluate_network,
+)
 from deeplabcut.pose_estimation_pytorch.apis.export import export_model
 from deeplabcut.pose_estimation_pytorch.apis.train import train_network
 from deeplabcut.pose_estimation_pytorch.apis.visualization import (
     extract_maps,
     extract_save_all_maps,
 )
+from deeplabcut.pose_estimation_pytorch.apis.utils import (
+    build_predictions_dataframe,
+    get_detector_inference_runner,
+    get_pose_inference_runner,
+)

From 976b67e759ec7eab29ea6708ea6a2aa1e22302d9 Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Fri, 20 Dec 2024 11:04:11 +0100
Subject: [PATCH 02/14] improve loader.ground_truth_bboxes

---
 .../pose_estimation_pytorch/apis/evaluate.py   |  5 ++++-
 .../apis/visualization.py                      |  4 +++-
 .../pose_estimation_pytorch/data/base.py       | 18 +++++++++++++++---
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/deeplabcut/pose_estimation_pytorch/apis/evaluate.py b/deeplabcut/pose_estimation_pytorch/apis/evaluate.py
index 79d213d892..cc8fe913c0 100755
--- a/deeplabcut/pose_estimation_pytorch/apis/evaluate.py
+++ b/deeplabcut/pose_estimation_pytorch/apis/evaluate.py
@@ -79,7 +79,10 @@ def predict(
             context = bbox_predictions
         else:
             ground_truth_bboxes = loader.ground_truth_bboxes(mode=mode)
-            context = [{"bboxes": ground_truth_bboxes[image]} for image in image_paths]
+            context = [
+                {"bboxes": ground_truth_bboxes[image]["bboxes"]}
+                for image in image_paths
+            ]
 
     images_with_context = image_paths
     if context is not None:
diff --git a/deeplabcut/pose_estimation_pytorch/apis/visualization.py b/deeplabcut/pose_estimation_pytorch/apis/visualization.py
index 0a790a2307..b923b2040b 100644
--- a/deeplabcut/pose_estimation_pytorch/apis/visualization.py
+++ b/deeplabcut/pose_estimation_pytorch/apis/visualization.py
@@ -550,7 +550,9 @@ def _get_context(
         bboxes_train = loader.ground_truth_bboxes(mode="train")
         bboxes_test = loader.ground_truth_bboxes(mode="test")
         bboxes = {**bboxes_train, **bboxes_test}
-        return [dict(bboxes=bboxes[str(img_path)]) for img_path in image_paths]
+        return [
+            dict(bboxes=bboxes[str(img_path)]["bboxes"]) for img_path in image_paths
+        ]
 
     detector_runner = utils.get_detector_inference_runner(
         model_config=loader.model_cfg,
diff --git a/deeplabcut/pose_estimation_pytorch/data/base.py b/deeplabcut/pose_estimation_pytorch/data/base.py
index 705df4a1ad..fc14818032 100644
--- a/deeplabcut/pose_estimation_pytorch/data/base.py
+++ b/deeplabcut/pose_estimation_pytorch/data/base.py
@@ -149,7 +149,7 @@ def ground_truth_keypoints(
 
         return ground_truth_dict
 
-    def ground_truth_bboxes(self, mode: str = "train") -> dict[str, np.ndarray]:
+    def ground_truth_bboxes(self, mode: str = "train") -> dict[str, dict]:
         """Creates a dictionary containing the ground truth bounding boxes
 
         Args:
@@ -158,7 +158,14 @@ def ground_truth_bboxes(self, mode: str = "train") -> dict[str, np.ndarray]:
         Returns:
             A dict mapping image paths to the ground truth annotations for the mode in
             the format:
-                {'image': bboxes with shape (num_individuals, xywh)}
+                {
+                    'path/to/image000.png': {
+                        "width": (int) the width of the image, in
+                        "height": (int) the height of the image, in pixels
+                        "bboxes": (np.ndarray) bboxes with shape (num_individuals, xywh)
+                    },
+                    'path/to/image000.png': {...},
+                }
         """
         if mode not in self._loaded_data:
             self._loaded_data[mode] = self.load_data(mode)
@@ -176,7 +183,12 @@ def ground_truth_bboxes(self, mode: str = "train") -> dict[str, np.ndarray]:
                 bboxes = np.zeros((0, 4))
             else:
                 bboxes = _compute_crop_bounds(np.stack(bboxes, axis=0), img_shape)
-            ground_truth_dict[image_path] = bboxes
+
+            ground_truth_dict[image_path] = dict(
+                width=image["width"],
+                height=image["height"],
+                bboxes=bboxes,
+            )
 
         return ground_truth_dict
 

From 703c115fcf84c4341073ea1bb9968c68fbcd1d66 Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Fri, 20 Dec 2024 13:55:11 +0100
Subject: [PATCH 03/14] improved imports

---
 deeplabcut/pose_estimation_pytorch/__init__.py      | 11 +++++++++++
 deeplabcut/pose_estimation_pytorch/apis/__init__.py |  1 +
 deeplabcut/pose_estimation_pytorch/data/__init__.py |  1 +
 3 files changed, 13 insertions(+)

diff --git a/deeplabcut/pose_estimation_pytorch/__init__.py b/deeplabcut/pose_estimation_pytorch/__init__.py
index a10d38940c..6d5d18c9ca 100644
--- a/deeplabcut/pose_estimation_pytorch/__init__.py
+++ b/deeplabcut/pose_estimation_pytorch/__init__.py
@@ -9,12 +9,23 @@
 # Licensed under GNU Lesser General Public License v3.0
 #
 from deeplabcut.pose_estimation_pytorch.apis import (
+    analyze_image_folder,
+    analyze_images,
     analyze_videos,
+    build_predictions_dataframe,
     convert_detections2tracklets,
+    evaluate,
     evaluate_network,
     extract_maps,
     extract_save_all_maps,
+    get_detector_inference_runner,
+    get_pose_inference_runner,
+    predict,
+    superanimal_analyze_images,
     train_network,
+    video_inference,
+    VideoIterator,
+    visualize_predictions,
 )
 from deeplabcut.pose_estimation_pytorch.config import (
     available_detectors,
diff --git a/deeplabcut/pose_estimation_pytorch/apis/__init__.py b/deeplabcut/pose_estimation_pytorch/apis/__init__.py
index 39456ccc4c..28390431dc 100644
--- a/deeplabcut/pose_estimation_pytorch/apis/__init__.py
+++ b/deeplabcut/pose_estimation_pytorch/apis/__init__.py
@@ -26,6 +26,7 @@
     predict,
     evaluate,
     evaluate_network,
+    visualize_predictions,
 )
 from deeplabcut.pose_estimation_pytorch.apis.export import export_model
 from deeplabcut.pose_estimation_pytorch.apis.train import train_network
diff --git a/deeplabcut/pose_estimation_pytorch/data/__init__.py b/deeplabcut/pose_estimation_pytorch/data/__init__.py
index a578b101d4..368926a462 100644
--- a/deeplabcut/pose_estimation_pytorch/data/__init__.py
+++ b/deeplabcut/pose_estimation_pytorch/data/__init__.py
@@ -10,6 +10,7 @@
 #
 from deeplabcut.pose_estimation_pytorch.data.base import Loader
 from deeplabcut.pose_estimation_pytorch.data.cocoloader import COCOLoader
+from deeplabcut.pose_estimation_pytorch.data.collate import COLLATE_FUNCTIONS
 from deeplabcut.pose_estimation_pytorch.data.dlcloader import DLCLoader
 from deeplabcut.pose_estimation_pytorch.data.dataset import (
     PoseDatasetParameters,

From d1f14465a453d385d0777862c24837a23cd4946c Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Fri, 20 Dec 2024 16:08:17 +0100
Subject: [PATCH 04/14] improvements to imports and analyze_image runner
 loading

---
 .../pose_estimation_pytorch/__init__.py       |  2 ++
 .../pose_estimation_pytorch/apis/__init__.py  |  1 +
 .../apis/analyze_images.py                    | 26 +++++++++----------
 3 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/deeplabcut/pose_estimation_pytorch/__init__.py b/deeplabcut/pose_estimation_pytorch/__init__.py
index 6d5d18c9ca..d4256cfba0 100644
--- a/deeplabcut/pose_estimation_pytorch/__init__.py
+++ b/deeplabcut/pose_estimation_pytorch/__init__.py
@@ -8,11 +8,13 @@
 #
 # Licensed under GNU Lesser General Public License v3.0
 #
+import deeplabcut.pose_estimation_pytorch.config as config
 from deeplabcut.pose_estimation_pytorch.apis import (
     analyze_image_folder,
     analyze_images,
     analyze_videos,
     build_predictions_dataframe,
+    create_labeled_images,
     convert_detections2tracklets,
     evaluate,
     evaluate_network,
diff --git a/deeplabcut/pose_estimation_pytorch/apis/__init__.py b/deeplabcut/pose_estimation_pytorch/apis/__init__.py
index 28390431dc..5851f1e8f5 100644
--- a/deeplabcut/pose_estimation_pytorch/apis/__init__.py
+++ b/deeplabcut/pose_estimation_pytorch/apis/__init__.py
@@ -31,6 +31,7 @@
 from deeplabcut.pose_estimation_pytorch.apis.export import export_model
 from deeplabcut.pose_estimation_pytorch.apis.train import train_network
 from deeplabcut.pose_estimation_pytorch.apis.visualization import (
+    create_labeled_images,
     extract_maps,
     extract_save_all_maps,
 )
diff --git a/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py b/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py
index a8ee75c881..7f0a913f4d 100644
--- a/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py
+++ b/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py
@@ -28,8 +28,9 @@
 from deeplabcut.core.engine import Engine
 from deeplabcut.modelzoo.utils import get_superanimal_colormaps
 from deeplabcut.pose_estimation_pytorch.apis.utils import (
-    get_inference_runners,
+    get_detector_inference_runner,
     get_model_snapshots,
+    get_pose_inference_runner,
     get_scorer_name,
     get_scorer_uid,
     parse_snapshot_index_for_analysis,
@@ -357,32 +358,29 @@ def analyze_image_folder(
             f" Please specify the `detector_path` parameter."
         )
 
-    bodyparts = model_cfg["metadata"]["bodyparts"]
-    unique_bodyparts = model_cfg["metadata"]["unique_bodyparts"]
-    individuals = model_cfg["metadata"]["individuals"]
     if max_individuals is None:
-        max_individuals = len(individuals)
+        max_individuals = len(model_cfg["metadata"]["individuals"])
 
     if device is None:
         device = resolve_device(model_cfg)
 
-    pose_runner, detector_runner = get_inference_runners(
+    pose_runner = get_pose_inference_runner(
         model_config=model_cfg,
         snapshot_path=snapshot_path,
-        max_individuals=max_individuals,
-        num_bodyparts=len(bodyparts),
-        num_unique_bodyparts=len(unique_bodyparts),
         device=device,
-        with_identity=False,
-        transform=None,
-        detector_path=detector_path,
-        detector_transform=None,
+        max_individuals=max_individuals,
     )
 
     image_paths = parse_images_and_image_folders(images)
     pose_inputs = image_paths
-    if detector_runner is not None:
+    if detector_path is not None:
         logging.info(f"Running object detection with {detector_path}")
+        detector_runner = get_detector_inference_runner(
+            model_config=model_cfg,
+            snapshot_path=detector_path,
+            device=device,
+            max_individuals=max_individuals,
+        )
 
         detector_image_paths = image_paths
         if progress_bar:

From 816a720b26615da8238c09d639de25d5087c0f9d Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Fri, 20 Dec 2024 16:27:56 +0100
Subject: [PATCH 05/14] improved video inference ux

---
 deeplabcut/pose_estimation_pytorch/README.md  |  2 -
 .../apis/analyze_videos.py                    | 50 +++++++++++++++----
 .../modelzoo/inference.py                     |  2 -
 3 files changed, 40 insertions(+), 14 deletions(-)

diff --git a/deeplabcut/pose_estimation_pytorch/README.md b/deeplabcut/pose_estimation_pytorch/README.md
index 0c3eec28a1..92ba3154f5 100644
--- a/deeplabcut/pose_estimation_pytorch/README.md
+++ b/deeplabcut/pose_estimation_pytorch/README.md
@@ -468,9 +468,7 @@ pose_runner, detector_runner = get_inference_runners(
 
 predictions = video_inference(
     video=video_path,
-    task=pose_task,
     pose_runner=pose_runner,
     detector_runner=detector_runner,
-    with_identity=False,
 )
 ```
diff --git a/deeplabcut/pose_estimation_pytorch/apis/analyze_videos.py b/deeplabcut/pose_estimation_pytorch/apis/analyze_videos.py
index f4ecc22856..6de6bd7520 100644
--- a/deeplabcut/pose_estimation_pytorch/apis/analyze_videos.py
+++ b/deeplabcut/pose_estimation_pytorch/apis/analyze_videos.py
@@ -88,7 +88,6 @@ def __next__(self) -> np.ndarray | tuple[str, dict[str, Any]]:
 
 def video_inference(
     video: str | Path | VideoIterator,
-    task: Task,
     pose_runner: InferenceRunner,
     detector_runner: InferenceRunner | None = None,
     cropping: list[int] | None = None,
@@ -99,10 +98,11 @@ def video_inference(
 
     Args:
         video: The video to analyze
-        task: The pose task to run (bottom-up or top-down)
         pose_runner: The pose runner to run inference with
-        detector_runner: When ``task==Task.TOP_DOWN``, the detector runner to obtain
-            bounding boxes for the video.
+        detector_runner: When the pose model is a top-down model, a detector runner can
+            be given to obtain bounding boxes for the video. If the pose model is a
+            top-down model and no detector_runner is given, the bounding boxes must
+            already be set in the VideoIterator (see examples).
         cropping: Optionally, video inference can be run on a cropped version of the
             video. To do so, pass a list containing 4 elements to specify which area
             of the video should be analyzed: ``[xmin, xmax, ymin, ymax]``.
@@ -119,6 +119,41 @@ def video_inference(
         Predictions for each frame in the video. If a shelf_manager is given, this list
         will be empty and the predictions will exclusively be stored in the file written
         by the shelf.
+
+    Examples:
+        Bottom-up video analysis:
+        >>> import deeplabcut.pose_estimation_pytorch as pep
+        >>> model_cfg = pep.config.read_config_as_dict("pytorch_config.yaml")
+        >>> runner = pep.get_pose_inference_runner(model_cfg, "snapshot.pt")
+        >>> video_predictions = pep.video_inference("video.mp4", runner)
+        >>>
+
+        Top-down video analysis:
+        >>> import deeplabcut.pose_estimation_pytorch as pep
+        >>> model_cfg = pep.config.read_config_as_dict("pytorch_config.yaml")
+        >>> runner = pep.get_pose_inference_runner(model_cfg, "snapshot.pt")
+        >>> d_runner = pep.get_pose_inference_runner(model_cfg, "snapshot-detector.pt")
+        >>> video_predictions = pep.video_inference("video.mp4", runner, d_runner)
+        >>>
+
+        Top-Down pose estimation with pre-computed bounding boxes:
+        >>> import numpy as np
+        >>> import deeplabcut.pose_estimation_pytorch as pep
+        >>>
+        >>> video_iterator = pep.VideoIterator("video.mp4")
+        >>> video_iterator.set_context([
+        >>>     { # frame 1 context
+        >>>         "bboxes": np.array([[12, 17, 4, 5]]),  # format (x0, y0, w, h)
+        >>>     },
+        >>>     { # frame 1 context
+        >>>         "bboxes": np.array([[12, 17, 4, 5], [18, 92, 54, 32]]),
+        >>>     },
+        >>>     ...
+        >>> ])
+        >>> model_cfg = pep.config.read_config_as_dict("pytorch_config.yaml")
+        >>> runner = pep.get_pose_inference_runner(model_cfg, "snapshot.pt")
+        >>> video_predictions = pep.video_inference(video_iterator, runner)
+        >>>
     """
     if not isinstance(video, VideoIterator):
         video = VideoIterator(str(video), cropping=cropping)
@@ -134,11 +169,7 @@ def video_inference(
         f"  resolution:             w={vid_w}, h={vid_h}\n"
     )
 
-    if task == Task.TOP_DOWN:
-        # Get bounding boxes for context
-        if detector_runner is None:
-            raise ValueError("Must use a detector for top-down video analysis")
-
+    if detector_runner is not None:
         print(f"Running detector with batch size {detector_runner.batch_size}")
         bbox_predictions = detector_runner.inference(images=tqdm(video))
         video.set_context(bbox_predictions)
@@ -407,7 +438,6 @@ def analyze_videos(
             predictions = video_inference(
                 video=video_iterator,
                 pose_runner=pose_runner,
-                task=pose_task,
                 detector_runner=detector_runner,
                 cropping=cropping,
                 shelf_writer=shelf_writer,
diff --git a/deeplabcut/pose_estimation_pytorch/modelzoo/inference.py b/deeplabcut/pose_estimation_pytorch/modelzoo/inference.py
index 361e1d5a6e..aa0869dc3b 100644
--- a/deeplabcut/pose_estimation_pytorch/modelzoo/inference.py
+++ b/deeplabcut/pose_estimation_pytorch/modelzoo/inference.py
@@ -108,7 +108,6 @@ def _video_inference_superanimal(
         detector_batch_size=detector_batch_size,
         detector_path=detector_snapshot_path,
     )
-    pose_task = Task(model_cfg.get("method", "BU"))
     results = {}
 
     if isinstance(video_paths, str):
@@ -138,7 +137,6 @@ def _video_inference_superanimal(
         video = VideoIterator(video_path, cropping=cropping)
         predictions = video_inference(
             video,
-            task=pose_task,
             pose_runner=pose_runner,
             detector_runner=detector_runner,
         )

From 40d7ab0eba2633c855bd1456803036a7306b4403 Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Fri, 20 Dec 2024 16:30:18 +0100
Subject: [PATCH 06/14] improved README

---
 deeplabcut/pose_estimation_pytorch/README.md | 32 +++++++++++---------
 1 file changed, 17 insertions(+), 15 deletions(-)

diff --git a/deeplabcut/pose_estimation_pytorch/README.md b/deeplabcut/pose_estimation_pytorch/README.md
index 92ba3154f5..fde5ddd871 100644
--- a/deeplabcut/pose_estimation_pytorch/README.md
+++ b/deeplabcut/pose_estimation_pytorch/README.md
@@ -420,17 +420,21 @@ train(
 
 ### Running Video Analysis outside a DeepLabCut Project
 
-DeepLabCut provides high-level APIs (via the GUI or the python package) to analyze your data. The usage of this API assumes the existance of a DLC project (with `config.yaml` file, etc.).
+DeepLabCut provides high-level APIs (via the GUI or the python package) to analyze your
+data. The usage of this API assumes the existence of a DLC project (with `config.yaml`
+file, etc.).
 
-Sometimes it might be more convenient to just run a model on your data via a low-level API. We also use this API under the hood, in particular for the Model Zoo. Check out the example below:
+Sometimes it might be more convenient to just run a model on your data via a low-level
+API. We also use this API under the hood, in particular for the Model Zoo. Check out the
+example below:
 
 ```python
 from pathlib import Path
 
+import deeplabcut.pose_estimation_pytorch as pep
 from deeplabcut.pose_estimation_pytorch.apis.analyze_videos import video_inference
 from deeplabcut.pose_estimation_pytorch.config import read_config_as_dict
 from deeplabcut.pose_estimation_pytorch.task import Task
-from deeplabcut.pose_estimation_pytorch.apis.utils import get_inference_runners
 
 train_dir = Path("/Users/Jaylen/my-dlc-models/train")
 pytorch_config_path = train_dir / "pytorch_config.yaml"
@@ -447,25 +451,23 @@ detector_batch_size = 8
 
 # read model configuration
 model_cfg = read_config_as_dict(pytorch_config_path)
-bodyparts = model_cfg["metadata"]["bodyparts"]
-unique_bodyparts = model_cfg["metadata"]["unique_bodyparts"]
-with_identity = model_cfg["metadata"].get("with_identity", False)
-
 pose_task = Task(model_cfg["method"])
-pose_runner, detector_runner = get_inference_runners(
+pose_runner = pep.get_pose_inference_runner(
     model_config=model_cfg,
     snapshot_path=snapshot_path,
     max_individuals=max_num_animals,
-    num_bodyparts=len(bodyparts),
-    num_unique_bodyparts=len(unique_bodyparts),
     batch_size=batch_size,
-    with_identity=with_identity,
-    transform=None,
-    detector_batch_size=detector_batch_size,
-    detector_path=detector_snapshot_path,
-    detector_transform=None,
 )
 
+detector_runner = None
+if pose_task == pep.Task.TOP_DOWN:
+    detector_runner = pep.get_detector_inference_runner(
+        model_config=model_cfg,
+        snapshot_path=detector_snapshot_path,
+        max_individuals=max_num_animals,
+        batch_size=detector_batch_size,
+    )
+
 predictions = video_inference(
     video=video_path,
     pose_runner=pose_runner,

From 5b913321d8b18f02c59140a488174a9c34d5209d Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Fri, 20 Dec 2024 16:46:42 +0100
Subject: [PATCH 07/14] fix typo

---
 deeplabcut/pose_estimation_pytorch/README.md | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/deeplabcut/pose_estimation_pytorch/README.md b/deeplabcut/pose_estimation_pytorch/README.md
index 659d6ce0f3..89511411da 100644
--- a/deeplabcut/pose_estimation_pytorch/README.md
+++ b/deeplabcut/pose_estimation_pytorch/README.md
@@ -421,11 +421,7 @@ train(
 ### Running Video Analysis outside a DeepLabCut Project
 
 DeepLabCut provides high-level APIs (via the GUI or the python package) to analyze your
-<<<<<<< HEAD
 data. The usage of this API assumes the existence of a DLC project (with `config.yaml`
-=======
-data. The usage of this API assumes the existance of a DLC project (with `config.yaml`
->>>>>>> pytorch_dlc
 file, etc.).
 
 Sometimes it might be more convenient to just run a model on your data via a low-level

From cd590c79aae8105f5d77ad35be54869e1852a5a5 Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Fri, 20 Dec 2024 17:00:25 +0100
Subject: [PATCH 08/14] improve imports

---
 .../pose_estimation_pytorch/__init__.py       | 22 ++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/deeplabcut/pose_estimation_pytorch/__init__.py b/deeplabcut/pose_estimation_pytorch/__init__.py
index d4256cfba0..3a28d7106e 100644
--- a/deeplabcut/pose_estimation_pytorch/__init__.py
+++ b/deeplabcut/pose_estimation_pytorch/__init__.py
@@ -33,13 +33,25 @@
     available_detectors,
     available_models,
 )
-from deeplabcut.pose_estimation_pytorch.data.base import Loader
-from deeplabcut.pose_estimation_pytorch.data.cocoloader import COCOLoader
-from deeplabcut.pose_estimation_pytorch.data.dataset import (
+from deeplabcut.pose_estimation_pytorch.data import (
+    build_transforms,
+    COCOLoader,
+    COLLATE_FUNCTIONS,
+    DLCLoader,
+    Loader,
     PoseDataset,
     PoseDatasetParameters,
 )
-from deeplabcut.pose_estimation_pytorch.data.dlcloader import DLCLoader
-from deeplabcut.pose_estimation_pytorch.runners.snapshots import TorchSnapshotManager
+from deeplabcut.pose_estimation_pytorch.runners import (
+    build_inference_runner,
+    build_training_runner,
+    DetectorInferenceRunner,
+    DetectorTrainingRunner,
+    InferenceRunner,
+    PoseInferenceRunner,
+    PoseTrainingRunner,
+    TorchSnapshotManager,
+    TrainingRunner,
+)
 from deeplabcut.pose_estimation_pytorch.task import Task
 from deeplabcut.pose_estimation_pytorch.utils import fix_seeds

From d8eed88141b5616d5d75042faee05b6b611508be Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Fri, 20 Dec 2024 17:04:22 +0100
Subject: [PATCH 09/14] improve imports in README

---
 deeplabcut/pose_estimation_pytorch/README.md | 55 +++++++-------------
 1 file changed, 20 insertions(+), 35 deletions(-)

diff --git a/deeplabcut/pose_estimation_pytorch/README.md b/deeplabcut/pose_estimation_pytorch/README.md
index 89511411da..73f19e1b88 100644
--- a/deeplabcut/pose_estimation_pytorch/README.md
+++ b/deeplabcut/pose_estimation_pytorch/README.md
@@ -64,20 +64,17 @@ disk.
 ```python
 from pathlib import Path
 
-from deeplabcut.pose_estimation_pytorch.config import (
-    make_pytorch_pose_config,
-    write_config,
-)
+import deeplabcut.pose_estimation_pytorch as pep
 
 project_cfg = { "Task": "mice", ... }  # the configuration for your DLC project
 pose_config_path = Path("/path/to/my/config/pytorch_cfg.yaml")
-model_cfg = make_pytorch_pose_config(
+model_cfg = pep.config.make_pytorch_pose_config(
     project_config=project_cfg,
     pose_config_path=pose_config_path,
     net_type="hrnet_w32",
     top_down=True,
+    save=True,
 )
-write_config(pose_config_path, model_cfg)
 ```
 
 #### Adding Models
@@ -238,9 +235,9 @@ dataset creation and test/train splitting. The `DLCLoader` class is used to load
 labeled data for a specific shuffle.
 
 ```python3
-from deeplabcut.pose_estimation_pytorch.data import DLCLoader
+import deeplabcut.pose_estimation_pytorch as pep
 
-loader = DLCLoader(
+loader = pep.DLCLoader(
     config="/path/to/my/project/config.yaml",
     trainset_index=0,
     shuffle=1,
@@ -265,23 +262,20 @@ images and keypoints to a tensor dataset for training and evaluation. You can ge
 an instance of training/test dataset with your `DLCLoader`:
 
 ```python3
-from deeplabcut.pose_estimation_pytorch.data import (
-    build_transforms,
-    DLCLoader,
-)
+import deeplabcut.pose_estimation_pytorch as pep
 
-loader = DLCLoader(
+loader = pep.DLCLoader(
     config="/path/to/my/project/config.yaml",
     trainset_index=0,
     shuffle=1,
 )
 train_dataset = loader.create_dataset(
-    transform=build_transforms(loader.model_cfg["data"]["train"]),
+    transform=pep.build_transforms(loader.model_cfg["data"]["train"]),
     mode="train",
     task=loader.pose_task,
 )
 valid_dataset = loader.create_dataset(
-    transform=build_transforms(loader.model_cfg["data"]["train"]),
+    transform=pep.build_transforms(loader.model_cfg["data"]["train"]),
     mode="test",
     task=loader.pose_task,
 )
@@ -320,15 +314,7 @@ configuration file (as described in [#model_configuration_files]).
 ```python3
 from pathlib import Path
 
-from deeplabcut.pose_estimation_pytorch.config import (
-    make_pytorch_pose_config,
-    write_config,
-)
-from deeplabcut.pose_estimation_pytorch.data import (
-    build_transforms,
-    COCOLoader,
-)
-from deeplabcut.pose_estimation_pytorch.task import Task
+import deeplabcut.pose_estimation_pytorch as pep
 
 # Specify project paths
 project_root = Path("/path/to/my/COCOProject")
@@ -336,14 +322,14 @@ train_json_filename = "train.json"
 test_json_filename = "test.json"
 
 # Parse information about the project
-train_dict = COCOLoader.load_json(project_root, filename=train_json_filename)
-max_num_individuals, bodyparts = COCOLoader.get_project_parameters(train_dict)
+train_dict = pep.COCOLoader.load_json(project_root, filename=train_json_filename)
+max_num_individuals, bodyparts = pep.COCOLoader.get_project_parameters(train_dict)
 
 # Generate a configuration file for your PyTorch model
 # In this case, it's for a Top-Down HRNet_w32
 experiment_path = project_root / "experiments" / "hrnet_w32"
 model_cfg_path = experiment_path / "train" / "pytorch_cfg.yaml"
-model_cfg = make_pytorch_pose_config(
+model_cfg = pep.config.make_pytorch_pose_config(
     project_config={
         "project_path": str(project_root.resolve()),
         "multianimalproject": max_num_individuals > 1,
@@ -355,23 +341,23 @@ model_cfg = make_pytorch_pose_config(
     pose_config_path=experiment_path,
     net_type="hrnet_w32",
     top_down=True,
+    save=True,
 )
-write_config(config_path=model_cfg_path, config=model_cfg)
 
 # Create the loader for the COCO dataset
-loader = COCOLoader(
+loader = pep.COCOLoader(
     project_root=project_root,
     model_config_path="/path/to/my/project/experiments/pytorch_config.yaml",
     train_json_filename=train_json_filename,
     test_json_filename=test_json_filename,
 )
 train_dataset = loader.create_dataset(
-    transform=build_transforms(loader.model_cfg["data"]["train"]),
+    transform=pep.build_transforms(loader.model_cfg["data"]["train"]),
     mode="train",
     task=loader.pose_task,
 )
 valid_dataset = loader.create_dataset(
-    transform=build_transforms(loader.model_cfg["data"]["train"]),
+    transform=pep.build_transforms(loader.model_cfg["data"]["train"]),
     mode="test",
     task=loader.pose_task,
 )
@@ -389,16 +375,15 @@ pretrained weights, and either train them or run inference with them.
 ```python
 from pathlib import Path
 
+import deeplabcut.pose_estimation_pytorch as pep
 from deeplabcut.pose_estimation_pytorch.apis.train import train
-from deeplabcut.pose_estimation_pytorch.data import COCOLoader
-from deeplabcut.pose_estimation_pytorch.task import Task
 
 # Specify project paths
 project_root = Path("/path/to/my/COCOProject")
 train_json_filename = "train.json"
 test_json_filename = "test.json"
 
-loader = COCOLoader(
+loader = pep.COCOLoader(
     project_root=project_root,
     model_config_path="/path/to/my/project/experiments/pytorch_config.yaml",
     train_json_filename=train_json_filename,
@@ -407,7 +392,7 @@ loader = COCOLoader(
 train(
     loader=loader,
     run_config=loader.model_cfg,
-    task=Task(loader.model_cfg["method"]),
+    task=pep.Task(loader.model_cfg["method"]),
     device="cuda:2",
     logger_config=dict(
         type="WandbLogger",

From 1af7a548c0352feb36e28ad3830c6ac0a7520b7b Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Mon, 6 Jan 2025 11:52:24 +0100
Subject: [PATCH 10/14] fix doc

---
 deeplabcut/pose_estimation_pytorch/data/base.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/deeplabcut/pose_estimation_pytorch/data/base.py b/deeplabcut/pose_estimation_pytorch/data/base.py
index fc14818032..e854a88120 100644
--- a/deeplabcut/pose_estimation_pytorch/data/base.py
+++ b/deeplabcut/pose_estimation_pytorch/data/base.py
@@ -160,7 +160,7 @@ def ground_truth_bboxes(self, mode: str = "train") -> dict[str, dict]:
             the format:
                 {
                     'path/to/image000.png': {
-                        "width": (int) the width of the image, in
+                        "width": (int) the width of the image, in pixels
                         "height": (int) the height of the image, in pixels
                         "bboxes": (np.ndarray) bboxes with shape (num_individuals, xywh)
                     },

From 9a5a3d37828fccd2d593392b1a954919245b5f57 Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Mon, 6 Jan 2025 13:44:13 +0100
Subject: [PATCH 11/14] pass load_head_weights to train method

---
 deeplabcut/pose_estimation_pytorch/apis/train.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/deeplabcut/pose_estimation_pytorch/apis/train.py b/deeplabcut/pose_estimation_pytorch/apis/train.py
index 41269d4924..a3f249c910 100644
--- a/deeplabcut/pose_estimation_pytorch/apis/train.py
+++ b/deeplabcut/pose_estimation_pytorch/apis/train.py
@@ -358,6 +358,7 @@ def train_network(
             logger_config=loader.model_cfg.get("logger"),
             snapshot_path=snapshot_path,
             max_snapshots_to_keep=max_snapshots_to_keep,
+            load_head_weights=load_head_weights,
         )
 
     destroy_file_logging()

From a426394ac082da4270bf57797e75b1a5090a0993 Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Mon, 6 Jan 2025 14:08:48 +0100
Subject: [PATCH 12/14] add train to pep init

---
 deeplabcut/pose_estimation_pytorch/__init__.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/deeplabcut/pose_estimation_pytorch/__init__.py b/deeplabcut/pose_estimation_pytorch/__init__.py
index 084633941d..d8582096af 100644
--- a/deeplabcut/pose_estimation_pytorch/__init__.py
+++ b/deeplabcut/pose_estimation_pytorch/__init__.py
@@ -24,6 +24,7 @@
     get_pose_inference_runner,
     predict,
     superanimal_analyze_images,
+    train,
     train_network,
     video_inference,
     VideoIterator,

From fc2cee5320285e31f05ac2db4709801e18ae6e22 Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Mon, 6 Jan 2025 15:46:31 +0100
Subject: [PATCH 13/14] add make_basic_project_config helper

---
 deeplabcut/pose_estimation_pytorch/README.md  | 14 ++--
 .../config/__init__.py                        |  1 +
 .../config/make_pose_config.py                | 69 +++++++++++++++++++
 .../config/test_make_pose_config.py           | 28 +++++++-
 4 files changed, 103 insertions(+), 9 deletions(-)

diff --git a/deeplabcut/pose_estimation_pytorch/README.md b/deeplabcut/pose_estimation_pytorch/README.md
index 73f19e1b88..0b907e7923 100644
--- a/deeplabcut/pose_estimation_pytorch/README.md
+++ b/deeplabcut/pose_estimation_pytorch/README.md
@@ -330,14 +330,12 @@ max_num_individuals, bodyparts = pep.COCOLoader.get_project_parameters(train_dic
 experiment_path = project_root / "experiments" / "hrnet_w32"
 model_cfg_path = experiment_path / "train" / "pytorch_cfg.yaml"
 model_cfg = pep.config.make_pytorch_pose_config(
-    project_config={
-        "project_path": str(project_root.resolve()),
-        "multianimalproject": max_num_individuals > 1,
-        "bodyparts": bodyparts,
-        "multianimalbodyparts": bodyparts,
-        "uniquebodyparts": [],
-        "individuals": [f"idv{i}" for i in range(max_num_individuals)],
-    },
+    project_config=pep.config.make_basic_project_config(
+        dataset_path=str(project_root.resolve()),
+        bodyparts=bodyparts,
+        max_individuals=max_num_individuals,
+        multi_animal=True,
+    ),
     pose_config_path=experiment_path,
     net_type="hrnet_w32",
     top_down=True,
diff --git a/deeplabcut/pose_estimation_pytorch/config/__init__.py b/deeplabcut/pose_estimation_pytorch/config/__init__.py
index 2cf1468e63..3e62f52750 100644
--- a/deeplabcut/pose_estimation_pytorch/config/__init__.py
+++ b/deeplabcut/pose_estimation_pytorch/config/__init__.py
@@ -9,6 +9,7 @@
 # Licensed under GNU Lesser General Public License v3.0
 #
 from deeplabcut.pose_estimation_pytorch.config.make_pose_config import (
+    make_basic_project_config,
     make_pytorch_pose_config,
     make_pytorch_test_config,
 )
diff --git a/deeplabcut/pose_estimation_pytorch/config/make_pose_config.py b/deeplabcut/pose_estimation_pytorch/config/make_pose_config.py
index 65cdff2d1d..25aa6b9050 100644
--- a/deeplabcut/pose_estimation_pytorch/config/make_pose_config.py
+++ b/deeplabcut/pose_estimation_pytorch/config/make_pose_config.py
@@ -217,6 +217,75 @@ def make_pytorch_test_config(
     return test_config
 
 
+def make_basic_project_config(
+    dataset_path: Path | str,
+    bodyparts: list[str],
+    max_individuals: int,
+    multi_animal: bool = True,
+) -> dict:
+    """Creates a basic configuration dict that can be used to create model configs.
+
+    This should be used to create the `project_config` given to
+    `make_pytorch_pose_config` for non-DeepLabCut projects (e.g. when creating a
+    configuration file for a model that will be trained on a COCO dataset).
+
+    Args:
+        dataset_path: The path to the dataset for which the config will be created.
+        bodyparts: The bodyparts labeled for individuals in the dataset.
+        max_individuals: The maximum number of individuals to detect in a single image.
+        multi_animal: Whether multiple animals can be present in an image.
+
+    Returns:
+        The created project configuration dict that can be given to
+        `make_pytorch_pose_config`.
+
+    Examples:
+        Creating a `pytorch_config` for a ResNet50 backbone with a part-affinity head (
+        as multi_animal=True and top_down=False)
+
+        >>> import deeplabcut.pose_estimation_pytorch as pep
+        >>> project_config = pep.config.make_basic_project_config(
+        >>>     dataset_path="/path/coco",
+        >>>     bodyparts=["nose", "left_eye", "right_eye"],
+        >>>     max_individuals=12,
+        >>>     multi_animal=True,
+        >>> )
+        >>> model_config = pep.config.make_pytorch_pose_config(
+        >>>     project_config=project_config,
+        >>>     pose_config_path="/path/coco/models/resnet50/pytorch_config.yaml",
+        >>>     net_type="resnet_50",
+        >>>     top_down=False,
+        >>>     save=True,
+        >>> )
+
+        Creating a `pytorch_config` for a ResNet50 backbone with a simple heatmap head
+        (as the project is single-animal):
+
+        >>> import deeplabcut.pose_estimation_pytorch as pep
+        >>> project_config = pep.config.make_basic_project_config(
+        >>>     dataset_path="/path/coco",
+        >>>     bodyparts=["nose", "left_eye", "right_eye"],
+        >>>     max_individuals=1,
+        >>>     multi_animal=False,
+        >>> )
+        >>> model_config = pep.config.make_pytorch_pose_config(
+        >>>     project_config=project_config,
+        >>>     pose_config_path="/path/coco/models/resnet50/pytorch_config.yaml",
+        >>>     net_type="resnet_50",
+        >>>     top_down=False,
+        >>>     save=True,
+        >>> )
+    """
+    return dict(
+        project_path=str(dataset_path),
+        multianimalproject=multi_animal,
+        bodyparts=bodyparts,
+        multianimalbodyparts=bodyparts,
+        uniquebodyparts=[],
+        individuals=[f"individual{i:03d}" for i in range(max_individuals)],
+    )
+
+
 def add_metadata(
     project_config: dict, config: dict, pose_config_path: str | Path
 ) -> dict:
diff --git a/tests/pose_estimation_pytorch/config/test_make_pose_config.py b/tests/pose_estimation_pytorch/config/test_make_pose_config.py
index aaa1f6566b..4208d96ab0 100644
--- a/tests/pose_estimation_pytorch/config/test_make_pose_config.py
+++ b/tests/pose_estimation_pytorch/config/test_make_pose_config.py
@@ -11,7 +11,11 @@
 """Tests the pre-processors"""
 import pytest
 
-from deeplabcut.pose_estimation_pytorch.config.make_pose_config import make_pytorch_pose_config
+import deeplabcut.utils.auxiliaryfunctions as af
+from deeplabcut.pose_estimation_pytorch.config.make_pose_config import (
+    make_basic_project_config,
+    make_pytorch_pose_config,
+)
 from deeplabcut.pose_estimation_pytorch.config.utils import pretty_print, update_config
 
 
@@ -417,3 +421,25 @@ def _make_project_config(
         project_config["bodyparts"] = bodyparts
 
     return project_config
+
+
+@pytest.mark.parametrize("bodyparts", [["nose"], ["nose", "ear", "eye"]])
+@pytest.mark.parametrize("max_idv", [1, 12, 20])
+@pytest.mark.parametrize("multi", [True, False])
+def test_make_basic_project_config(bodyparts: list[str], max_idv: int, multi: bool):
+    if not multi and max_idv > 1:
+        return
+
+    project_config = make_basic_project_config(
+        dataset_path="path/dataset",
+        bodyparts=bodyparts,
+        max_individuals=max_idv,
+        multi_animal=multi,
+    )
+
+    bpts = af.get_bodyparts(project_config)
+    assert bodyparts == bpts
+
+    individuals = project_config["individuals"]
+    assert len(individuals) == max_idv
+    assert len(set(individuals)) == max_idv

From ee03401ba35d7f0c2463c42c349438d0341ea3a5 Mon Sep 17 00:00:00 2001
From: Niels Poulsen <niels.poulsen@epfl.ch>
Date: Fri, 10 Jan 2025 10:08:35 +0100
Subject: [PATCH 14/14] import style: import deeplabcut.pose_estimation_pytorch
 as dlc_torch

---
 deeplabcut/pose_estimation_pytorch/README.md | 60 ++++++++++----------
 1 file changed, 30 insertions(+), 30 deletions(-)

diff --git a/deeplabcut/pose_estimation_pytorch/README.md b/deeplabcut/pose_estimation_pytorch/README.md
index 0b907e7923..831c618ca3 100644
--- a/deeplabcut/pose_estimation_pytorch/README.md
+++ b/deeplabcut/pose_estimation_pytorch/README.md
@@ -64,11 +64,11 @@ disk.
 ```python
 from pathlib import Path
 
-import deeplabcut.pose_estimation_pytorch as pep
+import deeplabcut.pose_estimation_pytorch as dlc_torch
 
 project_cfg = { "Task": "mice", ... }  # the configuration for your DLC project
 pose_config_path = Path("/path/to/my/config/pytorch_cfg.yaml")
-model_cfg = pep.config.make_pytorch_pose_config(
+model_cfg = dlc_torch.config.make_pytorch_pose_config(
     project_config=project_cfg,
     pose_config_path=pose_config_path,
     net_type="hrnet_w32",
@@ -235,9 +235,9 @@ dataset creation and test/train splitting. The `DLCLoader` class is used to load
 labeled data for a specific shuffle.
 
 ```python3
-import deeplabcut.pose_estimation_pytorch as pep
+import deeplabcut.pose_estimation_pytorch as dlc_torch
 
-loader = pep.DLCLoader(
+loader = dlc_torch.DLCLoader(
     config="/path/to/my/project/config.yaml",
     trainset_index=0,
     shuffle=1,
@@ -262,20 +262,20 @@ images and keypoints to a tensor dataset for training and evaluation. You can ge
 an instance of training/test dataset with your `DLCLoader`:
 
 ```python3
-import deeplabcut.pose_estimation_pytorch as pep
+import deeplabcut.pose_estimation_pytorch as dlc_torch
 
-loader = pep.DLCLoader(
+loader = dlc_torch.DLCLoader(
     config="/path/to/my/project/config.yaml",
     trainset_index=0,
     shuffle=1,
 )
 train_dataset = loader.create_dataset(
-    transform=pep.build_transforms(loader.model_cfg["data"]["train"]),
+    transform=dlc_torch.build_transforms(loader.model_cfg["data"]["train"]),
     mode="train",
     task=loader.pose_task,
 )
 valid_dataset = loader.create_dataset(
-    transform=pep.build_transforms(loader.model_cfg["data"]["train"]),
+    transform=dlc_torch.build_transforms(loader.model_cfg["data"]["train"]),
     mode="test",
     task=loader.pose_task,
 )
@@ -314,7 +314,7 @@ configuration file (as described in [#model_configuration_files]).
 ```python3
 from pathlib import Path
 
-import deeplabcut.pose_estimation_pytorch as pep
+import deeplabcut.pose_estimation_pytorch as dlc_torch
 
 # Specify project paths
 project_root = Path("/path/to/my/COCOProject")
@@ -322,15 +322,15 @@ train_json_filename = "train.json"
 test_json_filename = "test.json"
 
 # Parse information about the project
-train_dict = pep.COCOLoader.load_json(project_root, filename=train_json_filename)
-max_num_individuals, bodyparts = pep.COCOLoader.get_project_parameters(train_dict)
+train_dict = dlc_torch.COCOLoader.load_json(project_root, filename=train_json_filename)
+max_num_individuals, bodyparts = dlc_torch.COCOLoader.get_project_parameters(train_dict)
 
 # Generate a configuration file for your PyTorch model
 # In this case, it's for a Top-Down HRNet_w32
 experiment_path = project_root / "experiments" / "hrnet_w32"
 model_cfg_path = experiment_path / "train" / "pytorch_cfg.yaml"
-model_cfg = pep.config.make_pytorch_pose_config(
-    project_config=pep.config.make_basic_project_config(
+model_cfg = dlc_torch.config.make_pytorch_pose_config(
+    project_config=dlc_torch.config.make_basic_project_config(
         dataset_path=str(project_root.resolve()),
         bodyparts=bodyparts,
         max_individuals=max_num_individuals,
@@ -343,19 +343,19 @@ model_cfg = pep.config.make_pytorch_pose_config(
 )
 
 # Create the loader for the COCO dataset
-loader = pep.COCOLoader(
+loader = dlc_torch.COCOLoader(
     project_root=project_root,
     model_config_path="/path/to/my/project/experiments/pytorch_config.yaml",
     train_json_filename=train_json_filename,
     test_json_filename=test_json_filename,
 )
 train_dataset = loader.create_dataset(
-    transform=pep.build_transforms(loader.model_cfg["data"]["train"]),
+    transform=dlc_torch.build_transforms(loader.model_cfg["data"]["train"]),
     mode="train",
     task=loader.pose_task,
 )
 valid_dataset = loader.create_dataset(
-    transform=pep.build_transforms(loader.model_cfg["data"]["train"]),
+    transform=dlc_torch.build_transforms(loader.model_cfg["data"]["train"]),
     mode="test",
     task=loader.pose_task,
 )
@@ -373,7 +373,7 @@ pretrained weights, and either train them or run inference with them.
 ```python
 from pathlib import Path
 
-import deeplabcut.pose_estimation_pytorch as pep
+import deeplabcut.pose_estimation_pytorch as dlc_torch
 from deeplabcut.pose_estimation_pytorch.apis.train import train
 
 # Specify project paths
@@ -381,7 +381,7 @@ project_root = Path("/path/to/my/COCOProject")
 train_json_filename = "train.json"
 test_json_filename = "test.json"
 
-loader = pep.COCOLoader(
+loader = dlc_torch.COCOLoader(
     project_root=project_root,
     model_config_path="/path/to/my/project/experiments/pytorch_config.yaml",
     train_json_filename=train_json_filename,
@@ -390,7 +390,7 @@ loader = pep.COCOLoader(
 train(
     loader=loader,
     run_config=loader.model_cfg,
-    task=pep.Task(loader.model_cfg["method"]),
+    task=dlc_torch.Task(loader.model_cfg["method"]),
     device="cuda:2",
     logger_config=dict(
         type="WandbLogger",
@@ -414,7 +414,7 @@ example below:
 ```python
 from pathlib import Path
 
-import deeplabcut.pose_estimation_pytorch as pep
+import deeplabcut.pose_estimation_pytorch as dlc_torch
 
 train_dir = Path("/Users/Jaylen/my-dlc-models/train")
 pytorch_config_path = train_dir / "pytorch_config.yaml"
@@ -430,9 +430,9 @@ batch_size = 16
 detector_batch_size = 8
 
 # read model configuration
-model_cfg = pep.config.read_config_as_dict(pytorch_config_path)
-pose_task = pep.Task(model_cfg["method"])
-pose_runner = pep.get_pose_inference_runner(
+model_cfg = dlc_torch.config.read_config_as_dict(pytorch_config_path)
+pose_task = dlc_torch.Task(model_cfg["method"])
+pose_runner = dlc_torch.get_pose_inference_runner(
     model_config=model_cfg,
     snapshot_path=snapshot_path,
     max_individuals=max_num_animals,
@@ -440,15 +440,15 @@ pose_runner = pep.get_pose_inference_runner(
 )
 
 detector_runner = None
-if pose_task == pep.Task.TOP_DOWN:
-    detector_runner = pep.get_detector_inference_runner(
+if pose_task == dlc_torch.Task.TOP_DOWN:
+    detector_runner = dlc_torch.get_detector_inference_runner(
         model_config=model_cfg,
         snapshot_path=detector_snapshot_path,
         max_individuals=max_num_animals,
         batch_size=detector_batch_size,
     )
 
-predictions = pep.video_inference(
+predictions = dlc_torch.video_inference(
     video=video_path,
     pose_runner=pose_runner,
     detector_runner=detector_runner,
@@ -471,11 +471,11 @@ You can easily do so by writing a bit of custom code, as shown in the example be
 from pathlib import Path
 
 import numpy as np
-import deeplabcut.pose_estimation_pytorch as pep
+import deeplabcut.pose_estimation_pytorch as dlc_torch
 from tqdm import tqdm
 
 # create an iterator for your video
-video = pep.VideoIterator("/Users/Jayson/my-cool-video.mp4")
+video = dlc_torch.VideoIterator("/Users/Jayson/my-cool-video.mp4")
 
 # dummy bboxes - you can load yours from a file or in another way
 #  the bboxes should be in `xywh` format, i.e. (x_top_left, y_top_left, width, height)
@@ -495,8 +495,8 @@ video.set_context(bounding_boxes)
 max_individuals = np.max([len(context["bboxes"]) for context in bounding_boxes])
 
 # run inference!
-model_cfg = pep.config.read_config_as_dict("/Users/Jayson/pytorch_config.yaml")
-pose_runner = pep.get_pose_inference_runner(
+model_cfg = dlc_torch.config.read_config_as_dict("/Users/Jayson/pytorch_config.yaml")
+pose_runner = dlc_torch.get_pose_inference_runner(
     model_config=model_cfg,
     snapshot_path=Path("/Users/Jayson/model-snapshot.pt"),
     max_individuals=max_individuals,