DeepLabCut · n-poulsen · Oct 18, 2024 · Sep 18, 2024 · Oct 1, 2024 · Oct 1, 2024
diff --git a/deeplabcut/compat.py b/deeplabcut/compat.py
@@ -58,7 +58,7 @@ def get_available_aug_methods(engine: Engine) -> tuple[str, ...]:
 
 
 def train_network(
-    config: str,
+    config: str | Path,
     shuffle: int = 1,
     trainingsetindex: int = 0,
     max_snapshots_to_keep: int | None = None,
@@ -221,7 +221,7 @@ def train_network(
             max_snapshots_to_keep = 5
 
         return train_network(
-            config,
+            str(config),
             shuffle=shuffle,
             trainingsetindex=trainingsetindex,
             max_snapshots_to_keep=max_snapshots_to_keep,
@@ -317,7 +317,7 @@ def return_train_network_path(
 
 
 def evaluate_network(
-    config,
+    config: str | Path,
     Shuffles: Iterable[int] = (1,),
     trainingsetindex: int | str = 0,
     plotting: bool | str = False,
@@ -458,7 +458,7 @@ def evaluate_network(
     if engine == Engine.TF:
         from deeplabcut.pose_estimation_tensorflow import evaluate_network
         return evaluate_network(
-            config,
+            str(config),
             Shuffles=Shuffles,
             trainingsetindex=trainingsetindex,
             plotting=plotting,

diff --git a/deeplabcut/core/trackingutils.py b/deeplabcut/core/trackingutils.py
@@ -766,11 +766,11 @@ def fill_tracklets(tracklets, trackers, animals, imname):
         if tracklet_id not in tracklets:
             tracklets[tracklet_id] = {}
         if pred_id != -1:
-            tracklets[tracklet_id][imname] = animals[pred_id]
+            tracklets[tracklet_id][imname] = np.asarray(animals[pred_id])
         else:  # Resort to the tracker prediction
             xy = np.asarray(content[:-2])
             pred = np.insert(xy, range(2, len(xy) + 1, 2), 1)
-            tracklets[tracklet_id][imname] = pred
+            tracklets[tracklet_id][imname] = np.asarray(pred)
 
 
 def calc_bboxes_from_keypoints(data, slack=0, offset=0):

diff --git a/deeplabcut/core/weight_init.py b/deeplabcut/core/weight_init.py
@@ -11,63 +11,46 @@
 """Classes to configure how to initialize model weights"""
 from __future__ import annotations
 
+import warnings
 from dataclasses import dataclass
+from pathlib import Path
 
 import numpy as np
 
-import deeplabcut.modelzoo.utils as modelzoo_utils
-
 
 @dataclass
 class WeightInitialization:
-    """The dataset from which to initialize weights
-
-    To build a WeightInitialization instance for a project using the conversion table
-    specified in the project configuration file, use:
-
-        ```
-        from pathlib import Path
-        from deeplabcut.utils.auxiliaryfunctions import read_config
-
-        project_cfg = read_config("/path/to/my/project/config.yaml")
-        super_animal = "superanimal_quadruped"
-        weight_init = WeightInitialization.build(
-            cfg=project_cfg,
-            super_animal="superanimal_quadruped",
-            with_decoder=True,
-            memory_replay=True,
-        )
-        ```
+    """Configures weights initialization when transfer learning or fine-tuning models
 
     Args:
-        dataset: The dataset on which the model weights were trained. Must be one of the
-            SuperAnimal weights.
+        snapshot_path: The path to the snapshot used to initialize pose model weights
+            when training a model.
+        detector_snapshot_path: The path to the snapshot used to initialize detector
+            weights when training a model.
+        dataset: Optionally, the dataset on which the snapshots were trained. Required
+            when fine-tuning SuperAnimal models.
         with_decoder: Whether to load the decoder weights as well.
         memory_replay: Only when ``with_decoder=True``. Whether to train the model with
-            memory replay, so that it predicts all SuperAnimal bodyparts.
-        conversion_array: The mapping from SuperAnimal to project bodyparts. Required
-            when `with_decoder=True`.
+            memory replay, so that it predicts all SuperAnimal (or previous project)
+            bodyparts.
+        conversion_array: The mapping from SuperAnimal (or other project, on which the
+            weights were trained) to project bodyparts. Required when
+            `with_decoder=True`.
             An array [7, 0, 1] means the project has 3 bodyparts, where the 1st bodypart
             corresponds to the 8th bodypart in the pretrained model, the 2nd to the 1st
             and the 3rd to the 2nd (as arrays are 0-indexed).
         bodyparts: Optionally, the name of each bodypart entry in the conversion array.
-        customized_pose_checkpoint: A customized SuperAnimal pose checkpoint, as an
-            alternative to the Hugging Face one
-        customized_detector_checkpoint: A customized SuperAnimal detector
-            checkpoint, as an alternative to the Hugging Face one
     """
 
-    dataset: str
+    snapshot_path: Path
+    detector_snapshot_path: Path | None = None
+    dataset: str | None = None
     with_decoder: bool = False
     memory_replay: bool = False
     conversion_array: np.ndarray | None = None
     bodyparts: list[str] | None = None
-    customized_pose_checkpoint: str | None = None
-    customized_detector_checkpoint: str | None = None
 
     def __post_init__(self):
-        # check that the dataset exists; raises a ValueError if it doesn't
-        _ = modelzoo_utils.get_super_animal_project_cfg(self.dataset)
         if self.memory_replay and not self.with_decoder:
             raise ValueError(
                 "You cannot train a model with memory replay if you do not keep the "
@@ -97,51 +80,97 @@ def __post_init__(self):
 
     def to_dict(self) -> dict:
         """Returns: the weight initialization as a dict"""
-        data = {
-            "dataset": self.dataset,
-            "with_decoder": self.with_decoder,
-            "memory_replay": self.memory_replay,
-        }
+        data = dict()
+        if self.dataset is not None:
+            data["dataset"] = self.dataset
+
+        data["snapshot_path"] = str(self.snapshot_path)
+        if self.detector_snapshot_path is not None:
+            data["detector_snapshot_path"] = str(self.detector_snapshot_path)
+
+        data["with_decoder"] = self.with_decoder
+        data["memory_replay"] = self.memory_replay
 
         if self.conversion_array is not None:
             data["conversion_array"] = self.conversion_array.tolist()
-        if self.customized_pose_checkpoint is not None:
-            data["customized_pose_checkpoint"] = self.customized_pose_checkpoint
-        if self.customized_detector_checkpoint is not None:
-            data["customized_detector_checkpoint"] = self.customized_detector_checkpoint
+
+        if self.bodyparts is not None:
+            data["bodyparts"] = self.bodyparts
 
         return data
 
     @staticmethod
     def from_dict(data: dict) -> "WeightInitialization":
+        if "snapshot_path" not in data:
+            return WeightInitialization.from_dict_legacy(data)
+
+        detector_snapshot_path = data.get("detector_snapshot_path")
+        if detector_snapshot_path is not None:
+            detector_snapshot_path = Path(detector_snapshot_path)
+
         conversion_array = data.get("conversion_array")
         if conversion_array is not None:
+            conversion_array = np.array(conversion_array, dtype=int)
 
+        return WeightInitialization(
+            snapshot_path=Path(data["snapshot_path"]),
+            detector_snapshot_path=detector_snapshot_path,
+            dataset=data.get("dataset"),
+            with_decoder=data["with_decoder"],
+            memory_replay=data["memory_replay"],
+            conversion_array=conversion_array,
+            bodyparts=data.get("bodyparts"),
+        )
+
+    @staticmethod
+    def from_dict_legacy(data: dict) -> "WeightInitialization":
+        """Deals with weight initialization that were created before 3.0.0rc5"""
+        import deeplabcut.pose_estimation_pytorch.modelzoo.utils as utils
+
+        conversion_array = data.get("conversion_array")
+        if conversion_array is not None:
             conversion_array = np.array(conversion_array, dtype=int)
 
         return WeightInitialization(
-            dataset=data["dataset"],
+            snapshot_path=utils.get_super_animal_snapshot_path(
+                dataset=data["dataset"],
+                model_name="hrnet_w32",
+            ),
+            detector_snapshot_path=utils.get_super_animal_snapshot_path(
+                dataset=data["dataset"],
+                model_name="fasterrcnn_resnet50_fpn_v2",
+            ),
             with_decoder=data["with_decoder"],
             memory_replay=data["memory_replay"],
             conversion_array=conversion_array,
-            customized_pose_checkpoint=data.get("customized_pose_checkpoint"),
-            customized_detector_checkpoint=data.get("customized_detector_checkpoint"),
+            bodyparts=data.get("bodyparts"),
         )
 
     @staticmethod
     def build(
         cfg: dict,
         super_animal: str,
+        model_name: str = "hrnet_w32",
+        detector_name: str = "fasterrcnn_resnet50_fpn_v2",
         with_decoder: bool = False,
         memory_replay: bool = False,
         customized_pose_checkpoint: str | None = None,
         customized_detector_checkpoint: str | None = None,
     ) -> "WeightInitialization":
         """Builds a WeightInitialization for a project
 
+        `WeightInitialization.build` is deprecated and will be removed in a future
+        version of DeepLabCut. Please use `build_weight_init` from `deeplabcut.modelzoo`
+        instead.
+
         Args:
             cfg: The project's configuration.
             super_animal: The SuperAnimal model with which to initialize weights.
+            model_name: The name of the model architecture for which to load the weights
+                (defaults to "hrnet_w32" for backwards compatibility).
+            detector_name: The name of the detector architecture for which to load the
+                weights (defaults to "fasterrcnn_resnet50_fpn_v2" for backwards
+                compatibility).
             with_decoder: Whether to load the decoder weights as well. If this is true,
                 a conversion table must be specified for the given SuperAnimal in the
                 project configuration file. See
@@ -157,19 +186,21 @@ def build(
         Returns:
             The built WeightInitialization.
         """
-        conversion_array = None
-        bodyparts = None
-        if with_decoder:
-            conversion_table = modelzoo_utils.get_conversion_table(cfg, super_animal)
-            conversion_array = conversion_table.to_array()
-            bodyparts = conversion_table.converted_bodyparts()
-
-        return WeightInitialization(
-            dataset=super_animal,
-            with_decoder=with_decoder,
-            memory_replay=memory_replay,
-            conversion_array=conversion_array,
-            bodyparts=bodyparts,
-            customized_pose_checkpoint=customized_pose_checkpoint,
-            customized_detector_checkpoint=customized_detector_checkpoint,
+        from deeplabcut.modelzoo import build_weight_init
+        deprecation_warning = (
+            "The `WeightInitialization.build` is deprecated and will be removed in a "
+            "future version of DeepLabCut. Please use `build_weight_init` from "
+            "`deeplabcut.modelzoo` instead."
+        )
+        warnings.warn(deprecation_warning, DeprecationWarning)
+
+        return build_weight_init(
+            cfg,
+            super_animal,
+            model_name,
+            detector_name,
+            with_decoder,
+            memory_replay,
+            customized_pose_checkpoint,
+            customized_detector_checkpoint,
         )