Merge pull request #496 from roboflow/sb/rfdetr-keypoint-preview-deploy

Support deployment of RFDETRKeypointPreview weights

Author: mkaic

requirements.txt

@@ -3,7 +3,10 @@ idna>=3.7
 cycler
 kiwisolver>=1.3.1
 matplotlib
-numpy>=1.18.5,<2.5  # 2.5.0 ships type stubs using 3.12+ `type` syntax that mypy (pinned to 3.10) rejects
+# numpy 2.4 ships PEP 695 `type` statements in its stubs, which mypy rejects
+# under python_version=3.10 (see [tool.mypy] in pyproject.toml). Cap below 2.4,
+# matching rf-detr's typing constraint.
+numpy>=1.18.5,<2.4
 opencv-python-headless==4.10.0.84
 Pillow>=7.1.2
 # https://github.com/roboflow/roboflow-python/issues/390

roboflow/util/model_processor.py

@@ -107,6 +107,8 @@
     "rfdetr-seg-large": "RFDETRSegLarge",
     "rfdetr-seg-xlarge": "RFDETRSegXLarge",
     "rfdetr-seg-2xlarge": "RFDETRSeg2XLarge",
+    # Keypoint detection
+    "rfdetr-keypoint-preview": "RFDETRKeypointPreview",
 }
 
 SUPPORTED_RFDETR_TYPES = tuple(_RFDETR_MODEL_TYPE_TO_CLASS)
@@ -149,6 +151,8 @@
     "rfdetr-seg-large": 42,
     "rfdetr-seg-xlarge": 52,
     "rfdetr-seg-2xlarge": 64,
+    # Keypoint (576 / patch_size 12 = 48)
+    "rfdetr-keypoint-preview": 48,
 }
 
 
@@ -240,8 +244,13 @@ def task_of_model_type(model_type: str) -> str:
 
     Non-detect tasks double as the model_type suffix token
     (e.g. 'yolov11-seg' -> TASK_SEG). Plain 'yolov11' / 'rfdetr-base' -> TASK_DET.
+
+    Keypoint/pose models may spell the token as either 'pose' (Ultralytics) or
+    'keypoint' (rf-detr, e.g. 'rfdetr-keypoint-preview'); both map to TASK_POSE.
     """
     s = model_type.lower()
+    if "keypoint" in s:
+        return TASK_POSE
     for task in (TASK_SEM, TASK_SEG, TASK_POSE, TASK_CLS, TASK_OBB):
         if task in s:
             return task
@@ -813,27 +822,32 @@ def _process_yolo(
 def _detect_rfdetr_task(checkpoint: Any) -> str | None:
     """Detect the training task of an rf-detr checkpoint.
 
