feat: support rfdetr-keypoint-preview custom weights upload

CHANGELOG.md

@@ -6,6 +6,8 @@ All notable changes to this project will be documented in this file.
 
 ### Added
 
+- Weight upload support for `rfdetr-keypoint-preview` keypoint detection models
+  via `version.deploy()`, `workspace.deploy_model()`, and the `upload_model` CLI
 - Upload raw rf-detr PyTorch-Lightning checkpoints (e.g. `checkpoint_best_ema.pth`):
   `upload_model` detects them and rebuilds a deploy-ready bundle via rf-detr's
   `export_for_roboflow` (requires `rfdetr>=1.8.0`)

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,7 @@
     "rfdetr-seg-large": 42,
     "rfdetr-seg-xlarge": 52,
     "rfdetr-seg-2xlarge": 64,
+    "rfdetr-keypoint-preview": 48,
 }
 
 
@@ -239,9 +242,13 @@ def task_of_model_type(model_type: str) -> str:
     """Canonical task for a deploy model_type string.
 
     Non-detect tasks double as the model_type suffix token
-    (e.g. 'yolov11-seg' -> TASK_SEG). Plain 'yolov11' / 'rfdetr-base' -> TASK_DET.
+    (e.g. 'yolov11-seg' -> TASK_SEG). RF-DETR keypoint types spell the task out
+    instead ('rfdetr-keypoint-preview' -> TASK_POSE). Plain 'yolov11' /
+    'rfdetr-base' -> TASK_DET.
     """
     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 +820,28 @@ 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`.
+    Modern checkpoints (rf-detr v1.7+) store the Python class name at
+    `checkpoint["model_name"]` (e.g. 'RFDETRNano' vs 'RFDETRSegNano' vs
+    'RFDETRKeypointPreview'); older checkpoints — including those downloaded
+    from Roboflow — lack that field but carry the head flags in `args`
+    (`keypoint_head` / `segmentation_head`, both False on detection models).
     """
     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 configs also carry segmentation_head=False, so check the
+    # keypoint flag first.
+    if args.get("keypoint_head") is True:
+        return TASK_POSE
     segmentation_head = args.get("segmentation_head")
     if segmentation_head is True:
         return TASK_SEG
@@ -1050,8 +1058,8 @@ 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 (e.g. it rejects
+    # a keypoint checkpoint uploaded under a detection model_type and vice versa).
     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)
@@ -103,8 +104,6 @@ def test_detection_model_names(self):
             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.
         self.assertEqual(_detect_rfdetr_task({"model_name": "RFDETRKeypointPreview"}), TASK_POSE)
 
     def test_segmentation_head_fallback(self):
@@ -115,6 +114,14 @@ def test_segmentation_head_fallback(self):
         self.assertEqual(_detect_rfdetr_task({"args": {"segmentation_head": True}}), TASK_SEG)
         self.assertEqual(_detect_rfdetr_task({"args": {"segmentation_head": False}}), TASK_DET)
 
+    def test_keypoint_head_fallback(self):
+        # Keypoint training checkpoints from rf-detr-internal lack `model_name` and
+        # carry segmentation_head=False alongside keypoint_head=True; the keypoint
+        # flag must win so they are not misdetected as detection.
+        args = {"keypoint_head": True, "segmentation_head": False}
+        self.assertEqual(_detect_rfdetr_task({"args": args}), TASK_POSE)
+        self.assertEqual(_detect_rfdetr_task({"args": SimpleNamespace(**args)}), TASK_POSE)
+
     def test_model_name_preferred_over_args(self):
         # When both are present, model_name wins (matches rf-detr's loader).
         ckpt = {"model_name": "RFDETRNano", "args": SimpleNamespace(segmentation_head=True)}
@@ -186,6 +193,14 @@ def test_task_mismatch_is_a_value_error(self):
     def test_unknown_project_type_is_ignored(self):
         validate_model_type_for_project("yolov8", "some-new-type", "widgets")
 
+    def test_accepts_keypoint_model_for_keypoint_project(self):
+        validate_model_type_for_project("rfdetr-keypoint-preview", "keypoint-detection", "widgets")
+
+    def test_rejects_keypoint_model_for_detection_project(self):
+        with self.assertRaises(TaskMismatchError) as ctx:
+            validate_model_type_for_project("rfdetr-keypoint-preview", "object-detection", "widgets")
+        self.assertIn("task 'pose'", str(ctx.exception))
+
 
 class LegacyYoloArgsTest(unittest.TestCase):
     def test_reports_missing_batch_size(self):
@@ -307,6 +322,14 @@ def test_warns_but_allows_custom_resolution(self):
         self.assertEqual(resolved, "rfdetr-seg-nano")
         self.assertTrue(warnings and "matches no known RF-DETR variant" in warnings[0])
 
+    def test_keypoint_preview_grid_matches_without_warning(self):
+        # RFDETRKeypointPreviewConfig: resolution 576, patch_size 12 -> grid 48.
+        warnings: list = []
+        checkpoint = {"args": {"resolution": 576, "patch_size": 12}}
+        resolved = _resolve_rfdetr_variant("rfdetr-keypoint-preview", checkpoint, warnings)
+        self.assertEqual(resolved, "rfdetr-keypoint-preview")
+        self.assertEqual(warnings, [])
+
     def test_pe_size_derived_from_position_embeddings_tensor(self):
         class FakeTensor:
             shape = (1, 1025, 384)
@@ -500,6 +523,56 @@ def test_rfdetr_explicit_missing_filename_does_not_fall_back(self):
                 with self.assertRaises(MissingFileError):
                     package_custom_weights("rfdetr-base", str(model_dir), filename="checkpoint_epoch50.pth")
 
+    def test_rfdetr_keypoint_preview_full_flow(self):
+        # A keypoint training checkpoint (rf-detr-internal style: no model_name,
+        # keypoint_head=True in args) packages under 'rfdetr-keypoint-preview'.
+        with tempfile.TemporaryDirectory() as tmp:
+            model_dir = Path(tmp)
+            (model_dir / "weights.pt").write_bytes(b"checkpoint")
+            torch = _fake_torch(
+                {
+                    "args": {
+                        "keypoint_head": True,
+                        "segmentation_head": False,
+                        "resolution": 576,
+                        "patch_size": 12,
+                        "num_keypoints_per_class": [17],
+                        "class_names": ["person"],
+                    }
+                }
+            )
+            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")
+                self.assertEqual(bundle.warnings, ())
+                with zipfile.ZipFile(bundle.archive_path) as archive:
+                    self.assertIn("weights.pt", archive.namelist())
+                    class_names = archive.read("class_names.txt").decode().splitlines()
+                self.assertEqual(class_names, ["background_class83422", "person"])
+            finally:
+                bundle.cleanup()
+
+    def test_rfdetr_keypoint_checkpoint_rejected_for_detection_type(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            model_dir = Path(tmp)
+            (model_dir / "weights.pt").write_bytes(b"checkpoint")
+            torch = _fake_torch({"args": {"keypoint_head": True, "segmentation_head": False}})
+            with _import_patch({"torch": torch}):
+                with self.assertRaises(TaskMismatchError) as ctx:
+                    package_custom_weights("rfdetr-base", str(model_dir), filename="weights.pt")
+        self.assertIn("'pose'", str(ctx.exception))
+
+    def test_rfdetr_detection_checkpoint_rejected_for_keypoint_type(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            model_dir = Path(tmp)
+            (model_dir / "weights.pt").write_bytes(b"checkpoint")
+            torch = _fake_torch({"args": {"segmentation_head": False, "class_names": ["widget"]}})
+            with _import_patch({"torch": torch}):
+                with self.assertRaises(TaskMismatchError) as ctx:
+                    package_custom_weights("rfdetr-keypoint-preview", str(model_dir), filename="weights.pt")
+        self.assertIn("'det'", str(ctx.exception))
+
     def test_rfdetr_legacy_deploy_layout_does_not_self_copy(self):
         # Legacy deploy passes build_dir=model_path with a top-level weights.pt;
         # copying weights.pt onto itself would raise shutil.SameFileError.