Add specific export config for torchvision detectors

Author: deruyter92

dlclive/modelzoo/pytorch_model_zoo_export.py

@@ -32,10 +32,12 @@ def _load_model_weights(model_name: str, super_animal: str = super_animal) -> Or
         checkpoint: Path = download_super_animal_snapshot(dataset=super_animal, model_name=model_name)
         return torch.load(checkpoint, map_location="cpu", weights_only=True)["model"]
 
+    # Skip downloading the detector weights for humanbody models, as they are not on huggingface
+    skip_detector_download = (detector_name is None) or (super_animal == "superanimal_humanbody")
     export_dict = {
         "config": model_cfg,
         "pose": _load_model_weights(model_name),
-        "detector": _load_model_weights(detector_name) if detector_name is not None else None,
+        "detector": None if skip_detector_download else _load_model_weights(detector_name),
     }
     torch.save(export_dict, export_path)
 

dlclive/modelzoo/utils.py

@@ -12,6 +12,7 @@
 from ruamel.yaml import YAML
 
 from dlclive.modelzoo.resolve_config import update_config
+from dlclive.pose_estimation_pytorch.models.detectors.torchvision import SUPPORTED_TORCHVISION_DETECTORS
 
 _MODELZOO_PATH = Path(__file__).parent
 
@@ -131,12 +132,21 @@ def load_super_animal_config(
         model_config["method"] = "BU"
     else:
         model_config["method"] = "TD"
-        if super_animal != "superanimal_humanbody":
-            detector_cfg_path = get_super_animal_model_config_path(
-                model_name=detector_name
-            )
-            detector_cfg = read_config_as_dict(detector_cfg_path)
-            model_config["detector"] = detector_cfg
+        detector_cfg_path = get_super_animal_model_config_path(
+            model_name=detector_name
+        )
+        detector_cfg = read_config_as_dict(detector_cfg_path)
+        model_config["detector"] = detector_cfg
+        if super_animal == "superanimal_humanbody":
+            # Apply specific updates required to run the torchvision detector with pretrained weights
+            assert detector_name in SUPPORTED_TORCHVISION_DETECTORS
+            model_config["detector"]['model']= {
+                "type": "TorchvisionDetectorAdaptor",
+                "model": detector_name,
+                "weights": "COCO_V1",
+                "num_classes": None,
+                "box_score_thresh": 0.6,
+            }
     return model_config
 
 

dlclive/pose_estimation_pytorch/models/detectors/torchvision.py

@@ -16,6 +16,8 @@
 
 from dlclive.pose_estimation_pytorch.models.detectors.base import DETECTORS, BaseDetector
 
+SUPPORTED_TORCHVISION_DETECTORS = ["fasterrcnn_mobilenet_v3_large_fpn"]
+
 
 @DETECTORS.register_module
 class TorchvisionDetectorAdaptor(BaseDetector):