DeepLabCut · MMathisLab · Apr 14, 2025 · Feb 2, 2024 · Feb 2, 2024 · Feb 2, 2024
diff --git a/.codespellrc b/.codespellrc
@@ -1,5 +1,5 @@
 [codespell]
-skip = .git,*.pdf,*.svg,deeplabcut/pose_estimation_tensorflow/models/pretrained
+skip = .git,*.pdf,*.svg,*.ipynb,deeplabcut/pose_estimation_tensorflow/models/pretrained
 # MOT,SIE - legit acronyms
 # tThe - for \tThe. codespell is not good detecting those yet
 ignore-words-list = mot,sie,tthe,assertin,bu,td,ctd,wither
diff --git a/conda-environments/DEEPLABCUT.yaml b/conda-environments/DEEPLABCUT.yaml
@@ -35,4 +35,4 @@ dependencies:
   - pip:
     - torch
     - torchvision
-    - deeplabcut[gui,modelzoo,wandb]==3.0.0rc7
+    - deeplabcut[gui,modelzoo,wandb]==3.0.0rc8
diff --git a/deeplabcut/core/inferenceutils.py b/deeplabcut/core/inferenceutils.py
@@ -961,6 +961,9 @@ def calc_object_keypoint_similarity(
 
     k_squared = (2 * sigma) ** 2
     denom = 2 * scale_squared * k_squared
+    if isinstance(sigma, np.ndarray):
+        denom = denom[visible_gt]
+
     if symmetric_kpts is None:
         pred = xy_pred[visible_gt]
         pred[np.isnan(pred)] = np.inf

diff --git a/deeplabcut/core/metrics/api.py b/deeplabcut/core/metrics/api.py
@@ -24,7 +24,7 @@ def compute_metrics(
     unique_bodypart_poses: dict[str, np.ndarray] | None = None,
     pcutoff: float = -1,
     oks_bbox_margin: int = 0,
-    oks_sigma: float = 0.1,
+    oks_sigma: float | np.ndarray = 0.1,
     per_keypoint_rmse: bool = False,
     compute_detection_rmse: bool = True,
 ) -> dict:

diff --git a/deeplabcut/core/metrics/distance_metrics.py b/deeplabcut/core/metrics/distance_metrics.py
@@ -21,7 +21,7 @@
 def compute_oks_matrix(
     ground_truth: np.ndarray,
     predictions: np.ndarray,
-    oks_sigma: float,
+    oks_sigma: float | np.ndarray,
     oks_bbox_margin: float = 0.0,
 ) -> np.ndarray:
     """Computes the OKS score for each (prediction, gt) pair in an image
@@ -53,7 +53,7 @@ def compute_oks_matrix(
 def compute_oks(
     data: list[tuple[np.ndarray, np.ndarray]],
     oks_bbox_margin: float = 0.0,
-    oks_sigma: float = 0.1,
+    oks_sigma: float | np.ndarray = 0.1,
     oks_thresholds: np.ndarray | None = None,
     oks_recall_thresholds: np.ndarray | None = None,
 ) -> dict[str, float]:

diff --git a/deeplabcut/core/trackingutils.py b/deeplabcut/core/trackingutils.py
@@ -28,6 +28,7 @@
 
 TRACK_METHODS = {
     "box": "_bx",
+    "ctd": "_ctd",
     "skeleton": "_sk",
     "ellipse": "_el",
     "transformer": "_tr",

diff --git a/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py
@@ -119,6 +119,7 @@ def create_multianimaltraining_dataset(
     userfeedback: bool = True,
     weight_init: WeightInitialization | None = None,
     engine: Engine | None = None,
+    ctd_conditions: int | str | Path | tuple[int, str] | tuple[int, int] | None = None,
 ):
     """
     Creates a training dataset for multi-animal datasets. Labels from all the extracted
@@ -158,17 +159,38 @@ def create_multianimaltraining_dataset(
                 * ``efficientnet-b6``
             PyTorch (call ``deeplabcut.pose_estimation_pytorch.available_models()`` for
             a complete list)
-                * ``resnet_50``
-                * ``resnet_101``
+                * ``animaltokenpose_base``
+                * ``cspnext_m``
+                * ``cspnext_s``
+                * ``cspnext_x``
+                * ``ctd_coam_w32``
+                * ``ctd_coam_w48``
+                * ``ctd_prenet_hrnet_w32``
+                * ``ctd_prenet_hrnet_w48``
+                * ``ctd_prenet_rtmpose_m``
+                * ``ctd_prenet_rtmpose_x``
+                * ``ctd_prenet_rtmpose_x_human``
                 * ``dekr_w18``
                 * ``dekr_w32``
                 * ``dekr_w48``
-                * ``top_down_resnet_50``
-                * ``top_down_resnet_101``
+                * ``dlcrnet_stride16_ms5``
+                * ``dlcrnet_stride32_ms5``
+                * ``hrnet_w18``
+                * ``hrnet_w32``
+                * ``hrnet_w48``
+                * ``resnet_101``
+                * ``resnet_50``
+                * ``rtmpose_m``
+                * ``rtmpose_s``
+                * ``rtmpose_x``
+                * ``top_down_cspnext_m``
+                * ``top_down_cspnext_s``
+                * ``top_down_cspnext_x``
                 * ``top_down_hrnet_w18``
                 * ``top_down_hrnet_w32``
                 * ``top_down_hrnet_w48``
-                * ``animaltokenpose_base``
+                * ``top_down_resnet_101``
+                * ``top_down_resnet_50``
 
     detector_type: string, optional, default=None
         Only for the PyTorch engine.
@@ -233,6 +255,14 @@ def create_multianimaltraining_dataset(
         the value specified in the project configuration file. If no engine is specified
         for the project, defaults to ``deeplabcut.compat.DEFAULT_ENGINE``.
 
+    ctd_conditions: int | str | Path | tuple[int, str] | tuple[int, int] , optional, default = None,
+        If using a conditional-top-down (CTD) net_type, this argument needs to be specified.
+        It defines the conditions that will be used with the CTD model.
+        It can be either:
+            * A shuffle number (ctd_conditions: int), which must correspond to a bottom-up (BU) network type.
+            * A predictions file path (ctd_conditions: string | Path), which must correspond to a .json or .h5 predictions file.
+            * A shuffle number and a particular snapshot (ctd_conditions: tuple[int, str] | tuple[int, int]), which respectively correspond to a bottom-up (BU) network type and a particular snapshot name or index.
+
     Example
     --------
     >>> deeplabcut.create_multianimaltraining_dataset('/analysis/project/reaching-task/config.yaml',num_shuffles=1)
@@ -596,6 +626,7 @@ def create_multianimaltraining_dataset(
                         detector_type=detector_type,
                         weight_init=weight_init,
                         save=True,
+                        ctd_conditions=ctd_conditions,
                     )
 
                 make_pytorch_test_config(pytorch_cfg, path_test_config, save=True)

diff --git a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py
@@ -791,6 +791,7 @@ def create_training_dataset(
     superanimal_name="",
     weight_init: WeightInitialization | None = None,
     engine: Engine | None = None,
+    ctd_conditions: int | str | Path | tuple[int, str] | tuple[int, int] | None = None,
 ):
     """Creates a training dataset.
 
@@ -841,20 +842,41 @@ def create_training_dataset(
                 * ``efficientnet-b6``
             PyTorch (call ``deeplabcut.pose_estimation_pytorch.available_models()`` for
             a complete list)
-                * ``resnet_50``
-                * ``resnet_101``
-                * ``hrnet_w18``
-                * ``hrnet_w32``
-                * ``hrnet_w48``
+                * ``animaltokenpose_base``
+                * ``cspnext_m``
+                * ``cspnext_s``
+                * ``cspnext_x``
+                * ``ctd_coam_w32``
+                * ``ctd_coam_w48``
+                * ``ctd_prenet_cspnext_m``
+                * ``ctd_prenet_cspnext_x``
+                * ``ctd_prenet_rtmpose_x_human``
+                * ``ctd_prenet_hrnet_w32``
+                * ``ctd_prenet_hrnet_w48``
+                * ``ctd_prenet_rtmpose_m``
+                * ``ctd_prenet_rtmpose_x``
+                * ``ctd_prenet_rtmpose_x_human``
                 * ``dekr_w18``
                 * ``dekr_w32``
                 * ``dekr_w48``
-                * ``top_down_resnet_50``
-                * ``top_down_resnet_101``
+                * ``dlcrnet_stride16_ms5``
+                * ``dlcrnet_stride32_ms5``
+                * ``hrnet_w18``
+                * ``hrnet_w32``
+                * ``hrnet_w48``
+                * ``resnet_101``
+                * ``resnet_50``
+                * ``rtmpose_m``
+                * ``rtmpose_s``
+                * ``rtmpose_x``
+                * ``top_down_cspnext_m``
+                * ``top_down_cspnext_s``
+                * ``top_down_cspnext_x``
                 * ``top_down_hrnet_w18``
                 * ``top_down_hrnet_w32``
                 * ``top_down_hrnet_w48``
-                * ``animaltokenpose_base``
+                * ``top_down_resnet_101``
+                * ``top_down_resnet_50``
 
     detector_type: string, optional, default=None
         Only for the PyTorch engine.
@@ -900,6 +922,19 @@ def create_training_dataset(
         the value specified in the project configuration file. If no engine is specified
         for the project, defaults to ``deeplabcut.compat.DEFAULT_ENGINE``.
 
+    ctd_conditions: int | str | Path | tuple[int, str] | tuple[int, int] | None, default = None,
+        If using a conditional-top-down (CTD) net_type, this argument should be
+        specified. It defines the conditions that will be used with the CTD model.
+        It can be either:
+            * A shuffle number (ctd_conditions: int), which must correspond to a
+                bottom-up (BU) network type.
+            * A predictions file path (ctd_conditions: string | Path), which must
+                correspond to a .json or .h5 predictions file.
+            * A shuffle number and a particular snapshot
+                (ctd_conditions: tuple[int, str] | tuple[int, int]), which respectively
+                correspond to a bottom-up (BU) network type and a particular snapshot
+                name or index.
+
     Returns
     -------
     list(tuple) or None
@@ -985,6 +1020,7 @@ def create_training_dataset(
             userfeedback=userfeedback,
             engine=engine,
             weight_init=weight_init,
+            ctd_conditions=ctd_conditions,
         )
     else:
         scorer = cfg["scorer"]
@@ -1088,7 +1124,6 @@ def create_training_dataset(
 
         Shuffles = validate_shuffles(cfg, Shuffles, num_shuffles, userfeedback)
 
-        # print(trainIndices,testIndices, Shuffles, augmenter_type,net_type)
         if trainIndices is None and testIndices is None:
             splits = [
                 (
@@ -1305,6 +1340,7 @@ def create_training_dataset(
                             detector_type=detector_type,
                             weight_init=weight_init,
                             save=True,
+                            ctd_conditions=ctd_conditions,
                         )
 
                     make_pytorch_test_config(pytorch_cfg, path_test_config, save=True)
@@ -1599,6 +1635,7 @@ def create_training_dataset_from_existing_split(
     net_type: str | None = None,
     detector_type: str | None = None,
     augmenter_type: str | None = None,
+    ctd_conditions: int | str | Path | tuple[int, str] | tuple[int, int] | None = None,
     posecfg_template: dict | None = None,
     superanimal_name: str = "",
     weight_init: WeightInitialization | None = None,
@@ -1683,6 +1720,19 @@ def create_training_dataset_from_existing_split(
             engine is specified for the project, defaults to
             ``deeplabcut.compat.DEFAULT_ENGINE``.
 
+        ctd_conditions: int | str | Path | tuple[int, str] | tuple[int, int] | None, default = None,
+            If using a conditional-top-down (CTD) net_type, this argument should be
+            specified. It defines the conditions that will be used with the CTD model.
+            It can be either:
+                * A shuffle number (ctd_conditions: int), which must correspond to a
+                  bottom-up (BU) network type.
+                * A predictions file path (ctd_conditions: string | Path), which must
+                  correspond to a .json or .h5 predictions file.
+                * A shuffle number and a particular snapshot
+                  (ctd_conditions: tuple[int, str] | tuple[int, int]), which
+                  respectively correspond to a bottom-up (BU) network type and a
+                  particular snapshot name or index.
+
     Returns:
         If training dataset was successfully created, a list of tuples is returned.
         The first two elements in each tuple represent the training fraction and the
@@ -1734,6 +1784,7 @@ def create_training_dataset_from_existing_split(
         superanimal_name=superanimal_name,
         weight_init=weight_init,
         engine=engine,
+        ctd_conditions=ctd_conditions,
     )
 
 

diff --git a/deeplabcut/pose_estimation_pytorch/__init__.py b/deeplabcut/pose_estimation_pytorch/__init__.py
@@ -41,9 +41,13 @@
     COCOLoader,
     COLLATE_FUNCTIONS,
     DLCLoader,
+    GenerativeSampler,
+    GenSamplingConfig,
+    list_snapshots,
     Loader,
     PoseDataset,
     PoseDatasetParameters,
+    Snapshot,
 )
 from deeplabcut.pose_estimation_pytorch.runners import (
     build_inference_runner,