-    rf-detr currently only supports weight upload for detection and instance
-    segmentation. Modern checkpoints (rf-detr v1.7+) store the Python class
-    name at `checkpoint["model_name"]` (e.g. 'RFDETRNano' vs 'RFDETRSegNano');
-    older checkpoints — including those downloaded from Roboflow — lack that
-    field but always carry `args.segmentation_head: bool`.
+    rf-detr supports weight upload for detection, instance segmentation, and
+    keypoint detection. Modern checkpoints (rf-detr v1.7+) store the Python
+    class name at `checkpoint["model_name"]` (e.g. 'RFDETRNano' vs
+    'RFDETRSegNano' vs 'RFDETRKeypointPreview').
+
+    The deploy bundle written by rf-detr's `export_for_roboflow` only serialises
+    `{"model", "args"}` — it drops `model_name` — so detection must also work
+    from `args`: keypoint checkpoints carry a non-empty `args.num_keypoints_per_class`,
+    and detection/segmentation checkpoints carry `args.segmentation_head: bool`.
     """
     if not isinstance(checkpoint, dict):
         return None
     model_name = checkpoint.get("model_name")
     if isinstance(model_name, str):
         name = model_name.lower()
-        # Keypoint rf-detr checkpoints (e.g. 'RFDETRKeypointPreview') are not a
-        # supported upload type; classify them as pose so the task check rejects
-        # them instead of silently uploading a keypoint model as detection.
         if "keypoint" in name:
             return TASK_POSE
         return TASK_SEG if TASK_SEG in name else TASK_DET
     raw_args = checkpoint.get("args")
     if raw_args is None:
         return None
     args = _checkpoint_args_as_dict(raw_args)
+    # Keypoint checkpoints carry num_keypoints_per_class; classify them as pose so it agrees
+    # with task_of_model_type('rfdetr-keypoint-preview') == TASK_POSE and the upload proceeds.
+    if args.get("num_keypoints_per_class"):
+        return TASK_POSE
     segmentation_head = args.get("segmentation_head")
     if segmentation_head is True:
         return TASK_SEG
@@ -1050,8 +1064,9 @@ def _process_rfdetr(
     checkpoint_path = _find_rfdetr_checkpoint(model_path, filename, warnings)
     checkpoint = _load_checkpoint(torch, checkpoint_path, map_location="cpu")
 
-    # Task detection + mismatch runs for every checkpoint shape (it also rejects
-    # keypoint rf-detr, which is not a supported upload type).
+    # Task detection + mismatch runs for every checkpoint shape, so a checkpoint whose
+    # task disagrees with model_type (e.g. a keypoint checkpoint uploaded as 'rfdetr-base')
+    # is rejected instead of packaged under the wrong task.
     detected_task = _detect_rfdetr_task(checkpoint)
     if detected_task and detected_task != task_of_model_type(model_type):
         raise TaskMismatchError(

tests/util/test_model_processor.py

@@ -62,6 +62,7 @@ def test_segment(self):
 
     def test_pose(self):
         self.assertEqual(task_of_model_type("yolov11-pose"), TASK_POSE)
+        self.assertEqual(task_of_model_type("rfdetr-keypoint-preview"), TASK_POSE)
 
     def test_classify(self):
         self.assertEqual(task_of_model_type("yolov11-cls"), TASK_CLS)
@@ -102,11 +103,19 @@ def test_detection_model_names(self):
         for name in ("RFDETRNano", "RFDETRSmall", "RFDETRMedium", "RFDETRLarge", "RFDETRXLarge"):
             self.assertEqual(_detect_rfdetr_task({"model_name": name}), TASK_DET, name)
 
-    def test_keypoint_model_name_returns_pose(self):
-        # Keypoint checkpoints are unsupported; classifying them as pose lets the
-        # model_type task check reject them instead of uploading them as detection.
+    def test_keypoint_model_names(self):
         self.assertEqual(_detect_rfdetr_task({"model_name": "RFDETRKeypointPreview"}), TASK_POSE)
 
+    def test_keypoint_args_fallback(self):
+        # The deploy bundle from export_for_roboflow carries `args` but not
+        # `model_name`; a non-empty `num_keypoints_per_class` marks a keypoint model.
+        self.assertEqual(_detect_rfdetr_task({"args": SimpleNamespace(num_keypoints_per_class=[0, 17])}), TASK_POSE)
+        self.assertEqual(_detect_rfdetr_task({"args": {"num_keypoints_per_class": [0, 17]}}), TASK_POSE)
+        # Empty / absent keypoint schema must NOT be treated as a keypoint model.
+        self.assertEqual(
+            _detect_rfdetr_task({"args": {"num_keypoints_per_class": [], "segmentation_head": False}}), TASK_DET
+        )
+
     def test_segmentation_head_fallback(self):
         # Roboflow-hosted rf-detr .pt downloads lack `model_name` but always carry
         # `args.segmentation_head`. Cover both namespace and dict shapes.
@@ -454,6 +463,37 @@ def test_rfdetr_falls_back_to_discovered_checkpoint(self):
             finally:
                 bundle.cleanup()
 
+    def test_rfdetr_keypoint_exported_checkpoint_packages(self):
+        # The primary feature: an exported keypoint deploy-checkpoint (non-PTL shape,
+        # args carries class_names + num_keypoints_per_class) packages successfully as
+        # 'rfdetr-keypoint-preview'. num_keypoints_per_class marks it pose, which matches
+        # the model_type's task, so it passes the task check and copies weights.pt.
+        with tempfile.TemporaryDirectory() as tmp:
+            model_dir = Path(tmp)
+            (model_dir / "weights.pt").write_bytes(b"checkpoint")
+            torch = _fake_torch({"args": {"class_names": ["goal"], "num_keypoints_per_class": [0, 17]}})
+            with _import_patch({"torch": torch}):
+                bundle = package_custom_weights("rfdetr-keypoint-preview", str(model_dir), filename="weights.pt")
+            try:
+                self.assertEqual(bundle.model_type, "rfdetr-keypoint-preview")
+                with zipfile.ZipFile(bundle.archive_path) as archive:
+                    names = archive.namelist()
+                    self.assertIn("weights.pt", names)
+                    self.assertIn("class_names.txt", names)
+            finally:
+                bundle.cleanup()
+
+    def test_rfdetr_keypoint_checkpoint_rejected_as_detection_type(self):
+        # The same keypoint checkpoint uploaded under a detection model_type is a task
+        # mismatch (pose != detect) and must be rejected, not silently packaged.
+        with tempfile.TemporaryDirectory() as tmp:
+            model_dir = Path(tmp)
+            (model_dir / "weights.pt").write_bytes(b"checkpoint")
+            torch = _fake_torch({"args": {"class_names": ["goal"], "num_keypoints_per_class": [0, 17]}})
+            with _import_patch({"torch": torch}):
+                with self.assertRaises(TaskMismatchError):
+                    package_custom_weights("rfdetr-base", str(model_dir), filename="weights.pt")
+
     def test_rfdetr_without_any_checkpoint_raises(self):
         with tempfile.TemporaryDirectory() as tmp:
             torch = _fake_torch({})
@@ -824,6 +864,7 @@ class RfdetrModelTypeToClassTest(unittest.TestCase):
     def test_representative_mappings(self):
         self.assertEqual(_RFDETR_MODEL_TYPE_TO_CLASS["rfdetr-seg-medium"], "RFDETRSegMedium")
         self.assertEqual(_RFDETR_MODEL_TYPE_TO_CLASS["rfdetr-base"], "RFDETRBase")
+        self.assertEqual(_RFDETR_MODEL_TYPE_TO_CLASS["rfdetr-keypoint-preview"], "RFDETRKeypointPreview")
 
     def test_keys_are_rfdetr_types_and_values_are_class_names(self):
         for model_type, class_name in _RFDETR_MODEL_TYPE_TO_CLASS.items():
@@ -878,8 +919,10 @@ def __init__(self, *, pretrain_weights):
     module = SimpleNamespace()
     module.RFDETR = _RFDETR
     # The fallback resolves the subclass by name via _RFDETR_MODEL_TYPE_TO_CLASS,
-    # e.g. "rfdetr-seg-medium" -> getattr(rfdetr, "RFDETRSegMedium").
+    # e.g. "rfdetr-seg-medium" -> getattr(rfdetr, "RFDETRSegMedium") and
+    # "rfdetr-keypoint-preview" -> getattr(rfdetr, "RFDETRKeypointPreview").
     module.RFDETRSegMedium = _SizedModel
+    module.RFDETRKeypointPreview = _SizedModel
     if capabilities:
         _RFDETR.export_for_roboflow = _StubBundleModel.export_for_roboflow  # capability marker
     module._calls = calls
@@ -910,13 +953,13 @@ def test_returns_module_when_capable(self):
 class PackageRfdetrPtlTest(unittest.TestCase):
     """PyTorch-Lightning rf-detr checkpoints are rebuilt via rfdetr into build_dir."""
 
-    def _package(self, model_type, fake_rfdetr, *, segmentation_head=False):
+    def _package(self, model_type, fake_rfdetr, *, segmentation_head=False, num_keypoints_per_class=None):
         with tempfile.TemporaryDirectory() as model_dir:
             (Path(model_dir) / "checkpoint_best_ema.pth").write_bytes(b"raw-ptl")
-            ckpt = {
-                "pytorch-lightning_version": "2.1.0",
-                "args": {"segmentation_head": segmentation_head, "class_names": ["cat", "dog"]},
-            }
+            args = {"segmentation_head": segmentation_head, "class_names": ["cat", "dog"]}
+            if num_keypoints_per_class is not None:
+                args["num_keypoints_per_class"] = num_keypoints_per_class
+            ckpt = {"pytorch-lightning_version": "2.1.0", "args": args}
             torch = _fake_torch(ckpt)
             with _import_patch({"torch": torch}), mock.patch.dict(sys.modules, {"rfdetr": fake_rfdetr}):
                 bundle = package_custom_weights(model_type, model_dir, filename="checkpoint_best_ema.pth")
@@ -944,6 +987,28 @@ def test_from_checkpoint_valueerror_falls_back_to_model_type(self):
         self.assertEqual(fake._calls["fallback_constructed"], 1)
         self.assertIn("weights.pt", names)
 
+    def test_keypoint_from_checkpoint_success_produces_bundle(self):
+        # A keypoint PTL checkpoint (args.num_keypoints_per_class marks pose, matching
+        # the 'rfdetr-keypoint-preview' model_type) rebuilds via rfdetr.from_checkpoint.
+        fake = _make_fake_rfdetr()
+        bundle, names = self._package("rfdetr-keypoint-preview", fake, num_keypoints_per_class=[0, 17])
+        self.assertEqual(bundle.model_type, "rfdetr-keypoint-preview")
+        self.assertEqual(fake._calls["from_checkpoint"], 1)
+        self.assertEqual(fake._calls["fallback_constructed"], 0)
+        self.assertIn("weights.pt", names)
+        self.assertIn("class_names.txt", names)
+
+    def test_keypoint_from_checkpoint_valueerror_falls_back_to_model_type(self):
+        # When from_checkpoint can't infer the class, the fallback resolves the
+        # RFDETRKeypointPreview subclass from _RFDETR_MODEL_TYPE_TO_CLASS and rebuilds.
+        fake = _make_fake_rfdetr(from_checkpoint_raises=True)
+        bundle, names = self._package("rfdetr-keypoint-preview", fake, num_keypoints_per_class=[0, 17])
+        self.assertEqual(bundle.model_type, "rfdetr-keypoint-preview")
+        self.assertEqual(fake._calls["from_checkpoint"], 1)
+        self.assertEqual(fake._calls["fallback_constructed"], 1)
+        self.assertIn("weights.pt", names)
+        self.assertIn("class_names.txt", names)
+
     def test_ptl_path_raises_when_rfdetr_absent(self):
         with tempfile.TemporaryDirectory() as model_dir:
             (Path(model_dir) / "checkpoint_best_ema.pth").write_bytes(b"raw-ptl")