DeepLabCut · n-poulsen · Jan 10, 2025 · Dec 19, 2024 · Dec 19, 2024 · Dec 19, 2024
diff --git 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 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 = 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",
     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 dlc_torch
 
-loader = DLCLoader(
+loader = dlc_torch.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 dlc_torch
 
-loader = DLCLoader(
+loader = dlc_torch.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=dlc_torch.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=dlc_torch.build_transforms(loader.model_cfg["data"]["train"]),
     mode="test",
     task=loader.pose_task,
 )
@@ -320,58 +314,48 @@ 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 dlc_torch
 
 # Specify project paths
 project_root = Path("/path/to/my/COCOProject")
 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 = 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 = 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)],
-    },
+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,
+        multi_animal=True,
+    ),
     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 = 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=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=build_transforms(loader.model_cfg["data"]["train"]),
+    transform=dlc_torch.build_transforms(loader.model_cfg["data"]["train"]),
     mode="test",
     task=loader.pose_task,
 )
@@ -389,16 +373,15 @@ pretrained weights, and either train them or run inference with them.
 ```python
 from pathlib import Path
 
+import deeplabcut.pose_estimation_pytorch as dlc_torch
 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 = dlc_torch.COCOLoader(
     project_root=project_root,
     model_config_path="/path/to/my/project/experiments/pytorch_config.yaml",
     train_json_filename=train_json_filename,
@@ -407,7 +390,7 @@ loader = COCOLoader(
 train(
     loader=loader,
     run_config=loader.model_cfg,
-    task=Task(loader.model_cfg["method"]),
+    task=dlc_torch.Task(loader.model_cfg["method"]),
     device="cuda:2",
     logger_config=dict(
         type="WandbLogger",
@@ -421,7 +404,7 @@ 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`
+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
@@ -431,10 +414,7 @@ example below:
 ```python
 from pathlib import Path
 
-from deeplabcut.pose_estimation_pytorch import Task
-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.apis.utils import get_inference_runners
+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"
@@ -450,21 +430,26 @@ batch_size = 16
 detector_batch_size = 8
 
 # read model configuration
-model_cfg = read_config_as_dict(pytorch_config_path)
-
-pose_task = Task(model_cfg["method"])
-pose_runner, detector_runner = get_inference_runners(
+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,
     batch_size=batch_size,
-    detector_batch_size=detector_batch_size,
-    detector_path=detector_snapshot_path,
 )
 
-predictions = video_inference(
+detector_runner = None
+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 = dlc_torch.video_inference(
     video=video_path,
-    task=pose_task,
     pose_runner=pose_runner,
     detector_runner=detector_runner,
 )
@@ -486,13 +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
-from deeplabcut.pose_estimation_pytorch import get_inference_runners
-from deeplabcut.pose_estimation_pytorch.apis import VideoIterator
-from deeplabcut.pose_estimation_pytorch.config import read_config_as_dict
+import deeplabcut.pose_estimation_pytorch as dlc_torch
 from tqdm import tqdm
 
 # create an iterator for your video
-video = 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)
@@ -512,8 +495,8 @@ video.set_context(bounding_boxes)
 max_individuals = np.max([len(context["bboxes"]) for context in bounding_boxes])
 
 # run inference!
-model_cfg = read_config_as_dict("/Users/Jayson/pytorch_config.yaml")
-pose_runner, _ = get_inference_runners(
+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,

diff --git a/deeplabcut/pose_estimation_pytorch/__init__.py b/deeplabcut/pose_estimation_pytorch/__init__.py
@@ -8,31 +8,53 @@
 #
 # 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,
-    get_inference_runners,
+    evaluate,
     evaluate_network,
     extract_maps,
     extract_save_all_maps,
+    get_detector_inference_runner,
+    get_pose_inference_runner,
+    predict,
+    superanimal_analyze_images,
+    train,
     train_network,
+    video_inference,
+    VideoIterator,
+    visualize_predictions,
 )
 from deeplabcut.pose_estimation_pytorch.config import (
     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.base import (
+from deeplabcut.pose_estimation_pytorch.runners import (
+    build_inference_runner,
+    build_training_runner,
+    DetectorInferenceRunner,
+    DetectorTrainingRunner,
     get_load_weights_only,
+    InferenceRunner,
+    PoseInferenceRunner,
+    PoseTrainingRunner,
     set_load_weights_only,
+    TorchSnapshotManager,
+    TrainingRunner,
 )
-from deeplabcut.pose_estimation_pytorch.runners.snapshots import TorchSnapshotManager
 from deeplabcut.pose_estimation_pytorch.task import Task
 from deeplabcut.pose_estimation_pytorch.utils import fix_seeds
diff --git a/deeplabcut/pose_estimation_pytorch/apis/__init__.py b/deeplabcut/pose_estimation_pytorch/apis/__init__.py
@@ -12,6 +12,7 @@
 from deeplabcut.pose_estimation_pytorch.apis.analyze_images import (
     analyze_image_folder,
     analyze_images,
+    analyze_image_folder,
     superanimal_analyze_images,
 )
 from deeplabcut.pose_estimation_pytorch.apis.analyze_videos import (
@@ -23,8 +24,10 @@
     convert_detections2tracklets,
 )
 from deeplabcut.pose_estimation_pytorch.apis.evaluate import (
+    predict,
     evaluate,
     evaluate_network,
+    visualize_predictions,
 )
 from deeplabcut.pose_estimation_pytorch.apis.export import export_model
 from deeplabcut.pose_estimation_pytorch.apis.train import (
@@ -37,6 +40,12 @@
     get_pose_inference_runner,
 )
 from deeplabcut.pose_estimation_pytorch.apis.visualization import (
+    create_labeled_images,
     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,
+)