diff --git a/requirements.txt b/requirements.txt index 5258d57a..1a697b60 100644 --- a/requirements.txt +++ b/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 diff --git a/roboflow/util/model_processor.py b/roboflow/util/model_processor.py index 46389b35..48fc85be 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,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,20 +822,21 @@ 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 @@ -834,6 +844,10 @@ def _detect_rfdetr_task(checkpoint: Any) -> str | None: 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( diff --git a/tests/util/test_model_processor.py b/tests/util/test_model_processor.py index c28cfdd0..da39b96e 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) @@ -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")