From 3c6e7b1626e9e4149ec70f8dbd25db2e00bc9708 Mon Sep 17 00:00:00 2001 From: Kai Christensen Date: Thu, 16 Jul 2026 16:05:01 -0600 Subject: [PATCH] feat: support rfdetr-keypoint-preview custom weights upload Enables uploading RF-DETR Keypoint Preview models through version.deploy(), workspace.deploy_model(), and the upload_model CLI: - register rfdetr-keypoint-preview in _RFDETR_MODEL_TYPE_TO_CLASS (RFDETRKeypointPreview ships in rfdetr>=1.8.0, the already-required minimum for the PTL-checkpoint path) - map keypoint model types to TASK_POSE in task_of_model_type so project type validation accepts keypoint-detection projects - detect keypoint checkpoints lacking model_name via args.keypoint_head (keypoint configs also carry segmentation_head=False, which previously misdetected them as detection) - add the keypoint-preview position-encoding grid (576 // 12 = 48, mirroring RFDETRKeypointPreviewConfig) for the variant mislabel check Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 2 + roboflow/util/model_processor.py | 30 +++++++----- tests/util/test_model_processor.py | 77 +++++++++++++++++++++++++++++- 3 files changed, 96 insertions(+), 13 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 668cafe9..97b43acf 100644 --- a/CHANGELOG.md +++ b/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`) diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index 46389b35..283a8f37 100644 --- a/roboflow/util/model_processor.py +++ b/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,20 +820,17 @@ 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 @@ -834,6 +838,10 @@ def _detect_rfdetr_task(checkpoint: Any) -> str | None: 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( diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index c28cfdd0..72d7785f 100644 --- a/tests/util/test_model_processor.py +++ b/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.