chore: import upstream snapshot with attribution

tests/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/assets/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/assets/test_downloads.py

@@ -0,0 +1,422 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+from types import ModuleType
+from unittest.mock import Mock, patch
+
+import pytest
+
+from rfdetr.assets import ModelWeightAsset, ModelWeights
+from rfdetr.assets.model_weights import download_pretrain_weights
+from rfdetr.platform import _IS_RFDETR_PLUS_AVAILABLE
+
+
+# Module-level fixture for common file operation mocks
+@pytest.fixture
+def mock_file_operations():
+    """Mock file operations to avoid actual file I/O."""
+    with (
+        patch("rfdetr.assets.model_weights.os.path.exists") as mock_exists,
+        patch("rfdetr.assets.model_weights._download_file") as mock_download,
+        patch("rfdetr.assets.model_weights._validate_file_md5") as mock_validate,
+    ):
+        # Default: file doesn't exist
+        mock_exists.return_value = False
+        # Default: MD5 validation passes
+        mock_validate.return_value = True
+
+        yield {"exists": mock_exists, "download": mock_download, "validate": mock_validate}
+
+
+class TestDownloadPretrainWeights:
+    """Test download_pretrain_weights function with mocking for offline testing."""
+
+    def test_download_from_local_model_weights(self, mock_file_operations):
+        """Test downloading a model from local ModelWeights."""
+        download_pretrain_weights("rf-detr-base.pth")
+
+        # Should call download with correct URL and MD5
+        mock_file_operations["download"].assert_called_once()
+        call_kwargs = mock_file_operations["download"].call_args[1]
+
+        assert call_kwargs["filename"] == "rf-detr-base.pth"
+        assert "rf-detr-base-coco.pth" in call_kwargs["url"]
+        assert call_kwargs["expected_md5"] is not None  # Should have MD5 hash
+        assert len(call_kwargs["expected_md5"]) == 32  # Valid MD5 hash
+
+    @pytest.mark.skipif(not _IS_RFDETR_PLUS_AVAILABLE, reason="rf-detr-plus not installed - skip priority test")
+    def test_download_from_rfdetr_plus_when_available(self, mock_file_operations):
+        """Test that rf-detr-plus URL is used for models absent from local ModelWeights.
+
+        Verifies the priority contract: a model not registered in local ModelWeights but
+        present in rf-detr-plus triggers _download_file with the plus URL, not a fallback
+        or silent no-op.
+        """
+        plus_url = "https://plus.example.com/plus-only-model.pth"
+        plus_asset = ModelWeightAsset(filename="plus-only-model.pth", url=plus_url, md5_hash=None)
+
+        with (
+            patch("rfdetr.assets.model_weights.ModelWeights.from_filename", return_value=None),
+            patch("rfdetr_plus.assets.ModelWeights.from_filename", return_value=plus_asset),
+        ):
+            download_pretrain_weights("plus-only-model.pth")
+
+        mock_file_operations["download"].assert_called_once()
+        call_kwargs = mock_file_operations["download"].call_args[1]
+        assert call_kwargs["url"] == plus_url
+
+    def test_download_from_platform_models_fallback(self, mock_file_operations):
+        """Test falling back to PLATFORM_MODELS when model not in ModelWeights."""
+        # Mock PLATFORM_MODELS
+        mock_platform_module = Mock()
+        mock_platform_module.PLATFORM_MODELS = {"legacy-model.pth": "https://legacy.com/model.pth"}
+
+        with (
+            patch("rfdetr.assets.model_weights.ModelWeights.from_filename", return_value=None),
+            patch.dict("sys.modules", {"rfdetr.platform": Mock(), "rfdetr.platform.downloads": mock_platform_module}),
+        ):
+            download_pretrain_weights("legacy-model.pth")
+
+            # Should call download with legacy URL
+            mock_file_operations["download"].assert_called_once()
+            call_kwargs = mock_file_operations["download"].call_args[1]
+            assert call_kwargs["url"] == "https://legacy.com/model.pth"
+            assert call_kwargs["expected_md5"] is None  # Platform models don't have MD5
+
+    def test_private_platform_models_registry_imports_without_eager_fallback(self):
+        """Test that a private-only registry imports cleanly without evaluating PLATFORM_MODELS early."""
+        from importlib import reload
+
+        import rfdetr.platform as platform_pkg
+        import rfdetr.platform.downloads as platform_downloads
+
+        mock_rfdetr_plus_pkg = ModuleType("rfdetr_plus")
+        mock_rfdetr_plus_pkg.__path__ = []
+        mock_platforms_pkg = ModuleType("rfdetr_plus.models")
+        mock_platforms_pkg.__path__ = []
+        mock_downloads_module = ModuleType("rfdetr_plus.models.downloads")
+        mock_downloads_module._PLATFORM_MODELS = {"legacy-model.pth": "https://legacy.com/model.pth"}
+        mock_platforms_pkg.downloads = mock_downloads_module
+        mock_rfdetr_plus_pkg.models = mock_platforms_pkg
+
+        with (
+            patch.dict(
+                "sys.modules",
+                {
+                    "rfdetr_plus": mock_rfdetr_plus_pkg,
+                    "rfdetr_plus.models": mock_platforms_pkg,
+                    "rfdetr_plus.models.downloads": mock_downloads_module,
+                },
+            ),
+            patch.object(platform_pkg, "_IS_RFDETR_PLUS_AVAILABLE", True),
+        ):
+            reloaded = reload(platform_downloads)
+            assert reloaded.PLATFORM_MODELS == {"legacy-model.pth": "https://legacy.com/model.pth"}
+
+        reload(platform_downloads)
+
+    def test_file_exists_with_correct_md5(self, mock_file_operations):
+        """Test that download is skipped if file exists with correct MD5."""
+        mock_file_operations["exists"].return_value = True
+        mock_file_operations["validate"].return_value = True
+
+        download_pretrain_weights("rf-detr-base.pth")
+
+        # Should not download if file exists with correct hash
+        mock_file_operations["download"].assert_not_called()
+
+    def test_file_exists_with_incorrect_md5_warns_and_skips(self, mock_file_operations):
+        """Test that file is NOT re-downloaded when MD5 is incorrect and redownload=False.
+
+        This protects fine-tuned checkpoints that share the same filename as a registry model (e.g. rf-detr-nano.pth)
+        from being silently overwritten.
+        """
+        mock_file_operations["exists"].return_value = True
+        mock_file_operations["validate"].return_value = False  # Incorrect MD5
+
+        with patch("rfdetr.assets.model_weights.logger.warning") as mock_warning:
+            download_pretrain_weights("rf-detr-base.pth")
+
+        # Should NOT re-download — the user's file must be preserved
+        mock_file_operations["download"].assert_not_called()
+        mock_warning.assert_called_once()
+        warning_msg = mock_warning.call_args[0][0]
+        assert "incorrect MD5 hash" in warning_msg
+        assert "skipping re-download to avoid overwriting it" in warning_msg
+
+    def test_redownload_flag_forces_download(self, mock_file_operations):
+        """Test that redownload=True forces re-download even if file exists."""
+        mock_file_operations["exists"].return_value = True
+        mock_file_operations["validate"].return_value = True
+
+        download_pretrain_weights("rf-detr-base.pth", redownload=True)
+
+        # Should download despite file existing
+        mock_file_operations["download"].assert_called_once()
+
+    def test_redownload_flag_forces_download_despite_incorrect_md5(self, mock_file_operations):
+        """Test that redownload=True triggers download even when MD5 is incorrect.
+
+        Verifies the force-redownload path where the user explicitly wants to overwrite an existing file (e.g. a fine-
+        tuned checkpoint) with the original registry weights.
+        """
+        mock_file_operations["exists"].return_value = True
+        mock_file_operations["validate"].return_value = False  # Incorrect MD5
+
+        download_pretrain_weights("rf-detr-base.pth", redownload=True)
+
+        # Should download because redownload=True overrides the skip-on-existing-file guard
+        mock_file_operations["download"].assert_called_once()
+
+    def test_validate_md5_disabled(self, mock_file_operations):
+        """Test that MD5 validation can be disabled."""
+        download_pretrain_weights("rf-detr-base.pth", validate_md5=False)
+
+        # Should pass expected_md5=None when validation is disabled
+        call_kwargs = mock_file_operations["download"].call_args[1]
+        assert call_kwargs["expected_md5"] is None
+
+    @patch("rfdetr.assets.model_weights.ModelWeights.from_filename", return_value=None)
+    def test_nonexistent_model_returns_early(self, mock_from_filename, mock_file_operations):
+        """Test that function returns early for non-existent models."""
+        download_pretrain_weights("nonexistent-model.pth")
+
+        # Should not attempt download
+        mock_file_operations["download"].assert_not_called()
+
+    def test_model_without_md5_hash(self, mock_file_operations):
+        """Test downloading a model that has no MD5 hash."""
+        # Create a mock asset without MD5
+        mock_asset = ModelWeightAsset(filename="test-no-md5.pth", url="https://example.com/test.pth", md5_hash=None)
+
+        with patch("rfdetr.assets.model_weights.ModelWeights.from_filename", return_value=mock_asset):
+            download_pretrain_weights("test-no-md5.pth", validate_md5=True)
+
+        # Should pass None for expected_md5
+        call_kwargs = mock_file_operations["download"].call_args[1]
+        assert call_kwargs["expected_md5"] is None
+
+    def test_file_exists_no_md5_skips_download(self, mock_file_operations):
+        """Test that if file exists and no MD5 validation, download is skipped."""
+        mock_file_operations["exists"].return_value = True
+
+        # Create a mock asset without MD5
+        mock_asset = ModelWeightAsset(filename="test-no-md5.pth", url="https://example.com/test.pth", md5_hash=None)
+
+        with patch("rfdetr.assets.model_weights.ModelWeights.from_filename", return_value=mock_asset):
+            download_pretrain_weights("test-no-md5.pth")
+
+        # Should not download if file exists (no MD5 to validate)
+        mock_file_operations["download"].assert_not_called()
+
+
+class TestDownloadIntegration:
+    """Integration tests for the complete download flow."""
+
+    @pytest.mark.parametrize("model", [pytest.param(m, id=m.filename) for m in ModelWeights])
+    def test_all_models_have_valid_md5_format(self, model: ModelWeightAsset) -> None:
+        """Test that MD5 hashes are valid when present (prevent typos)."""
+        # MD5 should be None or valid 32-char hex string
+        if model.md5_hash is not None:
+            assert len(model.md5_hash) == 32, f"{model.filename} has invalid MD5 length: {len(model.md5_hash)}"
+            assert all(c in "0123456789abcdef" for c in model.md5_hash.lower()), (
+                f"{model.filename} has invalid MD5 characters"
+            )
+
+    def test_from_filename_bidirectional_lookup(self):
+        """Test that from_filename correctly maps back to enum values."""
+        from rfdetr.assets.model_weights import ModelWeights
+
+        # Pick a known model
+        original = ModelWeights.RF_DETR_BASE
+
+        # Look it up by filename
+        asset = ModelWeights.from_filename(original.filename)
+
+        # Should return the exact same asset
+        assert asset is not None
+        assert asset.filename == original.filename
+        assert asset.url == original.url
+        assert asset.md5_hash == original.md5_hash
+
+    @patch("rfdetr.assets.model_weights.os.path.exists")
+    @patch("rfdetr.assets.model_weights._validate_file_md5")
+    @patch("rfdetr.assets.model_weights._download_file")
+    def test_download_flow_for_real_model(self, mock_download, mock_validate, mock_exists):
+        """Test the complete download flow for a real model."""
+        mock_exists.return_value = False
+        mock_validate.return_value = True
+
+        # Download a real model (mocked network)
+        download_pretrain_weights("rf-detr-base.pth")
+
+        # Verify download was called with correct parameters
+        mock_download.assert_called_once()
+        call_kwargs = mock_download.call_args[1]
+
+        assert call_kwargs["filename"] == "rf-detr-base.pth"
+        assert "storage.googleapis.com/rfdetr" in call_kwargs["url"]
+        assert call_kwargs["expected_md5"] == "b4d3ce46099eaed50626ede388caf979"
+
+
+class TestDownloadErrorHandling:
+    """Test error handling in download mechanism."""
+
+    @patch("rfdetr.assets.model_weights.os.path.exists")
+    @patch("rfdetr.assets.model_weights._download_file")
+    def test_handles_missing_rfdetr_plus_gracefully(self, mock_download, mock_exists):
+        """Test that missing rf-detr-plus is handled gracefully."""
+        mock_exists.return_value = False
+
+        # Should not raise an error if rf-detr-plus is not installed
+        download_pretrain_weights("rf-detr-base.pth")
+
+        # Should still download from local ModelWeights
+        mock_download.assert_called_once()
+
+    @patch("rfdetr.assets.model_weights.ModelWeights.from_filename", return_value=None)
+    @patch("rfdetr.assets.model_weights._download_file")
+    @patch("rfdetr.assets.model_weights.os.path.exists")
+    def test_handles_missing_platform_models_gracefully(self, mock_exists, mock_download, mock_from_filename):
+        """Test that missing platform models is handled gracefully."""
+        mock_exists.return_value = False
+
+        # Should return early without raising error
+        download_pretrain_weights("unknown-model.pth")
+
+        # Should not attempt download
+        mock_download.assert_not_called()
+
+    @patch("rfdetr.assets.model_weights._validate_file_md5")
+    @patch("rfdetr.assets.model_weights._download_file")
+    @patch("rfdetr.assets.model_weights.os.path.exists")
+    @patch("rfdetr.assets.model_weights.logger")
+    def test_logs_info_messages(self, mock_logger, mock_exists, mock_download, mock_validate):
+        """Test that appropriate log messages are generated."""
+        mock_exists.return_value = False
+        mock_validate.return_value = True
+
+        download_pretrain_weights("rf-detr-base.pth")
+
+        # Should log download message
+        mock_logger.info.assert_called()
+        log_message = mock_logger.info.call_args[0][0]
+        assert "rf-detr-base.pth" in log_message
+
+    @patch("rfdetr.assets.model_weights._download_file")
+    @patch("rfdetr.assets.model_weights._validate_file_md5")
+    @patch("rfdetr.assets.model_weights.os.path.exists")
+    @patch("rfdetr.assets.model_weights.logger")
+    def test_logs_warning_on_incorrect_md5(self, mock_logger, mock_exists, mock_validate, mock_download):
+        """Test that warning is logged when MD5 is incorrect and no re-download occurs."""
+        mock_exists.return_value = True
+        mock_validate.return_value = False
+
+        download_pretrain_weights("rf-detr-base.pth")
+
+        # Should log warning about incorrect MD5
+        mock_logger.warning.assert_called()
+        warning_message = mock_logger.warning.call_args[0][0]
+        assert "incorrect MD5 hash" in warning_message
+
+        # Must NOT re-download — fine-tuned checkpoints should be preserved
+        mock_download.assert_not_called()
+
+    @patch("rfdetr.assets.model_weights._download_file")
+    @patch("rfdetr.assets.model_weights.os.path.exists")
+    def test_absolute_path_resolves_to_known_model(self, mock_exists, mock_download):
+        """Absolute paths like /content/rf-detr-base.pth must still match the registry.
+
+        Regression test: previously ModelWeights.from_filename received the full path instead of the basename, so it
+        returned None and the download was silently skipped.
+        """
+        mock_exists.return_value = False
+
+        download_pretrain_weights("/content/rf-detr-base.pth")
+
+        # Must have attempted a download — not silently returned
+        mock_download.assert_called_once()
+        call_kwargs = mock_download.call_args[1]
+        assert call_kwargs["filename"] == "/content/rf-detr-base.pth"
+        assert "rf-detr-base-coco.pth" in call_kwargs["url"]
+
+    @patch("rfdetr.assets.model_weights._download_file")
+    @patch("rfdetr.assets.model_weights.os.path.exists")
+    def test_nested_absolute_path_resolves_to_known_model(self, mock_exists, mock_download):
+        """Nested paths like /workspace/models/rf-detr-base.pth also resolve."""
+        mock_exists.return_value = False
+
+        download_pretrain_weights("/workspace/models/rf-detr-base.pth")
+
+        mock_download.assert_called_once()
+        call_kwargs = mock_download.call_args[1]
+        assert call_kwargs["filename"] == "/workspace/models/rf-detr-base.pth"
+
+
+# ---------------------------------------------------------------------------
+# maybe_download_pretrain_weights — RF_HOME cache-dir path resolution
+# ---------------------------------------------------------------------------
+
+
+class TestMaybeDownloadPretrainWeightsCacheDir:
+    """Verify that RFDETR.maybe_download_pretrain_weights resolves paths via RF_HOME."""
+
+    def _make_rfdetr(self, pretrain_weights):
+        """Return an RFDETR shell backed by a fully validated RFDETRBaseConfig.
+
+        Uses RFDETRBaseConfig (which supplies required field defaults) so the expand_path field validator on
+        pretrain_weights is exercised end-to-end.
+
+        Args:
+            pretrain_weights: Raw value to pass to RFDETRBaseConfig; the pydantic
+                validator resolves it before assigning to model_config.pretrain_weights.
+        """
+        from rfdetr.config import RFDETRBaseConfig
+        from rfdetr.detr import RFDETR
+
+        model = object.__new__(RFDETR)
+        model.model_config = RFDETRBaseConfig(pretrain_weights=pretrain_weights)
+        return model
+
+    def test_bare_filename_resolved_to_rf_home(self, monkeypatch, tmp_path):
+        """Bare filename (no directory separator) is joined with RF_HOME before download."""
+        monkeypatch.setenv("RF_HOME", str(tmp_path))
+        downloaded = []
+        monkeypatch.setattr("rfdetr.detr.download_pretrain_weights", lambda p, **kw: downloaded.append(p))
+
+        self._make_rfdetr("rf-detr-base.pth").maybe_download_pretrain_weights()
+
+        assert downloaded == [str(tmp_path / "rf-detr-base.pth")]
+
+    def test_path_with_directory_used_as_is(self, monkeypatch, tmp_path):
+        """Path containing a directory separator is not modified by RF_HOME."""
+        monkeypatch.setenv("RF_HOME", str(tmp_path / "should_not_be_used"))
+        downloaded = []
+        monkeypatch.setattr("rfdetr.detr.download_pretrain_weights", lambda p, **kw: downloaded.append(p))
+
+        explicit = str(tmp_path / "custom" / "my_weights.pth")
+        self._make_rfdetr(explicit).maybe_download_pretrain_weights()
+
+        assert downloaded == [explicit]
+
+    def test_none_pretrain_weights_skips_download(self, monkeypatch):
+        """None pretrain_weights returns without calling download."""
+        called = []
+        monkeypatch.setattr("rfdetr.detr.download_pretrain_weights", lambda *a, **kw: called.append(True))
+
+        self._make_rfdetr(None).maybe_download_pretrain_weights()
+
+        assert called == []
+
+    def test_cache_dir_created_when_absent(self, monkeypatch, tmp_path):
+        """RF_HOME directory is created if it does not already exist."""
+        cache_dir = tmp_path / "new_cache"
+        monkeypatch.setenv("RF_HOME", str(cache_dir))
+        monkeypatch.setattr("rfdetr.detr.download_pretrain_weights", lambda *a, **kw: None)
+
+        self._make_rfdetr("rf-detr-base.pth").maybe_download_pretrain_weights()
+
+        assert cache_dir.is_dir()

tests/assets/test_model_weights.py

@@ -0,0 +1,146 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+
+import os
+from unittest.mock import patch
+
+import pytest
+
+from rfdetr.assets import ModelWeightAsset, ModelWeights, ModelWeightsBase
+from rfdetr.assets.model_weights import get_model_cache_dir
+
+
+class TestGetModelCacheDir:
+    """Verify get_model_cache_dir reads RF_HOME with correct default."""
+
+    def test_default_when_rf_home_not_set(self, monkeypatch):
+        """Returns ~/.roboflow/models when RF_HOME env var is absent."""
+        monkeypatch.delenv("RF_HOME", raising=False)
+        expected = os.path.normpath(os.path.expanduser("~/.roboflow/models"))
+        assert get_model_cache_dir() == expected
+
+    def test_custom_rf_home_absolute_path(self, monkeypatch, tmp_path):
+        """Returns exact RF_HOME value when set to an absolute path."""
+        monkeypatch.setenv("RF_HOME", str(tmp_path))
+        assert get_model_cache_dir() == str(tmp_path)
+
+    def test_tilde_in_rf_home_is_expanded(self, monkeypatch):
+        """Tilde in RF_HOME is expanded; result contains no literal tilde."""
+        monkeypatch.setenv("RF_HOME", "~/custom_rfdetr_cache")
+        result = get_model_cache_dir()
+        assert result == os.path.normpath(os.path.expanduser("~/custom_rfdetr_cache"))
+        assert "~" not in result
+
+    def test_returns_string(self, monkeypatch):
+        """Return value is always a str, not a Path."""
+        monkeypatch.delenv("RF_HOME", raising=False)
+        assert isinstance(get_model_cache_dir(), str)
+
+
+def test_from_filename_found():
+    """Test from_filename with valid filename."""
+    asset = ModelWeights.from_filename("rf-detr-base.pth")
+
+    assert asset is not None
+    assert isinstance(asset, ModelWeightAsset)
+    assert asset.filename == "rf-detr-base.pth"
+    assert asset.url.startswith("http")
+    assert "rf-detr-base-coco.pth" in asset.url
+
+
+def test_from_filename_not_found():
+    """Test from_filename with invalid filename."""
+    asset = ModelWeights.from_filename("nonexistent-model.pth")
+    assert asset is None
+
+
+def test_get_url():
+    """Test get_url class method."""
+    url = ModelWeights.get_url("rf-detr-base.pth")
+
+    assert url is not None
+    assert isinstance(url, str)
+    assert url.startswith("http")
+    assert "rf-detr-base-coco.pth" in url
+
+
+def test_get_url_not_found():
+    """Test get_url with invalid filename."""
+    url = ModelWeights.get_url("nonexistent-model.pth")
+    assert url is None
+
+
+def test_get_md5():
+    """Test get_md5 class method."""
+    md5 = ModelWeights.get_md5("rf-detr-base.pth")
+
+    # MD5 may be None if not yet computed
+    assert md5 is None or isinstance(md5, str)
+
+    # If MD5 exists, verify format
+    if md5 is not None:
+        assert len(md5) == 32
+        assert all(c in "0123456789abcdef" for c in md5.lower())
+
+
+def test_list_models():
+    """Test list_models returns all model filenames."""
+    models = ModelWeights.list_models()
+
+    assert isinstance(models, list)
+    assert len(models) > 0
+    assert "rf-detr-base.pth" in models
+    assert "rf-detr-large.pth" in models
+
+    # All entries should be strings
+    assert all(isinstance(m, str) for m in models)
+
+
+@pytest.mark.parametrize("asset", [pytest.param(a, id=a.filename) for a in ModelWeights])
+def test_all_assets_have_valid_urls(asset: ModelWeightAsset) -> None:
+    """Test that all assets have valid URLs."""
+    assert asset.url.startswith("http")
+    assert len(asset.url) > 20  # Reasonable minimum URL length
+
+
+@pytest.mark.parametrize("asset", [pytest.param(a, id=a.filename) for a in ModelWeights])
+def test_all_assets_have_valid_filenames(asset: ModelWeightAsset) -> None:
+    """Test that all assets have valid filenames."""
+    assert len(asset.filename) > 0
+    assert asset.filename.endswith((".pth", ".pt"))
+
+
+def test_filenames_are_unique():
+    """Test that all filenames are unique (prevent accidental duplicates)."""
+    filenames = [asset.filename for asset in ModelWeights]
+    assert len(filenames) == len(set(filenames)), "Duplicate filenames detected"
+
+
+def test_model_weight_asset_optional_md5():
+    """Test that MD5 hash is optional (important for new models)."""
+    asset = ModelWeightAsset(filename="test-model.pth", url="https://example.com/test-model.pth")
+
+    assert asset.md5_hash is None, "MD5 hash should be optional"
+
+
+def test_model_weights_inherits_from_base():
+    """Test inheritance for compile-time safety contract."""
+    assert issubclass(ModelWeights, ModelWeightsBase), (
+        "ModelWeights must inherit from ModelWeightsBase for compatibility"
+    )
+
+
+def test_rfdetr_large_deprecated_emits_future_warning() -> None:
+    """RFDETRLargeDeprecated emits FutureWarning on instantiation via the pyDeprecate class decorator."""
+    from rfdetr.variants import RFDETRLargeDeprecated
+
+    # The proxy uses num_warns=1; reset counter so this test is order-independent
+    # regardless of whether another test already triggered the warning in this session.
+    RFDETRLargeDeprecated._cfg.warned = 0
+    with patch("rfdetr.detr.RFDETR.__init__", return_value=None):
+        with pytest.warns(FutureWarning):
+            RFDETRLargeDeprecated()

tests/benchmarks/conftest.py

@@ -0,0 +1,130 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+import socket
+from pathlib import Path
+
+import pytest
+
+from rfdetr.datasets._develop import (
+    _COCO_URLS,
+    _coco_val_images_complete,
+    _download_and_extract,
+    _download_lock,
+    _nonempty_file_exists,
+)
+from rfdetr.utilities.reproducibility import seed_all
+
+_PROJECT_ROOT = Path(__file__).resolve().parents[2]
+_DATA_DIR = _PROJECT_ROOT / "data"
+_COCO_HOST = "images.cocodataset.org"
+_COCO_PORT = 80
+
+
+def _is_online(host: str, port: int, timeout_s: float = 3.0) -> bool:
+    try:
+        with socket.create_connection((host, port), timeout=timeout_s):
+            return True
+    except OSError:
+        return False
+
+
+@pytest.fixture(scope="session")
+def download_coco_val() -> tuple[Path, Path]:
+    """Download COCO val2017 images and annotations if not already present.
+
+    Returns:
+        Tuple containing the images root directory and annotations file path.
+    """
+    if not _is_online(_COCO_HOST, _COCO_PORT):
+        pytest.skip("Offline environment, skipping COCO val2017 benchmark tests.")
+
+    images_root = _DATA_DIR / "val2017"
+    annotations_path = _DATA_DIR / "annotations" / "instances_val2017.json"
+
+    lock_path = _DATA_DIR / ".coco_download.lock"
+    with _download_lock(lock_path):
+        if not _coco_val_images_complete(images_root):
+            _download_and_extract(_COCO_URLS["val2017"], _DATA_DIR)
+        if not _nonempty_file_exists(annotations_path):
+            _download_and_extract(_COCO_URLS["annotations"], _DATA_DIR)
+
+    return images_root, annotations_path
+
+
+@pytest.fixture(scope="session")
+def download_coco_val_keypoints() -> tuple[Path, Path]:
+    """Prepare COCO val images plus person-keypoint annotations for benchmark tests."""
+    if not _is_online(_COCO_HOST, _COCO_PORT):
+        pytest.skip("Offline environment, skipping COCO keypoint benchmark tests.")
+
+    images_root = _DATA_DIR / "val2017"
+    keypoint_annotations = _DATA_DIR / "annotations" / "person_keypoints_val2017.json"
+
+    lock_path = _DATA_DIR / ".coco_keypoint_download.lock"
+    with _download_lock(lock_path):
+        if not images_root.exists():
+            _download_and_extract(_COCO_URLS["val2017"], _DATA_DIR)
+        if not keypoint_annotations.exists():
+            _download_and_extract(_COCO_URLS["annotations"], _DATA_DIR)
+
+    return images_root, keypoint_annotations
+
+
+@pytest.fixture(scope="session")
+def download_coco_train_val_keypoints() -> Path:
+    """Prepare full COCO train/val images plus person-keypoint annotations for release-qualification tests."""
+    if not _is_online(_COCO_HOST, _COCO_PORT):
+        pytest.skip("Offline environment, skipping full COCO keypoint training validation.")
+
+    lock_path = _DATA_DIR / ".coco_keypoint_train_val_download.lock"
+    with _download_lock(lock_path):
+        if not (_DATA_DIR / "train2017").exists():
+            _download_and_extract(_COCO_URLS["train2017"], _DATA_DIR)
+        if not (_DATA_DIR / "val2017").exists():
+            _download_and_extract(_COCO_URLS["val2017"], _DATA_DIR)
+        if (
+            not (_DATA_DIR / "annotations" / "person_keypoints_train2017.json").exists()
+            or not (_DATA_DIR / "annotations" / "person_keypoints_val2017.json").exists()
+        ):
+            _download_and_extract(_COCO_URLS["annotations"], _DATA_DIR)
+
+    return _DATA_DIR
+
+
+@pytest.fixture(autouse=True)
+def seed_everything(request: pytest.FixtureRequest) -> None:
+    """Reset random, numpy, torch, and CUDA seeds before each test.
+
+    Defaults to seed 7. Override per-test via indirect parametrize::
+
+        @pytest.mark.parametrize("seed_everything", [42], indirect=True)
+        def test_foo(seed_everything): ...
+
+    Args:
+        request: Pytest fixture request that may carry an overridden seed.
+    """
+    seed = request.param if hasattr(request, "param") else 7
+    seed_all(seed)
+
+
+def pytest_collection_modifyitems(config: pytest.Config, items: list[pytest.Item]) -> None:
+    """Reorder tests to prioritize long-running training test before xdist distribution.
+
+    This hook runs after collection but before xdist distributes tests to workers. By moving the training test to the
+    front, we ensure it gets scheduled early, maximizing parallel resource utilization.
+    """
+    training_tests = []
+    other_tests = []
+
+    for item in items:
+        # Prioritize the synthetic training convergence tests (detection + segmentation)
+        if "training" in item.nodeid:
+            training_tests.append(item)
+        else:
+            other_tests.append(item)
+
+    # Reorder: training tests first, then everything else
+    items[:] = training_tests + other_tests

tests/benchmarks/test_develop_downloads.py

@@ -0,0 +1,118 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for private developer download helpers."""
+
+import io
+import zipfile
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+from rfdetr.datasets._develop import (
+    _coco_val_images_complete,
+    _download_and_extract,
+    _download_lock,
+    _nonempty_file_exists,
+)
+
+
+class TestCocoValImagesComplete:
+    """Regression coverage for interrupted COCO val2017 image downloads."""
+
+    def test_missing_directory_is_incomplete(self, tmp_path: Path) -> None:
+        """A missing image directory must trigger a download."""
+        assert not _coco_val_images_complete(tmp_path / "val2017")
+
+    def test_empty_existing_directory_is_incomplete(self, tmp_path: Path) -> None:
+        """An empty ``val2017`` directory must not skip the image download."""
+        images_root = tmp_path / "val2017"
+        images_root.mkdir()
+
+        assert not _coco_val_images_complete(images_root)
+
+    @pytest.mark.parametrize(
+        "file_count,expected",
+        [
+            pytest.param(1, False, id="below_threshold_is_incomplete"),
+            pytest.param(2, True, id="at_threshold_is_complete"),
+            pytest.param(3, True, id="above_threshold_is_complete"),
+        ],
+    )
+    def test_file_count_threshold(
+        self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch, file_count: int, expected: bool
+    ) -> None:
+        """Directory completeness reflects the >= threshold semantics."""
+        import rfdetr.datasets._develop as _develop_mod
+
+        monkeypatch.setattr(_develop_mod, "_COCO_VAL_IMAGE_COUNT", 2)
+        images_root = tmp_path / "val2017"
+        images_root.mkdir()
+        for i in range(file_count):
+            (images_root / f"{i:012d}.jpg").write_bytes(b"jpeg")
+
+        assert _coco_val_images_complete(images_root) is expected
+
+
+class TestNonemptyFileExists:
+    """Regression coverage for annotation file integrity checks in benchmark downloads."""
+
+    def test_missing_file_is_incomplete(self, tmp_path: Path) -> None:
+        """A missing annotation file must trigger a download."""
+        annotations_path = tmp_path / "instances_val2017.json"
+
+        assert not _nonempty_file_exists(annotations_path)
+
+    def test_empty_file_is_incomplete(self, tmp_path: Path) -> None:
+        """An empty annotation file must trigger a re-download."""
+        annotations_path = tmp_path / "instances_val2017.json"
+        annotations_path.write_bytes(b"")
+
+        assert not _nonempty_file_exists(annotations_path)
+
+    def test_nonempty_file_is_complete(self, tmp_path: Path) -> None:
+        """A non-empty annotation file is accepted without re-download."""
+        annotations_path = tmp_path / "instances_val2017.json"
+        annotations_path.write_bytes(b"{}")
+
+        assert _nonempty_file_exists(annotations_path)
+
+
+class TestDownloadLock:
+    """Coverage for the cross-process file-lock context manager."""
+
+    def test_timeout_raises_when_lock_held(self, tmp_path: Path) -> None:
+        """TimeoutError is raised immediately when the lock file already exists and timeout_s=0."""
+        lock_path = tmp_path / "test.lock"
+        lock_path.touch()
+
+        with pytest.raises(TimeoutError):
+            with _download_lock(lock_path, timeout_s=0, poll_s=0):
+                pass
+
+
+class TestDownloadAndExtract:
+    """Coverage for the ZIP download-and-extract helper."""
+
+    def _make_zip(self, members: dict) -> bytes:
+        """Build an in-memory ZIP archive from a mapping of filename→content."""
+        buf = io.BytesIO()
+        with zipfile.ZipFile(buf, "w") as zf:
+            for name, content in members.items():
+                zf.writestr(name, content)
+        return buf.getvalue()
+
+    def test_path_traversal_raises_runtime_error(self, tmp_path: Path) -> None:
+        """A ZIP entry escaping dest_dir must raise RuntimeError (path-traversal guard)."""
+        zip_bytes = self._make_zip({"../evil.txt": "malicious"})
+        url = "http://example.com/test.zip"
+
+        def fake_urlretrieve(url: str, dest: str) -> None:
+            Path(dest).write_bytes(zip_bytes)
+
+        with patch("rfdetr.datasets._develop.urlretrieve", side_effect=fake_urlretrieve):
+            with pytest.raises(RuntimeError, match="Unsafe path detected"):
+                _download_and_extract(url, tmp_path)

tests/benchmarks/test_develop_urls.py

@@ -0,0 +1,26 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""No-network tests for private COCO developer helper URL selection."""
+
+from rfdetr.datasets._develop import _COCO_URLS, get_coco_download_url
+
+
+def test_coco_helper_train2017_url_selection() -> None:
+    """``train2017`` should resolve to the official COCO train archive URL."""
+    assert get_coco_download_url("train2017") == _COCO_URLS["train2017"]
+    assert get_coco_download_url("train2017").endswith("/train2017.zip")
+
+
+def test_coco_helper_val2017_url_selection() -> None:
+    """``val2017`` should resolve to the official COCO val archive URL."""
+    assert get_coco_download_url("val2017") == _COCO_URLS["val2017"]
+    assert get_coco_download_url("val2017").endswith("/val2017.zip")
+
+
+def test_coco_helper_annotations_url_selection() -> None:
+    """``annotations`` should resolve to the COCO train/val annotations archive URL."""
+    assert get_coco_download_url("annotations") == _COCO_URLS["annotations"]
+    assert get_coco_download_url("annotations").endswith("/annotations_trainval2017.zip")

tests/benchmarks/test_inference_coco.py

@@ -0,0 +1,660 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""COCO val2017 inference benchmarks asserting pretrained-weight accuracy on CPU and GPU.
+
+Each model family (detection, segmentation) is covered by **two independent code paths**:
+
+``RFDETR.predict()`` path (public API)
+    Loads images as PIL, calls ``RFDETR.predict()`` in batches, accumulates predictions into
+    ``torchmetrics.MeanAveragePrecision`` and a confidence-threshold sweep for macro-F1.  Exercises the
+    end-to-end public inference surface — preprocessing, backbone, decoder, postprocessing — without any
+    PTL machinery.  Tests: :func:`test_inference_detection_rfdetr_predict`,
+    :func:`test_inference_segmentation_rfdetr_predict`.
+
+PTL training-stack path (``Trainer.validate``)
+    Copies pretrained weights into :class:`~rfdetr.training.RFDETRModelModule`, runs ``Trainer.validate``
+    with a :class:`~rfdetr.training.RFDETRDataModule`, and reads ``val/mAP_50`` / ``val/F1`` from the
+    callback metrics.  Exercises ``validation_step``, ``on_after_batch_transfer``, and
+    :class:`~rfdetr.training.COCOEvalCallback` — the same code path used during training.  Tests:
+    :func:`test_inference_detection_ptl_predict`, :func:`test_inference_segmentation_ptl_predict`.
+
+Both paths run on CPU (nano models) and GPU (small and larger models, ``@pytest.mark.gpu``).
+
+API contract tests (return type, shape) live in ``tests/models/test_predict.py`` and do not require a COCO
+download.
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Optional, Sequence
+
+import numpy as np
+import PIL.Image
+import pytest
+import supervision as sv
+import torch
+from faster_coco_eval import COCO
+from pytorch_lightning import LightningModule
+from torchmetrics.detection import MeanAveragePrecision
+
+from rfdetr import (
+    RFDETRKeypointPreview,
+    RFDETRLarge,
+    RFDETRMedium,
+    RFDETRNano,
+    RFDETRSeg2XLarge,
+    RFDETRSegLarge,
+    RFDETRSegMedium,
+    RFDETRSegNano,
+    RFDETRSegSmall,
+    RFDETRSegXLarge,
+    RFDETRSmall,
+)
+from rfdetr.config import ModelConfig, TrainConfig
+from rfdetr.detr import RFDETR
+from rfdetr.evaluation.coco_eval import CocoEvaluator
+from rfdetr.evaluation.f1_sweep import sweep_confidence_thresholds
+from rfdetr.evaluation.matching import (
+    build_matching_data,
+    init_matching_accumulator,
+    merge_matching_data,
+)
+from rfdetr.training import RFDETRDataModule, RFDETRModelModule, build_trainer
+
+# All tests in this file download COCO val2017 (~1 GB); exclude from CPU CI with -m "not coco17".
+pytestmark = pytest.mark.coco17
+
+# ---------------------------------------------------------------------------
+# Shared helpers
+# ---------------------------------------------------------------------------
+
+
+def _bbox_dict(
+    boxes: "list[list[float]] | np.ndarray",
+    labels: "list[int] | np.ndarray",
+    scores: "list[float] | np.ndarray | None" = None,
+    iscrowd: "list[int] | np.ndarray | None" = None,
+) -> dict[str, torch.Tensor]:
+    """Build a torchmetrics-compatible bounding-box dict from raw list or array data.
+
+    Handles empty inputs transparently — an empty *boxes* list produces a ``(0, 4)`` tensor.
+
+    Args:
+        boxes: Bounding boxes in xyxy format, shape (N, 4).
+        labels: Integer class labels, length N.
+        scores: Per-detection confidence scores, length N.  Present in prediction dicts only.
+        iscrowd: Crowd flags (0/1), length N.  Present in target dicts only.
+
+    Returns:
+        Dict always containing ``boxes`` (N, 4) float32 and ``labels`` (N,) int64; optionally
+        ``scores`` (N,) float32 and/or ``iscrowd`` (N,) uint8.
+    """
+    result: dict[str, torch.Tensor] = {
+        "boxes": torch.tensor(boxes, dtype=torch.float32).reshape(-1, 4),
+        "labels": torch.tensor(labels, dtype=torch.int64),
+    }
+    if scores is not None:
+        result["scores"] = torch.tensor(scores, dtype=torch.float32)
+    if iscrowd is not None:
+        result["iscrowd"] = torch.tensor(iscrowd, dtype=torch.uint8)
+    return result
+
+
+def _coco_ann_to_target(coco_gt: "COCO", img_id: int) -> dict[str, torch.Tensor]:
+    """Build a torchmetrics target dict from COCO ground-truth annotations for one image.
+
+    Args:
+        coco_gt: Loaded ``pycocotools.coco.COCO`` object.
+        img_id: COCO image ID.
+
+    Returns:
+        Dict with ``boxes`` (M, 4) xyxy float, ``labels`` (M,) int64, ``iscrowd`` (M,) uint8.
+    """
+    anns = coco_gt.loadAnns(coco_gt.getAnnIds(imgIds=img_id))
+    gt_boxes: list[list[float]] = []
+    gt_labels: list[int] = []
+    iscrowd: list[int] = []
+    for ann in anns:
+        bx, by, bw, bh = ann["bbox"]
+        gt_boxes.append([bx, by, bx + bw, by + bh])
+        gt_labels.append(ann["category_id"])
+        iscrowd.append(int(ann.get("iscrowd", 0)))
+    return _bbox_dict(gt_boxes, gt_labels, iscrowd=iscrowd)
+
+
+def _score_rfdetr_predict(
+    rfdetr_obj: RFDETR,
+    images_root: Path,
+    annotations_path: Path,
+    num_samples: int,
+    batch_size: int,
+) -> tuple[float, float]:
+    """Run ``RFDETR.predict()`` on a COCO val subset and return ``(mAP@50, macro-F1)``.
+
+    Loads images from disk as PIL images, calls ``rfdetr_obj.predict()`` in batches, converts
+    :class:`~supervision.Detections` to torchmetrics format, and computes bbox mAP@50 via
+    ``MeanAveragePrecision`` and macro-F1 via a confidence-threshold sweep.
+
+    Args:
+        rfdetr_obj: Pretrained :class:`~rfdetr.detr.RFDETR` instance.
+        images_root: Directory containing COCO val images (``val2017/``).
+        annotations_path: Path to ``instances_val2017.json``.
+        num_samples: Number of images to evaluate (first N by sorted image ID).
+        batch_size: Number of images per ``predict()`` call.
+
+    Returns:
+        Tuple ``(mAP@50, macro_f1)`` computed over the evaluated subset.
+    """
+    coco_gt = COCO(str(annotations_path))
+    img_ids = sorted(coco_gt.getImgIds())[:num_samples]
+
+    map_metric = MeanAveragePrecision(
+        iou_type="bbox",
+        class_metrics=False,
+        max_detection_thresholds=[1, 10, 500],
+        backend="faster_coco_eval",
+    )
+    f1_local = init_matching_accumulator()
+
+    for start in range(0, len(img_ids), batch_size):
+        batch_ids = img_ids[start : start + batch_size]
+        images: list[PIL.Image.Image] = []
+        for img_id in batch_ids:
+            with PIL.Image.open(images_root / f"{img_id:012d}.jpg") as im:
+                images.append(im.convert("RGB"))
+        detections_batch = rfdetr_obj.predict(images, threshold=0.001, include_source_image=False)
+        if not isinstance(detections_batch, list):
+            detections_batch = [detections_batch]
+        preds = [_bbox_dict(det.xyxy, det.class_id, scores=det.confidence) for det in detections_batch]
+        targets = [_coco_ann_to_target(coco_gt, img_id) for img_id in batch_ids]
+
+        map_metric.update(preds, targets)
+        batch_matching = build_matching_data(preds, targets, iou_threshold=0.5, iou_type="bbox")
+        merge_matching_data(f1_local, batch_matching)
+
+    metrics = map_metric.compute()
+    map50 = float(metrics["map_50"])
+
+    f1_val = 0.0
+    if f1_local:
+        sorted_ids = sorted(f1_local.keys())
+        per_class_list = [f1_local[cid] for cid in sorted_ids]
+        classes_with_gt = [i for i, cid in enumerate(sorted_ids) if f1_local[cid]["total_gt"] > 0]
+        f1_results = sweep_confidence_thresholds(per_class_list, np.linspace(0, 1, 101), classes_with_gt)
+        best = max(f1_results, key=lambda x: x["macro_f1"])
+        f1_val = float(best["macro_f1"])
+
+    return map50, f1_val
+
+
+def _build_train_config(coco_root: Path, tmp_path: Path, batch_size: int) -> TrainConfig:
+    """Build a minimal :class:`~rfdetr.config.TrainConfig` for COCO inference runs.
+
+    Loggers and EMA are disabled; the config is only used for validation.
+
+    Args:
+        coco_root: Directory containing ``val2017/`` and ``annotations/``.
+        tmp_path: Temporary directory used as ``output_dir``.
+        batch_size: DataLoader batch size.
+
+    Returns:
+        Minimal :class:`~rfdetr.config.TrainConfig` suitable for validation.
+    """
+    return TrainConfig(
+        dataset_file="coco",
+        dataset_dir=str(coco_root),
+        output_dir=str(tmp_path),
+        batch_size=batch_size,
+        num_workers=0 if not torch.cuda.is_available() else min(os.cpu_count(), 4),
+        tensorboard=False,
+        wandb=False,
+        mlflow=False,
+        clearml=False,
+        use_ema=False,
+        run_test=False,
+        compute_val_loss=False,
+    )
+
+
+def _build_datamodule(
+    model_config: ModelConfig,
+    train_config: TrainConfig,
+    num_samples: Optional[int] = None,
+) -> RFDETRDataModule:
+    """Set up an :class:`~rfdetr.training.RFDETRDataModule` for validation.
+
+    Calls ``setup("validate")`` so ``_dataset_val`` is ready.  When *num_samples* is set the dataset is wrapped in a
+    :class:`torch.utils.data.Subset`.
+
+    Args:
+        model_config: Architecture config (``segmentation_head`` controls mask loading).
+        train_config: Training config.
+        num_samples: If set, truncate the val dataset to this many samples.
+
+    Returns:
+        Datamodule with ``_dataset_val`` populated.
+    """
+    dm = RFDETRDataModule(model_config, train_config)
+    dm.setup("validate")
+    if num_samples is not None:
+        dm._dataset_val = torch.utils.data.Subset(
+            dm._dataset_val,
+            list(range(min(num_samples, len(dm._dataset_val)))),
+        )
+    return dm
+
+
+def _build_ptl_module(rfdetr_obj: RFDETR, train_config: TrainConfig) -> RFDETRModelModule:
+    """Copy pretrained weights from *rfdetr_obj* into a fresh :class:`~rfdetr.training.RFDETRModelModule`.
+
+    Constructs the module with the same architecture (no pretrain download), loads weights from
+    ``rfdetr_obj.model.model``, and asserts PTL lineage and weight-copy correctness before returning.
+
+    Args:
+        rfdetr_obj: A pretrained :class:`~rfdetr.detr.RFDETR` instance.
+        train_config: Shared :class:`~rfdetr.config.TrainConfig` (must have a
+            valid ``output_dir``).
+
+    Returns:
+        Weight-synced :class:`~rfdetr.training.RFDETRModelModule` ready for ``Trainer.validate`` or ``Trainer.predict``.
+    """
+    module = RFDETRModelModule(rfdetr_obj.model_config, train_config)
+    module.model.load_state_dict(rfdetr_obj.model.model.state_dict())
+    module.model.eval()
+
+    assert isinstance(module, RFDETRModelModule), f"Expected RFDETRModelModule, got {type(module).__name__}"
+    assert isinstance(module, LightningModule), (
+        "module must be a pytorch_lightning.LightningModule — this confirms evaluation runs through the PTL stack"
+    )
+
+    _first_key = next(iter(rfdetr_obj.model.model.state_dict()))
+    assert torch.equal(
+        rfdetr_obj.model.model.state_dict()[_first_key].cpu(),
+        module.model.state_dict()[_first_key].cpu(),
+    ), f"Weight copy failed: '{_first_key}' differs between legacy model and PTL module"
+
+    return module
+
+
+def _select_fixed_person_images(
+    images_root: Path,
+    annotations_path: Path,
+    max_images: int = 8,
+) -> tuple[list[str], list[int]]:
+    """Load a deterministic subset of COCO person-keypoint validation images.
+
+    Args:
+        images_root: Directory containing COCO validation images.
+        annotations_path: COCO person-keypoints annotations JSON path.
+        max_images: Maximum number of keypoint-bearing images to load.
+
+    Returns:
+        RGB image paths and their corresponding COCO image IDs.
+
+    Raises:
+        RuntimeError: If no usable person-keypoint images are available.
+    """
+    with annotations_path.open(encoding="utf-8") as file:
+        payload = json.load(file)
+
+    image_id_to_name = {int(item["id"]): str(item["file_name"]) for item in payload["images"]}
+    person_image_ids = sorted(
+        {
+            int(annotation["image_id"])
+            for annotation in payload["annotations"]
+            if int(annotation.get("num_keypoints", 0)) > 0 and int(annotation.get("iscrowd", 0)) == 0
+        }
+    )
+    selected_ids = person_image_ids[:max_images]
+    if not selected_ids:
+        raise RuntimeError("No keypoint-bearing COCO validation images were found.")
+
+    image_paths: list[str] = []
+    for image_id in selected_ids:
+        image_path = images_root / image_id_to_name[image_id]
+        image_paths.append(str(image_path))
+
+    return image_paths, selected_ids
+
+
+def _predict_keypoint_preview_batches(
+    model: RFDETRKeypointPreview,
+    image_paths: Sequence[str],
+    batch_size: int,
+    threshold: float = 0.5,
+) -> list[sv.KeyPoints]:
+    """Run keypoint-preview inference in fixed-size batches.
+
+    Args:
+        model: Loaded keypoint-preview model.
+        image_paths: COCO image paths to evaluate.
+        batch_size: Number of RGB images to pass to each ``predict()`` call.
+        threshold: Minimum confidence score passed to ``RFDETRKeypointPreview.predict()``.
+
+    Returns:
+        Per-image keypoint detections in the same order as ``image_paths``.
+
+    Raises:
+        RuntimeError: If batched prediction unexpectedly returns a single detection object.
+    """
+    predictions: list[sv.KeyPoints] = []
+    for start_idx in range(0, len(image_paths), batch_size):
+        batch_paths = list(image_paths[start_idx : start_idx + batch_size])
+        batch_images: list[PIL.Image.Image] = []
+        for image_path in batch_paths:
+            with PIL.Image.open(image_path) as image:
+                batch_images.append(image.convert("RGB"))
+
+        batch_predictions = model.predict(batch_images, threshold=threshold, include_source_image=False)
+        if not isinstance(batch_predictions, list):
+            raise RuntimeError("Expected batched keypoint preview inference to return list[KeyPoints].")
+        predictions.extend(batch_predictions)
+    return predictions
+
+
+def _detections_to_coco_predictions(
+    detections_batch: list[sv.KeyPoints],
+    image_ids: list[int],
+) -> dict[int, dict[str, torch.Tensor]]:
+    """Convert batched supervision keypoints into the COCO evaluator format.
+
+    Args:
+        detections_batch: Per-image prediction batch returned by RF-DETR.
+        image_ids: COCO image IDs matching ``detections_batch`` order.
+
+    Returns:
+        COCO evaluator prediction dictionary keyed by image ID.
+    """
+    predictions: dict[int, dict[str, torch.Tensor]] = {}
+    for image_id, key_points in zip(image_ids, detections_batch):
+        xyxy = key_points.data.get("xyxy")
+        if xyxy is None or key_points.detection_confidence is None or key_points.class_id is None:
+            raise ValueError("Expected keypoint preview predictions to populate detection details.")
+        if key_points.keypoint_confidence is None:
+            raise ValueError("Expected keypoint preview predictions to populate per-keypoint confidence.")
+        keypoints = np.concatenate((key_points.xy, key_points.keypoint_confidence[:, :, np.newaxis]), axis=2)
+        predictions[image_id] = {
+            "boxes": torch.as_tensor(xyxy, dtype=torch.float32),
+            "scores": torch.as_tensor(key_points.detection_confidence, dtype=torch.float32),
+            "labels": torch.as_tensor(key_points.class_id, dtype=torch.int64),
+            "keypoints": torch.as_tensor(keypoints, dtype=torch.float32),
+        }
+    return predictions
+
+
+@pytest.fixture(scope="session")
+def keypoint_preview_predictions(
+    download_coco_val_keypoints: tuple[Path, Path],
+) -> tuple[list[sv.KeyPoints], list[int], Path]:
+    """Run one deterministic keypoint-preview inference pass for the COCO benchmark tests."""
+    images_root, annotations_path = download_coco_val_keypoints
+    image_paths, image_ids = _select_fixed_person_images(images_root, annotations_path)
+    model = RFDETRKeypointPreview(device="cuda" if torch.cuda.is_available() else "cpu")
+    predictions = _predict_keypoint_preview_batches(model, image_paths, batch_size=8)
+    return predictions, image_ids, annotations_path
+
+
+# ---------------------------------------------------------------------------
+# Inference — RFDETR.predict() (CPU nano) / Trainer.validate() (GPU)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("model_cls", "threshold_map", "threshold_f1", "num_samples", "batch_size"),
+    [
+        pytest.param(RFDETRNano, 0.66, 0.66, 200, 6, id="det-nano"),
+        pytest.param(RFDETRSmall, 0.72, 0.70, 500, 6, id="det-small", marks=pytest.mark.gpu),
+        pytest.param(RFDETRMedium, 0.73, 0.71, 500, 4, id="det-medium", marks=pytest.mark.gpu),
+        pytest.param(RFDETRLarge, 0.74, 0.72, 500, 2, id="det-large", marks=pytest.mark.gpu),
+    ],
+)
+def test_inference_detection_rfdetr_predict(
+    download_coco_val: tuple[Path, Path],
+    model_cls: type[RFDETR],
+    threshold_map: float,
+    threshold_f1: float,
+    num_samples: int,
+    batch_size: int,
+) -> None:
+    """Asserts mAP@50 and macro-F1 thresholds on COCO val for detection models via ``RFDETR.predict()``.
+
+    Loads a pretrained detection model, calls ``RFDETR.predict()`` in batches on *num_samples* COCO val images,
+    scores via ``torchmetrics.MeanAveragePrecision`` and a confidence-threshold sweep.  Runs on CPU (nano) and GPU
+    (small/medium/large) — GPU params use a smaller *num_samples* to stay within the CI timeout.
+
+    Args:
+        download_coco_val: Fixture providing ``(images_root, annotations_path)``.
+        model_cls: Detection model class to instantiate with pretrained weights.
+        threshold_map: Minimum bbox mAP@50 required.
+        threshold_f1: Minimum macro-F1 (best across confidence sweep) required.
+        num_samples: Number of COCO val images to evaluate.
+        batch_size: Number of images per batch.
+    """
+    device_str = "cuda" if torch.cuda.is_available() else "cpu"
+    images_root, annotations_path = download_coco_val
+
+    model = model_cls(device=device_str)
+    map_val, f1_val = _score_rfdetr_predict(model, images_root, annotations_path, num_samples, batch_size)
+
+    assert map_val >= threshold_map, f"mAP@50 {map_val:.4f} < {threshold_map}"
+    assert f1_val >= threshold_f1, f"F1 {f1_val:.4f} < {threshold_f1}"
+
+
+@pytest.mark.parametrize(
+    ("model_cls", "threshold_map", "threshold_f1", "num_samples", "batch_size"),
+    [
+        pytest.param(RFDETRSegNano, 0.63, 0.64, 200, 6, id="seg-nano"),
+        pytest.param(RFDETRSegSmall, 0.66, 0.67, 100, 6, id="seg-small", marks=pytest.mark.gpu),
+        pytest.param(RFDETRSegMedium, 0.68, 0.68, 100, 4, id="seg-medium", marks=pytest.mark.gpu),
+        pytest.param(RFDETRSegLarge, 0.70, 0.69, 100, 2, id="seg-large", marks=pytest.mark.gpu),
+        pytest.param(RFDETRSegXLarge, 0.72, 0.70, 100, 2, id="seg-xlarge", marks=pytest.mark.gpu),
+        pytest.param(RFDETRSeg2XLarge, 0.73, 0.71, 100, 2, id="seg-2xlarge", marks=pytest.mark.gpu),
+    ],
+)
+def test_inference_segmentation_rfdetr_predict(
+    download_coco_val: tuple[Path, Path],
+    model_cls: type[RFDETR],
+    threshold_map: float,
+    threshold_f1: float,
+    num_samples: int,
+    batch_size: int,
+) -> None:
+    """Asserts bbox mAP@50 and macro-F1 thresholds on COCO val for segmentation models via ``RFDETR.predict()``.
+
+    Loads a pretrained segmentation model, calls ``RFDETR.predict()`` in batches on *num_samples* COCO val images,
+    scores via ``torchmetrics.MeanAveragePrecision`` and a confidence-threshold sweep.  Masks are not required — only
+    bbox IoU is used for scoring.  Runs on CPU (nano) and GPU (small and larger variants).
+
+    Args:
+        download_coco_val: Fixture providing ``(images_root, annotations_path)``.
+        model_cls: Segmentation model class to instantiate with pretrained weights.
+        threshold_map: Minimum bbox mAP@50 required.
+        threshold_f1: Minimum macro-F1 (best across confidence sweep) required.
+        num_samples: Number of COCO val images to evaluate.
+        batch_size: Number of images per batch.
+    """
+    device_str = "cuda" if torch.cuda.is_available() else "cpu"
+    images_root, annotations_path = download_coco_val
+
+    model = model_cls(device=device_str)
+    map_val, f1_val = _score_rfdetr_predict(model, images_root, annotations_path, num_samples, batch_size)
+
+    assert map_val >= threshold_map, f"mAP@50 {map_val:.4f} < {threshold_map}"
+    assert f1_val >= threshold_f1, f"F1 {f1_val:.4f} < {threshold_f1}"
+
+
+@pytest.mark.coco17
+def test_keypoint_preview_pretrained_inference_thresholded(
+    keypoint_preview_predictions: tuple[list[sv.KeyPoints], list[int], Path],
+) -> None:
+    """Pretrained preview inference should emit thresholded person keypoints."""
+    predictions, _, _ = keypoint_preview_predictions
+    assert predictions, "Expected at least one inference result."
+
+    total_detections = 0
+    total_keypoint_sets = 0
+    confidences: list[np.ndarray] = []
+
+    for key_points in predictions:
+        total_detections += len(key_points)
+        assert key_points.detection_confidence is not None
+        confidences.append(key_points.detection_confidence)
+        assert key_points.keypoint_confidence is not None
+        assert key_points.xy.ndim == 3
+        assert key_points.xy.shape[1:] == (17, 2)
+        assert key_points.keypoint_confidence.shape == (len(key_points), 17)
+        assert np.isfinite(key_points.xy).all()
+        assert np.isfinite(key_points.keypoint_confidence).all()
+        total_keypoint_sets += key_points.xy.shape[0]
+
+    assert total_detections > 0, "Expected at least one detection above threshold=0.5."
+    assert total_keypoint_sets > 0, "Expected at least one emitted keypoint set."
+
+    all_confidences = np.concatenate(confidences) if confidences else np.array([], dtype=np.float32)
+    assert all_confidences.size > 0
+    assert float(np.mean(all_confidences)) >= 0.5
+
+
+@pytest.mark.gpu
+@pytest.mark.coco17
+@pytest.mark.parametrize(
+    ("threshold_keypoint_map", "num_samples", "batch_size"),
+    [
+        pytest.param(0.71, 500, 2, id="keypoint-preview"),
+    ],
+)
+def test_inference_keypoint_preview_rfdetr_predict(
+    download_coco_val_keypoints: tuple[Path, Path],
+    threshold_keypoint_map: float,
+    num_samples: int,
+    batch_size: int,
+) -> None:
+    """``RFDETRKeypointPreview.predict()`` meets the keypoint COCO AP threshold."""
+    images_root, annotations_path = download_coco_val_keypoints
+    image_paths, image_ids = _select_fixed_person_images(images_root, annotations_path, max_images=num_samples)
+    assert len(image_ids) >= num_samples, f"Expected at least {num_samples} keypoint-bearing images."
+
+    model = RFDETRKeypointPreview(device="cuda" if torch.cuda.is_available() else "cpu")
+    predictions = _predict_keypoint_preview_batches(model, image_paths, batch_size=batch_size, threshold=0.0)
+    coco_gt = COCO(str(annotations_path))
+    coco_gt.label2cat = {1: 1}
+    evaluator = CocoEvaluator(coco_gt, ["keypoints"])
+    evaluator.update(_detections_to_coco_predictions(predictions, image_ids))
+    evaluator.synchronize_between_processes()
+    evaluator.accumulate()
+
+    keypoint_ap_50_95 = float(evaluator.coco_eval["keypoints"].stats[0])
+    assert keypoint_ap_50_95 >= threshold_keypoint_map, (
+        f"keypoint AP@50:95 {keypoint_ap_50_95:.4f} < {threshold_keypoint_map}"
+    )
+
+
+# ---------------------------------------------------------------------------
+# Inference — Trainer.validate() via PTL stack (CPU + GPU, COCO val2017)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    ("model_cls", "threshold_map", "threshold_f1", "num_samples", "batch_size"),
+    [
+        pytest.param(RFDETRNano, 0.66, 0.66, 200, 6, id="det-nano"),
+        pytest.param(RFDETRSmall, 0.72, 0.70, 500, 6, id="det-small", marks=pytest.mark.gpu),
+        pytest.param(RFDETRMedium, 0.73, 0.71, 500, 4, id="det-medium", marks=pytest.mark.gpu),
+        pytest.param(RFDETRLarge, 0.74, 0.72, 500, 2, id="det-large", marks=pytest.mark.gpu),
+    ],
+)
+def test_inference_detection_ptl_predict(
+    tmp_path: Path,
+    download_coco_val: tuple[Path, Path],
+    model_cls: type[RFDETR],
+    threshold_map: float,
+    threshold_f1: float,
+    num_samples: int,
+    batch_size: int,
+) -> None:
+    """Asserts mAP@50 and macro-F1 thresholds on COCO val for detection models via the PTL training stack.
+
+    Loads a pretrained detection model, copies weights into a :class:`~rfdetr.training.RFDETRModelModule`, and asserts
+    mAP and F1 via ``Trainer.validate``.  Exercises the PTL validation loop (``validation_step`` + callbacks) rather
+    than the public ``RFDETR.predict()`` API.
+
+    Args:
+        tmp_path: Pytest-provided temporary directory.
+        download_coco_val: Fixture providing ``(images_root, annotations_path)``.
+        model_cls: Detection model class to instantiate with pretrained weights.
+        threshold_map: Minimum ``val/mAP_50`` required.
+        threshold_f1: Minimum ``val/F1`` (best macro-F1 across confidence sweep) required.
+        num_samples: Number of val samples used for ``Trainer.validate``.
+        batch_size: DataLoader batch size.
+    """
+    device_str = "cuda" if torch.cuda.is_available() else "cpu"
+    images_root, _ = download_coco_val
+    coco_root = images_root.parent
+    accelerator = "auto" if torch.cuda.is_available() else "cpu"
+
+    model = model_cls(device=device_str)
+    tc = _build_train_config(coco_root, tmp_path, batch_size)
+    module = _build_ptl_module(model, tc)
+    trainer = build_trainer(tc, model.model_config, accelerator=accelerator)
+
+    dm = _build_datamodule(model.model_config, tc, num_samples=num_samples)
+    (metrics,) = trainer.validate(module, datamodule=dm)
+    map_val = metrics["val/mAP_50"]
+    f1_val = metrics["val/F1"]
+    assert map_val >= threshold_map, f"mAP@50 {map_val:.4f} < {threshold_map}"
+    assert f1_val >= threshold_f1, f"F1 {f1_val:.4f} < {threshold_f1}"
+
+
+@pytest.mark.parametrize(
+    ("model_cls", "threshold_map", "threshold_f1", "num_samples", "batch_size"),
+    [
+        pytest.param(RFDETRSegNano, 0.63, 0.64, 200, 6, id="seg-nano"),
+        pytest.param(RFDETRSegSmall, 0.66, 0.67, 100, 6, id="seg-small", marks=pytest.mark.gpu),
+        pytest.param(RFDETRSegMedium, 0.68, 0.68, 100, 4, id="seg-medium", marks=pytest.mark.gpu),
+        pytest.param(RFDETRSegLarge, 0.70, 0.69, 100, 2, id="seg-large", marks=pytest.mark.gpu),
+        pytest.param(RFDETRSegXLarge, 0.72, 0.70, 100, 2, id="seg-xlarge", marks=pytest.mark.gpu),
+        pytest.param(RFDETRSeg2XLarge, 0.73, 0.71, 100, 2, id="seg-2xlarge", marks=pytest.mark.gpu),
+    ],
+)
+def test_inference_segmentation_ptl_predict(
+    tmp_path: Path,
+    download_coco_val: tuple[Path, Path],
+    model_cls: type[RFDETR],
+    threshold_map: float,
+    threshold_f1: float,
+    num_samples: int,
+    batch_size: int,
+) -> None:
+    """Asserts bbox mAP@50 and macro-F1 thresholds on COCO val for segmentation models via the PTL training stack.
+
+    Same structure as :func:`test_inference_detection_ptl_predict` but for segmentation variants.
+
+    Args:
+        tmp_path: Pytest-provided temporary directory.
+        download_coco_val: Fixture providing ``(images_root, annotations_path)``.
+        model_cls: Segmentation model class to instantiate with pretrained weights.
+        threshold_map: Minimum ``val/mAP_50`` (bbox) required.
+        threshold_f1: Minimum ``val/F1`` (best macro-F1 across confidence sweep) required.
+        num_samples: Number of val samples used for ``Trainer.validate``.
+        batch_size: DataLoader batch size.
+    """
+    device_str = "cuda" if torch.cuda.is_available() else "cpu"
+    images_root, _ = download_coco_val
+    coco_root = images_root.parent
+    accelerator = "auto" if torch.cuda.is_available() else "cpu"
+
+    model = model_cls(device=device_str)
+    tc = _build_train_config(coco_root, tmp_path, batch_size)
+    module = _build_ptl_module(model, tc)
+    trainer = build_trainer(tc, model.model_config, accelerator=accelerator)
+
+    dm = _build_datamodule(model.model_config, tc, num_samples=num_samples)
+    (metrics,) = trainer.validate(module, datamodule=dm)
+    map_val = metrics["val/mAP_50"]
+    f1_val = metrics["val/F1"]
+    assert map_val >= threshold_map, f"mAP@50 {map_val:.4f} < {threshold_map}"
+    assert f1_val >= threshold_f1, f"F1 {f1_val:.4f} < {threshold_f1}"

tests/benchmarks/test_training_coco.py

@@ -0,0 +1,243 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""COCO benchmark coverage for short keypoint-preview training on a deterministic subset."""
+
+from __future__ import annotations
+
+import json
+import shutil
+from pathlib import Path
+
+import pytest
+import torch
+from torch.utils.data import Subset
+
+from rfdetr import RFDETRKeypointPreview
+from rfdetr.config import KeypointTrainConfig
+from rfdetr.training import RFDETRDataModule, RFDETRModelModule, build_trainer
+from rfdetr.utilities.reproducibility import seed_all
+
+
+def _to_float(value: float | torch.Tensor) -> float:
+    return float(value.item()) if isinstance(value, torch.Tensor) else float(value)
+
+
+def _build_subset_annotations(
+    payload: dict,
+    image_ids: list[int],
+) -> dict:
+    image_id_set = set(image_ids)
+    images = [image for image in payload["images"] if int(image["id"]) in image_id_set]
+    annotations = [
+        annotation
+        for annotation in payload["annotations"]
+        if int(annotation["image_id"]) in image_id_set
+        and int(annotation.get("iscrowd", 0)) == 0
+        and int(annotation.get("num_keypoints", 0)) > 0
+    ]
+    categories = [category for category in payload["categories"] if int(category["id"]) == 1]
+    return {
+        "info": payload.get("info", {}),
+        "licenses": payload.get("licenses", []),
+        "images": images,
+        "annotations": annotations,
+        "categories": categories,
+    }
+
+
+def _build_coco_keypoint_subset_from_val(
+    *,
+    images_root: Path,
+    annotations_path: Path,
+    output_root: Path,
+    train_images: int,
+    val_images: int,
+) -> Path:
+    with annotations_path.open(encoding="utf-8") as file:
+        payload = json.load(file)
+
+    person_image_ids = sorted(
+        {
+            int(annotation["image_id"])
+            for annotation in payload["annotations"]
+            if int(annotation.get("iscrowd", 0)) == 0 and int(annotation.get("num_keypoints", 0)) > 0
+        }
+    )
+    required = train_images + val_images
+    if len(person_image_ids) < required:
+        raise RuntimeError(f"Need at least {required} keypoint images, found {len(person_image_ids)}.")
+
+    train_ids = person_image_ids[:train_images]
+    val_ids = person_image_ids[train_images : train_images + val_images]
+    image_by_id = {int(image["id"]): image for image in payload["images"]}
+
+    train_dir = output_root / "train2017"
+    val_dir = output_root / "val2017"
+    annotations_dir = output_root / "annotations"
+    train_dir.mkdir(parents=True, exist_ok=True)
+    val_dir.mkdir(parents=True, exist_ok=True)
+    annotations_dir.mkdir(parents=True, exist_ok=True)
+
+    for image_id in train_ids:
+        file_name = str(image_by_id[image_id]["file_name"])
+        shutil.copy2(images_root / file_name, train_dir / file_name)
+    for image_id in val_ids:
+        file_name = str(image_by_id[image_id]["file_name"])
+        shutil.copy2(images_root / file_name, val_dir / file_name)
+
+    train_payload = _build_subset_annotations(payload, train_ids)
+    val_payload = _build_subset_annotations(payload, val_ids)
+
+    train_annotations = annotations_dir / "person_keypoints_train2017.json"
+    val_annotations = annotations_dir / "person_keypoints_val2017.json"
+    train_annotations.write_text(json.dumps(train_payload), encoding="utf-8")
+    val_annotations.write_text(json.dumps(val_payload), encoding="utf-8")
+
+    return output_root
+
+
+def _build_subset_datamodule(
+    model: RFDETRKeypointPreview,
+    train_config: KeypointTrainConfig,
+    train_subset_size: int = 8,
+    val_subset_size: int = 4,
+) -> RFDETRDataModule:
+    datamodule = RFDETRDataModule(model.model_config, train_config)
+    datamodule.setup("fit")
+    if datamodule._dataset_train is None or datamodule._dataset_val is None:
+        raise RuntimeError("Expected both training and validation datasets to be initialized.")
+
+    train_count = min(train_subset_size, len(datamodule._dataset_train))
+    val_count = min(val_subset_size, len(datamodule._dataset_val))
+    datamodule._dataset_train = Subset(datamodule._dataset_train, list(range(train_count)))
+    datamodule._dataset_val = Subset(datamodule._dataset_val, list(range(val_count)))
+    return datamodule
+
+
+@pytest.mark.gpu
+@pytest.mark.coco17
+@pytest.mark.flaky(reruns=1, only_rerun="AssertionError")
+def test_keypoint_training_subset_reports_loss_and_metric(
+    tmp_path: Path,
+    download_coco_val_keypoints: tuple[Path, Path],
+) -> None:
+    """Short deterministic fine-tuning should report finite loss and keypoint AP on the fixed subset."""
+    seed_all(7)
+    images_root, annotations_path = download_coco_val_keypoints
+    subset_root = _build_coco_keypoint_subset_from_val(
+        images_root=images_root,
+        annotations_path=annotations_path,
+        output_root=tmp_path / "coco_keypoint_subset",
+        train_images=64,
+        val_images=16,
+    )
+    train_config = KeypointTrainConfig(
+        dataset_file="coco",
+        dataset_dir=str(subset_root),
+        output_dir=str(tmp_path / "train_output"),
+        epochs=1,
+        batch_size=1,
+        num_workers=0,
+        grad_accum_steps=4,
+        use_ema=False,
+        run_test=False,
+        compute_val_loss=True,
+        multi_scale=False,
+        expanded_scales=False,
+        do_random_resize_via_padding=False,
+        tensorboard=False,
+        wandb=False,
+        mlflow=False,
+        clearml=False,
+    )
+    model = RFDETRKeypointPreview()
+    datamodule = _build_subset_datamodule(
+        model,
+        train_config,
+        train_subset_size=8,
+        val_subset_size=4,
+    )
+
+    module = RFDETRModelModule(model.model_config, train_config)
+    module.model.load_state_dict(model.model.model.state_dict())
+
+    trainer = build_trainer(
+        train_config,
+        model.model_config,
+        accelerator="gpu",
+        limit_train_batches=8,
+        limit_val_batches=4,
+        num_sanity_val_steps=0,
+    )
+    (pre_metrics,) = trainer.validate(module, datamodule=datamodule)
+    pre_loss = _to_float(pre_metrics["val/loss"])
+    pre_map = _to_float(pre_metrics["val/keypoint_map_50_95"])
+    assert torch.isfinite(torch.tensor(pre_loss)), f"Expected finite pre-training val/loss, got {pre_loss:.6f}"
+    assert torch.isfinite(torch.tensor(pre_map)), f"Expected finite pre-training keypoint AP, got {pre_map:.6f}"
+    assert 0.0 <= pre_map <= 1.0, f"Expected pre-training keypoint AP in [0, 1], got {pre_map:.6f}"
+
+    trainer.fit(module, datamodule=datamodule)
+    (post_metrics,) = trainer.validate(module, datamodule=datamodule)
+    post_loss = _to_float(post_metrics["val/loss"])
+    post_map = _to_float(post_metrics["val/keypoint_map_50_95"])
+    assert torch.isfinite(torch.tensor(post_loss)), f"Expected finite post-training val/loss, got {post_loss:.6f}"
+    assert torch.isfinite(torch.tensor(post_map)), f"Expected finite post-training keypoint AP, got {post_map:.6f}"
+    assert 0.0 <= post_map <= 1.0, f"Expected post-training keypoint AP in [0, 1], got {post_map:.6f}"
+
+
+@pytest.mark.gpu
+@pytest.mark.coco17
+def test_keypoint_training_full_coco_release_qualification(
+    tmp_path: Path,
+    download_coco_val_keypoints: tuple[Path, Path],
+) -> None:
+    """Release smoke gate: train and validate keypoint preview on a bounded COCO subset."""
+    seed_all(7)
+    images_root, annotations_path = download_coco_val_keypoints
+    subset_root = _build_coco_keypoint_subset_from_val(
+        images_root=images_root,
+        annotations_path=annotations_path,
+        output_root=tmp_path / "full_coco_keypoint_subset",
+        train_images=8,
+        val_images=4,
+    )
+    train_config = KeypointTrainConfig(
+        dataset_file="coco",
+        dataset_dir=str(subset_root),
+        output_dir=str(tmp_path / "full_coco_keypoint_train"),
+        epochs=1,
+        batch_size=1,
+        num_workers=0,
+        grad_accum_steps=1,
+        use_ema=False,
+        run_test=False,
+        compute_val_loss=True,
+        tensorboard=False,
+        wandb=False,
+        mlflow=False,
+        clearml=False,
+    )
+    model = RFDETRKeypointPreview()
+    datamodule = RFDETRDataModule(model.model_config, train_config)
+    module = RFDETRModelModule(model.model_config, train_config)
+    module.model.load_state_dict(model.model.model.state_dict())
+
+    trainer = build_trainer(
+        train_config,
+        model.model_config,
+        accelerator="gpu",
+        limit_train_batches=1,
+        limit_val_batches=1,
+        num_sanity_val_steps=0,
+    )
+    trainer.fit(module, datamodule=datamodule)
+    (metrics,) = trainer.validate(module, datamodule=datamodule)
+
+    val_loss = _to_float(metrics["val/loss"])
+    keypoint_map = _to_float(metrics["val/keypoint_map_50_95"])
+    assert torch.isfinite(torch.tensor(val_loss)), f"Expected finite release val/loss, got {val_loss:.6f}"
+    assert torch.isfinite(torch.tensor(keypoint_map)), f"Expected finite release keypoint AP, got {keypoint_map:.6f}"
+    assert 0.0 <= keypoint_map <= 1.0, f"Expected release keypoint AP in [0, 1], got {keypoint_map:.6f}"

tests/benchmarks/test_training_synthetic.py

@@ -0,0 +1,322 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""End-to-end benchmarks for training convergence via the PTL stack.
+
+Smoke test (CPU-friendly):
+
+* :func:`test_train_fast_dev_run` — ``Trainer.fit`` completes without error on a synthetic dataset.
+
+Training convergence (GPU, synthetic dataset, no pretrained weights):
+
+* :func:`test_train_convergence_native_ptl` — ``RFDETRModelModule`` + ``Trainer.fit`` reaches ≥ 35 % mAP@50.
+* :func:`test_train_convergence_rfdetr_api` — ``RFDETR.train()`` reaches ≥ 35 % mAP@50.
+"""
+
+import json
+import os
+from pathlib import Path
+
+import pytest
+import torch
+from pytorch_lightning import LightningModule
+
+from rfdetr import RFDETRNano
+from rfdetr.config import RFDETRBaseConfig, RFDETRNanoConfig, RFDETRSegNanoConfig, SegmentationTrainConfig, TrainConfig
+from rfdetr.detr import RFDETR
+from rfdetr.training import RFDETRDataModule, RFDETRModelModule, build_trainer
+
+# ---------------------------------------------------------------------------
+# Shared helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_ptl_module_from(rfdetr_obj: RFDETR, dataset_dir: Path, output_dir: Path) -> RFDETRModelModule:
+    """Build an :class:`~rfdetr.training.RFDETRModelModule` from an RFDETR instance.
+
+    Creates the module with the same architecture as *rfdetr_obj*, copies its current weights, and asserts PTL lineage
+    before returning.
+
+    Args:
+        rfdetr_obj: A (possibly trained) :class:`~rfdetr.detr.RFDETR` instance.
+        dataset_dir: Dataset directory forwarded to :class:`~rfdetr.config.TrainConfig`.
+        output_dir: Output directory forwarded to :class:`~rfdetr.config.TrainConfig`.
+
+    Returns:
+        Weight-synced :class:`~rfdetr.training.RFDETRModelModule` in eval mode.
+    """
+    train_config = TrainConfig(
+        dataset_file="roboflow",
+        dataset_dir=str(dataset_dir),
+        output_dir=str(output_dir),
+    )
+    model_config = rfdetr_obj.model_config.model_copy(update={"pretrain_weights": None})
+    module = RFDETRModelModule(model_config, train_config)
+    module.model.load_state_dict(rfdetr_obj.model.model.state_dict())
+    module.model.eval()
+
+    assert isinstance(module, RFDETRModelModule), f"Expected RFDETRModelModule, got {type(module).__name__}"
+    assert isinstance(module, LightningModule), "Module must be a pytorch_lightning.LightningModule"
+    return module
+
+
+# ---------------------------------------------------------------------------
+# Smoke test (CPU-friendly, no GPU required)
+# ---------------------------------------------------------------------------
+
+
+def test_train_fast_dev_run(
+    tmp_path: Path,
+    synthetic_shape_dataset_dir: Path,
+) -> None:
+    """Smoke-test the full PTL stack on a real synthetic dataset with fast_dev_run.
+
+    Uses ``build_trainer(tc, mc, fast_dev_run=2)`` and ``trainer.fit(module, datamodule=datamodule)`` with a real model
+    and real data (no mocking).  Only asserts the pipeline runs without error; convergence is tested by the GPU-only
+    tests below.
+    """
+    output_dir = tmp_path / "output"
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    with open(synthetic_shape_dataset_dir / "train" / "_annotations.coco.json") as f:
+        num_classes = len(json.load(f)["categories"])
+
+    mc = RFDETRNanoConfig(num_classes=num_classes, pretrain_weights=None, amp=False)
+    tc = TrainConfig(
+        dataset_dir=str(synthetic_shape_dataset_dir),
+        output_dir=str(output_dir),
+        epochs=1,
+        batch_size=2,
+        num_workers=0,
+        use_ema=False,
+        run_test=False,
+        tensorboard=False,
+        multi_scale=False,
+        expanded_scales=False,
+        do_random_resize_via_padding=False,
+        drop_path=0.0,
+        grad_accum_steps=1,
+    )
+
+    module = RFDETRModelModule(mc, tc)
+    datamodule = RFDETRDataModule(mc, tc)
+    trainer = build_trainer(tc, mc, accelerator="auto", fast_dev_run=2)
+    trainer.fit(module, datamodule=datamodule)
+
+
+# ---------------------------------------------------------------------------
+# Training convergence (GPU, synthetic dataset)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.gpu
+@pytest.mark.flaky(reruns=1, only_rerun="AssertionError")
+def test_train_convergence_native_ptl(
+    tmp_path: Path,
+    synthetic_shape_dataset_dir: Path,
+) -> None:
+    """Native PTL stack converges: ``RFDETRModelModule`` + ``RFDETRDataModule`` + ``Trainer.fit``.
+
+    Uses ``Trainer.validate`` before and after ``Trainer.fit`` so only Lightning elements are exercised — no
+    ``engine.evaluate`` or legacy paths.
+
+    Assertions:
+        - ``val/mAP_50`` before training ≤ 5 %.
+        - ``val/mAP_50`` after 10 epochs ≥ 35 %.
+    """
+    output_dir = tmp_path / "train_output"
+    output_dir.mkdir(parents=True, exist_ok=True)
+    dataset_dir = synthetic_shape_dataset_dir
+
+    with open(dataset_dir / "train" / "_annotations.coco.json") as f:
+        num_classes = len(json.load(f)["categories"])
+
+    accelerator = "auto" if torch.cuda.is_available() else "cpu"
+
+    mc = RFDETRBaseConfig(num_classes=num_classes, pretrain_weights=None, amp=False)
+    tc = TrainConfig(
+        dataset_file="roboflow",
+        dataset_dir=str(dataset_dir),
+        output_dir=str(output_dir),
+        epochs=10,
+        batch_size=4,
+        grad_accum_steps=1,
+        num_workers=max(1, (os.cpu_count() or 1) // 2),
+        lr=1e-3,
+        warmup_epochs=1.0,
+        use_ema=True,
+        multi_scale=False,
+        run_test=False,
+        tensorboard=False,
+    )
+
+    module = RFDETRModelModule(mc, tc)
+    datamodule = RFDETRDataModule(mc, tc)
+
+    # Pre-training baseline — untrained model should have near-zero mAP.
+    pre_trainer = build_trainer(tc, mc, accelerator=accelerator)
+    pre_results = pre_trainer.validate(module, datamodule=datamodule)
+    map_before = pre_results[0]["val/mAP_50"]
+    assert map_before <= 0.05, f"Untrained val mAP {map_before:.3f} should be ≤ 5 %."
+
+    # Train via native PTL Trainer.fit.
+    trainer = build_trainer(tc, mc, accelerator=accelerator)
+    trainer.fit(module, datamodule=datamodule)
+
+    # Post-training validation — model should have converged.
+    post_results = trainer.validate(module, datamodule=datamodule)
+    map_after = post_results[0]["val/mAP_50"]
+    assert map_after >= 0.35, f"val mAP {map_after:.3f} should reach at least 0.35 after Trainer.fit."
+
+
+@pytest.mark.gpu
+@pytest.mark.flaky(reruns=1, only_rerun="AssertionError")
+def test_train_convergence_rfdetr_api(
+    tmp_path: Path,
+    synthetic_shape_dataset_dir: Path,
+) -> None:
+    """``RFDETR.train()`` entry-point converges on synthetic data.
+
+    Exercises the public ``model.train()`` API end-to-end.  Pre- and post-training mAP are measured via
+    ``Trainer.validate`` so the assertion is identical to :func:`test_train_convergence_native_ptl`.
+
+    Assertions:
+        - ``val/mAP_50`` before training ≤ 5 %.
+        - ``val/mAP_50`` after 10 epochs ≥ 35 %.
+    """
+    output_dir = tmp_path / "train_output"
+    output_dir.mkdir(parents=True, exist_ok=True)
+    dataset_dir = synthetic_shape_dataset_dir
+
+    with open(dataset_dir / "train" / "_annotations.coco.json") as f:
+        num_classes = len(json.load(f)["categories"])
+
+    accelerator = "auto" if torch.cuda.is_available() else "cpu"
+    device = None if torch.cuda.is_available() else "cpu"
+
+    model = RFDETRNano(num_classes=num_classes, pretrain_weights=None, amp=False)
+    # Use the model's own config so RFDETRDataModule uses the correct resolution.
+    # RFDETRNano (patch_size=16, num_windows=2) requires block_size=32 divisibility;
+    # its resolution=384 satisfies this, while RFDETRBaseConfig resolution=560 does not.
+    mc = model.model_config
+    tc = TrainConfig(
+        dataset_file="roboflow",
+        dataset_dir=str(dataset_dir),
+        output_dir=str(output_dir),
+        epochs=10,
+        batch_size=4,
+        grad_accum_steps=1,
+        num_workers=max(1, (os.cpu_count() or 1) // 2),
+        lr=1e-3,
+        warmup_epochs=1.0,
+        use_ema=True,
+        multi_scale=False,
+        run_test=False,
+        tensorboard=False,
+    )
+
+    datamodule = RFDETRDataModule(mc, tc)
+
+    # Pre-training baseline via a temporary PTL module.
+    pre_module = _make_ptl_module_from(model, dataset_dir, output_dir)
+    pre_trainer = build_trainer(tc, mc, accelerator=accelerator)
+    pre_results = pre_trainer.validate(pre_module, datamodule=datamodule)
+    map_before = pre_results[0]["val/mAP_50"]
+    assert map_before <= 0.05, f"Untrained val mAP {map_before:.3f} should be ≤ 5 %."
+
+    # Train via the public RFDETR.train() API.
+    train_kwargs = dict(
+        dataset_file="roboflow",
+        dataset_dir=str(dataset_dir),
+        output_dir=str(output_dir),
+        epochs=10,
+        batch_size=4,
+        grad_accum_steps=1,
+        num_workers=max(1, (os.cpu_count() or 1) // 2),
+        lr=1e-3,
+        warmup_epochs=1.0,
+        use_ema=True,
+        multi_scale=False,
+        run_test=False,
+        tensorboard=False,
+    )
+    if device is not None:
+        train_kwargs["device"] = device
+    model.train(**train_kwargs)
+
+    # Post-training: copy trained weights into a fresh module and validate.
+    post_module = _make_ptl_module_from(model, dataset_dir, output_dir)
+    post_trainer = build_trainer(tc, mc, accelerator=accelerator)
+    post_results = post_trainer.validate(post_module, datamodule=datamodule)
+    map_after = post_results[0]["val/mAP_50"]
+    assert map_after >= 0.35, f"val mAP {map_after:.3f} should reach at least 0.35 after RFDETR.train()."
+
+
+@pytest.mark.gpu
+@pytest.mark.flaky(reruns=1, only_rerun="AssertionError")
+def test_train_convergence_segmentation(
+    tmp_path: Path,
+    synthetic_shape_segmentation_dataset_dir: Path,
+) -> None:
+    """Segmentation PTL stack converges on synthetic polygon data.
+
+    Mirrors :func:`test_train_convergence_native_ptl` but uses :class:`~rfdetr.config.RFDETRSegNanoConfig` and
+    :class:`~rfdetr.config.SegmentationTrainConfig` with a dataset that includes COCO polygon annotations.
+
+    The mask mAP threshold is deliberately lower than the bbox threshold because segmentation convergence is harder
+    within the same epoch budget. Thresholds are calibrated conservatively: the goal is to verify that the segmentation
+    training pipeline is functional (loss flows, masks are loaded, both bbox and segm mAP improve) rather than to
+    validate final accuracy.
+
+    Assertions:
+        - ``val/mAP_50`` before training ≤ 5 %.
+        - ``val/mAP_50`` after 5 epochs ≥ 10 %.
+        - ``val/segm_mAP_50`` after 5 epochs ≥ 5 %.
+    """
+    output_dir = tmp_path / "train_output_seg"
+    output_dir.mkdir(parents=True, exist_ok=True)
+    dataset_dir = synthetic_shape_segmentation_dataset_dir
+
+    with open(dataset_dir / "train" / "_annotations.coco.json") as f:
+        num_classes = len(json.load(f)["categories"])
+
+    accelerator = "auto" if torch.cuda.is_available() else "cpu"
+
+    mc = RFDETRSegNanoConfig(num_classes=num_classes, pretrain_weights=None, amp=False)
+    tc = SegmentationTrainConfig(
+        dataset_file="roboflow",
+        dataset_dir=str(dataset_dir),
+        output_dir=str(output_dir),
+        epochs=5,
+        batch_size=4,
+        grad_accum_steps=1,
+        num_workers=max(1, (os.cpu_count() or 1) // 2),
+        lr=1e-3,
+        warmup_epochs=1.0,
+        use_ema=True,
+        multi_scale=False,
+        run_test=False,
+        tensorboard=False,
+    )
+
+    module = RFDETRModelModule(mc, tc)
+    datamodule = RFDETRDataModule(mc, tc)
+
+    # Pre-training baseline — untrained model should have near-zero mAP.
+    pre_trainer = build_trainer(tc, mc, accelerator=accelerator)
+    pre_results = pre_trainer.validate(module, datamodule=datamodule)
+    map_before = pre_results[0]["val/mAP_50"]
+    assert map_before <= 0.05, f"Untrained val bbox mAP {map_before:.3f} should be ≤ 5 %."
+
+    # Train via native PTL Trainer.fit.
+    trainer = build_trainer(tc, mc, accelerator=accelerator)
+    trainer.fit(module, datamodule=datamodule)
+
+    # Post-training validation — both bbox and mask mAP should have improved.
+    post_results = trainer.validate(module, datamodule=datamodule)
+    map_after = post_results[0]["val/mAP_50"]
+    segm_map_after = post_results[0]["val/segm_mAP_50"]
+    assert map_after >= 0.15, f"val bbox mAP {map_after:.3f} should reach at least 0.15 after Trainer.fit."
+    assert segm_map_after >= 0.05, f"val segm mAP {segm_map_after:.3f} should reach at least 0.05 after Trainer.fit."

tests/cli/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/cli/test_configs.py

@@ -0,0 +1,143 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for YAML config files in configs/ — PTL Ch4/T6.
+
+Verifies that every example YAML config file:
+  - exists on disk,
+  - parses as valid YAML,
+  - contains a ``model`` section with ``model_config`` and ``train_config``,
+  - references the expected model class_path, and
+  - segmentation configs use SegmentationTrainConfig.
+"""
+
+import pathlib
+
+import pytest
+import yaml
+
+CONFIGS_DIR = pathlib.Path(__file__).parent.parent.parent / "configs"
+
+DETECTION_CONFIGS = [
+    "rfdetr_nano",
+    "rfdetr_small",
+    "rfdetr_medium",
+    "rfdetr_base",
+    "rfdetr_large",
+]
+
+SEGMENTATION_CONFIGS = [
+    "rfdetr_seg_nano",
+    "rfdetr_seg_small",
+    "rfdetr_seg_medium",
+    "rfdetr_seg_large",
+    "rfdetr_seg_xlarge",
+    "rfdetr_seg_2xlarge",
+]
+
+ALL_CONFIGS = DETECTION_CONFIGS + SEGMENTATION_CONFIGS
+
+# Maps filename stem → expected model_config class_path.
+EXPECTED_MODEL_CLASS = {
+    "rfdetr_nano": "rfdetr.config.RFDETRNanoConfig",
+    "rfdetr_small": "rfdetr.config.RFDETRSmallConfig",
+    "rfdetr_medium": "rfdetr.config.RFDETRMediumConfig",
+    "rfdetr_base": "rfdetr.config.RFDETRBaseConfig",
+    "rfdetr_large": "rfdetr.config.RFDETRLargeConfig",
+    "rfdetr_seg_nano": "rfdetr.config.RFDETRSegNanoConfig",
+    "rfdetr_seg_small": "rfdetr.config.RFDETRSegSmallConfig",
+    "rfdetr_seg_medium": "rfdetr.config.RFDETRSegMediumConfig",
+    "rfdetr_seg_large": "rfdetr.config.RFDETRSegLargeConfig",
+    "rfdetr_seg_xlarge": "rfdetr.config.RFDETRSegXLargeConfig",
+    "rfdetr_seg_2xlarge": "rfdetr.config.RFDETRSeg2XLargeConfig",
+}
+
+
+def _load(name: str) -> dict:
+    """Parse a config file by stem name and return its dict."""
+    return yaml.safe_load((CONFIGS_DIR / f"{name}.yaml").read_text())
+
+
+# ---------------------------------------------------------------------------
+# File existence
+# ---------------------------------------------------------------------------
+
+
+class TestConfigFilesExist:
+    """Every expected YAML config file must be present on disk."""
+
+    @pytest.mark.parametrize("name", ALL_CONFIGS)
+    def test_config_file_exists(self, name):
+        """configs/{name}.yaml must exist."""
+        assert (CONFIGS_DIR / f"{name}.yaml").exists(), f"Missing config file: {name}.yaml"
+
+
+# ---------------------------------------------------------------------------
+# YAML validity
+# ---------------------------------------------------------------------------
+
+
+class TestConfigFilesValidYAML:
+    """Each file must be parseable as YAML and produce a mapping."""
+
+    @pytest.mark.parametrize("name", ALL_CONFIGS)
+    def test_config_is_valid_yaml(self, name):
+        """yaml.safe_load must succeed and return a dict."""
+        data = _load(name)
+        assert isinstance(data, dict), f"{name}.yaml did not parse to a dict"
+
+
+# ---------------------------------------------------------------------------
+# Structure
+# ---------------------------------------------------------------------------
+
+
+class TestConfigStructure:
+    """Each YAML must have a model section with model_config and train_config."""
+
+    @pytest.mark.parametrize("name", ALL_CONFIGS)
+    def test_has_model_section(self, name):
+        """Top-level 'model' key must be present."""
+        assert "model" in _load(name), f"{name}.yaml missing 'model' section"
+
+    @pytest.mark.parametrize("name", ALL_CONFIGS)
+    def test_has_model_config(self, name):
+        """model.model_config must be present."""
+        assert "model_config" in _load(name)["model"], f"{name}.yaml missing model.model_config"
+
+    @pytest.mark.parametrize("name", ALL_CONFIGS)
+    def test_has_train_config(self, name):
+        """model.train_config must be present."""
+        assert "train_config" in _load(name)["model"], f"{name}.yaml missing model.train_config"
+
+
+# ---------------------------------------------------------------------------
+# Class paths
+# ---------------------------------------------------------------------------
+
+
+class TestConfigClassPaths:
+    """model_config class_path must match the expected model variant."""
+
+    @pytest.mark.parametrize("name", ALL_CONFIGS)
+    def test_model_config_class_path(self, name):
+        """model.model_config.class_path must match the variant."""
+        got = _load(name)["model"]["model_config"]["class_path"]
+        want = EXPECTED_MODEL_CLASS[name]
+        assert got == want, f"{name}.yaml: expected class_path {want!r}, got {got!r}"
+
+    @pytest.mark.parametrize("name", SEGMENTATION_CONFIGS)
+    def test_seg_uses_segmentation_train_config(self, name):
+        """Segmentation configs must use SegmentationTrainConfig."""
+        got = _load(name)["model"]["train_config"]["class_path"]
+        assert got == "rfdetr.config.SegmentationTrainConfig", (
+            f"{name}.yaml: train_config must use SegmentationTrainConfig, got {got!r}"
+        )
+
+    @pytest.mark.parametrize("name", DETECTION_CONFIGS)
+    def test_det_uses_train_config(self, name):
+        """Detection configs must use TrainConfig (not a subclass)."""
+        got = _load(name)["model"]["train_config"]["class_path"]
+        assert got == "rfdetr.config.TrainConfig", f"{name}.yaml: train_config must use TrainConfig, got {got!r}"

tests/cli/test_legacy_cli.py

@@ -0,0 +1,29 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for the CLI entry point configuration."""
+
+import pathlib
+import re
+
+
+class TestEntryPoint:
+    """[project.scripts] in pyproject.toml uses the correct CLI entry point."""
+
+    def _read_entry_point(self) -> str:
+        """Return the rfdetr console_scripts value from pyproject.toml."""
+        root = pathlib.Path(__file__).parent.parent.parent
+        content = (root / "pyproject.toml").read_text()
+        m = re.search(r"\[project\.scripts\].*?rfdetr\s*=\s*\"([^\"]+)\"", content, re.DOTALL)
+        assert m, "rfdetr entry not found in [project.scripts]"
+        return m.group(1)
+
+    def test_entry_point_value(self):
+        """Rfdetr entry point must be rfdetr.cli:main."""
+        assert self._read_entry_point() == "rfdetr.cli:main"
+
+    def test_entry_point_not_legacy(self):
+        """Entry point must no longer reference rfdetr.cli.main:trainer."""
+        assert self._read_entry_point() != "rfdetr.cli.main:trainer"

tests/cli/test_smoke.py

@@ -0,0 +1,195 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""CLI smoke tests and YAML roundtrip tests — PTL Ch4/T7.
+
+Smoke tests run RFDETRCli in-process with args=['--help'] / ['fit', '--help'] / ['validate', '--help'] and assert
+SystemExit(0) — no subprocess needed.
+
+YAML roundtrip tests load each example config with yaml.safe_load, import the class_path, construct the config object
+with the YAML init_args, and verify every specified field survived the round-trip.
+"""
+
+import importlib
+import pathlib
+
+import pytest
+import yaml
+
+CONFIGS_DIR = pathlib.Path(__file__).parent.parent.parent / "configs"
+
+ALL_CONFIGS = [
+    "rfdetr_nano",
+    "rfdetr_small",
+    "rfdetr_medium",
+    "rfdetr_base",
+    "rfdetr_large",
+    "rfdetr_seg_nano",
+    "rfdetr_seg_small",
+    "rfdetr_seg_medium",
+    "rfdetr_seg_large",
+    "rfdetr_seg_xlarge",
+    "rfdetr_seg_2xlarge",
+]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _run_cli(*args: str) -> int:
+    """Run RFDETRCli in-process with the given args; return the SystemExit code."""
+    from rfdetr.training.cli import RFDETRCli
+    from rfdetr.training.module_data import RFDETRDataModule
+    from rfdetr.training.module_model import RFDETRModelModule
+
+    with pytest.raises(SystemExit) as exc_info:
+        RFDETRCli(RFDETRModelModule, RFDETRDataModule, args=list(args))
+    return exc_info.value.code
+
+
+def _load(name: str) -> dict:
+    return yaml.safe_load((CONFIGS_DIR / f"{name}.yaml").read_text())
+
+
+def _instantiate(class_path: str, init_args: dict) -> object:
+    """Import class_path and construct an instance with init_args."""
+    module_path, class_name = class_path.rsplit(".", 1)
+    cls = getattr(importlib.import_module(module_path), class_name)
+    return cls(**init_args)
+
+
+# ---------------------------------------------------------------------------
+# Smoke tests
+# ---------------------------------------------------------------------------
+
+
+class TestCLIEntrypoint:
+    """Module entrypoint and CLI import tests."""
+
+    def test_python_module_entrypoint_runs(self) -> None:
+        """Python -m rfdetr --help exits 0 and mentions rfdetr."""
+        import subprocess
+        import sys
+
+        result = subprocess.run(
+            [sys.executable, "-m", "rfdetr", "--help"],
+            capture_output=True,
+            text=True,
+            timeout=30,
+        )
+        assert result.returncode == 0
+        assert "rfdetr" in result.stdout.lower() or "rfdetr" in result.stderr.lower()
+
+    def test_cli_main_is_importable(self) -> None:
+        """rfdetr.cli module is importable and exposes a callable main."""
+        import importlib
+
+        mod = importlib.import_module("rfdetr.cli")
+        assert callable(getattr(mod, "main", None))
+
+
+class TestCLIHelp:
+    """Rfdetr --help and subcommand --help must exit 0."""
+
+    def test_top_level_help(self):
+        """Rfdetr --help exits with code 0."""
+        assert _run_cli("--help") == 0
+
+    def test_fit_help(self):
+        """Rfdetr fit --help exits with code 0."""
+        assert _run_cli("fit", "--help") == 0
+
+    def test_validate_help(self):
+        """Rfdetr validate --help exits with code 0."""
+        assert _run_cli("validate", "--help") == 0
+
+    def test_fit_help_exposes_model_config(self):
+        """Rfdetr fit --help output lists model.model_config arguments."""
+        import io
+        import sys
+
+        buf = io.StringIO()
+        old_stdout = sys.stdout
+        sys.stdout = buf
+        try:
+            _run_cli("fit", "--help")
+        finally:
+            sys.stdout = old_stdout
+        assert "model_config" in buf.getvalue()
+
+    def test_fit_help_exposes_train_config(self):
+        """Rfdetr fit --help output lists model.train_config arguments."""
+        import io
+        import sys
+
+        buf = io.StringIO()
+        old_stdout = sys.stdout
+        sys.stdout = buf
+        try:
+            _run_cli("fit", "--help")
+        finally:
+            sys.stdout = old_stdout
+        assert "train_config" in buf.getvalue()
+
+
+# ---------------------------------------------------------------------------
+# YAML roundtrip — model_config
+# ---------------------------------------------------------------------------
+
+
+class TestModelConfigRoundtrip:
+    """model_config init_args from YAML construct a valid config with matching values."""
+
+    @pytest.mark.parametrize("name", ALL_CONFIGS)
+    def test_model_config_fields_survive_roundtrip(self, name):
+        """Every field in model_config.init_args is preserved after instantiation."""
+        data = _load(name)
+        mc_section = data["model"]["model_config"]
+        class_path = mc_section["class_path"]
+        init_args = mc_section.get("init_args", {})
+
+        mc = _instantiate(class_path, init_args)
+        for field, value in init_args.items():
+            assert getattr(mc, field) == value, (
+                f"{name}.yaml: model_config.{field} expected {value!r}, got {getattr(mc, field)!r}"
+            )
+
+
+# ---------------------------------------------------------------------------
+# YAML roundtrip — train_config
+# ---------------------------------------------------------------------------
+
+
+class TestTrainConfigRoundtrip:
+    """train_config init_args from YAML construct a valid config with matching values."""
+
+    @pytest.mark.parametrize("name", ALL_CONFIGS)
+    def test_train_config_fields_survive_roundtrip(self, name, tmp_path):
+        """Every field in train_config.init_args is preserved after instantiation.
+
+        dataset_dir is rewritten to tmp_path so path expansion doesn't fail on the placeholder /data/coco value.
+        TrainConfig.expand_paths() converts relative paths containing separators to absolute, so both sides are
+        normalised with os.path.abspath before comparison.
+        """
+        import os
+
+        data = _load(name)
+        tc_section = data["model"]["train_config"]
+        class_path = tc_section["class_path"]
+        init_args = dict(tc_section.get("init_args", {}))
+
+        # Substitute placeholder dataset_dir with a real tmp_path.
+        init_args["dataset_dir"] = str(tmp_path)
+
+        tc = _instantiate(class_path, init_args)
+        for field, value in init_args.items():
+            actual = getattr(tc, field)
+            # TrainConfig.expand_paths() resolves relative path strings to
+            # absolute, so normalise both sides for string path fields.
+            if isinstance(value, str) and (os.sep in value or "/" in value):
+                value = os.path.abspath(value)
+            assert actual == value, f"{name}.yaml: train_config.{field} expected {value!r}, got {actual!r}"

tests/conftest.py

@@ -0,0 +1,151 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+import shutil
+from pathlib import Path
+from typing import Any, Generator
+
+import pytest
+
+# NOTE: Model weights (rf-detr-*.pth) download to the CWD (typically tests/ when running pytest).
+# This is a known limitation. Route: change pretrain_weights default to
+# platformdirs.user_cache_dir("rfdetr") in a future PR.
+from rfdetr.datasets.synthetic import DatasetSplitRatios, generate_coco_dataset
+from rfdetr.utilities.reproducibility import seed_all
+
+
+@pytest.fixture(scope="session", autouse=True)
+def _prewarm_dinov2_cache() -> None:
+    """Download DINOv2 backbone weights once per test session.
+
+    HuggingFace hub uses file-level locking internally, so concurrent xdist
+    workers block on each other rather than issuing duplicate network requests.
+    After the first worker finishes, all others read from the local disk cache.
+
+    Examples:
+        This fixture is autouse — no explicit reference needed in tests.
+    """
+    import os
+
+    if os.environ.get("RFDETR_SKIP_DINOV2_PREWARM") == "1":
+        return
+
+    from huggingface_hub import snapshot_download
+
+    try:
+        snapshot_download(
+            "facebook/dinov2-with-registers-base",
+            ignore_patterns=["*.msgpack", "flax_model*", "tf_model*", "rust_model*"],
+        )
+    except Exception as exc:
+        pytest.skip(f"Skipping DINOv2 cache prewarm (snapshot_download failed): {exc}")
+
+
+@pytest.fixture(autouse=True)
+def reset_random_seeds() -> None:
+    """Reset all RNG sources before every test for reproducibility."""
+    seed_all()
+
+
+@pytest.fixture
+def reset_build_namespace_warning_state() -> Generator[None, Any, None]:
+    """Reset ``build_namespace`` deprecation call counters before each test.
+
+    ``@deprecated(..., num_warns=1)`` emits only once per process by default. This fixture makes warning assertions
+    deterministic regardless of test order.
+    """
+    from rfdetr._namespace import build_namespace
+
+    state = build_namespace._state
+    snapshot = (state.called, state.warned_calls, dict(state.warned_args))
+
+    state.called = 0
+    state.warned_calls = 0
+    state.warned_args = {}
+    try:
+        yield
+    finally:
+        state.called = snapshot[0]
+        state.warned_calls = snapshot[1]
+        state.warned_args = snapshot[2]
+
+
+@pytest.fixture(scope="session")
+def synthetic_shape_dataset_dir(tmp_path_factory: pytest.TempPathFactory) -> Generator[Path, Any, None]:
+    """Build a synthetic COCO-style dataset on disk and clean it up after tests.
+
+    Args:
+        tmp_path_factory: Pytest factory for temporary directories.
+
+    Yields:
+        Path to the synthetic dataset directory.
+    """
+    seed_all()
+    dataset_dir = tmp_path_factory.mktemp("synthetic_dataset")
+    generate_coco_dataset(
+        output_dir=str(dataset_dir),
+        num_images=100,
+        img_size=224,
+        class_mode="shape",
+        min_objects=3,
+        max_objects=7,
+        split_ratios=DatasetSplitRatios(train=0.8, val=0.2, test=0.0),
+    )
+    val_dir = dataset_dir / "val"
+    valid_dir = dataset_dir / "valid"
+    if val_dir.exists() and not valid_dir.exists():
+        val_dir.rename(valid_dir)
+    test_dir = dataset_dir / "test"
+    if not test_dir.exists():
+        test_dir.mkdir(parents=True, exist_ok=True)
+        (test_dir / "_annotations.coco.json").write_text((valid_dir / "_annotations.coco.json").read_text())
+        # Ensure test split has corresponding images referenced by the annotations
+        for item in valid_dir.iterdir():
+            if item.is_file() and item.name != "_annotations.coco.json":
+                shutil.copy2(item, test_dir / item.name)
+    yield dataset_dir
+    shutil.rmtree(dataset_dir)
+
+
+@pytest.fixture(scope="session")
+def synthetic_shape_segmentation_dataset_dir(
+    tmp_path_factory: pytest.TempPathFactory,
+) -> Generator[Path, Any, None]:
+    """Build a synthetic COCO-style dataset with polygon segmentation annotations.
+
+    Same layout as :func:`synthetic_shape_dataset_dir` but every annotation includes a ``segmentation`` polygon field so
+    the dataset can be used to train or evaluate segmentation models.
+
+    Args:
+        tmp_path_factory: Pytest factory for temporary directories.
+
+    Yields:
+        Path to the synthetic segmentation dataset directory.
+    """
+    seed_all()
+    dataset_dir = tmp_path_factory.mktemp("synthetic_seg_dataset")
+    generate_coco_dataset(
+        output_dir=str(dataset_dir),
+        num_images=100,
+        img_size=224,
+        class_mode="shape",
+        min_objects=3,
+        max_objects=7,
+        split_ratios=DatasetSplitRatios(train=0.8, val=0.2, test=0.0),
+        with_segmentation=True,
+    )
+    val_dir = dataset_dir / "val"
+    valid_dir = dataset_dir / "valid"
+    if val_dir.exists() and not valid_dir.exists():
+        val_dir.rename(valid_dir)
+    test_dir = dataset_dir / "test"
+    if not test_dir.exists():
+        test_dir.mkdir(parents=True, exist_ok=True)
+        (test_dir / "_annotations.coco.json").write_text((valid_dir / "_annotations.coco.json").read_text())
+        for item in valid_dir.iterdir():
+            if item.is_file() and item.name != "_annotations.coco.json":
+                shutil.copy2(item, test_dir / item.name)
+    yield dataset_dir
+    shutil.rmtree(dataset_dir)

tests/datasets/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/datasets/test_coco.py

@@ -0,0 +1,928 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests for COCO dataset handling.
+
+Tests cover:
+- Sparse COCO category ID remapping in ``ConvertCoco``
+- ``_load_classes`` hierarchy detection (GitHub #609)
+"""
+
+import json
+import types
+from pathlib import Path
+from typing import Dict, List
+
+import pytest
+import torch
+from PIL import Image
+
+from rfdetr.datasets._keypoint_schema import infer_coco_keypoint_schema
+from rfdetr.datasets.coco import CocoDetection, ConvertCoco, build_coco, build_roboflow_from_coco
+from rfdetr.detr import RFDETR
+
+# Minimal image shared across all tests
+_IMAGE = Image.new("RGB", (100, 100))
+
+# Sparse COCO-style category IDs (as in the real COCO dataset: 1-90 with gaps)
+# e.g. COCO skips IDs 12, 26, 29, 30, 45, 66, 68, 69, 71, 83, 91
+_SPARSE_CAT_IDS = [1, 2, 3, 7, 8]  # sparse, non-zero-indexed
+
+_ANNOTATIONS = [
+    {"bbox": [10, 10, 30, 30], "category_id": 1, "area": 900, "iscrowd": 0},
+    {"bbox": [50, 50, 20, 20], "category_id": 7, "area": 400, "iscrowd": 0},
+]
+
+_CAT2LABEL = {cat_id: i for i, cat_id in enumerate(sorted(_SPARSE_CAT_IDS))}
+# {1: 0, 2: 1, 3: 2, 7: 3, 8: 4}
+
+
+def _make_target(annotations=_ANNOTATIONS):
+    return {"image_id": 1, "annotations": annotations}
+
+
+class TestConvertCocoWithoutMapping:
+    """Without cat2label, sparse IDs pass through unchanged — demonstrating the bug."""
+
+    def test_labels_are_raw_category_ids(self):
+        converter = ConvertCoco(cat2label=None)
+        _, target = converter(_IMAGE, _make_target())
+        # Raw COCO IDs — NOT safe to use as indices into an 80-class tensor
+        assert target["labels"].tolist() == [1, 7]
+
+    def test_raw_ids_would_exceed_num_classes(self):
+        """Illustrates why raw IDs cause CUDA out-of-bounds with num_classes=80."""
+        converter = ConvertCoco(cat2label=None)
+        _, target = converter(_IMAGE, _make_target())
+        num_classes = len(_SPARSE_CAT_IDS)  # 5 — same as model would see
+        assert any(lbl >= num_classes for lbl in target["labels"].tolist()), (
+            "At least one raw category_id should exceed num_classes, "
+            "triggering an out-of-bounds index in the matcher/loss."
+        )
+
+
+class TestConvertCocoWithMapping:
+    """With cat2label, sparse IDs are remapped to contiguous 0-indexed labels."""
+
+    def test_labels_are_remapped_to_zero_indexed(self):
+        converter = ConvertCoco(cat2label=_CAT2LABEL)
+        _, target = converter(_IMAGE, _make_target())
+        # category_id 1 → 0, category_id 7 → 3
+        assert target["labels"].tolist() == [0, 3]
+
+    def test_all_labels_within_num_classes(self):
+        converter = ConvertCoco(cat2label=_CAT2LABEL)
+        _, target = converter(_IMAGE, _make_target())
+        num_classes = len(_SPARSE_CAT_IDS)
+        assert all(lbl < num_classes for lbl in target["labels"].tolist())
+
+    def test_keypoints_retain_instances_with_all_invisible_keypoints(self) -> None:
+        """Instances with all-invisible keypoints must be retained for box/class supervision."""
+        converter = ConvertCoco(include_keypoints=True, num_keypoints_per_class=[17])
+        visible_keypoints = [0.0, 0.0, 0.0] * 17
+        visible_keypoints[2] = 2.0
+        unlabeled_keypoints = [0.0, 0.0, 0.0] * 17
+
+        _, target = converter(
+            _IMAGE,
+            _make_target(
+                [
+                    {
+                        "id": 1,
+                        "image_id": 1,
+                        "category_id": 1,
+                        "bbox": [10.0, 10.0, 20.0, 20.0],
+                        "area": 400.0,
+                        "iscrowd": 0,
+                        "keypoints": unlabeled_keypoints,
+                    },
+                    {
+                        "id": 2,
+                        "image_id": 1,
+                        "category_id": 1,
+                        "bbox": [30.0, 30.0, 20.0, 20.0],
+                        "area": 400.0,
+                        "iscrowd": 0,
+                        "keypoints": visible_keypoints,
+                    },
+                ]
+            ),
+        )
+
+        assert target["boxes"].shape == (2, 4)
+        assert target["labels"].tolist() == [1, 1]
+        assert target["keypoints"].shape == (2, 17, 3)
+        assert target["keypoints"][1, 0, 2].item() == 2.0
+
+    def test_roboflow_zero_indexed_is_identity(self):
+        """Roboflow datasets already use 0-indexed IDs — mapping must be identity."""
+        roboflow_cat2label = {0: 0, 1: 1, 2: 2}
+        annotations = [
+            {"bbox": [10, 10, 30, 30], "category_id": 0, "area": 900, "iscrowd": 0},
+            {"bbox": [50, 50, 20, 20], "category_id": 2, "area": 400, "iscrowd": 0},
+        ]
+        converter = ConvertCoco(cat2label=roboflow_cat2label)
+        _, target = converter(_IMAGE, _make_target(annotations))
+        assert target["labels"].tolist() == [0, 2]
+
+    def test_label_tensor_dtype(self):
+        converter = ConvertCoco(cat2label=_CAT2LABEL)
+        _, target = converter(_IMAGE, _make_target())
+        assert target["labels"].dtype == torch.int64
+
+
+def _write_coco_json(path: Path, categories: List[Dict]) -> None:
+    """Write a minimal valid COCO annotation file."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    data = {"images": [], "annotations": [], "categories": categories}
+    path.write_text(json.dumps(data))
+
+
+def _write_roboflow_keypoint_coco(path: Path, *, category_id: int = 0) -> None:
+    """Write a minimal Roboflow-style COCO keypoint split."""
+    path.parent.mkdir(parents=True, exist_ok=True)
+    image_path = path.parent / "person.png"
+    Image.new("RGB", (64, 48), color=(255, 255, 255)).save(image_path)
+    keypoint_names = [
+        "nose",
+        "left_eye",
+        "right_eye",
+        "left_ear",
+        "right_ear",
+        "left_shoulder",
+        "right_shoulder",
+        "left_elbow",
+        "right_elbow",
+        "left_wrist",
+        "right_wrist",
+        "left_hip",
+        "right_hip",
+        "left_knee",
+        "right_knee",
+        "left_ankle",
+        "right_ankle",
+    ]
+    keypoints = []
+    for idx in range(len(keypoint_names)):
+        keypoints.extend([10 + idx, 20 + idx, 2])
+    data = {
+        "images": [{"id": 1, "file_name": image_path.name, "width": 64, "height": 48}],
+        "annotations": [
+            {
+                "id": 1,
+                "image_id": 1,
+                "category_id": category_id,
+                "bbox": [8, 18, 24, 24],
+                "area": 576,
+                "iscrowd": 0,
+                "num_keypoints": len(keypoint_names),
+                "keypoints": keypoints,
+            }
+        ],
+        "categories": [
+            {
+                "id": category_id,
+                "name": "person",
+                "supercategory": "person",
+                "keypoints": keypoint_names,
+                "skeleton": [],
+            }
+        ],
+    }
+    path.write_text(json.dumps(data), encoding="utf-8")
+
+
+class TestLoadClassesHierarchy:
+    """Regression tests for ``_load_classes`` supercategory filtering (#609).
+
+    When all categories have ``supercategory: "none"`` (flat COCO datasets), ``_load_classes`` previously returned an
+    empty list. It should only filter when a Roboflow hierarchical export is detected.
+    """
+
+    def test_roboflow_hierarchy_filters_parent(self, tmp_path: Path) -> None:
+        """Roboflow exports include a parent node — only leaf categories kept."""
+        categories = [
+            {"id": 0, "name": "annotations", "supercategory": "none"},
+            {"id": 1, "name": "dog", "supercategory": "annotations"},
+            {"id": 2, "name": "cat", "supercategory": "annotations"},
+        ]
+        _write_coco_json(tmp_path / "train" / "_annotations.coco.json", categories)
+        result = RFDETR._load_classes(str(tmp_path))
+        assert result == ["dog", "cat"]
+
+    def test_flat_none_supercategory_keeps_all(self, tmp_path: Path) -> None:
+        """Flat datasets where every category has supercategory 'none' (#609)."""
+        categories = [
+            {"id": 1, "name": "dog", "supercategory": "none"},
+            {"id": 2, "name": "cat", "supercategory": "none"},
+        ]
+        _write_coco_json(tmp_path / "train" / "_annotations.coco.json", categories)
+        result = RFDETR._load_classes(str(tmp_path))
+        assert result == ["dog", "cat"]
+
+    def test_mixed_supercategories_keeps_all(self, tmp_path: Path) -> None:
+        """Mix of 'none' and non-'none' supercategories where no category is a parent of another.
+
+        'animal' appears as a supercategory but is not itself a category name, so ``has_children`` is empty and all
+        categories pass the ``name not in has_children`` filter — both 'dog' and 'cat' are returned.
+        """
+        categories = [
+            {"id": 1, "name": "dog", "supercategory": "none"},
+            {"id": 2, "name": "cat", "supercategory": "animal"},
+        ]
+        _write_coco_json(tmp_path / "train" / "_annotations.coco.json", categories)
+        result = RFDETR._load_classes(str(tmp_path))
+        assert result == ["dog", "cat"]
+
+    def test_category_named_none_does_not_empty_list(self, tmp_path: Path) -> None:
+        """If a category is literally named 'none' and all supercategories are placeholders, the loader must return all
+        class names instead of []."""
+
+        categories = [
+            {"id": 1, "name": "none", "supercategory": "none"},
+            {"id": 2, "name": "dog", "supercategory": "none"},
+            {"id": 3, "name": "cat", "supercategory": "none"},
+        ]
+        _write_coco_json(tmp_path / "train" / "_annotations.coco.json", categories)
+        result = RFDETR._load_classes(str(tmp_path))
+        assert result == ["none", "dog", "cat"]
+
+    def test_mixed_hierarchy_leaf_and_standalone_forwarding(self, tmp_path: Path) -> None:
+        """Mixed hierarchy: only leaf classes + standalone top-level categories should be forwarded.
+
+        Parent/grouping nodes are dropped.
+        """
+        categories = [
+            {"id": 1, "name": "animals", "supercategory": "none"},
+            {"id": 2, "name": "mammal", "supercategory": "animals"},
+            {"id": 3, "name": "dog", "supercategory": "mammal"},
+            {"id": 4, "name": "cat", "supercategory": "mammal"},
+            {"id": 5, "name": "bird", "supercategory": "animals"},
+            {"id": 6, "name": "eagle", "supercategory": "bird"},
+            {"id": 7, "name": "pigeon", "supercategory": "bird"},
+            {"id": 8, "name": "objects", "supercategory": "none"},
+            {"id": 9, "name": "vehicle", "supercategory": "objects"},
+            {"id": 10, "name": "car", "supercategory": "vehicle"},
+            {"id": 11, "name": "truck", "supercategory": "vehicle"},
+            {"id": 12, "name": "appliance", "supercategory": "objects"},
+            {"id": 13, "name": "toaster", "supercategory": "appliance"},
+            {"id": 14, "name": "microwave", "supercategory": "appliance"},
+            {"id": 15, "name": "person", "supercategory": "none"},
+        ]
+        _write_coco_json(tmp_path / "train" / "_annotations.coco.json", categories)
+        result = RFDETR._load_classes(str(tmp_path))
+        expected = [
+            "dog",
+            "cat",
+            "eagle",
+            "pigeon",
+            "car",
+            "truck",
+            "toaster",
+            "microwave",
+            "person",
+        ]
+        assert result == expected
+
+    def test_placeholder_values_treated_as_no_parent(self, tmp_path: Path) -> None:
+        """Placeholders like None, '', and 'null' should be treated the same as 'none'."""
+        categories = [
+            {"id": 1, "name": "dog", "supercategory": None},
+            {"id": 2, "name": "cat", "supercategory": ""},
+            {"id": 3, "name": "elephant", "supercategory": "null"},
+        ]
+        _write_coco_json(tmp_path / "train" / "_annotations.coco.json", categories)
+        result = RFDETR._load_classes(str(tmp_path))
+        assert result == ["dog", "cat", "elephant"]
+
+    def test_unsorted_category_ids_return_id_sorted_class_order(self, tmp_path: Path) -> None:
+        """Returned class names must follow category-ID order for stable index mapping."""
+        categories = [
+            {"id": 30, "name": "truck", "supercategory": "vehicle"},
+            {"id": 10, "name": "vehicle", "supercategory": "none"},
+            {"id": 20, "name": "car", "supercategory": "vehicle"},
+            {"id": 40, "name": "person", "supercategory": "none"},
+        ]
+        _write_coco_json(tmp_path / "train" / "_annotations.coco.json", categories)
+        result = RFDETR._load_classes(str(tmp_path))
+        assert result == ["car", "truck", "person"]
+
+
+class TestRoboflowCocoKeypointFormat:
+    """Roboflow COCO keypoint datasets should align labels with the keypoint schema."""
+
+    def _make_args(self, dataset_dir: Path) -> types.SimpleNamespace:
+        """Return minimal args consumed by ``build_roboflow_from_coco`` in keypoint mode."""
+        return types.SimpleNamespace(
+            dataset_dir=str(dataset_dir),
+            square_resize_div_64=False,
+            segmentation_head=False,
+            multi_scale=False,
+            expanded_scales=False,
+            do_random_resize_via_padding=False,
+            patch_size=16,
+            num_windows=4,
+            use_grouppose_keypoints=True,
+            num_keypoints_per_class=[17],
+            aug_config={},
+            augmentation_backend="cpu",
+        )
+
+    def test_keypoint_category_maps_to_active_schema_slot(self, tmp_path: Path) -> None:
+        """A one-class Roboflow keypoint dataset maps person to label 0 for the `[17]` preview schema."""
+        _write_roboflow_keypoint_coco(tmp_path / "train" / "_annotations.coco.json", category_id=0)
+
+        dataset = build_roboflow_from_coco("train", self._make_args(tmp_path), resolution=64)
+        _, target = dataset[0]
+
+        assert target["labels"].tolist() == [0]
+        assert target["keypoints"].shape == (1, 17, 3)
+        assert dataset.cat2label == {0: 0}
+        assert dataset.label2cat == {0: 0}
+        assert dataset.coco.label2cat == {0: 0}
+
+    def test_standard_coco_cat_id_maps_to_active_schema_slot(self, tmp_path: Path) -> None:
+        """Standard COCO person (cat_id=1) maps to slot 0 under the active-first [17] schema."""
+        _write_roboflow_keypoint_coco(tmp_path / "train" / "_annotations.coco.json", category_id=1)
+
+        dataset = build_roboflow_from_coco("train", self._make_args(tmp_path), resolution=64)
+
+        assert dataset.cat2label == {1: 0}
+
+    def test_keypoint_coco_without_keypoint_schema_raises(self, tmp_path: Path) -> None:
+        """Keypoint mode should fail clearly if a COCO dataset has no keypoint metadata or annotations."""
+        _write_coco_json(
+            tmp_path / "train" / "_annotations.coco.json",
+            [{"id": 0, "name": "person", "supercategory": "none"}],
+        )
+
+        with pytest.raises(ValueError, match="Keypoint COCO dataset"):
+            build_roboflow_from_coco("train", self._make_args(tmp_path), resolution=64)
+
+
+class TestInferCocoKeypointSchema:
+    """COCO keypoint schema inference."""
+
+    def test_reads_category_keypoint_metadata(self, tmp_path: Path) -> None:
+        """Category keypoint names define the per-class keypoint count."""
+        _write_roboflow_keypoint_coco(tmp_path / "train" / "_annotations.coco.json", category_id=0)
+
+        schema = infer_coco_keypoint_schema(tmp_path / "train" / "_annotations.coco.json")
+
+        assert schema.class_names == ["person"]
+        assert schema.num_keypoints_per_class == [17]
+        assert len(schema.keypoint_oks_sigmas) == 17
+
+    def test_falls_back_to_annotation_keypoint_vectors(self, tmp_path: Path) -> None:
+        """Annotation vectors can define keypoint count when category names are absent."""
+        annotation_path = tmp_path / "train" / "_annotations.coco.json"
+        annotation_path.parent.mkdir(parents=True, exist_ok=True)
+        annotation_path.write_text(
+            json.dumps(
+                {
+                    "images": [],
+                    "annotations": [
+                        {
+                            "id": 1,
+                            "image_id": 1,
+                            "category_id": 0,
+                            "bbox": [0, 0, 10, 10],
+                            "area": 100,
+                            "iscrowd": 0,
+                            "keypoints": [1, 2, 2, 3, 4, 2],
+                        }
+                    ],
+                    "categories": [{"id": 0, "name": "person", "supercategory": "none"}],
+                }
+            ),
+            encoding="utf-8",
+        )
+
+        schema = infer_coco_keypoint_schema(annotation_path)
+
+        assert schema.num_keypoints_per_class == [2]
+
+
+# ---------------------------------------------------------------------------
+# TestBuildO365RawGpuBackend — validates that build_o365_raw emits a WARNING
+# and passes gpu_postprocess when augmentation_backend != 'cpu'.
+# ---------------------------------------------------------------------------
+
+
+class TestBuildO365RawGpuBackend:
+    """build_o365_raw warns and wires gpu_postprocess for non-cpu backends."""
+
+    class _FakeArgs:
+        """Minimal args stub for build_o365_raw."""
+
+        def __init__(self, augmentation_backend="cpu", square_resize_div_64=False):
+            self.augmentation_backend = augmentation_backend
+            self.square_resize_div_64 = square_resize_div_64
+            self.multi_scale = False
+            self.expanded_scales = False
+            self.dataset_dir = "/nonexistent/o365"
+            self.coco_path = "/nonexistent/o365"
+
+    def _call_build_o365_raw(self, augmentation_backend, square_resize_div_64=False):
+        """Call build_o365_raw with mocked CocoDetection and transform builders."""
+        from unittest.mock import MagicMock, patch
+
+        from rfdetr.datasets.o365 import build_o365_raw
+
+        args = self._FakeArgs(augmentation_backend=augmentation_backend, square_resize_div_64=square_resize_div_64)
+        fake_dataset = MagicMock()
+
+        with (
+            patch("rfdetr.datasets.o365.CocoDetection", return_value=fake_dataset),
+            patch("rfdetr.datasets.o365.make_coco_transforms") as mock_transform,
+            patch("rfdetr.datasets.o365.make_coco_transforms_square_div_64") as mock_sq_transform,
+        ):
+            mock_transform.return_value = MagicMock()
+            mock_sq_transform.return_value = MagicMock()
+            result = build_o365_raw("train", args, resolution=640)
+            return result, mock_transform, mock_sq_transform
+
+    def test_cpu_backend_no_warning(self):
+        """Cpu backend does not call logger.warning with O365 content."""
+        from unittest.mock import patch
+
+        with patch("rfdetr.datasets.o365.logger") as mock_logger:
+            self._call_build_o365_raw("cpu")
+        o365_warns = [c for c in mock_logger.warning.call_args_list if "O365" in str(c)]
+        assert len(o365_warns) == 0, "cpu backend must not warn about O365 GPU augmentation"
+
+    def test_auto_backend_emits_warning(self):
+        """Auto + CUDA + kornia available: logger.warning about O365 Phase 1 limitation."""
+        import sys
+        from unittest.mock import MagicMock, patch
+
+        with (
+            patch("rfdetr.datasets.kornia_transforms._has_cuda_device", return_value=True),
+            patch.dict(sys.modules, {"kornia": MagicMock(), "kornia.augmentation": MagicMock()}),
+            patch("rfdetr.datasets.o365.logger") as mock_logger,
+        ):
+            self._call_build_o365_raw("auto")
+        o365_warns = [c for c in mock_logger.warning.call_args_list if "O365" in str(c)]
+        assert len(o365_warns) >= 1, "auto backend must warn about O365 GPU aug limitation"
+
+    def test_auto_backend_no_cuda_no_warning(self):
+        """Auto + no CUDA: resolves to cpu, no O365 warning emitted."""
+        from unittest.mock import patch
+
+        with (
+            patch("rfdetr.datasets.kornia_transforms._has_cuda_device", return_value=False),
+            patch("rfdetr.datasets.o365.logger") as mock_logger,
+        ):
+            self._call_build_o365_raw("auto")
+        o365_warns = [c for c in mock_logger.warning.call_args_list if "O365" in str(c)]
+        assert len(o365_warns) == 0, "auto + no CUDA must not warn about O365 GPU aug"
+
+    def test_gpu_postprocess_false_for_cpu_backend(self):
+        """Cpu backend passes gpu_postprocess=False (or omits it) to make_coco_transforms."""
+        _, mock_transform, _ = self._call_build_o365_raw("cpu")
+        call_kwargs = mock_transform.call_args.kwargs if mock_transform.call_args else {}
+        assert call_kwargs.get("gpu_postprocess", False) is False
+
+    def test_gpu_postprocess_true_for_auto_backend(self):
+        """Auto + CUDA + kornia available: gpu_postprocess=True passed to make_coco_transforms."""
+        import sys
+        from unittest.mock import MagicMock, patch
+
+        with (
+            patch("rfdetr.datasets.kornia_transforms._has_cuda_device", return_value=True),
+            patch.dict(sys.modules, {"kornia": MagicMock(), "kornia.augmentation": MagicMock()}),
+        ):
+            _, mock_transform, _ = self._call_build_o365_raw("auto")
+        call_kwargs = mock_transform.call_args.kwargs if mock_transform.call_args else {}
+        assert call_kwargs.get("gpu_postprocess", False) is True
+
+    def test_gpu_postprocess_false_for_auto_no_cuda(self):
+        """Auto + no CUDA: gpu_postprocess=False so CPU Normalize is retained."""
+        from unittest.mock import patch
+
+        with patch("rfdetr.datasets.kornia_transforms._has_cuda_device", return_value=False):
+            _, mock_transform, _ = self._call_build_o365_raw("auto")
+        call_kwargs = mock_transform.call_args.kwargs if mock_transform.call_args else {}
+        assert call_kwargs.get("gpu_postprocess", False) is False, "auto + no CUDA must not strip CPU Normalize"
+
+    def test_square_resize_uses_square_transform(self):
+        """square_resize_div_64=True delegates to make_coco_transforms_square_div_64."""
+        _, mock_transform, mock_sq_transform = self._call_build_o365_raw("cpu", square_resize_div_64=True)
+        mock_sq_transform.assert_called_once()
+        mock_transform.assert_not_called()
+
+    def test_gpu_backend_no_cuda_raises_runtime_error(self):
+        """Gpu backend must fail fast when CUDA is unavailable."""
+        from unittest.mock import patch
+
+        with (
+            patch("rfdetr.datasets.kornia_transforms._has_cuda_device", return_value=False),
+            pytest.raises(RuntimeError, match="CUDA"),
+        ):
+            self._call_build_o365_raw("gpu")
+
+    def test_gpu_backend_no_kornia_raises_import_error(self):
+        """Gpu backend must raise with install hint when kornia is missing."""
+        from unittest.mock import patch
+
+        original_import = __builtins__.__import__ if hasattr(__builtins__, "__import__") else __import__
+
+        def _mock_import(name, *args, **kwargs):
+            if name == "kornia" or name.startswith("kornia."):
+                raise ImportError("No module named 'kornia'")
+            return original_import(name, *args, **kwargs)
+
+        with (
+            patch("rfdetr.datasets.kornia_transforms._has_cuda_device", return_value=True),
+            patch("builtins.__import__", side_effect=_mock_import),
+            pytest.raises(ImportError, match="rfdetr\\[kornia\\]"),
+        ):
+            self._call_build_o365_raw("gpu")
+
+
+class TestBuildRoboflowFromCocoBackendResolution:
+    """Roboflow COCO builder should resolve backend for gpu_postprocess consistently."""
+
+    def test_auto_no_cuda_keeps_cpu_normalize(self):
+        """Auto + no CUDA must set gpu_postprocess=False."""
+        from unittest.mock import MagicMock, patch
+
+        from rfdetr.datasets.coco import build_roboflow_from_coco
+
+        args = types.SimpleNamespace(
+            dataset_dir="/fake/dataset",
+            augmentation_backend="auto",
+            square_resize_div_64=False,
+            segmentation_head=False,
+            multi_scale=False,
+            expanded_scales=False,
+            do_random_resize_via_padding=False,
+            patch_size=16,
+            num_windows=4,
+            aug_config=None,
+        )
+        with (
+            patch("rfdetr.datasets.coco.Path") as mock_path,
+            patch("rfdetr.datasets.coco.make_coco_transforms") as mock_transforms,
+            patch("rfdetr.datasets.coco.CocoDetection", return_value=MagicMock()),
+            patch("rfdetr.datasets.kornia_transforms._has_cuda_device", return_value=False),
+        ):
+            mock_path.return_value.exists.return_value = True
+            mock_transforms.return_value = MagicMock()
+            build_roboflow_from_coco("train", args, resolution=640)
+        assert mock_transforms.call_args.kwargs["gpu_postprocess"] is False
+
+    @pytest.mark.parametrize(
+        ("square_resize_div_64", "transform_factory"),
+        [
+            pytest.param(False, "make_coco_transforms", id="standard_resize"),
+            pytest.param(True, "make_coco_transforms_square_div_64", id="square_resize"),
+        ],
+    )
+    def test_keypoint_flip_pairs_forwarded_to_transforms(
+        self,
+        tmp_path: Path,
+        square_resize_div_64: bool,
+        transform_factory: str,
+    ) -> None:
+        """Roboflow keypoint datasets must pass flip pairs to CPU augmentation transforms."""
+        from unittest.mock import MagicMock, patch
+
+        from rfdetr.datasets.coco import build_roboflow_from_coco
+
+        args = types.SimpleNamespace(
+            dataset_dir=str(tmp_path),
+            augmentation_backend="cpu",
+            square_resize_div_64=square_resize_div_64,
+            segmentation_head=False,
+            multi_scale=False,
+            expanded_scales=False,
+            do_random_resize_via_padding=False,
+            patch_size=16,
+            num_windows=4,
+            use_grouppose_keypoints=True,
+            num_keypoints_per_class=[0, 4],
+            keypoint_flip_pairs=[0, 1, 2, 3],
+            aug_config={},
+        )
+
+        with (
+            patch(f"rfdetr.datasets.coco.{transform_factory}") as mock_transforms,
+            patch("rfdetr.datasets.coco.CocoDetection") as mock_coco,
+        ):
+            mock_transforms.return_value = MagicMock()
+            mock_coco.return_value = MagicMock()
+
+            build_roboflow_from_coco("train", args, resolution=640)
+
+        assert mock_transforms.call_args.kwargs["keypoint_flip_pairs"] == [0, 1, 2, 3]
+
+
+class TestBuilderGpuPostprocess:
+    """Verify Roboflow COCO builder sets gpu_postprocess for segmentation models."""
+
+    @pytest.mark.parametrize(
+        "segmentation_head, augmentation_backend, resolved_backend, expected_gpu_postprocess",
+        [
+            pytest.param(False, "cpu", "cpu", False, id="cpu_backend_no_seg"),
+            pytest.param(True, "cpu", "cpu", False, id="cpu_backend_with_seg"),
+            pytest.param(False, "gpu", "gpu", True, id="gpu_backend_no_seg"),
+            pytest.param(True, "gpu", "gpu", True, id="gpu_backend_with_seg"),
+            pytest.param(True, "auto", "gpu", True, id="auto_resolved_gpu_with_seg"),
+            pytest.param(True, "auto", "cpu", False, id="auto_resolved_cpu_with_seg"),
+        ],
+    )
+    def test_gpu_postprocess_flag(
+        self,
+        tmp_path,
+        segmentation_head,
+        augmentation_backend,
+        resolved_backend,
+        expected_gpu_postprocess,
+    ):
+        """Build Roboflow COCO datasets and assert the GPU postprocess flag passed to transforms."""
+        from unittest.mock import MagicMock, patch
+
+        from rfdetr.datasets.coco import build_roboflow_from_coco
+
+        annotations_dir = tmp_path / "train"
+        annotations_dir.mkdir()
+        (annotations_dir / "_annotations.coco.json").write_text(
+            json.dumps({"images": [], "annotations": [], "categories": []}),
+            encoding="utf-8",
+        )
+        args = types.SimpleNamespace(
+            dataset_dir=str(tmp_path),
+            segmentation_head=segmentation_head,
+            augmentation_backend=augmentation_backend,
+            square_resize_div_64=False,
+            multi_scale=False,
+            expanded_scales=False,
+            do_random_resize_via_padding=False,
+            patch_size=16,
+            num_windows=4,
+            aug_config=None,
+        )
+
+        with (
+            patch("rfdetr.datasets.coco._resolve_runtime_augmentation_backend", return_value=resolved_backend),
+            patch("rfdetr.datasets.coco.make_coco_transforms") as mock_transforms,
+            patch("rfdetr.datasets.coco.CocoDetection") as mock_coco,
+        ):
+            mock_transforms.return_value = MagicMock()
+            mock_coco.return_value = MagicMock()
+
+            build_roboflow_from_coco("train", args, resolution=640)
+
+        call_kwargs = mock_transforms.call_args.kwargs if mock_transforms.call_args else mock_transforms.call_args[1]
+        assert call_kwargs["gpu_postprocess"] is expected_gpu_postprocess
+
+
+def _make_keypoint_annotation(
+    *,
+    category_id: int = 1,
+    bbox: List[float] | None = None,
+    area: float = 80.0,
+    keypoints: List[float] | None = None,
+) -> Dict[str, object]:
+    """Build a minimal keypoint annotation used in keypoint conversion tests."""
+    return {
+        "bbox": bbox if bbox is not None else [10.0, 5.0, 8.0, 10.0],
+        "category_id": category_id,
+        "area": area,
+        "iscrowd": 0,
+        "keypoints": keypoints if keypoints is not None else [1.0, 2.0, 2.0] * 17,
+    }
+
+
+def _make_coco_builder_args(tmp_path: Path, *, use_grouppose_keypoints: bool) -> types.SimpleNamespace:
+    """Return a namespace with all fields consumed by ``build_coco``."""
+    return types.SimpleNamespace(
+        dataset_dir=None,
+        coco_path=str(tmp_path),
+        square_resize_div_64=False,
+        segmentation_head=False,
+        multi_scale=False,
+        expanded_scales=False,
+        do_random_resize_via_padding=False,
+        patch_size=16,
+        num_windows=4,
+        # Empty aug_config disables augmentation — these tests verify annotation routing, not aug.
+        aug_config={},
+        augmentation_backend="cpu",
+        use_grouppose_keypoints=use_grouppose_keypoints,
+        num_keypoints_per_class=[17] if use_grouppose_keypoints else [],
+        keypoint_flip_pairs=[],
+    )
+
+
+class TestConvertCocoKeypoints:
+    """ConvertCoco keypoint-mode coverage."""
+
+    def test_keypoint_target_includes_keypoints(self) -> None:
+        """Keypoint-enabled conversion should emit keypoints in ``[N, K, 3]`` format."""
+        converter = ConvertCoco(
+            include_masks=False,
+            include_keypoints=True,
+            cat2label=None,
+            num_keypoints_per_class=[17],
+        )
+
+        _, target = converter(
+            _IMAGE,
+            {"image_id": 42, "annotations": [_make_keypoint_annotation()]},
+        )
+
+        assert target["keypoints"].shape == (1, 17, 3)
+        assert target["keypoints"].dtype == torch.float32
+        assert target["labels"].tolist() == [1]
+
+    def test_person_category_stays_raw_coco_id(self) -> None:
+        """COCO person category ``1`` remains raw when no category remapping is supplied."""
+        converter = ConvertCoco(
+            include_masks=False,
+            include_keypoints=True,
+            cat2label=None,
+            num_keypoints_per_class=[17],
+        )
+        _, target = converter(
+            _IMAGE,
+            {"image_id": 7, "annotations": [_make_keypoint_annotation(category_id=1)]},
+        )
+
+        assert target["labels"].shape == (1,)
+        assert target["labels"].item() == 1
+
+    def test_num_keypoints_zero_annotation_retains_instance_for_box_supervision(self) -> None:
+        """All-zero-visibility keypoints must not drop the instance; box/class targets are still valid."""
+        converter = ConvertCoco(
+            include_masks=False,
+            include_keypoints=True,
+            cat2label=None,
+            num_keypoints_per_class=[17],
+        )
+        _, target = converter(
+            _IMAGE,
+            {"image_id": 3, "annotations": [_make_keypoint_annotation(keypoints=[0.0] * (17 * 3))]},
+        )
+
+        assert target["boxes"].shape == (1, 4)
+        assert target["labels"].shape == (1,)
+        assert target["keypoints"].shape == (1, 17, 3)
+        assert torch.count_nonzero(target["keypoints"]) == 0
+
+    def test_empty_image_uses_schema_max_shape(self) -> None:
+        """Empty images should emit ``(0, max(num_keypoints_per_class), 3)`` keypoint tensors."""
+        converter = ConvertCoco(
+            include_masks=False,
+            include_keypoints=True,
+            cat2label={1: 0},
+            num_keypoints_per_class=[2, 1],
+        )
+        _, target = converter(_IMAGE, {"image_id": 99, "annotations": []})
+
+        assert target["keypoints"].shape == (0, 2, 3)
+
+    def test_multiclass_keypoints_use_schema_max_shape(self) -> None:
+        """Multi-class keypoint targets should be padded to Kmax, not schema sum."""
+        converter = ConvertCoco(
+            include_masks=False,
+            include_keypoints=True,
+            cat2label=None,
+            num_keypoints_per_class=[2, 1],
+        )
+        _, target = converter(
+            _IMAGE,
+            {
+                "image_id": 100,
+                "annotations": [
+                    _make_keypoint_annotation(category_id=0, keypoints=[1.0, 2.0, 2.0, 3.0, 4.0, 2.0]),
+                    _make_keypoint_annotation(category_id=1, keypoints=[5.0, 6.0, 2.0]),
+                ],
+            },
+        )
+
+        assert target["labels"].tolist() == [0, 1]
+        assert target["keypoints"].shape == (2, 2, 3)
+        torch.testing.assert_close(
+            target["keypoints"][0],
+            torch.tensor([[1.0, 2.0, 2.0], [3.0, 4.0, 2.0]], dtype=torch.float32),
+            rtol=1e-4,
+            atol=1e-6,
+        )
+        torch.testing.assert_close(
+            target["keypoints"][1],
+            torch.tensor([[5.0, 6.0, 2.0], [0.0, 0.0, 0.0]], dtype=torch.float32),
+            rtol=1e-4,
+            atol=1e-6,
+        )
+
+
+class TestBuildCocoKeypointMode:
+    """COCO builder mode switch for person keypoints."""
+
+    def test_keypoint_mode_uses_person_keypoints_annotations(self, tmp_path: Path) -> None:
+        """Keypoint mode should switch train annotations to ``person_keypoints_train2017.json``."""
+        args = _make_coco_builder_args(tmp_path, use_grouppose_keypoints=True)
+
+        from unittest.mock import patch
+
+        with (
+            patch("rfdetr.datasets.coco.make_coco_transforms", return_value=lambda image, target: (image, target)),
+            patch("rfdetr.datasets.coco.CocoDetection", return_value=object()) as mock_dataset,
+        ):
+            build_coco("train", args, resolution=640)
+
+        _, kwargs = mock_dataset.call_args
+        ann_file = Path(mock_dataset.call_args.args[1])
+        assert ann_file.parent.name == "annotations"
+        assert ann_file.name == "person_keypoints_train2017.json"
+        assert kwargs["include_keypoints"] is True
+        assert kwargs["remap_category_ids"] is True
+
+    def test_default_mode_uses_instances_annotations_with_raw_coco_ids(self, tmp_path: Path) -> None:
+        """Default COCO detection mode should keep raw sparse category IDs for pretrained checkpoints."""
+        from unittest.mock import patch
+
+        args = _make_coco_builder_args(tmp_path, use_grouppose_keypoints=False)
+        with (
+            patch("rfdetr.datasets.coco.make_coco_transforms", return_value=lambda image, target: (image, target)),
+            patch("rfdetr.datasets.coco.CocoDetection", return_value=object()) as mock_dataset,
+        ):
+            build_coco("train", args, resolution=640)
+
+        _, kwargs = mock_dataset.call_args
+        ann_file = Path(mock_dataset.call_args.args[1])
+        assert ann_file.parent.name == "annotations"
+        assert ann_file.name == "instances_train2017.json"
+        assert kwargs["include_keypoints"] is False
+        assert kwargs["remap_category_ids"] is False
+
+
+class TestBuildKeypointCat2Label:
+    """Unit tests for ``_build_keypoint_cat2label`` schema alignment."""
+
+    def _person_coco(self, cat_id: int = 1) -> types.SimpleNamespace:
+        """Return a minimal COCO-like object with a single keypoint-bearing person category."""
+        return types.SimpleNamespace(
+            cats={cat_id: {"name": "person", "keypoints": ["nose"] * 17}},
+            anns={},
+        )
+
+    def test_legacy_bgfirst_schema_maps_person_to_slot_1(self) -> None:
+        """Legacy [0, 17] schema maps person (cat_id=1) to slot 1, not slot 0."""
+        from rfdetr.datasets.coco import _build_keypoint_cat2label
+
+        result = _build_keypoint_cat2label(self._person_coco(cat_id=1), num_keypoints_per_class=[0, 17])
+
+        assert result == {1: 1}
+
+    def test_mixed_detection_and_keypoint_categories(self) -> None:
+        """Non-keypoint categories fill free slots after keypoint categories are assigned."""
+        from rfdetr.datasets.coco import _build_keypoint_cat2label
+
+        coco = types.SimpleNamespace(
+            cats={
+                1: {"name": "person", "keypoints": ["nose"] * 17},
+                3: {"name": "car"},
+            },
+            anns={},
+        )
+        result = _build_keypoint_cat2label(coco, num_keypoints_per_class=[17])
+
+        assert result == {1: 0, 3: 1}
+
+
+class TestCocoDetectionZeroAnnotations:
+    """CocoDetection correctly handles images with no annotations."""
+
+    def test_zero_annotation_sample_yields_empty_boxes_and_labels(self, tmp_path: Path) -> None:
+        """An image with no annotations yields boxes (0, 4) float32 and labels (0,) int64 tensors."""
+        img_dir = tmp_path / "images"
+        img_dir.mkdir()
+        Image.new("RGB", (100, 100)).save(img_dir / "img1.jpg")
+        Image.new("RGB", (100, 100)).save(img_dir / "img2.jpg")
+        ann_file = tmp_path / "annotations.json"
+        ann_file.write_text(
+            json.dumps(
+                {
+                    "images": [
+                        {"id": 1, "file_name": "img1.jpg", "width": 100, "height": 100},
+                        {"id": 2, "file_name": "img2.jpg", "width": 100, "height": 100},
+                    ],
+                    "annotations": [
+                        {"id": 1, "image_id": 1, "category_id": 1, "bbox": [10, 10, 30, 30], "area": 900, "iscrowd": 0}
+                    ],
+                    "categories": [{"id": 1, "name": "cat", "supercategory": "animal"}],
+                }
+            )
+        )
+        dataset = CocoDetection(img_dir, ann_file, transforms=None)
+        zero_ann_idx = dataset.ids.index(2)
+        _, target = dataset[zero_ann_idx]
+        assert target["boxes"].shape == torch.Size([0, 4])
+        assert target["labels"].shape == torch.Size([0])
+        assert target["boxes"].dtype == torch.float32
+        assert target["labels"].dtype == torch.int64

tests/datasets/test_coco_resize_config.py

@@ -0,0 +1,267 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Characterization tests for _build_train_resize_config."""
+
+import pytest
+
+from rfdetr.datasets.coco import _build_train_resize_config
+
+
+class TestBuildTrainResizeConfigStructure:
+    """Top-level structure is always a single-element list wrapping a OneOf."""
+
+    @pytest.mark.parametrize(
+        "scales,square",
+        [
+            pytest.param([640], True, id="square-single"),
+            pytest.param([480, 640], True, id="square-multi"),
+            pytest.param([640], False, id="nonsquare-single"),
+            pytest.param([480, 640], False, id="nonsquare-multi"),
+        ],
+    )
+    def test_returns_single_element_list(self, scales, square):
+        result = _build_train_resize_config(scales, square=square)
+        assert isinstance(result, list)
+        assert len(result) == 1
+
+    @pytest.mark.parametrize(
+        "scales,square",
+        [
+            pytest.param([640], True, id="square-single"),
+            pytest.param([480, 640], True, id="square-multi"),
+            pytest.param([640], False, id="nonsquare-single"),
+            pytest.param([480, 640], False, id="nonsquare-multi"),
+        ],
+    )
+    def test_top_level_is_oneof_with_two_branches(self, scales, square):
+        result = _build_train_resize_config(scales, square=square)
+        entry = result[0]
+        assert "OneOf" in entry
+        oneof = entry["OneOf"]
+        assert len(oneof["transforms"]) == 2
+
+
+class TestBuildTrainResizeConfigSquareSingleScale:
+    """Square=True, single scale — OneOf[Resize] + Sequential[..., OneOf[RandomSizedCrop]]."""
+
+    def test_option_a_is_oneof_wrapping_single_resize(self):
+        result = _build_train_resize_config([640], square=True)
+        option_a = result[0]["OneOf"]["transforms"][0]
+        assert option_a == {
+            "OneOf": {
+                "transforms": [{"Resize": {"height": 640, "width": 640}}],
+            }
+        }
+
+    def test_option_b_is_sequential_with_oneof_crop(self):
+        result = _build_train_resize_config([640], square=True)
+        option_b = result[0]["OneOf"]["transforms"][1]
+        assert option_b == {
+            "Sequential": {
+                "transforms": [
+                    {"SmallestMaxSize": {"max_size": [400, 500, 600]}},
+                    {
+                        "OneOf": {
+                            "transforms": [
+                                {"RandomSizedCrop": {"min_max_height": [384, 600], "height": 640, "width": 640}},
+                            ],
+                        }
+                    },
+                ]
+            }
+        }
+
+    def test_uses_correct_scale_value(self):
+        result = _build_train_resize_config([480], square=True)
+        option_a = result[0]["OneOf"]["transforms"][0]
+        assert option_a == {
+            "OneOf": {
+                "transforms": [{"Resize": {"height": 480, "width": 480}}],
+            }
+        }
+
+
+class TestBuildTrainResizeConfigSquareMultiScale:
+    """Square=True, multiple scales — OneOf[Resize] + Sequential[..., OneOf[RandomSizedCrop]]."""
+
+    def test_option_a_is_oneof_of_resizes(self):
+        result = _build_train_resize_config([480, 640], square=True)
+        option_a = result[0]["OneOf"]["transforms"][0]
+        assert option_a == {
+            "OneOf": {
+                "transforms": [
+                    {"Resize": {"height": 480, "width": 480}},
+                    {"Resize": {"height": 640, "width": 640}},
+                ],
+            }
+        }
+
+    def test_option_b_is_sequential_with_oneof_crop(self):
+        result = _build_train_resize_config([480, 640], square=True)
+        option_b = result[0]["OneOf"]["transforms"][1]
+        assert option_b == {
+            "Sequential": {
+                "transforms": [
+                    {"SmallestMaxSize": {"max_size": [400, 500, 600]}},
+                    {
+                        "OneOf": {
+                            "transforms": [
+                                {"RandomSizedCrop": {"min_max_height": [384, 600], "height": 480, "width": 480}},
+                                {"RandomSizedCrop": {"min_max_height": [384, 600], "height": 640, "width": 640}},
+                            ],
+                        }
+                    },
+                ]
+            }
+        }
+
+    def test_three_scales_produce_three_resize_options(self):
+        result = _build_train_resize_config([384, 512, 640], square=True)
+        option_a = result[0]["OneOf"]["transforms"][0]
+        assert len(option_a["OneOf"]["transforms"]) == 3
+
+
+class TestBuildTrainResizeConfigNonSquareSingleScale:
+    """Square=False, single scale — SmallestMaxSize uses scalar, default cap 1333."""
+
+    def test_option_a_uses_scalar_size(self):
+        result = _build_train_resize_config([640], square=False)
+        option_a = result[0]["OneOf"]["transforms"][0]
+        assert option_a == {
+            "Sequential": {
+                "transforms": [
+                    {"SmallestMaxSize": {"max_size": 640}},
+                    {"LongestMaxSize": {"max_size": 1333}},
+                ]
+            }
+        }
+
+    def test_option_b_uses_scalar_size(self):
+        result = _build_train_resize_config([640], square=False)
+        option_b = result[0]["OneOf"]["transforms"][1]
+        assert option_b == {
+            "Sequential": {
+                "transforms": [
+                    {"SmallestMaxSize": {"max_size": [400, 500, 600]}},
+                    {
+                        "OneOf": {
+                            "transforms": [
+                                {"RandomSizedCrop": {"min_max_height": [384, 600], "height": 640, "width": 640}},
+                            ]
+                        }
+                    },
+                ]
+            }
+        }
+
+    def test_custom_max_size(self):
+        result = _build_train_resize_config([640], square=False, max_size=800)
+        option_a = result[0]["OneOf"]["transforms"][0]
+        assert option_a["Sequential"]["transforms"][1] == {"LongestMaxSize": {"max_size": 800}}
+
+
+class TestBuildTrainResizeConfigNonSquareMultiScale:
+    """Square=False, multiple scales — SmallestMaxSize uses list directly."""
+
+    def test_option_a_uses_list_size(self):
+        result = _build_train_resize_config([480, 640], square=False)
+        option_a = result[0]["OneOf"]["transforms"][0]
+        assert option_a == {
+            "Sequential": {
+                "transforms": [
+                    {"SmallestMaxSize": {"max_size": [480, 640]}},
+                    {"LongestMaxSize": {"max_size": 1333}},
+                ]
+            }
+        }
+
+    def test_option_b_uses_list_size(self):
+        result = _build_train_resize_config([480, 640], square=False)
+        option_b = result[0]["OneOf"]["transforms"][1]
+        assert option_b == {
+            "Sequential": {
+                "transforms": [
+                    {"SmallestMaxSize": {"max_size": [400, 500, 600]}},
+                    {
+                        "OneOf": {
+                            "transforms": [
+                                {"RandomSizedCrop": {"min_max_height": [384, 600], "height": 480, "width": 480}},
+                                {"RandomSizedCrop": {"min_max_height": [384, 600], "height": 640, "width": 640}},
+                            ]
+                        }
+                    },
+                ]
+            }
+        }
+
+    def test_custom_max_size_applies_to_option_a_only(self):
+        """max_size caps option_a's long side; option_b now resizes the crop directly to the target (no cap step)."""
+        result = _build_train_resize_config([480, 640], square=False, max_size=1000)
+        option_a = result[0]["OneOf"]["transforms"][0]
+        option_b_steps = result[0]["OneOf"]["transforms"][1]["Sequential"]["transforms"]
+        assert option_a["Sequential"]["transforms"][1] == {"LongestMaxSize": {"max_size": 1000}}
+        assert not any("LongestMaxSize" in step for step in option_b_steps)
+
+
+class TestBuildTrainResizeConfigNonSquareScaleJitter:
+    """Non-square option_b must keep RandomSizedCrop (scale jitter), not a fixed RandomCrop.
+
+    Regression tests for https://github.com/roboflow/rf-detr/issues/1018 — PR #752 replaced RandomSizeCrop(384, 600)
+    with a fixed RandomCrop(384, 384), silently removing scale jitter from the non-square training pipeline.
+
+    The ``fix-resize-crop`` branch keeps RandomSizedCrop and removes the wasteful fixed-384 intermediate hop: the crop
+    now resizes directly to the target scale (per-scale ``OneOf``, mirroring the square path). ``min_max_height`` uses
+    ``[384, 600]`` to match the full SmallestMaxSize range — when the image short side is 400, albumentations clamps
+    the crop to the image height (a full-image crop), which is the original DETR recipe behaviour and preserves
+    zoom-out diversity across the SmallestMaxSize range.
+    """
+
+    @pytest.mark.parametrize(
+        "scales",
+        [
+            pytest.param([640], id="nonsquare-single"),
+            pytest.param([480, 640], id="nonsquare-multi"),
+        ],
+    )
+    def test_option_b_crop_step_uses_random_sized_crop(self, scales):
+        """Non-square option_b crop must use RandomSizedCrop, never fixed RandomCrop (issue #1018)."""
+        result = _build_train_resize_config(scales, square=False)
+        option_b = result[0]["OneOf"]["transforms"][1]
+        crop_step = option_b["Sequential"]["transforms"][1]
+        crop_variants = crop_step["OneOf"]["transforms"]
+        assert crop_variants and all(
+            "RandomSizedCrop" in entry and "RandomCrop" not in entry for entry in crop_variants
+        )
+
+    @pytest.mark.parametrize(
+        "scales",
+        [
+            pytest.param([640], id="nonsquare-single"),
+            pytest.param([480, 640], id="nonsquare-multi"),
+        ],
+    )
+    def test_option_b_crop_uses_full_scale_jitter_range(self, scales):
+        """RandomSizedCrop min_max_height matches SmallestMaxSize range [384, 600] for full zoom-out diversity."""
+        result = _build_train_resize_config(scales, square=False)
+        option_b = result[0]["OneOf"]["transforms"][1]
+        crop_variants = option_b["Sequential"]["transforms"][1]["OneOf"]["transforms"]
+        assert all(entry["RandomSizedCrop"]["min_max_height"] == [384, 600] for entry in crop_variants)
+
+    @pytest.mark.parametrize(
+        "scales,square",
+        [
+            pytest.param([640], True, id="square-single"),
+            pytest.param([480, 640], True, id="square-multi"),
+        ],
+    )
+    def test_square_option_b_unchanged(self, scales, square):
+        """Square path must still use RandomSizedCrop parameterized by scale."""
+        result = _build_train_resize_config(scales, square=square)
+        option_b = result[0]["OneOf"]["transforms"][1]
+        inner_transforms = option_b["Sequential"]["transforms"][1]["OneOf"]["transforms"]
+        for entry in inner_transforms:
+            assert "RandomSizedCrop" in entry
+            assert entry["RandomSizedCrop"]["min_max_height"] == [384, 600]

tests/datasets/test_coco_rle.py

@@ -0,0 +1,280 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for native RLE annotation support in the COCO dataset pipeline.
+
+Verifies that :func:`convert_coco_poly_to_mask` and :class:`ConvertCoco` correctly handle compressed RLE, uncompressed
+RLE, and polygon segmentation formats — including mixed annotations within the same image.
+"""
+
+import numpy as np
+import pycocotools.mask as mask_util
+import pytest
+import torch
+from PIL import Image
+
+from rfdetr.datasets.coco import ConvertCoco, _is_rle, convert_coco_poly_to_mask
+
+# Shared test dimensions
+_H, _W = 100, 100
+_IMAGE = Image.new("RGB", (_W, _H))
+
+
+def _make_reference_mask() -> np.ndarray:
+    """Create a deterministic 100x100 binary mask with a rectangular region."""
+    mask = np.zeros((_H, _W), dtype=np.uint8)
+    mask[20:50, 30:70] = 1
+    return mask
+
+
+def _encode_compressed_rle(mask: np.ndarray) -> dict:
+    """Encode a binary mask to compressed RLE with string counts (COCO JSON format)."""
+    rle = mask_util.encode(np.asfortranarray(mask))
+    # COCO JSON stores counts as a UTF-8 string, not bytes
+    rle["counts"] = rle["counts"].decode("utf-8") if isinstance(rle["counts"], bytes) else rle["counts"]
+    rle["size"] = list(rle["size"])
+    return rle
+
+
+def _encode_uncompressed_rle(mask: np.ndarray) -> dict:
+    """Encode a binary mask to uncompressed RLE with integer counts."""
+    flat = mask.flatten(order="F")
+    counts = []
+    current_val = 0
+    run_length = 0
+    for pixel in flat:
+        if pixel == current_val:
+            run_length += 1
+        else:
+            counts.append(run_length)
+            current_val = pixel
+            run_length = 1
+    counts.append(run_length)
+    return {"counts": counts, "size": [_H, _W]}
+
+
+def _make_polygon(mask: np.ndarray) -> list:
+    """Create a polygon annotation from a rectangular mask region."""
+    # Simple rectangle polygon matching the mask region [20:50, 30:70]
+    return [[30, 20, 70, 20, 70, 50, 30, 50]]
+
+
+class TestIsRle:
+    """Tests for the ``_is_rle`` helper."""
+
+    def test_compressed_rle_detected(self) -> None:
+        assert _is_rle({"counts": "abc", "size": [100, 100]}) is True
+
+    def test_uncompressed_rle_detected(self) -> None:
+        assert _is_rle({"counts": [0, 5, 10], "size": [100, 100]}) is True
+
+    def test_bytes_counts_detected(self) -> None:
+        assert _is_rle({"counts": b"abc", "size": [100, 100]}) is True
+
+    def test_polygon_not_detected(self) -> None:
+        assert _is_rle([[30, 20, 70, 20, 70, 50, 30, 50]]) is False
+
+    def test_empty_list_not_detected(self) -> None:
+        assert _is_rle([]) is False
+
+    def test_none_not_detected(self) -> None:
+        assert _is_rle(None) is False
+
+
+class TestConvertCocoPolyToMaskRle:
+    """Tests for RLE support in ``convert_coco_poly_to_mask``."""
+
+    def test_compressed_rle_decodes_correctly(self) -> None:
+        """Compressed RLE (string counts) should decode to the expected mask."""
+        ref_mask = _make_reference_mask()
+        rle = _encode_compressed_rle(ref_mask)
+
+        result = convert_coco_poly_to_mask([rle], _H, _W)
+
+        assert result.shape == (1, _H, _W)
+        assert result.dtype == torch.uint8
+        assert torch.equal(result[0], torch.as_tensor(ref_mask, dtype=torch.uint8))
+
+    def test_uncompressed_rle_decodes_correctly(self) -> None:
+        """Uncompressed RLE (int-list counts) should decode to the expected mask."""
+        ref_mask = _make_reference_mask()
+        uncompressed = _encode_uncompressed_rle(ref_mask)
+
+        result = convert_coco_poly_to_mask([uncompressed], _H, _W)
+
+        assert result.shape == (1, _H, _W)
+        assert result.dtype == torch.uint8
+        assert torch.equal(result[0], torch.as_tensor(ref_mask, dtype=torch.uint8))
+
+    def test_polygon_still_works(self) -> None:
+        """Polygon annotations should continue to work as before."""
+        polygon = _make_polygon(_make_reference_mask())
+
+        result = convert_coco_poly_to_mask([polygon], _H, _W)
+
+        assert result.shape == (1, _H, _W)
+        assert result.dtype == torch.uint8
+        # The polygon covers the same rectangular region
+        assert result[0, 30, 50] == 1  # inside the region
+        assert result[0, 0, 0] == 0  # outside
+
+    def test_compressed_rle_matches_polygon(self) -> None:
+        """Compressed RLE and polygon for the same region should produce identical masks."""
+        polygon = _make_polygon(_make_reference_mask())
+        poly_masks = convert_coco_poly_to_mask([polygon], _H, _W)
+
+        # Encode the polygon result as RLE, then decode via our path
+        ref_np = poly_masks[0].numpy()
+        rle = _encode_compressed_rle(ref_np)
+        rle_masks = convert_coco_poly_to_mask([rle], _H, _W)
+
+        assert torch.equal(poly_masks, rle_masks)
+
+    def test_mixed_polygon_and_rle(self) -> None:
+        """An image can have both polygon and RLE annotations across instances."""
+        ref_mask = _make_reference_mask()
+        polygon = _make_polygon(ref_mask)
+        rle = _encode_compressed_rle(ref_mask)
+
+        result = convert_coco_poly_to_mask([polygon, rle], _H, _W)
+
+        assert result.shape == (2, _H, _W)
+        # Both should produce the same mask
+        assert torch.equal(result[0], result[1])
+
+    def test_empty_segmentation_unchanged(self) -> None:
+        """Empty segmentation should produce a zero mask."""
+        result = convert_coco_poly_to_mask([[]], _H, _W)
+        assert result.shape == (1, _H, _W)
+        assert result.sum() == 0
+
+    def test_none_segmentation_unchanged(self) -> None:
+        """None segmentation should produce a zero mask."""
+        result = convert_coco_poly_to_mask([None], _H, _W)
+        assert result.shape == (1, _H, _W)
+        assert result.sum() == 0
+
+    def test_empty_list_returns_zero_tensor(self) -> None:
+        """No segmentations at all should return (0, H, W) tensor."""
+        result = convert_coco_poly_to_mask([], _H, _W)
+        assert result.shape == (0, _H, _W)
+
+    def test_rle_size_mismatch_behavior(self) -> None:
+        """Compressed RLE with mismatched embedded size should raise a decode error."""
+        ref_mask = _make_reference_mask()
+        rle = _encode_compressed_rle(ref_mask)
+        rle["size"] = [50, 50]
+
+        # Observed behavior: pycocotools rejects mismatched RLE metadata during decode.
+        with pytest.raises(ValueError, match="Invalid RLE mask representation"):
+            convert_coco_poly_to_mask([rle], _H, _W)
+
+    def test_compressed_rle_bytes_counts_decode(self) -> None:
+        """Compressed RLE with bytes counts should decode correctly."""
+        ref_mask = _make_reference_mask()
+        rle = mask_util.encode(np.asfortranarray(ref_mask))
+        rle["counts"] = rle["counts"].encode("utf-8") if isinstance(rle["counts"], str) else rle["counts"]
+        rle["size"] = list(rle["size"])
+
+        result = convert_coco_poly_to_mask([rle], _H, _W)
+
+        assert result.shape == (1, _H, _W)
+        assert result[0, 30, 50] == 1
+        assert result[0, 0, 0] == 0
+
+    def test_malformed_rle_counts_none_raises_value_error(self) -> None:
+        """Malformed RLE with counts=None should raise ValueError."""
+        with pytest.raises(ValueError, match="unsupported counts type"):
+            convert_coco_poly_to_mask([{"counts": None, "size": [_H, _W]}], _H, _W)
+
+
+class TestConvertCocoClassWithRle:
+    """Tests that ``ConvertCoco`` correctly passes RLE annotations through."""
+
+    def _make_annotation(self, segmentation: object, category_id: int = 0) -> dict:
+        return {
+            "bbox": [30, 20, 40, 30],
+            "category_id": category_id,
+            "area": 1200,
+            "iscrowd": 0,
+            "segmentation": segmentation,
+        }
+
+    def _make_target(self, annotations: list) -> dict:
+        return {"image_id": 1, "annotations": annotations}
+
+    def test_rle_masks_included_in_target(self) -> None:
+        """ConvertCoco with include_masks=True should handle RLE segmentations."""
+        ref_mask = _make_reference_mask()
+        rle = _encode_compressed_rle(ref_mask)
+        anno = self._make_annotation(rle)
+
+        converter = ConvertCoco(include_masks=True)
+        _, target = converter(_IMAGE, self._make_target([anno]))
+
+        assert "masks" in target
+        assert target["masks"].shape == (1, _H, _W)
+        assert target["masks"].dtype == torch.bool
+        assert target["masks"][0].any()
+
+    def test_polygon_masks_still_work(self) -> None:
+        """ConvertCoco should still handle polygon segmentations."""
+        polygon = _make_polygon(_make_reference_mask())
+        anno = self._make_annotation(polygon)
+
+        converter = ConvertCoco(include_masks=True)
+        _, target = converter(_IMAGE, self._make_target([anno]))
+
+        assert "masks" in target
+        assert target["masks"].shape == (1, _H, _W)
+        assert target["masks"].dtype == torch.bool
+
+    def test_mixed_rle_and_polygon_in_same_image(self) -> None:
+        """An image with both polygon and RLE annotations across instances."""
+        ref_mask = _make_reference_mask()
+        rle_anno = self._make_annotation(_encode_compressed_rle(ref_mask), category_id=0)
+        poly_anno = self._make_annotation(_make_polygon(ref_mask), category_id=1)
+
+        converter = ConvertCoco(include_masks=True)
+        _, target = converter(_IMAGE, self._make_target([rle_anno, poly_anno]))
+
+        assert target["masks"].shape == (2, _H, _W)
+        assert target["labels"].tolist() == [0, 1]
+
+    def test_no_masks_without_flag(self) -> None:
+        """RLE annotations should not produce masks when include_masks=False."""
+        rle = _encode_compressed_rle(_make_reference_mask())
+        anno = self._make_annotation(rle)
+
+        converter = ConvertCoco(include_masks=False)
+        _, target = converter(_IMAGE, self._make_target([anno]))
+
+        assert "masks" not in target
+
+
+class TestMalformedRle:
+    """Documents _is_rle behaviour for structurally malformed inputs.
+
+    Before this PR a bare ``except:`` in the polygon path silently swallowed any pycocotools error.  These tests confirm
+    that ``_is_rle`` is a *structural* check only (it does not validate values inside the dict) and that dicts missing
+    required keys are correctly classified as non-RLE so they are routed through the polygon path — where pycocotools
+    will either handle them or raise a descriptive error rather than silently falling back.
+    """
+
+    def test_missing_size_key_is_not_rle(self) -> None:
+        """Dict with 'counts' but no 'size' is not treated as RLE."""
+        assert _is_rle({"counts": [1, 2, 3]}) is False
+
+    def test_missing_counts_key_is_not_rle(self) -> None:
+        """Dict with 'size' but no 'counts' is not treated as RLE."""
+        assert _is_rle({"size": [100, 100]}) is False
+
+    def test_counts_none_is_classified_as_rle(self) -> None:
+        """_is_rle is a structural check: presence of both keys suffices regardless of value types."""
+        assert _is_rle({"counts": None, "size": [_H, _W]}) is True
+
+    def test_size_mismatch_is_still_classified_as_rle(self) -> None:
+        """Dicts with both keys are RLE even when the embedded size mismatches the image dimensions."""
+        assert _is_rle({"counts": [1, 2], "size": [50, 50]}) is True

tests/datasets/test_keypoint_schema.py

@@ -0,0 +1,435 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for private COCO keypoint schema inference helpers."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+
+from rfdetr.datasets._keypoint_schema import (
+    CocoKeypointSchema,
+    YoloKeypointSchema,
+    _infer_keypoint_flip_pairs_from_names,
+    _merge_category_keypoint_flip_pairs,
+    infer_coco_keypoint_schema,
+    infer_yolo_keypoint_schema,
+)
+
+
+def _write_coco_annotations(
+    path: Path,
+    *,
+    categories: list[dict],
+    annotations: list[dict] | None = None,
+) -> None:
+    """Write a minimal COCO annotation file.
+
+    Args:
+        path: Destination JSON path.
+        categories: COCO category objects.
+        annotations: Optional COCO annotation objects.
+
+    Returns:
+        ``None``.
+
+    Raises:
+        OSError: If the file cannot be written.
+
+    Example:
+        >>> import tempfile
+        >>> output = Path(tempfile.mkdtemp()) / "annotations.json"
+        >>> _write_coco_annotations(output, categories=[{"id": 0, "name": "person"}])
+        >>> output.exists()
+        True
+    """
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(
+        json.dumps({"images": [], "annotations": annotations or [], "categories": categories}),
+        encoding="utf-8",
+    )
+
+
+def test_infer_coco_keypoint_schema_uses_declared_category_keypoints(tmp_path: Path) -> None:
+    """Declared category keypoints should produce a category-aligned keypoint schema."""
+    annotation_path = tmp_path / "annotations.json"
+    _write_coco_annotations(
+        annotation_path,
+        categories=[{"id": 0, "name": "person", "keypoints": ["nose", "left_eye"], "skeleton": []}],
+    )
+
+    schema = infer_coco_keypoint_schema(annotation_path)
+
+    assert schema == CocoKeypointSchema(
+        class_names=["person"],
+        num_keypoints_per_class=[2],
+        keypoint_oks_sigmas=[0.1, 0.1],
+    )
+
+
+def test_infer_coco_keypoint_schema_infers_left_right_flip_pairs(tmp_path: Path) -> None:
+    """COCO category keypoint names should infer horizontal flip pairs when unambiguous."""
+    annotation_path = tmp_path / "annotations.json"
+    _write_coco_annotations(
+        annotation_path,
+        categories=[
+            {
+                "id": 0,
+                "name": "person",
+                "keypoints": ["nose", "left_eye", "right_eye", "left_wrist", "right_wrist"],
+                "skeleton": [],
+            }
+        ],
+    )
+
+    schema = infer_coco_keypoint_schema(annotation_path)
+
+    assert schema.keypoint_flip_pairs == [1, 2, 3, 4]
+
+
+def test_infer_coco_keypoint_schema_does_not_invent_missing_mirror_pairs(tmp_path: Path) -> None:
+    """A left/right token without its counterpart should keep the keypoint slots unswapped."""
+    annotation_path = tmp_path / "annotations.json"
+    _write_coco_annotations(
+        annotation_path,
+        categories=[
+            {
+                "id": 0,
+                "name": "person",
+                "keypoints": ["nose", "left_eye", "left_wrist"],
+                "skeleton": [],
+            }
+        ],
+    )
+
+    schema = infer_coco_keypoint_schema(annotation_path)
+
+    assert schema.num_keypoints_per_class == [3]
+    assert schema.keypoint_flip_pairs == []
+
+
+def test_infer_coco_keypoint_schema_drops_pairs_when_keypoint_categories_disagree(tmp_path: Path) -> None:
+    """A global flip-pair list is unsafe when keypoint classes use different slot layouts."""
+    annotation_path = tmp_path / "annotations.json"
+    _write_coco_annotations(
+        annotation_path,
+        categories=[
+            {
+                "id": 0,
+                "name": "standing_person",
+                "keypoints": ["left_eye", "right_eye", "nose"],
+                "skeleton": [],
+            },
+            {
+                "id": 1,
+                "name": "seated_person",
+                "keypoints": ["nose", "left_eye", "right_eye"],
+                "skeleton": [],
+            },
+        ],
+    )
+
+    schema = infer_coco_keypoint_schema(annotation_path)
+
+    assert schema.num_keypoints_per_class == [3, 3]
+    assert schema.keypoint_flip_pairs == []
+
+
+def test_infer_coco_keypoint_schema_uses_annotation_keypoints_when_category_metadata_is_missing(
+    tmp_path: Path,
+) -> None:
+    """Annotation keypoint arrays should define the count when category metadata is absent."""
+    annotation_path = tmp_path / "annotations.json"
+    _write_coco_annotations(
+        annotation_path,
+        categories=[{"id": 7, "name": "pose"}],
+        annotations=[{"id": 1, "image_id": 1, "category_id": 7, "keypoints": [1, 2, 2, 3, 4, 1]}],
+    )
+
+    schema = infer_coco_keypoint_schema(annotation_path, keypoint_oks_sigma=0.1)
+
+    assert schema.class_names == ["pose"]
+    assert schema.num_keypoints_per_class == [2]
+    assert schema.keypoint_oks_sigmas == [0.1, 0.1]
+
+
+def test_infer_coco_keypoint_schema_places_detection_only_categories_in_free_slots(tmp_path: Path) -> None:
+    """Detection-only categories should stay category-aligned with zero keypoint counts."""
+    annotation_path = tmp_path / "annotations.json"
+    _write_coco_annotations(
+        annotation_path,
+        categories=[
+            {"id": 0, "name": "person", "keypoints": ["nose", "left_eye"]},
+            {"id": 1, "name": "helmet"},
+            {"id": 2, "name": "vest"},
+        ],
+    )
+
+    schema = infer_coco_keypoint_schema(annotation_path)
+
+    assert schema.class_names == ["person", "helmet", "vest"]
+    assert schema.num_keypoints_per_class == [2, 0, 0]
+    assert schema.keypoint_oks_sigmas == [0.1, 0.1]
+
+
+def test_infer_coco_keypoint_schema_supports_multiple_keypoint_categories_with_same_count(tmp_path: Path) -> None:
+    """Multiple keypoint classes with the same keypoint count should stay category-aligned."""
+    annotation_path = tmp_path / "annotations.json"
+    _write_coco_annotations(
+        annotation_path,
+        categories=[
+            {"id": 3, "name": "adult", "keypoints": ["head", "foot"]},
+            {"id": 9, "name": "child", "keypoints": ["head", "foot"]},
+        ],
+    )
+
+    schema = infer_coco_keypoint_schema(annotation_path)
+
+    assert schema.class_names == ["adult", "child"]
+    assert schema.num_keypoints_per_class == [2, 2]
+    assert schema.keypoint_oks_sigmas == [0.1, 0.1]
+
+
+def test_infer_coco_keypoint_schema_rejects_missing_keypoints(tmp_path: Path) -> None:
+    """Detection-only COCO files should fail fast instead of silently training without keypoints."""
+    annotation_path = tmp_path / "annotations.json"
+    _write_coco_annotations(annotation_path, categories=[{"id": 0, "name": "person"}])
+
+    with pytest.raises(ValueError, match="has no keypoint metadata"):
+        infer_coco_keypoint_schema(annotation_path)
+
+
+def test_infer_coco_keypoint_schema_supports_mixed_keypoint_counts(tmp_path: Path) -> None:
+    """Different keypoint counts are represented per class and padded later by the dataset."""
+    annotation_path = tmp_path / "annotations.json"
+    _write_coco_annotations(
+        annotation_path,
+        categories=[
+            {"id": 0, "name": "person", "keypoints": ["nose"]},
+            {"id": 1, "name": "animal", "keypoints": ["head", "tail"]},
+        ],
+    )
+
+    schema = infer_coco_keypoint_schema(annotation_path)
+
+    assert schema.class_names == ["person", "animal"]
+    assert schema.num_keypoints_per_class == [1, 2]
+    assert schema.keypoint_oks_sigmas == [0.1, 0.1]
+
+
+def test_infer_coco_keypoint_schema_rejects_malformed_annotation_keypoint_length(tmp_path: Path) -> None:
+    """COCO keypoint arrays must be flattened ``x, y, visibility`` triples."""
+    annotation_path = tmp_path / "annotations.json"
+    _write_coco_annotations(
+        annotation_path,
+        categories=[{"id": 0, "name": "person"}],
+        annotations=[{"id": 1, "image_id": 1, "category_id": 0, "keypoints": [1, 2]}],
+    )
+
+    with pytest.raises(ValueError, match="length divisible by 3"):
+        infer_coco_keypoint_schema(annotation_path)
+
+
+def test_infer_coco_keypoint_schema_raises_file_not_found(tmp_path: Path) -> None:
+    """Non-existent annotation file should raise FileNotFoundError before any parsing."""
+    missing = tmp_path / "does_not_exist.json"
+
+    with pytest.raises(FileNotFoundError):
+        infer_coco_keypoint_schema(missing)
+
+
+def test_infer_coco_keypoint_schema_raises_key_error_for_missing_categories_key(tmp_path: Path) -> None:
+    """JSON file missing the 'categories' key should raise KeyError."""
+    annotation_path = tmp_path / "no_categories.json"
+    annotation_path.write_text('{"images": [], "annotations": []}', encoding="utf-8")
+
+    with pytest.raises(KeyError):
+        infer_coco_keypoint_schema(annotation_path)
+
+
+def test_infer_coco_keypoint_schema_raises_value_error_for_list_root(tmp_path: Path) -> None:
+    """JSON file whose root is a list (not an object) should raise ValueError."""
+    annotation_path = tmp_path / "list_root.json"
+    annotation_path.write_text("[]", encoding="utf-8")
+
+    with pytest.raises(ValueError):
+        infer_coco_keypoint_schema(annotation_path)
+
+
+@pytest.mark.parametrize(
+    "counts,expected",
+    [
+        pytest.param([0, 17, 25], [17, 25], id="leading-zero-filtered"),
+        pytest.param([0, 0], [], id="all-zero-returns-empty"),
+        pytest.param([5, 17], [5, 17], id="all-nonzero-returned-unchanged"),
+        pytest.param([], [], id="empty-input-returns-empty"),
+    ],
+)
+def test_active_keypoint_counts_filters_zeros(counts: list[int], expected: list[int]) -> None:
+    """active_keypoint_counts should return only positive counts in schema order."""
+    from rfdetr.datasets._keypoint_schema import active_keypoint_counts
+
+    result = active_keypoint_counts(counts)
+
+    assert result == expected, f"active_keypoint_counts({counts!r}) = {result!r}, expected {expected!r}"
+
+
+def test_infer_yolo_keypoint_schema_reads_pose_yaml_metadata(tmp_path: Path) -> None:
+    """YOLO pose YAML should define class names, keypoint count, names, and flip pairs."""
+    data_file = tmp_path / "data.yaml"
+    data_file.write_text(
+        "names:\n  0: person\nkpt_shape: [2, 3]\nflip_idx: [0, 1]\nkpt_names:\n  0:\n    - left_eye\n    - right_eye\n",
+        encoding="utf-8",
+    )
+
+    schema = infer_yolo_keypoint_schema(data_file)
+
+    assert schema == YoloKeypointSchema(
+        class_names=["person"],
+        num_keypoints_per_class=[2],
+        keypoint_oks_sigmas=[0.1, 0.1],
+        keypoint_names=["left_eye", "right_eye"],
+        flip_idx=[0, 1],
+        keypoint_dim=3,
+    )
+
+
+def test_infer_yolo_keypoint_schema_rejects_detection_yaml(tmp_path: Path) -> None:
+    """Detection-only YOLO YAML should fail fast in keypoint schema inference."""
+    data_file = tmp_path / "data.yaml"
+    data_file.write_text("names:\n  - person\n", encoding="utf-8")
+
+    with pytest.raises(ValueError, match="kpt_shape"):
+        infer_yolo_keypoint_schema(data_file)
+
+
+@pytest.mark.parametrize("kpt_shape", ["[17, 1]", "[0, 3]", "[17]", "[17, 4]"])
+def test_infer_yolo_keypoint_schema_rejects_invalid_kpt_shape(tmp_path: Path, kpt_shape: str) -> None:
+    """YOLO pose kpt_shape must be [positive_count, 2_or_3]."""
+    data_file = tmp_path / "data.yaml"
+    data_file.write_text(f"names:\n  - person\nkpt_shape: {kpt_shape}\n", encoding="utf-8")
+
+    with pytest.raises(ValueError, match="kpt_shape"):
+        infer_yolo_keypoint_schema(data_file)
+
+
+@pytest.mark.parametrize(
+    "flip_idx_text, expected_match",
+    [
+        pytest.param("[0, 5]", "permutation", id="out_of_range"),
+        pytest.param("[0, 0]", "permutation", id="duplicate"),
+        pytest.param("[0]", "integer indexes", id="wrong_length"),
+    ],
+)
+def test_infer_yolo_keypoint_schema_rejects_invalid_flip_idx(
+    tmp_path: Path, flip_idx_text: str, expected_match: str
+) -> None:
+    """flip_idx must be a valid permutation of 0..N-1 matching kpt_shape count."""
+    data_file = tmp_path / "data.yaml"
+    data_file.write_text(
+        f"names:\n  0: person\nkpt_shape: [2, 3]\nflip_idx: {flip_idx_text}\n",
+        encoding="utf-8",
+    )
+    with pytest.raises(ValueError, match=expected_match):
+        infer_yolo_keypoint_schema(data_file)
+
+
+# ---------------------------------------------------------------------------
+# _infer_keypoint_flip_pairs_from_names edge cases
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "names,expected",
+    [
+        pytest.param([], [], id="empty-input"),
+        pytest.param(["left_eye"], [], id="single-directional-no-mirror"),
+        pytest.param(["Left-Eye", "left_eye"], [], id="duplicate-normalized-names"),
+        pytest.param(["left_right_wrist"], [], id="two-directional-tokens-ambiguous"),
+        pytest.param(["nose", "left_eye", "right_eye"], [1, 2], id="standard-coco-pair"),
+    ],
+)
+def test_infer_keypoint_flip_pairs_from_names_edge_cases(names: list[str], expected: list[int]) -> None:
+    """Edge-case inputs should return the expected flat pair list without raising."""
+    assert _infer_keypoint_flip_pairs_from_names(names) == expected
+
+
+# ---------------------------------------------------------------------------
+# _merge_category_keypoint_flip_pairs success path
+# ---------------------------------------------------------------------------
+
+
+def test_merge_category_keypoint_flip_pairs_returns_common_pairs_when_all_agree() -> None:
+    """All-agreeing categories should return the shared pair list."""
+    assert _merge_category_keypoint_flip_pairs([[1, 2], [1, 2], [1, 2]]) == [1, 2]
+
+
+def test_merge_category_keypoint_flip_pairs_single_category() -> None:
+    """Single-category input should return that category's pairs unchanged."""
+    assert _merge_category_keypoint_flip_pairs([[3, 5, 1, 2]]) == [3, 5, 1, 2]
+
+
+# ---------------------------------------------------------------------------
+# YoloKeypointSchema.keypoint_flip_pairs
+# ---------------------------------------------------------------------------
+
+
+def test_infer_yolo_keypoint_schema_populates_keypoint_flip_pairs(tmp_path: Path) -> None:
+    """YOLO flip_idx should produce an equivalent keypoint_flip_pairs list on the schema."""
+    data_file = tmp_path / "data.yaml"
+    data_file.write_text(
+        "names:\n  0: person\nkpt_shape: [3, 3]\nflip_idx: [0, 2, 1]\n",
+        encoding="utf-8",
+    )
+
+    schema = infer_yolo_keypoint_schema(data_file)
+
+    assert schema.flip_idx == [0, 2, 1]
+    assert schema.keypoint_flip_pairs == [1, 2]
+
+
+def test_infer_yolo_keypoint_schema_empty_flip_idx_gives_empty_pairs(tmp_path: Path) -> None:
+    """Missing flip_idx should result in an empty keypoint_flip_pairs list."""
+    data_file = tmp_path / "data.yaml"
+    data_file.write_text(
+        "names:\n  0: person\nkpt_shape: [2, 3]\n",
+        encoding="utf-8",
+    )
+
+    schema = infer_yolo_keypoint_schema(data_file)
+
+    assert schema.flip_idx == []
+    assert schema.keypoint_flip_pairs == []
+
+
+# ---------------------------------------------------------------------------
+# Public re-export from rfdetr.datasets
+# ---------------------------------------------------------------------------
+
+
+def test_infer_coco_keypoint_schema_importable_from_datasets_package(tmp_path: Path) -> None:
+    """infer_coco_keypoint_schema should be importable from the public rfdetr.datasets package."""
+    from rfdetr.datasets import infer_coco_keypoint_schema as public_fn
+
+    annotation_path = tmp_path / "annotations.json"
+    _write_coco_annotations(
+        annotation_path,
+        categories=[{"id": 0, "name": "person", "keypoints": ["nose", "left_eye", "right_eye"], "skeleton": []}],
+    )
+    schema = public_fn(annotation_path)
+    assert schema.keypoint_flip_pairs == [1, 2]
+
+
+def test_infer_yolo_keypoint_schema_importable_from_datasets_package(tmp_path: Path) -> None:
+    """infer_yolo_keypoint_schema should be importable from the public rfdetr.datasets package."""
+    from rfdetr.datasets import infer_yolo_keypoint_schema as public_fn
+
+    data_file = tmp_path / "data.yaml"
+    data_file.write_text("names:\n  0: person\nkpt_shape: [1, 3]\n", encoding="utf-8")
+    schema = public_fn(data_file)
+    assert schema.num_keypoints_per_class == [1]

tests/datasets/test_kornia_transforms.py

@@ -0,0 +1,732 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for Kornia GPU augmentation pipeline builder and bbox utilities.
+
+All tests in this module are CPU-compatible — Kornia operates on CPU tensors identically to GPU tensors, so no
+``@pytest.mark.gpu`` is needed.
+"""
+
+import pytest
+import torch
+
+from rfdetr.datasets.aug_configs import (
+    AUG_AERIAL,
+    AUG_AGGRESSIVE,
+    AUG_CONSERVATIVE,
+    AUG_INDUSTRIAL,
+)
+
+# ---------------------------------------------------------------------------
+# TestBuildKorniaPipeline — validates the factory that translates aug_config
+# dicts into a Kornia AugmentationSequential pipeline.
+# ---------------------------------------------------------------------------
+
+
+class TestBuildKorniaPipeline:
+    """build_kornia_pipeline returns a valid pipeline for every preset and rejects unknown transform keys with a clear
+    error."""
+
+    @pytest.fixture(autouse=True)
+    def _require_kornia(self):
+        pytest.importorskip("kornia")
+
+    @pytest.mark.parametrize(
+        "config,config_name",
+        [
+            pytest.param(AUG_CONSERVATIVE, "AUG_CONSERVATIVE", id="conservative"),
+            pytest.param(AUG_AGGRESSIVE, "AUG_AGGRESSIVE", id="aggressive"),
+            pytest.param(AUG_AERIAL, "AUG_AERIAL", id="aerial"),
+            pytest.param(AUG_INDUSTRIAL, "AUG_INDUSTRIAL", id="industrial"),
+        ],
+    )
+    def test_each_preset_config(self, config, config_name):
+        """Each named preset builds a pipeline without errors."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        pipeline = build_kornia_pipeline(config, 560)
+        assert pipeline is not None, f"build_kornia_pipeline({config_name}, 560) must return a non-None pipeline"
+
+    def test_unknown_key_raises_value_error(self):
+        """An unrecognised transform key raises ValueError immediately."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        with pytest.raises(ValueError, match="FooBarTransform"):
+            build_kornia_pipeline({"FooBarTransform": {"p": 0.5}}, 560)
+
+    def test_empty_config_returns_pipeline(self):
+        """An empty config dict returns a valid (no-op) pipeline, not None."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        pipeline = build_kornia_pipeline({}, 560)
+        assert pipeline is not None, "Empty config must still return a pipeline object"
+
+    def test_known_plus_unknown_raises(self):
+        """Mixing a valid key with an unknown key still raises ValueError."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        mixed = {"HorizontalFlip": {"p": 0.5}, "BogusTransform": {"p": 0.3}}
+        with pytest.raises(ValueError, match="BogusTransform"):
+            build_kornia_pipeline(mixed, 560)
+
+    def test_hflip_disabled_for_keypoint_pipeline(self):
+        """Keypoint-mode Kornia augmentation drops hflip transforms with a warning."""
+        from unittest import mock
+
+        from rfdetr.datasets import kornia_transforms
+
+        config = {"HorizontalFlip": {"p": 0.5}, "VerticalFlip": {"p": 0.5}}
+        mock_warning = mock.patch.object(kornia_transforms.logger, "warning")
+
+        with mock_warning as warning:
+            pipeline = kornia_transforms.build_kornia_pipeline(config, 560, include_keypoints=True)
+
+        transform_names = [child.__class__.__name__ for child in pipeline.children()]
+        assert "RandomHorizontalFlip" not in transform_names
+        assert "RandomVerticalFlip" in transform_names
+        assert warning.called
+        assert "HorizontalFlip" in str(warning.call_args)
+
+
+# ---------------------------------------------------------------------------
+# TestCollateBoxes — validates packing of variable-length per-image boxes
+# into a zero-padded [B, N_max, 4] tensor with a boolean validity mask.
+# ---------------------------------------------------------------------------
+
+
+class TestCollateBoxes:
+    """collate_boxes packs variable-length boxes into [B, N_max, 4] with mask."""
+
+    @pytest.fixture(autouse=True)
+    def _require_kornia(self):
+        pytest.importorskip("kornia")
+
+    def _make_targets(self, box_counts):
+        """Build a list of target dicts with the given per-image box counts.
+
+        Each box is a valid xyxy rectangle within a 100x100 image.
+        """
+        targets = []
+        for n in box_counts:
+            boxes = (
+                torch.tensor([[10.0, 10.0, 50.0, 50.0]] * n, dtype=torch.float32)
+                if n > 0
+                else torch.zeros(0, 4, dtype=torch.float32)
+            )
+            targets.append({"boxes": boxes})
+        return targets
+
+    def test_normal_batch(self):
+        """Batch of 2 images: output shape is [2, N_max, 4] with valid mask [2, N_max]."""
+        from rfdetr.datasets.kornia_transforms import collate_boxes
+
+        targets = self._make_targets([2, 3])
+        boxes_padded, valid = collate_boxes(targets, torch.device("cpu"))
+
+        assert boxes_padded.shape == (2, 3, 4), f"Expected shape (2, 3, 4), got {boxes_padded.shape}"
+        assert valid.shape == (2, 3), f"Expected valid shape (2, 3), got {valid.shape}"
+        assert valid.dtype == torch.bool
+
+    def test_b_zero(self):
+        """Empty target list produces shape [0, 0, 4] and valid [0, 0]."""
+        from rfdetr.datasets.kornia_transforms import collate_boxes
+
+        boxes_padded, valid = collate_boxes([], torch.device("cpu"))
+
+        assert boxes_padded.shape == (0, 0, 4), f"Expected (0, 0, 4) for empty batch, got {boxes_padded.shape}"
+        assert valid.shape == (0, 0), f"Expected valid (0, 0) for empty batch, got {valid.shape}"
+
+    def test_n_zero_per_image(self):
+        """One image with 0 boxes: shape [1, 0, 4], valid all-False."""
+        from rfdetr.datasets.kornia_transforms import collate_boxes
+
+        targets = self._make_targets([0])
+        boxes_padded, valid = collate_boxes(targets, torch.device("cpu"))
+
+        assert boxes_padded.shape == (1, 0, 4), f"Expected (1, 0, 4), got {boxes_padded.shape}"
+        assert valid.shape == (1, 0), f"Expected (1, 0), got {valid.shape}"
+
+    def test_single_image(self):
+        """B=1 with 3 boxes: output shape is [1, 3, 4]."""
+        from rfdetr.datasets.kornia_transforms import collate_boxes
+
+        targets = self._make_targets([3])
+        boxes_padded, valid = collate_boxes(targets, torch.device("cpu"))
+
+        assert boxes_padded.shape == (1, 3, 4)
+        assert valid.shape == (1, 3)
+
+    def test_valid_mask_matches_box_count(self):
+        """The valid mask has True for real boxes and False for padding."""
+        from rfdetr.datasets.kornia_transforms import collate_boxes
+
+        targets = self._make_targets([1, 3])
+        _, valid = collate_boxes(targets, torch.device("cpu"))
+
+        # Image 0: 1 real box, 2 padding → [True, False, False]
+        assert valid[0].tolist() == [True, False, False], f"Image 0 valid mask wrong: {valid[0].tolist()}"
+        # Image 1: 3 real boxes, 0 padding → [True, True, True]
+        assert valid[1].tolist() == [True, True, True], f"Image 1 valid mask wrong: {valid[1].tolist()}"
+
+
+# ---------------------------------------------------------------------------
+# TestUnpackBoxes — validates the inverse: writing augmented boxes back into
+# per-image target dicts with clamping, zero-area removal, and label sync.
+# ---------------------------------------------------------------------------
+
+
+class TestUnpackBoxes:
+    """unpack_boxes writes augmented boxes back and removes zero-area entries."""
+
+    @pytest.fixture(autouse=True)
+    def _require_kornia(self):
+        pytest.importorskip("kornia")
+
+    def _make_inputs(
+        self,
+        boxes_aug,
+        valid_mask,
+        original_targets,
+        image_height=100,
+        image_width=100,
+    ):
+        """Return tensors suitable for unpack_boxes."""
+        boxes_tensor = torch.tensor(boxes_aug, dtype=torch.float32)
+        valid_tensor = torch.tensor(valid_mask, dtype=torch.bool)
+        return boxes_tensor, valid_tensor, original_targets, image_height, image_width
+
+    def test_all_boxes_removed_after_aug(self):
+        """When all augmented boxes are zero-area, output targets have empty boxes."""
+        from rfdetr.datasets.kornia_transforms import unpack_boxes
+
+        # B=1, N=2: both boxes are zero-area (x1==x2 or y1==y2)
+        boxes_aug = [[[10.0, 10.0, 10.0, 10.0], [20.0, 20.0, 20.0, 20.0]]]
+        valid = [[True, True]]
+        targets = [
+            {
+                "boxes": torch.tensor([[10.0, 10.0, 50.0, 50.0], [20.0, 20.0, 60.0, 60.0]]),
+                "labels": torch.tensor([1, 2]),
+                "area": torch.tensor([1600.0, 1600.0]),
+                "iscrowd": torch.tensor([0, 0]),
+            }
+        ]
+        boxes_t, valid_t, tgts, image_height, image_width = self._make_inputs(boxes_aug, valid, targets)
+        result = unpack_boxes(boxes_t, valid_t, tgts, image_height, image_width)
+
+        assert result[0]["boxes"].shape[0] == 0, (
+            f"Expected 0 boxes after zero-area removal, got {result[0]['boxes'].shape[0]}"
+        )
+        assert result[0]["labels"].shape[0] == 0
+
+    def test_partial_removal(self):
+        """Some boxes survive, some removed; labels/area/iscrowd synced."""
+        from rfdetr.datasets.kornia_transforms import unpack_boxes
+
+        # Box 0: valid, non-zero area; Box 1: zero-area
+        boxes_aug = [[[10.0, 10.0, 50.0, 50.0], [30.0, 30.0, 30.0, 30.0]]]
+        valid = [[True, True]]
+        targets = [
+            {
+                "boxes": torch.tensor([[10.0, 10.0, 50.0, 50.0], [30.0, 30.0, 70.0, 70.0]]),
+                "labels": torch.tensor([1, 2]),
+                "area": torch.tensor([1600.0, 1600.0]),
+                "iscrowd": torch.tensor([0, 1]),
+            }
+        ]
+        boxes_t, valid_t, tgts, image_height, image_width = self._make_inputs(boxes_aug, valid, targets)
+        result = unpack_boxes(boxes_t, valid_t, tgts, image_height, image_width)
+
+        assert result[0]["boxes"].shape[0] == 1, f"Expected 1 surviving box, got {result[0]['boxes'].shape[0]}"
+        assert result[0]["labels"].tolist() == [1]
+
+    def test_labels_area_iscrowd_sync(self):
+        """When boxes are removed, labels/area/iscrowd entries are also removed."""
+        from rfdetr.datasets.kornia_transforms import unpack_boxes
+
+        # Box 0: zero-area (removed), Box 1: valid
+        boxes_aug = [[[5.0, 5.0, 5.0, 5.0], [10.0, 10.0, 40.0, 40.0]]]
+        valid = [[True, True]]
+        targets = [
+            {
+                "boxes": torch.tensor([[5.0, 5.0, 30.0, 30.0], [10.0, 10.0, 40.0, 40.0]]),
+                "labels": torch.tensor([7, 9]),
+                "area": torch.tensor([625.0, 900.0]),
+                "iscrowd": torch.tensor([0, 1]),
+            }
+        ]
+        boxes_t, valid_t, tgts, image_height, image_width = self._make_inputs(boxes_aug, valid, targets)
+        result = unpack_boxes(boxes_t, valid_t, tgts, image_height, image_width)
+
+        assert result[0]["labels"].tolist() == [9], (
+            f"Expected label [9] after removal of box 0, got {result[0]['labels'].tolist()}"
+        )
+        assert result[0]["area"].shape[0] == 1
+        assert result[0]["iscrowd"].tolist() == [1]
+
+    def test_boxes_clamped_to_image_bounds(self):
+        """Boxes outside [0,W]x[0,H] are clamped to image bounds."""
+        from rfdetr.datasets.kornia_transforms import unpack_boxes
+
+        # Box extends beyond 100x100 image
+        boxes_aug = [[[-10.0, -5.0, 120.0, 110.0]]]
+        valid = [[True]]
+        targets = [
+            {
+                "boxes": torch.tensor([[0.0, 0.0, 90.0, 90.0]]),
+                "labels": torch.tensor([1]),
+                "area": torch.tensor([8100.0]),
+                "iscrowd": torch.tensor([0]),
+            }
+        ]
+        image_height, image_width = 100, 100
+        boxes_t, valid_t, tgts, image_height, image_width = self._make_inputs(
+            boxes_aug,
+            valid,
+            targets,
+            image_height,
+            image_width,
+        )
+        result = unpack_boxes(boxes_t, valid_t, tgts, image_height, image_width)
+
+        result_boxes = result[0]["boxes"]
+        assert result_boxes.shape[0] == 1, "Clamped box should survive (non-zero area)"
+        # Verify clamping: x1>=0, y1>=0, x2<=W, y2<=H
+        assert result_boxes[0, 0].item() >= 0.0, "x1 not clamped to >= 0"
+        assert result_boxes[0, 1].item() >= 0.0, "y1 not clamped to >= 0"
+        assert result_boxes[0, 2].item() <= image_width, f"x2 not clamped to <= {image_width}"
+        assert result_boxes[0, 3].item() <= image_height, f"y2 not clamped to <= {image_height}"
+
+
+# ---------------------------------------------------------------------------
+# TestRotateFactory — validates the Rotate parameter translation from
+# Albumentations-style limit (scalar or tuple) to Kornia RandomRotation.
+# ---------------------------------------------------------------------------
+
+
+class TestRotateFactory:
+    """Rotate factory translates limit (scalar or tuple) to K.RandomRotation(degrees=...)."""
+
+    @pytest.fixture(autouse=True)
+    def _require_kornia(self):
+        pytest.importorskip("kornia")
+
+    def test_limit_as_scalar(self):
+        """Rotate(limit=45) produces K.RandomRotation(degrees=(-45, 45))."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        # Build a pipeline with just Rotate(limit=45)
+        pipeline = build_kornia_pipeline({"Rotate": {"limit": 45, "p": 1.0}}, 560)
+        assert pipeline is not None
+
+        # Inspect the pipeline's children to find the RandomRotation and check degrees
+        import kornia.augmentation as kornia_augmentation
+
+        rotation_augs = [
+            child for child in pipeline.children() if isinstance(child, kornia_augmentation.RandomRotation)
+        ]
+        assert len(rotation_augs) == 1, f"Expected exactly 1 RandomRotation, found {len(rotation_augs)}"
+        degrees = rotation_augs[0].flags["degrees"]
+        # degrees should be a tensor representing (-45, 45)
+        assert float(degrees[0]) == pytest.approx(-45.0, abs=0.1)
+        assert float(degrees[1]) == pytest.approx(45.0, abs=0.1)
+
+    def test_limit_as_tuple(self):
+        """Rotate(limit=(90, 90)) produces K.RandomRotation(degrees=(90, 90))."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        pipeline = build_kornia_pipeline({"Rotate": {"limit": (90, 90), "p": 1.0}}, 560)
+        assert pipeline is not None
+
+        import kornia.augmentation as kornia_augmentation
+
+        rotation_augs = [
+            child for child in pipeline.children() if isinstance(child, kornia_augmentation.RandomRotation)
+        ]
+        assert len(rotation_augs) == 1
+        degrees = rotation_augs[0].flags["degrees"]
+        assert float(degrees[0]) == pytest.approx(90.0, abs=0.1)
+        assert float(degrees[1]) == pytest.approx(90.0, abs=0.1)
+
+    def test_flags_include_degrees(self):
+        """Rotate factory keeps a legacy degrees entry in Kornia flags for compatibility."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        pipeline = build_kornia_pipeline({"Rotate": {"limit": 30, "p": 1.0}}, 560)
+        assert pipeline is not None
+
+        import kornia.augmentation as kornia_augmentation
+
+        rotation_augs = [
+            child for child in pipeline.children() if isinstance(child, kornia_augmentation.RandomRotation)
+        ]
+        assert len(rotation_augs) == 1
+        assert "degrees" in rotation_augs[0].flags
+        assert rotation_augs[0].flags["degrees"] == (-30, 30)
+
+
+# ---------------------------------------------------------------------------
+# TestGpuPostprocessFlag — validates that make_coco_transforms respects the
+# gpu_postprocess flag to omit augmentation and normalization from CPU path.
+# ---------------------------------------------------------------------------
+
+
+class TestGpuPostprocessFlag:
+    """gpu_postprocess flag controls whether aug + normalize appear in CPU pipeline."""
+
+    def test_gpu_postprocess_true_omits_aug_and_normalize_from_train(self):
+        """gpu_postprocess=True: train pipeline has no Normalize; fewer AlbumentationsWrappers (no aug_wrappers)."""
+        from rfdetr.datasets.coco import make_coco_transforms
+        from rfdetr.datasets.transforms import AlbumentationsWrapper, Normalize
+
+        pipeline_gpu = make_coco_transforms("train", 560, gpu_postprocess=True)
+        pipeline_cpu = make_coco_transforms("train", 560, gpu_postprocess=False)
+
+        steps_gpu = pipeline_gpu.transforms
+        steps_cpu = pipeline_cpu.transforms
+
+        normalize_gpu = [s for s in steps_gpu if isinstance(s, Normalize)]
+        assert len(normalize_gpu) == 0, "gpu_postprocess=True must omit Normalize from train pipeline"
+
+        # Resize wrappers (AlbumentationsWrapper) remain; aug wrappers are removed.
+        # Default AUG_CONFIG adds 1 aug wrapper, so gpu version must have fewer wrappers.
+        n_alb_gpu = sum(isinstance(s, AlbumentationsWrapper) for s in steps_gpu)
+        n_alb_cpu = sum(isinstance(s, AlbumentationsWrapper) for s in steps_cpu)
+        assert n_alb_gpu < n_alb_cpu, "gpu_postprocess=True must remove aug AlbumentationsWrappers from train pipeline"
+
+    def test_gpu_postprocess_false_includes_aug_and_normalize_from_train(self):
+        """gpu_postprocess=False (default): train pipeline includes Normalize."""
+        from rfdetr.datasets.coco import make_coco_transforms
+        from rfdetr.datasets.transforms import Normalize
+
+        pipeline = make_coco_transforms("train", 560, gpu_postprocess=False)
+        steps = pipeline.transforms
+
+        normalize_steps = [s for s in steps if isinstance(s, Normalize)]
+        assert len(normalize_steps) > 0, "gpu_postprocess=False must include Normalize in train pipeline"
+
+    def test_val_path_unaffected_by_gpu_postprocess(self):
+        """Val pipeline is unchanged regardless of gpu_postprocess value."""
+        from rfdetr.datasets.coco import make_coco_transforms
+        from rfdetr.datasets.transforms import Normalize
+
+        pipeline_default = make_coco_transforms("val", 560, gpu_postprocess=False)
+        pipeline_gpu = make_coco_transforms("val", 560, gpu_postprocess=True)
+
+        # Both should have Normalize (val is never stripped)
+        norm_default = [s for s in pipeline_default.transforms if isinstance(s, Normalize)]
+        norm_gpu = [s for s in pipeline_gpu.transforms if isinstance(s, Normalize)]
+
+        assert len(norm_default) > 0, "Val pipeline (default) must include Normalize"
+        assert len(norm_gpu) > 0, "Val pipeline (gpu_postprocess=True) must include Normalize"
+
+        # Same number of pipeline steps
+        assert len(pipeline_default.transforms) == len(pipeline_gpu.transforms), (
+            "Val pipeline step count must be identical regardless of gpu_postprocess"
+        )
+
+
+# ---------------------------------------------------------------------------
+# TestGaussianBlurMinKernel — validates that blur_limit < 3 is clamped so
+# Kornia never receives an invalid kernel_size < 3.
+# ---------------------------------------------------------------------------
+
+
+class TestGaussianBlurMinKernel:
+    """_make_gaussian_blur enforces kernel_size >= 3 regardless of blur_limit."""
+
+    @pytest.fixture(autouse=True)
+    def _require_kornia(self):
+        pytest.importorskip("kornia")
+
+    @pytest.mark.parametrize(
+        "blur_limit",
+        [pytest.param(1, id="blur_limit_1"), pytest.param(2, id="blur_limit_2")],
+    )
+    def test_small_blur_limit_produces_valid_kernel(self, blur_limit):
+        """blur_limit below 3 must be clamped so the resulting kernel_size >= 3."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        # Should not raise; previously blur_limit=1 produced kernel_size=(3,1)
+        pipeline = build_kornia_pipeline({"GaussianBlur": {"blur_limit": blur_limit, "p": 1.0}}, 560)
+        assert pipeline is not None
+
+        import kornia.augmentation as kornia_augmentation
+
+        blur_augs = [c for c in pipeline.children() if isinstance(c, kornia_augmentation.RandomGaussianBlur)]
+        assert len(blur_augs) == 1
+        ks = blur_augs[0].flags["kernel_size"]
+        assert int(ks[0]) >= 3, f"kernel_size[0]={int(ks[0])} must be >= 3"
+        assert int(ks[1]) >= 3, f"kernel_size[1]={int(ks[1])} must be >= 3"
+
+    def test_blur_limit_3_unchanged(self):
+        """blur_limit=3 (default) passes through without modification."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        pipeline = build_kornia_pipeline({"GaussianBlur": {"blur_limit": 3, "p": 1.0}}, 560)
+        import kornia.augmentation as kornia_augmentation
+
+        blur_augs = [c for c in pipeline.children() if isinstance(c, kornia_augmentation.RandomGaussianBlur)]
+        ks = blur_augs[0].flags["kernel_size"]
+        assert int(ks[0]) == 3
+        assert int(ks[1]) == 3
+
+
+# ---------------------------------------------------------------------------
+# TestKorniaPipelineForwardPass — validates that a built pipeline produces
+# output of the correct shape and dtype on CPU tensors.
+# ---------------------------------------------------------------------------
+
+
+class TestKorniaPipelineForwardPass:
+    """build_kornia_pipeline output passes through without shape/dtype errors."""
+
+    @pytest.fixture(autouse=True)
+    def _require_kornia(self):
+        pytest.importorskip("kornia")
+
+    def test_forward_pass_shape_and_dtype(self):
+        """Pipeline output images have same shape as input; boxes shape is [B, N, 4]."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        pipeline = build_kornia_pipeline({"HorizontalFlip": {"p": 1.0}}, resolution=64)
+
+        batch_size, channels, image_height, image_width = 2, 3, 64, 64
+        img = torch.rand(batch_size, channels, image_height, image_width)
+        boxes = torch.tensor([[[0.0, 0.0, 32.0, 32.0]], [[10.0, 10.0, 50.0, 50.0]]], dtype=torch.float32)
+
+        img_out, boxes_out = pipeline(img, boxes)
+
+        assert img_out.shape == (batch_size, channels, image_height, image_width), (
+            f"Image shape changed: {img_out.shape}"
+        )
+        assert img_out.dtype == torch.float32
+        assert boxes_out.shape == (batch_size, 1, 4), f"Boxes shape wrong: {boxes_out.shape}"
+
+    def test_forward_pass_empty_boxes(self):
+        """Pipeline handles a batch where N_max=0 (no boxes) without error."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        pipeline = build_kornia_pipeline({"HorizontalFlip": {"p": 1.0}}, resolution=32)
+
+        batch_size, channels, image_height, image_width = 2, 3, 32, 32
+        img = torch.rand(batch_size, channels, image_height, image_width)
+        # [B, 0, 4] — no boxes
+        boxes = torch.zeros(batch_size, 0, 4, dtype=torch.float32)
+
+        img_out, boxes_out = pipeline(img, boxes)
+
+        assert img_out.shape == (batch_size, channels, image_height, image_width)
+        assert boxes_out.shape == (batch_size, 0, 4)
+
+
+# ---------------------------------------------------------------------------
+# TestCollateMasks — validates packing of variable-length per-image masks
+# into a zero-padded [B, N_max, H, W] float32 tensor.
+# ---------------------------------------------------------------------------
+
+
+class TestCollateMasks:
+    """collate_masks packs [N_i, H, W] instance masks into [B, N_max, H, W]."""
+
+    def _make_targets_with_masks(self, mask_counts, h=16, w=16):
+        """Build target dicts with boolean mask tensors for given instance counts."""
+        targets = []
+        for n in mask_counts:
+            masks = torch.ones(n, h, w, dtype=torch.bool) if n > 0 else torch.zeros(0, h, w, dtype=torch.bool)
+            targets.append({"masks": masks, "boxes": torch.zeros(n, 4)})
+        return targets
+
+    def test_normal_batch(self):
+        """Batch of [2 masks, 3 masks] → shape [2, 3, H, W] float32."""
+        from rfdetr.datasets.kornia_transforms import collate_masks
+
+        targets = self._make_targets_with_masks([2, 3])
+        masks_padded = collate_masks(targets, torch.device("cpu"), n_max=3, image_height=16, image_width=16)
+
+        assert masks_padded.shape == (2, 3, 16, 16), f"Expected (2, 3, 16, 16), got {masks_padded.shape}"
+        assert masks_padded.dtype == torch.float32, f"Expected float32, got {masks_padded.dtype}"
+
+    def test_padding_is_zero(self):
+        """Padded slots (beyond real instance count) are filled with zeros."""
+        from rfdetr.datasets.kornia_transforms import collate_masks
+
+        targets = self._make_targets_with_masks([1, 3])  # image 0 padded to 3
+        masks_padded = collate_masks(targets, torch.device("cpu"), n_max=3, image_height=16, image_width=16)
+
+        # Image 0: slot 0 real (ones), slots 1-2 zero-padded
+        assert masks_padded[0, 0].min() == pytest.approx(1.0), "Real mask slot must be all ones"
+        assert masks_padded[0, 1].max() == pytest.approx(0.0), "Padded slot 1 must be all zeros"
+        assert masks_padded[0, 2].max() == pytest.approx(0.0), "Padded slot 2 must be all zeros"
+
+    def test_n_max_zero_returns_empty(self):
+        """n_max=0 → shape [B, 0, H, W]."""
+        from rfdetr.datasets.kornia_transforms import collate_masks
+
+        targets = self._make_targets_with_masks([0, 0])
+        masks_padded = collate_masks(targets, torch.device("cpu"), n_max=0, image_height=16, image_width=16)
+
+        assert masks_padded.shape == (2, 0, 16, 16), f"Expected (2, 0, 16, 16), got {masks_padded.shape}"
+
+    def test_empty_target_list(self):
+        """Empty target list → shape [0, 0, H, W]."""
+        from rfdetr.datasets.kornia_transforms import collate_masks
+
+        masks_padded = collate_masks([], torch.device("cpu"), n_max=0, image_height=16, image_width=16)
+
+        assert masks_padded.shape == (0, 0, 16, 16), f"Expected (0, 0, 16, 16), got {masks_padded.shape}"
+
+    def test_targets_without_masks_key(self):
+        """Targets without 'masks' key produce all-zero rows."""
+        from rfdetr.datasets.kornia_transforms import collate_masks
+
+        targets = [{"boxes": torch.zeros(2, 4)}, {"boxes": torch.zeros(1, 4)}]
+        masks_padded = collate_masks(targets, torch.device("cpu"), n_max=2, image_height=8, image_width=8)
+
+        assert masks_padded.shape == (2, 2, 8, 8)
+        assert masks_padded.max() == pytest.approx(0.0), "Targets without masks key must produce all-zero output"
+
+
+# ---------------------------------------------------------------------------
+# TestBuildKorniaPipelineWithMasks — validates that with_masks=True produces
+# a pipeline with mask data_key included.
+# ---------------------------------------------------------------------------
+
+
+class TestBuildKorniaPipelineWithMasks:
+    """build_kornia_pipeline(with_masks=True) includes mask in data_keys."""
+
+    @pytest.fixture(autouse=True)
+    def _require_kornia(self):
+        """Skip when Kornia is unavailable (optional extra not installed in CPU CI)."""
+        pytest.importorskip("kornia")
+
+    def test_with_masks_false_is_default(self):
+        """with_masks defaults to False; pipeline returns (img, boxes) on call."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        pipeline = build_kornia_pipeline({"HorizontalFlip": {"p": 1.0}}, resolution=32)
+        img = torch.rand(1, 3, 32, 32)
+        boxes = torch.tensor([[[0.0, 0.0, 16.0, 16.0]]])
+        result = pipeline(img, boxes)
+        assert len(result) == 2, f"Detection pipeline must return 2 values, got {len(result)}"
+
+    def test_with_masks_true_returns_three_values(self):
+        """with_masks=True: pipeline(img, boxes, masks) returns (img, boxes, masks)."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        pipeline = build_kornia_pipeline({"HorizontalFlip": {"p": 1.0}}, resolution=32, with_masks=True)
+        img = torch.rand(1, 3, 32, 32)
+        boxes = torch.tensor([[[0.0, 0.0, 16.0, 16.0]]])
+        masks = torch.ones(1, 1, 32, 32, dtype=torch.float32)
+        result = pipeline(img, boxes, masks)
+        assert len(result) == 3, f"Segmentation pipeline must return 3 values, got {len(result)}"
+
+    def test_with_masks_true_preserves_mask_shape(self):
+        """Mask shape [B, N, H, W] is preserved after pipeline pass."""
+        from rfdetr.datasets.kornia_transforms import build_kornia_pipeline
+
+        pipeline = build_kornia_pipeline({"HorizontalFlip": {"p": 0.0}}, resolution=32, with_masks=True)
+        img = torch.rand(2, 3, 32, 32)
+        boxes = torch.tensor([[[0.0, 0.0, 16.0, 16.0]], [[8.0, 8.0, 24.0, 24.0]]])
+        masks = torch.ones(2, 1, 32, 32, dtype=torch.float32)
+        _, _, masks_aug = pipeline(img, boxes, masks)
+        assert masks_aug.shape == (2, 1, 32, 32), f"Mask shape must be preserved: {masks_aug.shape}"
+
+
+# ---------------------------------------------------------------------------
+# TestUnpackBoxesWithMasks — validates that unpack_boxes propagates the same
+# keep filter to masks when masks_aug is provided.
+# ---------------------------------------------------------------------------
+
+
+class TestUnpackBoxesWithMasks:
+    """unpack_boxes with masks_aug keeps/removes masks in sync with boxes."""
+
+    def test_masks_filtered_same_as_boxes(self):
+        """Box removed → corresponding mask also removed from output."""
+        from rfdetr.datasets.kornia_transforms import unpack_boxes
+
+        # B=1, N=2: box 0 valid, box 1 zero-area (will be removed)
+        boxes_aug = torch.tensor([[[5.0, 5.0, 25.0, 25.0], [30.0, 30.0, 30.0, 30.0]]])
+        valid = torch.tensor([[True, True]])
+        targets = [
+            {
+                "boxes": torch.tensor([[5.0, 5.0, 25.0, 25.0], [30.0, 30.0, 60.0, 60.0]]),
+                "labels": torch.tensor([1, 2]),
+            }
+        ]
+        # 2 masks: instance 0 = all ones, instance 1 = all twos (distinguishable)
+        masks_aug = torch.zeros(1, 2, 8, 8, dtype=torch.float32)
+        masks_aug[0, 0] = 1.0
+        masks_aug[0, 1] = 1.0  # will be removed with box 1
+
+        result = unpack_boxes(boxes_aug, valid, targets, 100, 100, masks_aug=masks_aug)
+
+        assert "masks" in result[0], "masks key must be present in output target"
+        assert result[0]["masks"].shape[0] == 1, f"Expected 1 surviving mask, got {result[0]['masks'].shape[0]}"
+
+    def test_masks_converted_to_bool(self):
+        """Float masks > 0.5 threshold converted to bool in output."""
+        from rfdetr.datasets.kornia_transforms import unpack_boxes
+
+        boxes_aug = torch.tensor([[[5.0, 5.0, 25.0, 25.0]]])
+        valid = torch.tensor([[True]])
+        targets = [{"boxes": torch.tensor([[5.0, 5.0, 25.0, 25.0]]), "labels": torch.tensor([1])}]
+        masks_aug = torch.full((1, 1, 8, 8), 0.8, dtype=torch.float32)  # float, all 0.8
+
+        result = unpack_boxes(boxes_aug, valid, targets, 100, 100, masks_aug=masks_aug)
+
+        assert result[0]["masks"].dtype == torch.bool, f"masks must be bool, got {result[0]['masks'].dtype}"
+        assert result[0]["masks"].all(), "All values > 0.5 should be True after thresholding"
+
+    def test_no_masks_aug_leaves_masks_key_unchanged(self):
+        """When masks_aug=None, existing masks key in target is preserved as-is."""
+        from rfdetr.datasets.kornia_transforms import unpack_boxes
+
+        boxes_aug = torch.tensor([[[5.0, 5.0, 25.0, 25.0]]])
+        valid = torch.tensor([[True]])
+        original_mask = torch.ones(1, 8, 8, dtype=torch.bool)
+        targets = [
+            {
+                "boxes": torch.tensor([[5.0, 5.0, 25.0, 25.0]]),
+                "labels": torch.tensor([1]),
+                "masks": original_mask,
+            }
+        ]
+
+        result = unpack_boxes(boxes_aug, valid, targets, 100, 100, masks_aug=None)
+
+        assert "masks" in result[0], "masks key must still be present when masks_aug=None"
+        assert result[0]["masks"] is original_mask, "Original masks object must be preserved unchanged"
+
+
+class TestGaussNoiseStdRangeWarning:
+    """_make_gauss_noise warns when the configured std range is non-degenerate (GPU uses a fixed upper-bound std)."""
+
+    @pytest.fixture(autouse=True)
+    def _require_kornia(self):
+        pytest.importorskip("kornia")
+
+    def test_warns_for_unequal_std_range(self):
+        """A non-degenerate std_range emits a divergence warning at build time."""
+        from unittest import mock
+
+        from rfdetr.datasets import kornia_transforms
+
+        with mock.patch.object(kornia_transforms.logger, "warning") as mock_warning:
+            kornia_transforms._make_gauss_noise({"std_range": (0.01, 0.05), "p": 0.5})
+
+        mock_warning.assert_called_once()
+
+    def test_no_warning_for_degenerate_std_range(self):
+        """An equal-bound std_range matches the CPU path exactly and stays silent."""
+        from unittest import mock
+
+        from rfdetr.datasets import kornia_transforms
+
+        with mock.patch.object(kornia_transforms.logger, "warning") as mock_warning:
+            kornia_transforms._make_gauss_noise({"std_range": (0.05, 0.05), "p": 0.5})
+
+        mock_warning.assert_not_called()

tests/datasets/test_o365.py

@@ -0,0 +1,20 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for the Object365 dataset module."""
+
+import PIL.Image
+
+
+def test_o365_import_keeps_finite_decompression_bomb_guard() -> None:
+    """Importing ``o365`` must not disable PIL's decompression-bomb guard process-wide."""
+    from rfdetr.datasets import o365  # noqa: F401
+
+    assert PIL.Image.MAX_IMAGE_PIXELS is not None, (
+        "o365 must set a finite MAX_IMAGE_PIXELS cap, not None (which disables the guard globally)"
+    )
+    assert PIL.Image.MAX_IMAGE_PIXELS >= 178_956_970, (
+        "the cap must stay above PIL's default so legitimate large O365 images still load"
+    )

tests/datasets/test_save_grids.py

@@ -0,0 +1,62 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for DatasetGridSaver — verifies that annotated grid images are written without OpenCV layout errors across all
+supported OpenCV versions."""
+
+from pathlib import Path
+
+import numpy as np
+import torch
+from PIL import Image
+from torch.utils.data import DataLoader
+
+
+class _FakeDataset:
+    """Minimal dataset returning a single synthetic image + target."""
+
+    def __init__(self, num_samples: int = 4) -> None:
+        self.num_samples = num_samples
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def __getitem__(self, idx):
+        # CHW float tensor in ImageNet-normalised range
+        image = torch.zeros(3, 224, 224)
+        target = {
+            "size": torch.tensor([224, 224]),
+            "boxes": torch.tensor([[0.25, 0.25, 0.5, 0.5], [0.6, 0.6, 0.2, 0.2]]),
+            "labels": torch.tensor([0, 1]),
+        }
+        return image, target
+
+
+def _collate(batch):
+    from rfdetr.utilities import nested_tensor_from_tensor_list
+
+    images, targets = zip(*batch)
+    # NestedTensor expected by DatasetGridSaver
+    nested = nested_tensor_from_tensor_list(list(images))
+    return nested, list(targets)
+
+
+def test_save_grid_writes_files(tmp_path: Path) -> None:
+    """DatasetGridSaver must write JPEG grid files without raising OpenCV errors."""
+    from rfdetr.datasets.save_grids import DatasetGridSaver
+
+    dataset = _FakeDataset(num_samples=4)
+    loader = DataLoader(dataset, batch_size=2, collate_fn=_collate)
+
+    saver = DatasetGridSaver(loader, tmp_path, max_batches=2, dataset_type="train")
+    saver.save_grid()
+
+    grids = list(tmp_path.glob("train_batch*_grid.jpg"))
+    assert len(grids) == 2, f"Expected 2 grid files, got {len(grids)}"
+    for grid_path in grids:
+        with Image.open(grid_path) as pil_img:
+            img = np.array(pil_img)
+        assert img.ndim == 3
+        assert img.shape[2] == 3

tests/datasets/test_synthetic.py

@@ -0,0 +1,601 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+import json
+
+import numpy as np
+import pytest
+import supervision as sv
+
+from rfdetr.datasets.synthetic import (
+    DEFAULT_SPLIT_RATIOS,
+    SYNTHETIC_SHAPES,
+    DatasetSplitRatios,
+    _calculate_polygon_area,
+    _write_coco_json,
+    calculate_boundary_overlap,
+    draw_synthetic_shape,
+    generate_coco_dataset,
+    generate_synthetic_sample,
+)
+
+
+class TestCalculateBoundaryOverlap:
+    @pytest.mark.parametrize(
+        "bbox,expected_overlap",
+        [
+            pytest.param(np.array([40.0, 40.0, 60.0, 60.0]), 0.0, id="fully_inside"),
+            pytest.param(np.array([-10.0, 40.0, 10.0, 60.0]), 0.5, id="half_outside_horizontally"),
+            pytest.param(np.array([110.0, 40.0, 130.0, 60.0]), 1.0, id="fully_outside"),
+            pytest.param(np.array([0.0, 0.0, 50.0, 50.0]), 0.0, id="exactly_at_boundary"),
+            pytest.param(np.array([50.0, 50.0, 100.0, 100.0]), 0.0, id="exactly_at_max_boundary"),
+        ],
+    )
+    def test_overlap_values(self, bbox, expected_overlap):
+        result = calculate_boundary_overlap(bbox, img_size=100)
+        assert result == pytest.approx(expected_overlap)
+
+
+class TestDrawSyntheticShape:
+    @pytest.mark.parametrize(
+        "shape,color",
+        [
+            pytest.param("square", sv.Color.RED, id="square_red"),
+            pytest.param("triangle", sv.Color.GREEN, id="triangle_green"),
+            pytest.param("circle", sv.Color.BLUE, id="circle_blue"),
+        ],
+    )
+    def test_pixels_are_modified(self, shape, color):
+        img = np.zeros((100, 100, 3), dtype=np.uint8)
+        img_modified, polygon = draw_synthetic_shape(img.copy(), shape, color, (50, 50), 20)
+        assert not np.array_equal(img, img_modified)
+        assert len(polygon) >= 6
+        assert len(polygon) % 2 == 0
+
+    @pytest.mark.parametrize(
+        "shape,cx,cy,size",
+        [
+            pytest.param("square", 50, 50, 20, id="square"),
+            pytest.param("triangle", 50, 50, 20, id="triangle"),
+            pytest.param("circle", 50, 50, 20, id="circle"),
+        ],
+    )
+    def test_polygon_min_points(self, shape, cx, cy, size):
+        """Returned polygon must have at least 3 points (6 values) for COCO."""
+        img = np.zeros((100, 100, 3), dtype=np.uint8)
+        _, poly = draw_synthetic_shape(img, shape, sv.Color.WHITE, (cx, cy), size)
+        assert len(poly) >= 6, f"{shape} polygon has fewer than 6 values: {poly}"
+        assert len(poly) % 2 == 0, f"{shape} polygon has an odd number of values: {poly}"
+
+    @pytest.mark.parametrize(
+        "shape,cx,cy,size,expected_n_coords",
+        [
+            pytest.param("square", 50, 50, 20, 8, id="square_4pts"),
+            pytest.param("triangle", 50, 50, 20, 6, id="triangle_3pts"),
+            pytest.param("circle", 50, 50, 20, 64, id="circle_32pts"),
+        ],
+    )
+    def test_polygon_coord_count(self, shape, cx, cy, size, expected_n_coords):
+        """Each shape must return the expected number of flat coordinate values."""
+        img = np.zeros((100, 100, 3), dtype=np.uint8)
+        _, poly = draw_synthetic_shape(img, shape, sv.Color.WHITE, (cx, cy), size)
+        assert len(poly) == expected_n_coords
+
+    def test_square_polygon_matches_bbox(self):
+        """Square polygon corners must align with the drawn rectangle bounds."""
+        cx, cy, size = 60, 40, 30
+        img = np.zeros((100, 100, 3), dtype=np.uint8)
+        _, poly = draw_synthetic_shape(img, "square", sv.Color.WHITE, (cx, cy), size)
+        hs = size // 2
+        expected = [
+            float(cx - hs),
+            float(cy - hs),
+            float(cx - hs + size),
+            float(cy - hs),
+            float(cx - hs + size),
+            float(cy - hs + size),
+            float(cx - hs),
+            float(cy - hs + size),
+        ]
+        assert poly == pytest.approx(expected)
+
+    def test_unknown_shape_returns_empty_polygon(self):
+        """An unrecognised shape name must return an empty polygon without crashing."""
+        img = np.zeros((100, 100, 3), dtype=np.uint8)
+        _, poly = draw_synthetic_shape(img, "hexagon", sv.Color.WHITE, (50, 50), 20)
+        assert poly == []
+
+
+class TestGenerateSyntheticSample:
+    @pytest.mark.parametrize(
+        "img_size,min_objects,max_objects,class_mode",
+        [
+            pytest.param(100, 1, 3, "shape", id="small_shape_mode"),
+            pytest.param(200, 2, 5, "color", id="medium_color_mode"),
+            pytest.param(100, 1, 1, "shape", id="single_object"),
+            pytest.param(100, 0, 0, "shape", id="zero_objects"),
+        ],
+    )
+    def test_output_shape_and_detection_count(self, img_size, min_objects, max_objects, class_mode):
+        img, detections = generate_synthetic_sample(
+            img_size=img_size, min_objects=min_objects, max_objects=max_objects, class_mode=class_mode
+        )
+        assert img.shape == (img_size, img_size, 3)
+        assert min_objects <= len(detections) <= max_objects
+        assert hasattr(detections, "xyxy")
+        assert hasattr(detections, "class_id")
+
+    def test_polygon_data_present(self):
+        """detections.data must contain a 'polygons' array with one entry per detection."""
+        _, detections = generate_synthetic_sample(img_size=100, min_objects=2, max_objects=4, class_mode="shape")
+        assert "polygons" in detections.data
+        assert len(detections.data["polygons"]) == len(detections)
+
+    def test_polygon_data_non_empty(self):
+        """Each stored polygon must be a non-empty list of floats."""
+        _, detections = generate_synthetic_sample(img_size=100, min_objects=1, max_objects=3, class_mode="shape")
+        for poly in detections.data["polygons"]:
+            assert isinstance(poly, list)
+            assert len(poly) >= 6
+
+    def test_zero_objects_polygon_data(self):
+        """With zero objects the polygon data array must be present but empty."""
+        _, detections = generate_synthetic_sample(img_size=100, min_objects=0, max_objects=0, class_mode="shape")
+        assert "polygons" in detections.data
+        assert len(detections.data["polygons"]) == 0
+
+    def test_polygon_bbox_consistency(self):
+        """detections.xyxy must match the min/max of the corresponding polygon."""
+        _, detections = generate_synthetic_sample(img_size=200, min_objects=3, max_objects=5, class_mode="shape")
+        for i in range(len(detections)):
+            poly = detections.data["polygons"][i]
+            poly_array = np.asarray(poly, dtype=float).reshape(-1, 2)
+            expected_x_min = float(np.min(poly_array[:, 0]))
+            expected_y_min = float(np.min(poly_array[:, 1]))
+            expected_x_max = float(np.max(poly_array[:, 0]))
+            expected_y_max = float(np.max(poly_array[:, 1]))
+            x_min, y_min, x_max, y_max = detections.xyxy[i]
+            assert x_min == pytest.approx(expected_x_min), f"detection {i} x_min mismatch"
+            assert y_min == pytest.approx(expected_y_min), f"detection {i} y_min mismatch"
+            assert x_max == pytest.approx(expected_x_max), f"detection {i} x_max mismatch"
+            assert y_max == pytest.approx(expected_y_max), f"detection {i} y_max mismatch"
+
+
+class TestGenerateCocoDataset:
+    @pytest.mark.parametrize(
+        "num_images,img_size,class_mode,split_ratios,expected_splits",
+        [
+            # Test with dictionary (legacy support)
+            pytest.param(
+                5,
+                100,
+                "shape",
+                {"train": 0.6, "val": 0.2, "test": 0.2},
+                ["train", "val", "test"],
+                id="shape_mode_all_splits_dict",
+            ),
+            pytest.param(
+                3,
+                64,
+                "color",
+                {"train": 0.5, "val": 0.5},
+                ["train", "val"],
+                id="color_mode_two_splits_dict",
+            ),
+            pytest.param(
+                2,
+                128,
+                "shape",
+                {"train": 1.0},
+                ["train"],
+                id="single_split_only_dict",
+            ),
+            # Test with DatasetSplitRatios dataclass
+            pytest.param(
+                4,
+                100,
+                "shape",
+                DatasetSplitRatios(train=0.7, val=0.2, test=0.1),
+                ["train", "val", "test"],
+                id="split_ratios_dataclass",
+            ),
+            pytest.param(
+                3,
+                64,
+                "color",
+                DatasetSplitRatios(train=0.8, val=0.2, test=0.0),
+                ["train", "val"],
+                id="split_ratios_no_test",
+            ),
+            # Test with tuple
+            pytest.param(
+                4,
+                100,
+                "shape",
+                (0.7, 0.2, 0.1),
+                ["train", "val", "test"],
+                id="split_ratios_tuple_three",
+            ),
+            pytest.param(
+                3,
+                64,
+                "color",
+                (0.8, 0.2),
+                ["train", "val"],
+                id="split_ratios_tuple_two",
+            ),
+            # Test with default
+            pytest.param(
+                10,
+                64,
+                "shape",
+                DEFAULT_SPLIT_RATIOS,
+                ["train", "val", "test"],
+                id="split_ratios_default",
+            ),
+        ],
+    )
+    def test_splits_created(self, num_images, img_size, class_mode, split_ratios, expected_splits, tmp_path):
+        output_dir = tmp_path / "test_dataset"
+        generate_coco_dataset(
+            output_dir=str(output_dir),
+            num_images=num_images,
+            img_size=img_size,
+            class_mode=class_mode,
+            split_ratios=split_ratios,
+        )
+
+        assert output_dir.exists()
+        for split in expected_splits:
+            split_dir = output_dir / split
+            assert split_dir.exists()
+            assert (split_dir / "_annotations.coco.json").exists()
+
+            with open(split_dir / "_annotations.coco.json") as f:
+                data = json.load(f)
+            assert "images" in data
+            assert "annotations" in data
+            assert "categories" in data
+            for img_info in data["images"]:
+                assert (split_dir / img_info["file_name"]).exists()
+
+    @pytest.mark.parametrize(
+        "num_images,split_ratios",
+        [
+            pytest.param(10, (0.33, 0.33, 0.34), id="truncating_ratios"),
+            pytest.param(7, (0.7, 0.2, 0.1), id="standard_ratios"),
+            pytest.param(5, (0.8, 0.2), id="two_split"),
+        ],
+    )
+    def test_split_image_count_equals_total(self, num_images, split_ratios, tmp_path):
+        """Total images assigned across all splits must equal num_images."""
+        output_dir = tmp_path / "test_dataset"
+        generate_coco_dataset(
+            output_dir=str(output_dir),
+            num_images=num_images,
+            img_size=64,
+            class_mode="shape",
+            split_ratios=split_ratios,
+        )
+        total_images = 0
+        for split_dir in output_dir.iterdir():
+            ann_file = split_dir / "_annotations.coco.json"
+            if ann_file.exists():
+                with open(ann_file) as fh:
+                    total_images += len(json.load(fh)["images"])
+        assert total_images == num_images
+
+    @pytest.mark.parametrize(
+        "split_ratios,error_message",
+        [
+            pytest.param(
+                (1.1, -0.1),
+                "Split ratios must be non-negative",
+                id="tuple_negative_ratio",
+            ),
+            pytest.param(
+                {"train": 1.1, "val": -0.1},
+                "Split ratios must be non-negative",
+                id="dict_negative_ratio",
+            ),
+            pytest.param(
+                (0.5, 0.3),
+                "Split ratios must sum to 1.0",
+                id="tuple_invalid_sum",
+            ),
+        ],
+    )
+    def test_invalid_split_ratios(self, split_ratios, error_message, tmp_path):
+        output_dir = tmp_path / "test_dataset"
+        with pytest.raises(ValueError, match=error_message):
+            generate_coco_dataset(
+                output_dir=str(output_dir),
+                num_images=5,
+                img_size=100,
+                class_mode="shape",
+                split_ratios=split_ratios,
+            )
+
+
+class TestGenerateCocoDatasetWithSegmentation:
+    def test_write_coco_json_raises_when_polygons_key_missing(self, tmp_path):
+        """with_segmentation=True must raise if detections.data has no 'polygons' key."""
+        annotations_path = tmp_path / "_annotations.coco.json"
+        detections = sv.Detections(
+            xyxy=np.array([[0.0, 0.0, 10.0, 10.0]], dtype=float),
+            class_id=np.array([0], dtype=int),
+            data={},  # intentionally no "polygons" key
+        )
+        with pytest.raises(ValueError, match="no 'polygons' found"):
+            _write_coco_json(
+                annotations_path=annotations_path,
+                classes=["shape"],
+                file_paths=["/tmp/synthetic.png"],
+                detections_list=[detections],
+                img_size=64,
+                with_segmentation=True,
+            )
+
+    def test_write_coco_json_raises_for_mismatched_inputs(self, tmp_path):
+        """Mismatched file/detection list lengths must raise to avoid silent truncation."""
+        annotations_path = tmp_path / "_annotations.coco.json"
+        detections = sv.Detections(
+            xyxy=np.empty((0, 4), dtype=float),
+            class_id=np.empty((0,), dtype=int),
+            data={"polygons": np.empty(0, dtype=object)},
+        )
+
+        with pytest.raises(ValueError, match="file_paths and detections_list must have the same length"):
+            _write_coco_json(
+                annotations_path=annotations_path,
+                classes=["shape"],
+                file_paths=["/tmp/a.png", "/tmp/b.png"],
+                detections_list=[detections],
+                img_size=64,
+            )
+
+    def test_creates_files(self, tmp_path):
+        """with_segmentation=True must create the same directory/file structure as the default."""
+        output_dir = tmp_path / "seg_dataset"
+        generate_coco_dataset(
+            output_dir=str(output_dir),
+            num_images=4,
+            img_size=64,
+            class_mode="shape",
+            split_ratios={"train": 0.75, "val": 0.25},
+            with_segmentation=True,
+        )
+        for split in ("train", "val"):
+            assert (output_dir / split / "_annotations.coco.json").exists()
+
+    def test_json_structure(self, tmp_path):
+        """COCO JSON produced with segmentation must have the required top-level keys."""
+        output_dir = tmp_path / "seg_dataset"
+        generate_coco_dataset(
+            output_dir=str(output_dir),
+            num_images=4,
+            img_size=64,
+            class_mode="shape",
+            split_ratios={"train": 1.0},
+            with_segmentation=True,
+        )
+        with open(output_dir / "train" / "_annotations.coco.json") as fh:
+            data = json.load(fh)
+        assert "images" in data
+        assert "annotations" in data
+        assert "categories" in data
+
+    def test_has_polygon_field(self, tmp_path):
+        """Every annotation must have a non-empty segmentation polygon."""
+        output_dir = tmp_path / "seg_dataset"
+        generate_coco_dataset(
+            output_dir=str(output_dir),
+            num_images=3,
+            img_size=64,
+            class_mode="shape",
+            min_objects=1,
+            max_objects=2,
+            split_ratios={"train": 1.0},
+            with_segmentation=True,
+        )
+        with open(output_dir / "train" / "_annotations.coco.json") as fh:
+            data = json.load(fh)
+        assert len(data["annotations"]) > 0, "Expected at least one annotation"
+        for ann in data["annotations"]:
+            assert "segmentation" in ann
+            assert isinstance(ann["segmentation"], list)
+            assert len(ann["segmentation"]) == 1, "Expected exactly one polygon per annotation"
+            assert len(ann["segmentation"][0]) >= 6, "Polygon must have at least 3 points"
+
+    def test_area_uses_polygon_when_segmentation_enabled(self, tmp_path):
+        """COCO area must match polygon area when segmentation annotations are present."""
+        annotations_path = tmp_path / "_annotations.coco.json"
+        polygon_data = np.empty(1, dtype=object)
+        polygon_data[0] = [0.0, 0.0, 10.0, 0.0, 0.0, 10.0]  # Right triangle area = 50
+        detections = sv.Detections(
+            xyxy=np.array([[0.0, 0.0, 10.0, 10.0]], dtype=float),
+            class_id=np.array([0], dtype=int),
+            data={"polygons": polygon_data},
+        )
+
+        _write_coco_json(
+            annotations_path=annotations_path,
+            classes=["shape"],
+            file_paths=["/tmp/synthetic.png"],
+            detections_list=[detections],
+            img_size=64,
+            with_segmentation=True,
+        )
+
+        with open(annotations_path) as fh:
+            data = json.load(fh)
+
+        assert len(data["annotations"]) == 1
+        assert data["annotations"][0]["area"] == pytest.approx(50.0)
+
+    def test_sparse_category_ids(self, tmp_path):
+        """Category IDs must use sparse 1-based encoding (1, 3, 5, …)."""
+        output_dir = tmp_path / "seg_dataset"
+        generate_coco_dataset(
+            output_dir=str(output_dir),
+            num_images=4,
+            img_size=64,
+            class_mode="shape",
+            split_ratios={"train": 1.0},
+            with_segmentation=True,
+        )
+        with open(output_dir / "train" / "_annotations.coco.json") as fh:
+            data = json.load(fh)
+        cat_ids = {c["id"] for c in data["categories"]}
+        expected_ids = {idx * 2 + 1 for idx in range(len(SYNTHETIC_SHAPES))}
+        assert cat_ids == expected_ids
+        ann_cat_ids = {a["category_id"] for a in data["annotations"]}
+        assert ann_cat_ids.issubset(expected_ids)
+
+    def test_images_exist(self, tmp_path):
+        """All images referenced in the JSON must exist on disk."""
+        output_dir = tmp_path / "seg_dataset"
+        generate_coco_dataset(
+            output_dir=str(output_dir),
+            num_images=3,
+            img_size=64,
+            class_mode="shape",
+            split_ratios={"train": 1.0},
+            with_segmentation=True,
+        )
+        split_dir = output_dir / "train"
+        with open(split_dir / "_annotations.coco.json") as fh:
+            data = json.load(fh)
+        for img_info in data["images"]:
+            assert (split_dir / img_info["file_name"]).exists()
+
+    def test_empty_polygon_falls_back_to_empty_segmentation(self, tmp_path):
+        """An empty polygon entry silently falls back to ``segmentation=[]``.
+
+        The ``len(polygon_data) < len(detections)`` guard only checks array length, not contents.  An element that is an
+        empty list passes the guard and takes the ``else`` branch producing ``segmentation=[]``. This test documents the
+        existing silent-fallback behaviour.
+        """
+        annotations_path = tmp_path / "_annotations.coco.json"
+        polygon_data = np.empty(1, dtype=object)
+        polygon_data[0] = []  # empty polygon — passes length guard
+        detections = sv.Detections(
+            xyxy=np.array([[0.0, 0.0, 10.0, 10.0]], dtype=float),
+            class_id=np.array([0], dtype=int),
+            data={"polygons": polygon_data},
+        )
+        _write_coco_json(
+            annotations_path=annotations_path,
+            classes=["shape"],
+            file_paths=["/tmp/synthetic.png"],
+            detections_list=[detections],
+            img_size=64,
+            with_segmentation=True,
+        )
+        with open(annotations_path) as fh:
+            data = json.load(fh)
+        assert data["annotations"][0]["segmentation"] == []
+
+
+class TestCalculatePolygonArea:
+    @pytest.mark.parametrize(
+        "polygon,expected_area",
+        [
+            pytest.param(
+                [0.0, 0.0, 10.0, 0.0, 0.0, 10.0],
+                50.0,
+                id="right_triangle",
+            ),
+            pytest.param(
+                [0.0, 0.0, 10.0, 0.0, 10.0, 10.0, 0.0, 10.0],
+                100.0,
+                id="unit_square_10x10",
+            ),
+            pytest.param(
+                [0.0, 0.0, 5.0, 0.0, 10.0, 0.0],
+                0.0,
+                id="collinear_points_degenerate",
+            ),
+            pytest.param(
+                [0.0, 0.0, 1.0, 1.0],
+                0.0,
+                id="fewer_than_3_points",
+            ),
+            pytest.param(
+                [],
+                0.0,
+                id="empty_polygon",
+            ),
+        ],
+    )
+    def test_area(self, polygon, expected_area):
+        assert _calculate_polygon_area(polygon) == pytest.approx(expected_area)
+
+
+class TestDrawSyntheticShapeEdgeCases:
+    def test_square_polygon_respects_half_size_and_image_bounds_for_odd_size(self):
+        """For odd sizes, the square polygon should:
+
+        * Have all vertices within the image bounds.
+        * Be horizontally contained within ``cx ± size / 2``.
+        """
+        cx, cy, size = 50, 50, 21
+        img = np.zeros((100, 100, 3), dtype=np.uint8)
+        _, poly = draw_synthetic_shape(img, "square", sv.Color.WHITE, (cx, cy), size)
+
+        half_size = size / 2.0
+        xs = [poly[i] for i in range(0, len(poly), 2)]
+        ys = [poly[i] for i in range(1, len(poly), 2)]
+
+        # All vertices must be inside the image
+        assert min(xs) >= 0.0
+        assert max(xs) <= float(img.shape[1])
+        assert min(ys) >= 0.0
+        assert max(ys) <= float(img.shape[0])
+
+        # Horizontal extent should not exceed the intended half-size around cx
+        assert min(xs) >= cx - half_size - 1.0
+        assert max(xs) <= cx + half_size + 1.0
+
+    def test_triangle_vertices_within_half_size_and_image_bounds(self):
+        """Triangle vertices should:
+
+        * Have all vertices within the image bounds.
+        * Be vertically contained within ``cy ± size / 2`` so the apex does not
+          extend beyond the intended half-size boundary.
+        """
+        cx, cy, size = 50, 50, 20
+        img = np.zeros((100, 100, 3), dtype=np.uint8)
+        _, poly = draw_synthetic_shape(img, "triangle", sv.Color.WHITE, (cx, cy), size)
+
+        half_size = size / 2.0
+        xs = [poly[i] for i in range(0, len(poly), 2)]
+        ys = [poly[i] for i in range(1, len(poly), 2)]
+
+        # All vertices must be inside the image
+        assert min(xs) >= 0.0
+        assert max(xs) <= float(img.shape[1])
+        assert min(ys) >= 0.0
+        assert max(ys) <= float(img.shape[0])
+
+        # Vertical extent should not exceed the intended half-size around cy
+        assert min(ys) >= cy - half_size - 1.0
+        assert max(ys) <= cy + half_size + 1.0
+
+    @pytest.mark.parametrize(
+        "shape,size,expected_n_coords",
+        [
+            pytest.param("square", 0, 8, id="square_size_0"),
+            pytest.param("square", 1, 8, id="square_size_1"),
+            pytest.param("circle", 0, 64, id="circle_size_0"),
+            pytest.param("circle", 1, 64, id="circle_size_1"),
+        ],
+    )
+    def test_degenerate_size_returns_polygon_without_crashing(self, shape, size, expected_n_coords):
+        """draw_synthetic_shape with size=0 or size=1 must not raise and must return the expected number of flat
+        coordinate values."""
+        img = np.zeros((100, 100, 3), dtype=np.uint8)
+        _, poly = draw_synthetic_shape(img, shape, sv.Color.WHITE, (50, 50), size)
+        assert len(poly) == expected_n_coords

tests/evaluation/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/evaluation/test_coco_eval.py

@@ -0,0 +1,501 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests for the local COCO evaluator wrapper."""
+
+import json
+from pathlib import Path
+from unittest.mock import patch
+
+import numpy as np
+import pycocotools.coco as pycoco
+import pytest
+import torch
+from faster_coco_eval import COCO
+
+from rfdetr.evaluation import coco_eval as coco_eval_module
+from rfdetr.evaluation.coco_eval import CocoEvaluator
+
+
+def _write_person_keypoint_coco(path: Path, *, include_num_keypoints: bool = True, keypoint_count: int = 17) -> None:
+    """Write a minimal COCO keypoint annotation file."""
+    if keypoint_count == 17:
+        keypoints = [
+            "nose",
+            "left_eye",
+            "right_eye",
+            "left_ear",
+            "right_ear",
+            "left_shoulder",
+            "right_shoulder",
+            "left_elbow",
+            "right_elbow",
+            "left_wrist",
+            "right_wrist",
+            "left_hip",
+            "right_hip",
+            "left_knee",
+            "right_knee",
+            "left_ankle",
+            "right_ankle",
+        ]
+    else:
+        keypoints = [f"point_{idx}" for idx in range(keypoint_count)]
+    coords = []
+    for idx in range(len(keypoints)):
+        coords.extend([20.0 + idx, 30.0 + idx, 2.0])
+    annotation = {
+        "id": 1,
+        "image_id": 1,
+        "category_id": 1,
+        "bbox": [10.0, 20.0, 50.0, 60.0],
+        "area": 3000.0,
+        "iscrowd": 0,
+        "keypoints": coords,
+    }
+    if include_num_keypoints:
+        annotation["num_keypoints"] = len(keypoints)
+    payload = {
+        "images": [{"id": 1, "width": 100, "height": 100, "file_name": "image.jpg"}],
+        "annotations": [annotation],
+        "categories": [
+            {
+                "id": 1,
+                "name": "person",
+                "supercategory": "person",
+                "keypoints": keypoints,
+                "skeleton": [],
+            }
+        ],
+    }
+    path.write_text(json.dumps(payload), encoding="utf-8")
+
+
+def _write_mixed_keypoint_coco(path: Path) -> None:
+    """Write a COCO keypoint file with two categories using different keypoint counts."""
+    categories = [
+        {
+            "id": 1,
+            "name": "dart",
+            "supercategory": "object",
+            "keypoints": [f"dart_{idx}" for idx in range(4)],
+            "skeleton": [],
+        },
+        {
+            "id": 2,
+            "name": "person",
+            "supercategory": "person",
+            "keypoints": [f"person_{idx}" for idx in range(21)],
+            "skeleton": [],
+        },
+    ]
+    annotations = []
+    for annotation_id, (category_id, keypoint_count, x0, y0) in enumerate(
+        [(1, 4, 10.0, 20.0), (2, 21, 50.0, 60.0)],
+        start=1,
+    ):
+        keypoints = []
+        for idx in range(keypoint_count):
+            keypoints.extend([x0 + idx, y0 + idx, 2.0])
+        annotations.append(
+            {
+                "id": annotation_id,
+                "image_id": 1,
+                "category_id": category_id,
+                "bbox": [x0, y0, 20.0, 20.0],
+                "area": 400.0,
+                "iscrowd": 0,
+                "keypoints": keypoints,
+                "num_keypoints": keypoint_count,
+            }
+        )
+
+    payload = {
+        "images": [{"id": 1, "width": 100, "height": 100, "file_name": "image.jpg"}],
+        "annotations": annotations,
+        "categories": categories,
+    }
+    path.write_text(json.dumps(payload), encoding="utf-8")
+
+
+def test_coco_evaluator_keypoints_uses_faster_evaluate_without_deprecated_evaluate_img(tmp_path: Path) -> None:
+    """Keypoint evaluation should not call faster-coco-eval's deprecated ``evaluateImg`` shim."""
+    annotation_path = tmp_path / "person_keypoints_val2017.json"
+    _write_person_keypoint_coco(annotation_path)
+    coco_gt = COCO(str(annotation_path))
+    coco_gt.label2cat = {0: 1}
+    evaluator = CocoEvaluator(coco_gt, ["keypoints"])
+    keypoints = np.asarray(coco_gt.anns[1]["keypoints"], dtype=np.float32).reshape(1, 17, 3)
+
+    evaluator.update(
+        {
+            1: {
+                "boxes": torch.tensor([[10.0, 20.0, 60.0, 80.0]], dtype=torch.float32),
+                "scores": torch.tensor([0.99], dtype=torch.float32),
+                "labels": torch.tensor([0], dtype=torch.int64),
+                "keypoints": torch.as_tensor(keypoints, dtype=torch.float32),
+            }
+        }
+    )
+    evaluator.synchronize_between_processes()
+    evaluator.accumulate()
+
+    stats = evaluator.coco_eval["keypoints"].stats
+    assert np.isfinite(stats[0])
+
+
+def test_coco_evaluator_keypoints_log_summary_false_suppresses_summary_rows(tmp_path: Path) -> None:
+    """Keypoint accumulation should compute stats without AP/AR logger spam when summaries are disabled."""
+    annotation_path = tmp_path / "person_keypoints_val2017.json"
+    _write_person_keypoint_coco(annotation_path)
+    coco_gt = COCO(str(annotation_path))
+    coco_gt.label2cat = {0: 1}
+    evaluator = CocoEvaluator(coco_gt, ["keypoints"], log_summary=False)
+    keypoints = np.asarray(coco_gt.anns[1]["keypoints"], dtype=np.float32).reshape(1, 17, 3)
+
+    evaluator.update(
+        {
+            1: {
+                "boxes": torch.tensor([[10.0, 20.0, 60.0, 80.0]], dtype=torch.float32),
+                "scores": torch.tensor([0.99], dtype=torch.float32),
+                "labels": torch.tensor([0], dtype=torch.int64),
+                "keypoints": torch.as_tensor(keypoints, dtype=torch.float32),
+            }
+        }
+    )
+
+    with patch.object(coco_eval_module.logger, "info") as info:
+        evaluator.synchronize_between_processes()
+        evaluator.accumulate()
+
+    info.assert_not_called()
+    stats = evaluator.coco_eval["keypoints"].stats
+    assert np.isfinite(stats[0])
+
+
+def test_coco_evaluator_keypoints_accepts_pycocotools_coco_api(tmp_path: Path) -> None:
+    """Keypoint evaluation should accept COCO APIs returned by torchvision datasets."""
+    annotation_path = tmp_path / "person_keypoints_val2017.json"
+    _write_person_keypoint_coco(annotation_path)
+    coco_gt = pycoco.COCO(str(annotation_path))
+    coco_gt.label2cat = {0: 1}
+
+    evaluator = CocoEvaluator(coco_gt, ["keypoints"])
+    keypoints = np.asarray(coco_gt.anns[1]["keypoints"], dtype=np.float32).reshape(1, 17, 3)
+    evaluator.update(
+        {
+            1: {
+                "boxes": torch.tensor([[10.0, 20.0, 60.0, 80.0]], dtype=torch.float32),
+                "scores": torch.tensor([0.99], dtype=torch.float32),
+                "labels": torch.tensor([0], dtype=torch.int64),
+                "keypoints": torch.as_tensor(keypoints, dtype=torch.float32),
+            }
+        }
+    )
+    evaluator.synchronize_between_processes()
+    evaluator.accumulate()
+
+    stats = evaluator.coco_eval["keypoints"].stats
+    assert np.isfinite(stats[0])
+
+
+def test_coco_evaluator_keypoints_infers_custom_oks_sigmas(tmp_path: Path) -> None:
+    """Custom keypoint-count datasets should not use COCO's fixed 17-keypoint OKS sigmas."""
+    annotation_path = tmp_path / "custom_keypoints_val.json"
+    _write_person_keypoint_coco(annotation_path, keypoint_count=25)
+    coco_gt = COCO(str(annotation_path))
+    coco_gt.label2cat = {0: 1}
+    evaluator = CocoEvaluator(coco_gt, ["keypoints"])
+    keypoints = np.asarray(coco_gt.anns[1]["keypoints"], dtype=np.float32).reshape(1, 25, 3)
+
+    evaluator.update(
+        {
+            1: {
+                "boxes": torch.tensor([[10.0, 20.0, 60.0, 80.0]], dtype=torch.float32),
+                "scores": torch.tensor([0.99], dtype=torch.float32),
+                "labels": torch.tensor([0], dtype=torch.int64),
+                "keypoints": torch.as_tensor(keypoints, dtype=torch.float32),
+            }
+        }
+    )
+    evaluator.synchronize_between_processes()
+    evaluator.accumulate()
+
+    stats = evaluator.coco_eval["keypoints"].stats
+    assert np.isfinite(stats[0])
+
+
+def test_coco_evaluator_warns_once_per_custom_keypoint_count(tmp_path: Path) -> None:
+    """Repeated evaluator construction should not spam the same custom OKS fallback warning."""
+    annotation_path = tmp_path / "custom_keypoints_val.json"
+    _write_person_keypoint_coco(annotation_path, keypoint_count=25)
+    coco_gt = COCO(str(annotation_path))
+
+    coco_eval_module._WARNED_CUSTOM_KEYPOINT_OKS_COUNTS.clear()
+    try:
+        with patch.object(coco_eval_module.logger, "warning") as warning:
+            CocoEvaluator(coco_gt, ["keypoints"])
+            CocoEvaluator(coco_gt, ["keypoints"])
+    finally:
+        coco_eval_module._WARNED_CUSTOM_KEYPOINT_OKS_COUNTS.clear()
+
+    warning.assert_called_once()
+
+
+def test_coco_evaluator_rejects_mismatched_custom_oks_sigmas(tmp_path: Path) -> None:
+    """Explicit OKS sigmas must match the dataset keypoint count."""
+    annotation_path = tmp_path / "custom_keypoints_val.json"
+    _write_person_keypoint_coco(annotation_path, keypoint_count=25)
+    coco_gt = COCO(str(annotation_path))
+
+    with pytest.raises(ValueError, match="keypoint_oks_sigmas length 17 does not match dataset keypoint count 25"):
+        CocoEvaluator(coco_gt, ["keypoints"], keypoint_oks_sigmas=[0.05] * 17)
+
+
+def test_coco_evaluator_keypoints_handles_mixed_counts_and_multi_instance_image(tmp_path: Path) -> None:
+    """Mixed keypoint-count categories should evaluate by group instead of being skipped."""
+    annotation_path = tmp_path / "mixed_keypoints_val.json"
+    _write_mixed_keypoint_coco(annotation_path)
+    coco_gt = COCO(str(annotation_path))
+    coco_gt.label2cat = {0: 1, 1: 2}
+    evaluator = CocoEvaluator(coco_gt, ["keypoints"], keypoint_oks_sigmas=[0.05] * 21)
+
+    padded_keypoints = np.zeros((2, 21, 3), dtype=np.float32)
+    for detection_idx, annotation in enumerate(coco_gt.dataset["annotations"]):
+        keypoints = np.asarray(annotation["keypoints"], dtype=np.float32).reshape(-1, 3)
+        padded_keypoints[detection_idx, : keypoints.shape[0]] = keypoints
+
+    evaluator.update(
+        {
+            1: {
+                "boxes": torch.tensor([[10.0, 20.0, 30.0, 40.0], [50.0, 60.0, 70.0, 80.0]], dtype=torch.float32),
+                "scores": torch.tensor([0.99, 0.98], dtype=torch.float32),
+                "labels": torch.tensor([0, 1], dtype=torch.int64),
+                "keypoints": torch.as_tensor(padded_keypoints, dtype=torch.float32),
+            }
+        }
+    )
+
+    results = evaluator.coco_results["keypoints"]
+    assert len(results) == 2
+    assert len(results[0]["keypoints"]) == 4 * 3
+    assert len(results[1]["keypoints"]) == 21 * 3
+
+    evaluator.synchronize_between_processes()
+    evaluator.accumulate()
+
+    grouped_eval = evaluator.coco_eval["keypoints"]
+    assert len(grouped_eval.evals) == 2
+    stats = grouped_eval.stats
+    assert stats.shape == (10,)
+    assert np.isfinite(stats[0])
+
+
+def test_coco_evaluator_backfills_missing_num_keypoints(tmp_path: Path) -> None:
+    """Keypoint GT without `num_keypoints` should not be ignored during OKS evaluation."""
+    annotation_path = tmp_path / "person_keypoints_val2017.json"
+    _write_person_keypoint_coco(annotation_path, include_num_keypoints=False)
+    coco_gt = COCO(str(annotation_path))
+    assert "num_keypoints" not in coco_gt.anns[1]
+
+    evaluator = CocoEvaluator(coco_gt, ["keypoints"])
+
+    assert evaluator.coco_gt.anns[1]["num_keypoints"] == 17
+
+
+def test_coco_evaluator_handles_empty_keypoint_predictions(tmp_path: Path) -> None:
+    """Keypoint evaluation should handle images with no detections."""
+    annotation_path = tmp_path / "person_keypoints_val2017.json"
+    _write_person_keypoint_coco(annotation_path)
+    coco_gt = COCO(str(annotation_path))
+    evaluator = CocoEvaluator(coco_gt, ["keypoints"])
+
+    evaluator.update(
+        {
+            1: {
+                "boxes": torch.zeros((0, 4), dtype=torch.float32),
+                "scores": torch.zeros((0,), dtype=torch.float32),
+                "labels": torch.zeros((0,), dtype=torch.int64),
+                "keypoints": torch.zeros((0, 17, 3), dtype=torch.float32),
+            }
+        }
+    )
+    evaluator.synchronize_between_processes()
+    evaluator.accumulate()
+
+    stats = evaluator.coco_eval["keypoints"].stats
+    assert stats.shape == (10,)
+
+
+class TestSynchronizeBetweenProcesses:
+    """synchronize_between_processes() deduplicates DT when DDP padding repeats image_ids."""
+
+    def _make_evaluator(self, tmp_path: Path) -> CocoEvaluator:
+        """Return a single-annotation evaluator with label2cat identity mapping."""
+        annotation_path = tmp_path / "kp.json"
+        _write_person_keypoint_coco(annotation_path)
+        coco_gt = COCO(str(annotation_path))
+        coco_gt.label2cat = {0: 1}
+        return CocoEvaluator(coco_gt, ["keypoints"])
+
+    def _pred(self, image_id: int, score: float = 0.99) -> dict:
+        """Single detection prediction dict for image_id."""
+        kp = np.zeros((1, 17, 3), dtype=np.float32)
+        return {
+            image_id: {
+                "boxes": torch.tensor([[10.0, 20.0, 60.0, 80.0]]),
+                "scores": torch.tensor([score]),
+                "labels": torch.tensor([0], dtype=torch.long),
+                "keypoints": torch.as_tensor(kp),
+            }
+        }
+
+    def test_single_gpu_no_dedup_needed(self, tmp_path: Path) -> None:
+        """Single-GPU path (world_size=1): all_gather returns one-element list; all results preserved."""
+        ev = self._make_evaluator(tmp_path)
+        ev.update(self._pred(1))
+
+        with patch("rfdetr.evaluation.coco_eval.all_gather", side_effect=lambda x: [x]):
+            ev.synchronize_between_processes()
+
+        assert ev.img_ids == [1]
+        assert len(ev.coco_results["keypoints"]) == 1
+
+    def test_no_overlap_across_ranks_all_results_kept(self, tmp_path: Path) -> None:
+        """When image_ids are disjoint across ranks, all predictions are preserved."""
+        ev = self._make_evaluator(tmp_path)
+        # Simulate rank 0 has already called update() with image_id=1
+        ev.img_ids = [1]
+        ev.coco_results["keypoints"] = [{"image_id": 1, "category_id": 1, "keypoints": [], "score": 0.9}]
+
+        # all_gather returns rank-0 list + rank-1 list (no overlap)
+        rank1_ids = [2]
+        rank1_results = [{"image_id": 2, "category_id": 1, "keypoints": [], "score": 0.8}]
+        call_count = [0]
+
+        def _all_gather(x: list) -> list:
+            call_count[0] += 1
+            if call_count[0] == 1:
+                return [x, rank1_ids]
+            return [x, rank1_results]
+
+        with patch("rfdetr.evaluation.coco_eval.all_gather", side_effect=_all_gather):
+            ev.synchronize_between_processes()
+
+        assert sorted(ev.img_ids) == [1, 2]
+        image_ids_in_results = [r["image_id"] for r in ev.coco_results["keypoints"]]
+        assert sorted(image_ids_in_results) == [1, 2]
+
+    def test_ddp_padding_duplicate_image_id_deduped(self, tmp_path: Path) -> None:
+        """DDP DistributedSampler padding: same image_id on two ranks → only rank-0 results kept."""
+        ev = self._make_evaluator(tmp_path)
+        # image_id=1 on BOTH ranks (padding), image_id=2 only on rank-1
+        rank0_ids = [1]
+        rank1_ids = [1, 2]
+        rank0_results = [{"image_id": 1, "category_id": 1, "keypoints": [0.9], "score": 0.9}]
+        rank1_results = [
+            {"image_id": 1, "category_id": 1, "keypoints": [0.9], "score": 0.9},  # duplicate
+            {"image_id": 2, "category_id": 1, "keypoints": [0.5], "score": 0.8},
+        ]
+        call_count = [0]
+
+        def _all_gather(x: list) -> list:
+            call_count[0] += 1
+            if call_count[0] == 1:
+                return [rank0_ids, rank1_ids]
+            return [rank0_results, rank1_results]
+
+        with patch("rfdetr.evaluation.coco_eval.all_gather", side_effect=_all_gather):
+            ev.synchronize_between_processes()
+
+        assert sorted(ev.img_ids) == [1, 2]
+        image_ids_in_results = [r["image_id"] for r in ev.coco_results["keypoints"]]
+        # image_id=1 from rank-0 only (not duplicated), image_id=2 from rank-1
+        assert image_ids_in_results.count(1) == 1, "image_id=1 must appear exactly once (no DDP duplicate)"
+        assert image_ids_in_results.count(2) == 1
+
+    def test_ddp_padding_rank0_predictions_chosen_over_rank1(self, tmp_path: Path) -> None:
+        """When image_id appears on rank-0 and rank-1, rank-0's prediction is kept (first-wins)."""
+        ev = self._make_evaluator(tmp_path)
+        rank0_ids = [1]
+        rank1_ids = [1]
+        rank0_results = [{"image_id": 1, "category_id": 1, "keypoints": [], "score": 0.9}]
+        rank1_results = [{"image_id": 1, "category_id": 1, "keypoints": [], "score": 0.5}]
+        call_count = [0]
+
+        def _all_gather(x: list) -> list:
+            call_count[0] += 1
+            if call_count[0] == 1:
+                return [rank0_ids, rank1_ids]
+            return [rank0_results, rank1_results]
+
+        with patch("rfdetr.evaluation.coco_eval.all_gather", side_effect=_all_gather):
+            ev.synchronize_between_processes()
+
+        assert len(ev.coco_results["keypoints"]) == 1
+        assert ev.coco_results["keypoints"][0]["score"] == pytest.approx(0.9), "rank-0 prediction must win"
+
+    def test_multiple_detections_same_image_all_kept(self, tmp_path: Path) -> None:
+        """Multiple DT per image (multi-instance) on the owning rank are all preserved."""
+        ev = self._make_evaluator(tmp_path)
+        # rank-0 has 3 detections for image_id=1 (3 distinct instances)
+        rank0_ids = [1]
+        rank0_results = [
+            {"image_id": 1, "category_id": 1, "keypoints": [], "score": 0.9},
+            {"image_id": 1, "category_id": 1, "keypoints": [], "score": 0.8},
+            {"image_id": 1, "category_id": 1, "keypoints": [], "score": 0.7},
+        ]
+        call_count = [0]
+
+        def _all_gather(x: list) -> list:
+            call_count[0] += 1
+            if call_count[0] == 1:
+                return [rank0_ids]
+            return [rank0_results]
+
+        with patch("rfdetr.evaluation.coco_eval.all_gather", side_effect=_all_gather):
+            ev.synchronize_between_processes()
+
+        assert len(ev.coco_results["keypoints"]) == 3, "all 3 per-image detections must be kept"
+
+
+def test_coco_evaluator_skips_unmapped_labels_when_label2cat_is_present(tmp_path: Path) -> None:
+    """A non-identity label2cat map should not fall back to raw category IDs for unmapped labels."""
+    annotation_path = tmp_path / "person_keypoints_val2017.json"
+    _write_person_keypoint_coco(annotation_path)
+    coco_gt = COCO(str(annotation_path))
+    coco_gt.label2cat = {1: 1}
+    evaluator = CocoEvaluator(coco_gt, ["keypoints"])
+    keypoints = np.asarray(coco_gt.anns[1]["keypoints"], dtype=np.float32).reshape(1, 17, 3)
+
+    evaluator.update(
+        {
+            1: {
+                "boxes": torch.tensor(
+                    [[10.0, 20.0, 60.0, 80.0], [10.0, 20.0, 60.0, 80.0]],
+                    dtype=torch.float32,
+                ),
+                "scores": torch.tensor([0.99, 0.98], dtype=torch.float32),
+                "labels": torch.tensor([0, 1], dtype=torch.int64),
+                "keypoints": torch.as_tensor(np.concatenate([keypoints, keypoints], axis=0), dtype=torch.float32),
+            }
+        }
+    )
+
+    results = evaluator.coco_results["keypoints"]
+    assert len(results) == 1
+    assert results[0]["category_id"] == 1
+
+
+def test_patched_pycocotools_summarize_raises_on_unknown_iou_type() -> None:
+    """patched_pycocotools_summarize raises ValueError for an unrecognised iouType."""
+    from unittest.mock import MagicMock
+
+    mock_eval = MagicMock()
+    mock_eval.eval = {"precision": np.zeros((1, 1, 1, 1, 1)), "recall": np.zeros((1, 1, 1, 1))}
+    mock_eval.params.iouType = "custom"
+
+    with pytest.raises(ValueError, match="Unknown iou type custom"):
+        coco_eval_module.patched_pycocotools_summarize(mock_eval)

tests/evaluation/test_evaluation.py

@@ -0,0 +1,471 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+from unittest.mock import patch
+
+import numpy as np
+import pytest
+import torch
+
+from rfdetr.evaluation.matching import (
+    _compute_mask_iou,
+    _match_single_class,
+    build_matching_data,
+    distributed_merge_matching_data,
+    init_matching_accumulator,
+    merge_matching_data,
+)
+
+# ---------------------------------------------------------------------------
+# _compute_mask_iou
+# ---------------------------------------------------------------------------
+
+
+class TestComputeMaskIou:
+    """Unit tests for the private _compute_mask_iou helper."""
+
+    @staticmethod
+    def _bool_mask(h: int, w: int, rows: slice, cols: slice) -> torch.Tensor:
+        """Return a [1, h, w] boolean mask with the specified region set to True."""
+        m = torch.zeros(h, w, dtype=torch.bool)
+        m[rows, cols] = True
+        return m.unsqueeze(0)
+
+    def test_identical_masks_give_iou_one(self) -> None:
+        """Masks that are identical should produce IoU of exactly 1.0."""
+        mask = self._bool_mask(4, 4, slice(0, 2), slice(0, 2))  # [1, 4, 4]
+        result = _compute_mask_iou(mask, mask)
+        assert result.shape == (1, 1)
+        assert float(result[0, 0]) == pytest.approx(1.0)
+
+    def test_disjoint_masks_give_iou_zero(self) -> None:
+        """Non-overlapping masks should produce IoU of 0.0."""
+        pred = self._bool_mask(4, 4, slice(0, 2), slice(0, 2))
+        gt = self._bool_mask(4, 4, slice(2, 4), slice(2, 4))
+        result = _compute_mask_iou(pred, gt)
+        assert float(result[0, 0]) == pytest.approx(0.0)
+
+    def test_known_partial_overlap(self) -> None:
+        """50% row overlap on a 4x4 grid: inter=4, union=12, IoU=1/3."""
+        pred = torch.zeros(1, 4, 4, dtype=torch.bool)
+        pred[0, :2, :] = True  # rows 0-1: 8 px
+        gt = torch.zeros(1, 4, 4, dtype=torch.bool)
+        gt[0, 1:3, :] = True  # rows 1-2: 8 px — 4 px of overlap at row 1
+        result = _compute_mask_iou(pred, gt)
+        assert float(result[0, 0]) == pytest.approx(4.0 / 12.0)
+
+    def test_empty_masks_return_zero_without_error(self) -> None:
+        """All-zero masks must yield IoU 0.0 (no divide-by-zero)."""
+        pred = torch.zeros(1, 4, 4, dtype=torch.bool)
+        gt = torch.zeros(1, 4, 4, dtype=torch.bool)
+        result = _compute_mask_iou(pred, gt)
+        assert float(result[0, 0]) == pytest.approx(0.0)
+
+    def test_output_shape_is_n_by_m(self) -> None:
+        """Output shape must be [N, M] for N predictions and M ground truths."""
+        pred = torch.zeros(3, 4, 4, dtype=torch.bool)
+        gt = torch.zeros(5, 4, 4, dtype=torch.bool)
+        result = _compute_mask_iou(pred, gt)
+        assert result.shape == (3, 5)
+
+
+# ---------------------------------------------------------------------------
+# _match_single_class
+# ---------------------------------------------------------------------------
+
+
+class TestMatchSingleClass:
+    """Unit tests for the private _match_single_class helper."""
+
+    @staticmethod
+    def _box(*coords: float) -> torch.Tensor:
+        """Return a [1, 4] float32 box tensor from (x1, y1, x2, y2)."""
+        return torch.tensor([list(coords)], dtype=torch.float32)
+
+    @staticmethod
+    def _boxes(*rows: list[float]) -> torch.Tensor:
+        """Return an [N, 4] float32 tensor from a sequence of [x1,y1,x2,y2] rows."""
+        return torch.tensor(list(rows), dtype=torch.float32)
+
+    def _run(
+        self,
+        pred_scores: torch.Tensor,
+        pred_items: torch.Tensor,
+        gt_items: torch.Tensor,
+        gt_crowd: torch.Tensor | None = None,
+        iou_threshold: float = 0.5,
+        iou_type: str = "bbox",
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, int]:
+        if gt_crowd is None:
+            gt_crowd = torch.zeros(len(gt_items), dtype=torch.bool)
+        return _match_single_class(pred_scores, pred_items, gt_items, gt_crowd, iou_threshold, iou_type)
+
+    def test_perfect_overlap_is_tp(self) -> None:
+        """A prediction that perfectly overlaps the GT box is a true positive."""
+        scores = torch.tensor([0.9])
+        box = self._box(0, 0, 10, 10)
+        _, matches, ignore, total_gt = self._run(scores, box, box)
+        assert matches[0] == 1
+        assert not ignore[0]
+        assert total_gt == 1
+
+    def test_disjoint_box_is_fp(self) -> None:
+        """A prediction with no overlap with the GT box is a false positive."""
+        scores = torch.tensor([0.9])
+        pred = self._box(0, 0, 10, 10)
+        gt = self._box(50, 50, 60, 60)
+        _, matches, ignore, total_gt = self._run(scores, pred, gt)
+        assert matches[0] == 0
+        assert not ignore[0]
+        assert total_gt == 1
+
+    def test_iou_below_threshold_is_fp(self) -> None:
+        """A detection with IoU < threshold must be marked as FP."""
+        scores = torch.tensor([0.9])
+        pred = self._box(0, 0, 5, 10)  # area = 50
+        gt = self._box(6, 0, 10, 10)  # area = 40 — no overlap
+        _, matches, _, _ = self._run(scores, pred, gt, iou_threshold=0.5)
+        assert matches[0] == 0
+
+    def test_greedy_matching_higher_score_wins(self) -> None:
+        """When two predictions compete for one GT, the higher-score pred wins."""
+        # Sorted descending: [0.9, 0.5] -> first gets TP, second gets FP.
+        scores = torch.tensor([0.5, 0.9])
+        preds = self._boxes([0, 0, 10, 10], [0, 0, 10, 10])
+        gt = self._box(0, 0, 10, 10)
+        scores_out, matches, _, _ = self._run(scores, preds, gt)
+        assert list(scores_out) == pytest.approx([0.9, 0.5])
+        assert matches[0] == 1  # highest score -> TP
+        assert matches[1] == 0  # lower score -> FP
+
+    def test_crowd_gt_match_is_ignored_not_fp(self) -> None:
+        """A detection matched to a crowd GT is ignored, not a false positive."""
+        scores = torch.tensor([0.9])
+        box = self._box(0, 0, 10, 10)
+        gt_crowd = torch.tensor([True])
+        _, matches, ignore, total_gt = self._run(scores, box, box, gt_crowd=gt_crowd)
+        assert matches[0] == 0  # not TP
+        assert ignore[0]  # ignored -> not counted as FP
+        assert total_gt == 0  # crowd GT excluded from denominator
+
+    def test_non_crowd_gt_counts_in_total_gt(self) -> None:
+        """Non-crowd GTs are counted in total_gt."""
+        scores = torch.tensor([0.9])
+        box = self._box(0, 0, 10, 10)
+        gt_crowd = torch.tensor([False])
+        _, _, _, total_gt = self._run(scores, box, box, gt_crowd=gt_crowd)
+        assert total_gt == 1
+
+    def test_mixed_crowd_only_non_crowd_in_total_gt(self) -> None:
+        """Only non-crowd instances contribute to total_gt."""
+        scores = torch.tensor([0.9])
+        pred = self._box(0, 0, 5, 5)  # overlaps neither GT significantly
+        gt_boxes = self._boxes([0, 0, 10, 10], [20, 20, 30, 30])
+        gt_crowd = torch.tensor([False, True])  # second GT is crowd
+        _, _, _, total_gt = self._run(scores, pred, gt_boxes, gt_crowd=gt_crowd)
+        assert total_gt == 1
+
+    def test_scores_returned_in_descending_order(self) -> None:
+        """Output scores must be sorted in descending order."""
+        scores = torch.tensor([0.3, 0.9, 0.6])
+        preds = self._boxes([0, 0, 10, 10], [20, 20, 30, 30], [40, 40, 50, 50])
+        gt = self._box(20, 20, 30, 30)
+        scores_out, _, _, _ = self._run(scores, preds, gt)
+        assert list(scores_out) == pytest.approx([0.9, 0.6, 0.3])
+
+    def test_segm_iou_type_identical_masks_is_tp(self) -> None:
+        """Identical masks with iou_type='segm' should yield a TP."""
+        mask = torch.ones(1, 4, 4, dtype=torch.bool)
+        scores = torch.tensor([0.9])
+        gt_crowd = torch.tensor([False])
+        _, matches, _, total_gt = _match_single_class(scores, mask, mask, gt_crowd, 0.5, "segm")
+        assert matches[0] == 1
+        assert total_gt == 1
+
+
+# ---------------------------------------------------------------------------
+# build_matching_data
+# ---------------------------------------------------------------------------
+
+
+class TestBuildMatchingData:
+    """Unit tests for build_matching_data()."""
+
+    @staticmethod
+    def _make_pred(
+        boxes: list,
+        scores: list,
+        labels: list,
+        masks: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        d: dict[str, torch.Tensor] = {
+            "boxes": torch.tensor(boxes, dtype=torch.float32).reshape(-1, 4),
+            "scores": torch.tensor(scores, dtype=torch.float32),
+            "labels": torch.tensor(labels, dtype=torch.int64),
+        }
+        if masks is not None:
+            d["masks"] = masks
+        return d
+
+    @staticmethod
+    def _make_target(
+        boxes: list,
+        labels: list,
+        iscrowd: list | None = None,
+        masks: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        d: dict[str, torch.Tensor] = {
+            "boxes": torch.tensor(boxes, dtype=torch.float32).reshape(-1, 4),
+            "labels": torch.tensor(labels, dtype=torch.int64),
+        }
+        if iscrowd is not None:
+            d["iscrowd"] = torch.tensor(iscrowd, dtype=torch.int64)
+        if masks is not None:
+            d["masks"] = masks
+        return d
+
+    def test_output_has_required_keys(self) -> None:
+        """Every class entry must contain scores, matches, ignore, total_gt."""
+        pred = self._make_pred([[0, 0, 10, 10]], [0.9], [0])
+        target = self._make_target([[0, 0, 10, 10]], [0])
+        result = build_matching_data([pred], [target])
+        assert 0 in result
+        assert set(result[0].keys()) == {"scores", "matches", "ignore", "total_gt"}
+
+    def test_perfect_detection_is_tp(self) -> None:
+        """A pred box identical to the GT box must be a TP."""
+        pred = self._make_pred([[0, 0, 10, 10]], [0.9], [0])
+        target = self._make_target([[0, 0, 10, 10]], [0])
+        result = build_matching_data([pred], [target])
+        assert result[0]["matches"][0] == 1
+        assert result[0]["total_gt"] == 1
+
+    def test_disjoint_box_is_fp(self) -> None:
+        """A pred box with no overlap against any GT must be a FP."""
+        pred = self._make_pred([[0, 0, 10, 10]], [0.9], [0])
+        target = self._make_target([[50, 50, 60, 60]], [0])
+        result = build_matching_data([pred], [target])
+        assert result[0]["matches"][0] == 0
+        assert result[0]["total_gt"] == 1
+
+    def test_no_predictions_records_total_gt_only(self) -> None:
+        """With no preds for a class, total_gt is recorded but scores list is empty."""
+        pred = self._make_pred([], [], [])
+        target = self._make_target([[0, 0, 10, 10]], [0])
+        result = build_matching_data([pred], [target])
+        assert result[0]["total_gt"] == 1
+        assert len(result[0]["scores"]) == 0
+
+    def test_no_gts_all_predictions_are_fp(self) -> None:
+        """With no GTs for a class, all predictions are FP and total_gt is 0."""
+        pred = self._make_pred([[0, 0, 10, 10]], [0.9], [0])
+        target = self._make_target([], [])
+        result = build_matching_data([pred], [target])
+        assert result[0]["matches"][0] == 0
+        assert result[0]["total_gt"] == 0
+
+    def test_multi_class_results_are_separated(self) -> None:
+        """Two classes in the same image must be tracked independently."""
+        pred = self._make_pred([[0, 0, 10, 10], [20, 20, 30, 30]], [0.9, 0.8], [0, 1])
+        target = self._make_target([[0, 0, 10, 10], [20, 20, 30, 30]], [0, 1])
+        result = build_matching_data([pred], [target])
+        assert result[0]["matches"][0] == 1
+        assert result[1]["matches"][0] == 1
+        assert result[0]["total_gt"] == 1
+        assert result[1]["total_gt"] == 1
+
+    def test_multi_image_batch_accumulates(self) -> None:
+        """Two-image batch must concatenate scores and sum total_gt."""
+        pred1 = self._make_pred([[0, 0, 10, 10]], [0.9], [0])
+        target1 = self._make_target([[0, 0, 10, 10]], [0])
+        pred2 = self._make_pred([[50, 50, 60, 60]], [0.8], [0])
+        target2 = self._make_target([[50, 50, 60, 60]], [0])
+        result = build_matching_data([pred1, pred2], [target1, target2])
+        assert len(result[0]["scores"]) == 2
+        assert result[0]["total_gt"] == 2
+
+    def test_crowd_gt_excluded_from_total_and_detection_ignored(self) -> None:
+        """A pred matched to a crowd GT must be ignored; crowd GT not counted."""
+        pred = self._make_pred([[0, 0, 10, 10]], [0.9], [0])
+        target = self._make_target([[0, 0, 10, 10]], [0], iscrowd=[1])
+        result = build_matching_data([pred], [target])
+        assert result[0]["total_gt"] == 0
+        assert result[0]["ignore"][0]
+        assert result[0]["matches"][0] == 0
+
+    def test_mixed_crowd_non_crowd_gts(self) -> None:
+        """Pred matched to non-crowd GT is TP; crowd GT not counted in total_gt."""
+        pred = self._make_pred([[0, 0, 10, 10]], [0.9], [0])
+        target = self._make_target([[0, 0, 10, 10], [20, 20, 30, 30]], [0, 0], iscrowd=[0, 1])
+        result = build_matching_data([pred], [target])
+        assert result[0]["total_gt"] == 1
+        assert result[0]["matches"][0] == 1
+        assert not result[0]["ignore"][0]
+
+    def test_segmentation_iou_type_identical_masks(self) -> None:
+        """iou_type='segm' path with identical masks must yield a TP."""
+        mask = torch.ones(1, 8, 8, dtype=torch.bool)
+        pred = {
+            "boxes": torch.zeros(1, 4),
+            "scores": torch.tensor([0.9]),
+            "labels": torch.tensor([0]),
+            "masks": mask,
+        }
+        target = {
+            "boxes": torch.zeros(1, 4),
+            "labels": torch.tensor([0]),
+            "masks": mask,
+        }
+        result = build_matching_data([pred], [target], iou_type="segm")
+        assert result[0]["matches"][0] == 1
+        assert result[0]["total_gt"] == 1
+
+    def test_segmentation_missing_masks_raises_value_error(self) -> None:
+        """iou_type='segm' without masks must raise ValueError."""
+        pred = self._make_pred([[0, 0, 10, 10]], [0.9], [0])
+        target = self._make_target([[0, 0, 10, 10]], [0])
+        with pytest.raises(ValueError, match="masks"):
+            build_matching_data([pred], [target], iou_type="segm")
+
+    def test_class_only_in_predictions_is_tracked_as_fp(self) -> None:
+        """A class seen only in predictions (no GT) must appear in output as FP."""
+        pred = self._make_pred([[0, 0, 10, 10]], [0.9], [99])
+        target = self._make_target([[0, 0, 10, 10]], [0])
+        result = build_matching_data([pred], [target])
+        assert 99 in result
+        assert result[99]["total_gt"] == 0
+        assert result[99]["matches"][0] == 0
+
+
+# ---------------------------------------------------------------------------
+# Helper shared by TestMergeMatchingData and TestDistributedMergeMatchingData
+# (used by multiple classes, so module-level rather than a staticmethod)
+# ---------------------------------------------------------------------------
+
+
+def _make_matching_entry(
+    scores: list,
+    matches: list,
+    ignore: list,
+    total_gt: int,
+) -> dict:
+    """Return a compact matching dict as produced by ``build_matching_data()``."""
+    return {
+        "scores": np.array(scores, dtype=np.float32),
+        "matches": np.array(matches, dtype=np.int64),
+        "ignore": np.array(ignore, dtype=bool),
+        "total_gt": total_gt,
+    }
+
+
+class TestInitMatchingAccumulator:
+    """init_matching_accumulator() returns a correct empty accumulator."""
+
+    def test_returns_empty_dict(self) -> None:
+        """Returns an empty dict."""
+        assert init_matching_accumulator() == {}
+
+    def test_returned_dict_is_mutable_via_merge(self) -> None:
+        """The returned dict can be populated by merge_matching_data."""
+        acc = init_matching_accumulator()
+        merge_matching_data(acc, {0: _make_matching_entry([0.9], [1], [False], 1)})
+        assert 0 in acc
+
+
+class TestMergeMatchingData:
+    """merge_matching_data() correctly accumulates per-class matching dicts."""
+
+    def test_empty_accumulator_copies_new_data(self) -> None:
+        """First merge populates the accumulator with the batch data."""
+        data = _make_matching_entry([0.9, 0.8], [1, 0], [False, False], 1)
+        acc = merge_matching_data({}, {0: data})
+        np.testing.assert_allclose(acc[0]["scores"], [0.9, 0.8], rtol=1e-6)
+        np.testing.assert_array_equal(acc[0]["matches"], [1, 0])
+        assert acc[0]["total_gt"] == 1
+
+    def test_second_merge_concatenates_arrays_and_sums_total_gt(self) -> None:
+        """Merging a second batch appends scores/matches/ignore and sums total_gt."""
+        acc: dict = {}
+        merge_matching_data(acc, {0: _make_matching_entry([0.9], [1], [False], 2)})
+        merge_matching_data(acc, {0: _make_matching_entry([0.7], [0], [False], 3)})
+        np.testing.assert_allclose(acc[0]["scores"], [0.9, 0.7], rtol=1e-6)
+        np.testing.assert_array_equal(acc[0]["matches"], [1, 0])
+        assert acc[0]["total_gt"] == 5
+
+    def test_new_class_added_independently(self) -> None:
+        """A class not yet in the accumulator is added without touching others."""
+        acc = {0: _make_matching_entry([0.9], [1], [False], 1)}
+        merge_matching_data(acc, {1: _make_matching_entry([0.5], [0], [False], 2)})
+        assert acc[0]["total_gt"] == 1
+        assert acc[1]["total_gt"] == 2
+
+    def test_returns_same_accumulator_object(self) -> None:
+        """merge_matching_data returns the same dict it was given (in-place)."""
+        acc: dict = {}
+        result = merge_matching_data(acc, {})
+        assert result is acc
+
+    def test_no_op_when_new_data_is_empty(self) -> None:
+        """Merging an empty batch leaves the accumulator unchanged."""
+        acc = {0: _make_matching_entry([0.9], [1], [False], 1)}
+        merge_matching_data(acc, {})
+        assert len(acc) == 1
+        assert acc[0]["total_gt"] == 1
+
+    def test_copied_arrays_are_independent_of_source(self) -> None:
+        """Mutating the source entry after the first merge must not corrupt acc."""
+        data = _make_matching_entry([0.9], [1], [False], 1)
+        acc: dict = {}
+        merge_matching_data(acc, {0: data})
+        data["scores"][0] = 0.0
+        assert acc[0]["scores"][0] == pytest.approx(0.9)
+
+    def test_multiple_classes_in_single_batch_all_added(self) -> None:
+        """All classes present in a single batch are merged into the accumulator."""
+        batch = {
+            0: _make_matching_entry([0.9], [1], [False], 1),
+            1: _make_matching_entry([0.8], [0], [False], 2),
+        }
+        acc = merge_matching_data({}, batch)
+        assert set(acc.keys()) == {0, 1}
+        assert acc[0]["total_gt"] == 1
+        assert acc[1]["total_gt"] == 2
+
+
+class TestDistributedMergeMatchingData:
+    """distributed_merge_matching_data() gathers and merges across DDP ranks."""
+
+    def test_single_rank_returns_same_content(self) -> None:
+        """In single-process mode (world_size=1), data passes through unchanged."""
+        local_data = {0: _make_matching_entry([0.9], [1], [False], 1)}
+        result = distributed_merge_matching_data(local_data)
+        np.testing.assert_allclose(result[0]["scores"], [0.9], rtol=1e-6)
+        assert result[0]["total_gt"] == 1
+
+    def test_two_ranks_disjoint_classes(self) -> None:
+        """Two ranks with disjoint classes -> merged result contains both."""
+        rank0 = {0: _make_matching_entry([0.9], [1], [False], 1)}
+        rank1 = {1: _make_matching_entry([0.7], [0], [False], 2)}
+        with patch("rfdetr.evaluation.matching.all_gather", return_value=[rank0, rank1]):
+            result = distributed_merge_matching_data(rank0)
+        assert set(result.keys()) == {0, 1}
+        assert result[0]["total_gt"] == 1
+        assert result[1]["total_gt"] == 2
+
+    def test_two_ranks_overlapping_class_concatenates(self) -> None:
+        """Two ranks sharing class 0 -> arrays concatenated, total_gt summed."""
+        rank0 = {0: _make_matching_entry([0.9], [1], [False], 2)}
+        rank1 = {0: _make_matching_entry([0.7, 0.5], [0, 1], [False, False], 3)}
+        with patch("rfdetr.evaluation.matching.all_gather", return_value=[rank0, rank1]):
+            result = distributed_merge_matching_data(rank0)
+        np.testing.assert_allclose(result[0]["scores"], [0.9, 0.7, 0.5], rtol=1e-6)
+        assert result[0]["total_gt"] == 5
+
+    def test_returns_new_dict_not_input(self) -> None:
+        """Result is a new dict, not a reference to the local input."""
+        local_data = {0: _make_matching_entry([0.9], [1], [False], 1)}
+        result = distributed_merge_matching_data(local_data)
+        assert result is not local_data

tests/evaluation/test_keypoint_oks.py

@@ -0,0 +1,504 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit and integration tests for MetricKeypointOKS.
+
+Integration tests (class TestOKSValues) build a minimal COCO ground-truth object and feed known predictions so that
+expected mAP values can be derived by hand without running model inference.  They are the first line of defence against
+silent metric-computation regressions.
+"""
+
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+import torch
+from faster_coco_eval import COCO
+
+from rfdetr.evaluation.keypoint_oks import MetricKeypointOKS
+
+# ---------------------------------------------------------------------------
+# Shared helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_coco_gt() -> MagicMock:
+    """Return a minimal COCO ground-truth mock (for unit tests that patch evaluator)."""
+    return MagicMock(name="coco_gt")
+
+
+def _make_predictions(image_id: int = 1, num_dets: int = 1, num_keypoints: int = 3) -> dict:
+    """Return a single-image prediction dict with zero-valued tensors."""
+    return {
+        image_id: {
+            "boxes": torch.zeros(num_dets, 4),
+            "scores": torch.ones(num_dets),
+            "labels": torch.zeros(num_dets, dtype=torch.long),
+            "keypoints": torch.zeros(num_dets, num_keypoints, 3),
+        }
+    }
+
+
+def _make_evaluator_mock(stats: list[float]) -> MagicMock:
+    """Return a CocoEvaluator mock that returns the given stats array.
+
+    Stats list must have exactly 10 elements matching _summarizeKps() output shape.
+    """
+    assert len(stats) == 10, f"_make_evaluator_mock: expected 10 stats, got {len(stats)}"
+    evaluator = MagicMock(name="evaluator")
+    evaluator.coco_eval = {"keypoints": MagicMock(stats=np.array(stats, dtype=np.float32))}
+    return evaluator
+
+
+def _build_coco_gt(
+    num_keypoints: int,
+    gt_keypoints: list[float],
+    area: float = 2500.0,
+    bbox: list[float] | None = None,
+) -> COCO:
+    """Build a minimal COCO GT object with a single annotation.
+
+    Args:
+        num_keypoints: Number of keypoints per instance.
+        gt_keypoints: Flat COCO keypoint list [x0,y0,v0, x1,y1,v1, ...].
+            Visibility values should be 2 (labelled+visible).
+        area: Ground-truth object area for OKS normalisation.
+        bbox: Ground-truth bounding box [x, y, w, h].  Defaults to [0,0,100,100].
+
+    Returns:
+        A fully-indexed ``faster_coco_eval.COCO`` object.
+    """
+    if bbox is None:
+        bbox = [0.0, 0.0, 100.0, 100.0]
+    kp_names = [f"kp{i}" for i in range(num_keypoints)]
+    dataset = {
+        "images": [{"id": 1, "width": 100, "height": 100, "file_name": "img.jpg"}],
+        "categories": [{"id": 1, "name": "obj", "keypoints": kp_names, "skeleton": []}],
+        "annotations": [
+            {
+                "id": 1,
+                "image_id": 1,
+                "category_id": 1,
+                "keypoints": gt_keypoints,
+                "num_keypoints": num_keypoints,
+                "area": area,
+                "bbox": bbox,
+                "iscrowd": 0,
+            }
+        ],
+    }
+    coco_gt = COCO()
+    coco_gt.dataset = dataset
+    coco_gt.createIndex()
+    return coco_gt
+
+
+def _make_keypoint_prediction(
+    image_id: int,
+    kp_xy: list[tuple[float, float]],
+    category_id: int = 1,
+    score: float = 0.99,
+    box: list[float] | None = None,
+) -> dict:
+    """Build a per-image prediction dict for MetricKeypointOKS.update().
+
+    Args:
+        image_id: COCO image ID.
+        kp_xy: List of (x, y) coordinates — one per keypoint.  Visibility is
+            set to 1.0 for all.
+        category_id: Category label (raw COCO ID used here, no remapping).
+        score: Detection confidence score.
+        box: Bounding box [x1, y1, x2, y2] in pixel coords.  Defaults to the
+            full 100x100 image.
+
+    Returns:
+        Dict mapping ``image_id`` to a prediction dict accepted by
+        :meth:`~rfdetr.evaluation.keypoint_oks.MetricKeypointOKS.update`.
+    """
+    if box is None:
+        box = [0.0, 0.0, 100.0, 100.0]
+    num_kp = len(kp_xy)
+    kp_tensor = torch.zeros(1, num_kp, 3)
+    for i, (x, y) in enumerate(kp_xy):
+        kp_tensor[0, i, 0] = x
+        kp_tensor[0, i, 1] = y
+        kp_tensor[0, i, 2] = 1.0
+    return {
+        image_id: {
+            "boxes": torch.tensor([box], dtype=torch.float32),
+            "scores": torch.tensor([score]),
+            "labels": torch.tensor([category_id], dtype=torch.long),
+            "keypoints": kp_tensor,
+        }
+    }
+
+
+# ---------------------------------------------------------------------------
+# Unit tests (evaluator is mocked)
+# ---------------------------------------------------------------------------
+
+
+class TestHasUpdates:
+    """has_updates reflects whether any batch has been accumulated."""
+
+    def test_false_on_construction(self) -> None:
+        """Fresh metric reports no updates."""
+        metric = MetricKeypointOKS(_make_coco_gt())
+        assert metric.has_updates is False
+
+    def test_true_after_update(self) -> None:
+        """has_updates becomes True after any update() call."""
+        metric = MetricKeypointOKS(_make_coco_gt())
+        metric.update({1: {}})
+        assert metric.has_updates is True
+
+    def test_false_after_reset(self) -> None:
+        """has_updates returns False after reset() clears all batches."""
+        metric = MetricKeypointOKS(_make_coco_gt())
+        metric.update({1: {}})
+        metric.reset()
+        assert metric.has_updates is False
+
+
+class TestReset:
+    """Reset() clears all accumulated batches."""
+
+    def test_clears_all_batches(self) -> None:
+        """Reset() empties internal _batches list."""
+        metric = MetricKeypointOKS(_make_coco_gt())
+        metric.update(_make_predictions(image_id=1))
+        metric.update(_make_predictions(image_id=2))
+        metric.reset()
+        assert metric._batches == []
+
+    def test_idempotent_on_empty_state(self) -> None:
+        """Reset() on empty metric does not raise."""
+        metric = MetricKeypointOKS(_make_coco_gt())
+        metric.reset()
+        assert metric.has_updates is False
+
+
+class TestUpdate:
+    """Update() appends batches without merging or overwriting."""
+
+    def test_each_call_appends_one_batch(self) -> None:
+        """Two update() calls produce two entries in _batches."""
+        metric = MetricKeypointOKS(_make_coco_gt())
+        metric.update(_make_predictions(image_id=1))
+        metric.update(_make_predictions(image_id=2))
+        assert len(metric._batches) == 2
+
+    def test_same_image_id_in_two_batches_both_preserved(self) -> None:
+        """Predictions for the same image_id in separate batches are NOT overwritten."""
+        metric = MetricKeypointOKS(_make_coco_gt())
+        metric.update({5: {"scores": torch.tensor([0.9])}})
+        metric.update({5: {"scores": torch.tensor([0.3])}})
+        # Both batches must be preserved — not merged/overwritten
+        assert len(metric._batches) == 2
+        assert float(metric._batches[0][5]["scores"][0]) == pytest.approx(0.9)
+        assert float(metric._batches[1][5]["scores"][0]) == pytest.approx(0.3)
+
+    def test_empty_prediction_dict_appended_as_batch(self) -> None:
+        """Empty dict marks an image with no detections and is preserved as a batch."""
+        metric = MetricKeypointOKS(_make_coco_gt())
+        metric.update({42: {}})
+        assert len(metric._batches) == 1
+        assert metric._batches[0] == {42: {}}
+
+
+class TestCompute:
+    """Compute() delegates to CocoEvaluator and returns correct stat dict."""
+
+    def test_returns_correct_stat_keys(self) -> None:
+        """Compute() returns dict with map, map_50, map_75, mar keys."""
+        evaluator = _make_evaluator_mock([0.5, 0.7, 0.4, -1.0, -1.0, 0.6, -1.0, -1.0, -1.0, -1.0])
+        metric = MetricKeypointOKS(_make_coco_gt())
+        with patch("rfdetr.evaluation.keypoint_oks.CocoEvaluator", return_value=evaluator):
+            result = metric.compute()
+        assert set(result.keys()) == {"map", "map_50", "map_75", "mar"}
+
+    def test_maps_stats_indices_to_dict_keys(self) -> None:
+        """Compute() maps stats[0,1,2,5] to map, map_50, map_75, mar."""
+        evaluator = _make_evaluator_mock([0.42, 0.72, 0.31, -1.0, -1.0, 0.55, -1.0, -1.0, -1.0, -1.0])
+        metric = MetricKeypointOKS(_make_coco_gt())
+        with patch("rfdetr.evaluation.keypoint_oks.CocoEvaluator", return_value=evaluator):
+            result = metric.compute()
+        assert result["map"] == pytest.approx(0.42)
+        assert result["map_50"] == pytest.approx(0.72)
+        assert result["map_75"] == pytest.approx(0.31)
+        assert result["mar"] == pytest.approx(0.55)
+
+    def test_raises_on_wrong_stats_shape(self) -> None:
+        """Compute() raises AssertionError when stats array is not shape (10,).
+
+        _summarizeKps() always returns (10,); a shape mismatch signals a pycocotools contract violation (e.g. wrong
+        faster_coco_eval version) and must not silently produce incorrect metric values via index-out-of-bounds sentinel
+        fallback.
+        """
+        evaluator = MagicMock(name="evaluator")
+        evaluator.coco_eval = {"keypoints": MagicMock(stats=np.array([0.3], dtype=np.float32))}
+        metric = MetricKeypointOKS(_make_coco_gt())
+        with patch("rfdetr.evaluation.keypoint_oks.CocoEvaluator", return_value=evaluator):
+            with pytest.raises(AssertionError, match="Expected coco keypoint stats shape"):
+                metric.compute()
+
+    def test_calls_synchronize_and_accumulate(self) -> None:
+        """Compute() calls synchronize_between_processes() and accumulate() on the evaluator."""
+        evaluator = _make_evaluator_mock([0.5, 0.7, 0.4, -1.0, -1.0, 0.6, -1.0, -1.0, -1.0, -1.0])
+        metric = MetricKeypointOKS(_make_coco_gt())
+        with patch("rfdetr.evaluation.keypoint_oks.CocoEvaluator", return_value=evaluator):
+            metric.compute()
+        evaluator.synchronize_between_processes.assert_called_once()
+        evaluator.accumulate.assert_called_once()
+
+    def test_constructs_evaluator_with_metric_params(self) -> None:
+        """Compute() passes max_dets and keypoint_oks_sigmas to CocoEvaluator."""
+        coco_gt = _make_coco_gt()
+        evaluator = _make_evaluator_mock([0.5, 0.7, 0.4, -1.0, -1.0, 0.6, -1.0, -1.0, -1.0, -1.0])
+        metric = MetricKeypointOKS(coco_gt, keypoint_oks_sigmas=[0.05, 0.1], max_dets=100)
+        with patch("rfdetr.evaluation.keypoint_oks.CocoEvaluator", return_value=evaluator) as cls:
+            metric.compute()
+        cls.assert_called_once_with(
+            coco_gt,
+            ["keypoints"],
+            max_dets=100,
+            keypoint_oks_sigmas=[0.05, 0.1],
+            log_summary=False,
+        )
+
+    def test_replays_each_batch_as_separate_evaluator_update(self) -> None:
+        """Compute() calls evaluator.update() once per accumulated batch."""
+        evaluator = _make_evaluator_mock([0.5, 0.7, 0.4, -1.0, -1.0, 0.6, -1.0, -1.0, -1.0, -1.0])
+        metric = MetricKeypointOKS(_make_coco_gt())
+        metric.update(_make_predictions(image_id=1))
+        metric.update(_make_predictions(image_id=2))
+        with patch("rfdetr.evaluation.keypoint_oks.CocoEvaluator", return_value=evaluator):
+            metric.compute()
+        assert evaluator.update.call_count == 2
+
+    def test_forwards_correct_image_id_to_evaluator(self) -> None:
+        """Compute() passes predictions with the correct image_id to the evaluator."""
+        evaluator = _make_evaluator_mock([0.5, 0.7, 0.4, -1.0, -1.0, 0.6, -1.0, -1.0, -1.0, -1.0])
+        metric = MetricKeypointOKS(_make_coco_gt())
+        metric.update(_make_predictions(image_id=7))
+        with patch("rfdetr.evaluation.keypoint_oks.CocoEvaluator", return_value=evaluator):
+            metric.compute()
+        passed_preds = evaluator.update.call_args.args[0]
+        assert 7 in passed_preds
+
+    def test_skips_evaluator_update_when_no_predictions(self) -> None:
+        """Compute() does not call evaluator.update() when no batches accumulated."""
+        evaluator = _make_evaluator_mock([0.5, 0.7, 0.4, -1.0, -1.0, 0.6, -1.0, -1.0, -1.0, -1.0])
+        metric = MetricKeypointOKS(_make_coco_gt())
+        with patch("rfdetr.evaluation.keypoint_oks.CocoEvaluator", return_value=evaluator):
+            metric.compute()
+        evaluator.update.assert_not_called()
+
+    def test_compute_is_idempotent(self) -> None:
+        """Two compute() calls with identical batches return the same stats.
+
+        Proves the shared coco_gt reference is not mutated between calls.
+        """
+        evaluator = _make_evaluator_mock([0.42, 0.72, 0.31, -1.0, -1.0, 0.55, -1.0, -1.0, -1.0, -1.0])
+        metric = MetricKeypointOKS(_make_coco_gt())
+        metric.update(_make_predictions(image_id=1))
+        with patch("rfdetr.evaluation.keypoint_oks.CocoEvaluator", return_value=evaluator):
+            result_a = metric.compute()
+            result_b = metric.compute()
+        assert result_a == result_b
+
+    @pytest.mark.parametrize(
+        "sigmas",
+        [
+            pytest.param(None, id="no_sigmas"),
+            pytest.param([0.05] * 17, id="17kp_sigmas"),
+            pytest.param([0.05] * 4, id="4kp_sigmas"),
+        ],
+    )
+    def test_compute_accepts_arbitrary_keypoint_counts(self, sigmas: list[float] | None) -> None:
+        """Compute() passes any keypoint_oks_sigmas length to CocoEvaluator without restriction."""
+        evaluator = _make_evaluator_mock([0.5, 0.7, 0.4, -1.0, -1.0, 0.6, -1.0, -1.0, -1.0, -1.0])
+        metric = MetricKeypointOKS(_make_coco_gt(), keypoint_oks_sigmas=sigmas)
+        with patch("rfdetr.evaluation.keypoint_oks.CocoEvaluator", return_value=evaluator) as cls:
+            metric.compute()
+        assert cls.call_args.kwargs["keypoint_oks_sigmas"] == sigmas
+
+
+# ---------------------------------------------------------------------------
+# Integration tests — real COCO GT, visually validatable expected values
+# ---------------------------------------------------------------------------
+
+
+class TestOKSValues:
+    """End-to-end OKS mAP computation against a real CocoEvaluator.
+
+    Each test uses a single annotation and a single prediction so that
+    expected mAP can be derived by hand:
+
+      OKS = exp(-d² / (8 * σ² * s²))
+
+    where d = Euclidean pixel distance, s = sqrt(GT area), σ = OKS sigma.
+    The 8× factor comes from pycocotools: vars = (2σ)², e = d² / (vars * s² * 2).
+
+    All tests use 1 keypoint, GT at (50, 50), area = 2500.0 (s = 50.0),
+    and sigma = 0.05 (default for custom keypoint counts via
+    _DEFAULT_CUSTOM_KEYPOINT_OKS_SIGMA).
+    """
+
+    _GT_KP_X = 50.0
+    _GT_KP_Y = 50.0
+    _AREA = 2500.0  # s = 50.0
+    _SIGMA = 0.05  # default custom OKS sigma
+
+    def _make_gt(self) -> COCO:
+        """Build a 1-image, 1-annotation, 1-keypoint COCO GT."""
+        return _build_coco_gt(
+            num_keypoints=1,
+            gt_keypoints=[self._GT_KP_X, self._GT_KP_Y, 2],
+            area=self._AREA,
+        )
+
+    def test_perfect_prediction_gives_map_one(self) -> None:
+        """Prediction exactly at GT keypoint location must yield mAP@50 = 1.0.
+
+        d = 0  →  OKS = exp(0) = 1.0  →  all IoU thresholds pass  →  mAP = 1.0.
+        """
+        coco_gt = self._make_gt()
+        metric = MetricKeypointOKS(coco_gt, keypoint_oks_sigmas=[self._SIGMA])
+        metric.update(_make_keypoint_prediction(1, [(self._GT_KP_X, self._GT_KP_Y)]))
+        result = metric.compute()
+        assert result["map_50"] == pytest.approx(1.0, abs=1e-3)
+        assert result["map"] == pytest.approx(1.0, abs=1e-3)
+
+    def test_no_predictions_gives_map_zero(self) -> None:
+        """Empty prediction set must yield mAP = 0.0 (no true positives)."""
+        coco_gt = self._make_gt()
+        metric = MetricKeypointOKS(coco_gt, keypoint_oks_sigmas=[self._SIGMA])
+        metric.update({1: {}})
+        result = metric.compute()
+        assert result["map_50"] == pytest.approx(0.0, abs=1e-3)
+        assert result["map"] == pytest.approx(0.0, abs=1e-3)
+
+    def test_far_prediction_gives_map_near_zero(self) -> None:
+        """Prediction far from GT must yield near-zero mAP.
+
+        d = sqrt(50² + 50²) ≈ 70.7, s = 50, σ = 0.05.
+        pycocotools formula: OKS = exp(-d² / (8 * σ² * s²))
+          = exp(-5000 / (8 * 0.0025 * 2500))
+          = exp(-5000 / 50)
+          = exp(-100) ≈ 0  →  mAP@50 = 0.0.
+        """
+        coco_gt = self._make_gt()
+        metric = MetricKeypointOKS(coco_gt, keypoint_oks_sigmas=[self._SIGMA])
+        metric.update(_make_keypoint_prediction(1, [(0.0, 0.0)]))
+        result = metric.compute()
+        assert result["map_50"] == pytest.approx(0.0, abs=1e-3)
+
+    def test_known_oks_threshold_boundary(self) -> None:
+        """Prediction at OKS ≈ 0.6 passes @50 but fails @75.
+
+        pycocotools (and faster_coco_eval) compute OKS as::
+
+          vars = (sigma * 2)²
+          e    = d² / (vars * s² * 2) = d² / (8 * sigma² * s²)
+          OKS  = exp(-e)
+
+        Solving for d where OKS = 0.6 (clearly between the 0.5 and 0.75 thresholds)::
+
+          d² = -ln(0.6) * 8 * sigma² * s²
+             = 0.511 * 8 * 0.0025 * 2500 = 25.55
+          d  = sqrt(25.55) ≈ 5.05 pixels.
+
+        Displacement along x-axis only: predict at (50 + 5.05, 50).
+        OKS ≈ 0.6 → passes @50 (threshold 0.5), fails @75 (threshold 0.75).
+        mAP@50 should be 1.0, mAP@75 should be 0.0.
+        """
+        sigma = self._SIGMA
+        s = np.sqrt(self._AREA)
+        oks_target = 0.6  # between 0.50 and 0.75 — clear boundary
+        # pycocotools formula: OKS = exp(-d² / (8 * sigma² * s²))
+        d = np.sqrt(-np.log(oks_target) * 8 * sigma**2 * s**2)
+        pred_x = float(self._GT_KP_X + d)
+
+        coco_gt = self._make_gt()
+        metric = MetricKeypointOKS(coco_gt, keypoint_oks_sigmas=[sigma])
+        metric.update(_make_keypoint_prediction(1, [(pred_x, self._GT_KP_Y)]))
+        result = metric.compute()
+        assert result["map_50"] == pytest.approx(1.0, abs=1e-3), "OKS=0.6 should pass @50 threshold"
+        assert result["map_75"] == pytest.approx(0.0, abs=1e-3), "OKS=0.6 should fail @75 threshold"
+
+    def test_metric_stable_across_identical_epochs(self) -> None:
+        """Identical predictions fed across three separate compute() cycles give identical mAP.
+
+        This proves no GT mutation, no accumulation bleed, and correct reset between epochs.  A metric that peaks-then-
+        decreases under frozen predictions would fail here.
+        """
+        coco_gt = self._make_gt()
+        metric = MetricKeypointOKS(coco_gt, keypoint_oks_sigmas=[self._SIGMA])
+        results = []
+        for _ in range(3):
+            metric.reset()
+            metric.update(_make_keypoint_prediction(1, [(self._GT_KP_X, self._GT_KP_Y)]))
+            results.append(metric.compute())
+        assert results[0]["map_50"] == pytest.approx(results[1]["map_50"], abs=1e-6)
+        assert results[1]["map_50"] == pytest.approx(results[2]["map_50"], abs=1e-6)
+
+    def test_multi_batch_same_result_as_single_batch(self) -> None:
+        """Splitting predictions across two update() calls gives same mAP as one call.
+
+        Two images, each predicted correctly.  Whether fed in one batch or two, mAP@50 must equal 1.0 — verifies that
+        per-batch append semantics are equivalent to batched evaluation.
+        """
+        kp_names = ["kp0"]
+        dataset = {
+            "images": [
+                {"id": 1, "width": 100, "height": 100, "file_name": "a.jpg"},
+                {"id": 2, "width": 100, "height": 100, "file_name": "b.jpg"},
+            ],
+            "categories": [{"id": 1, "name": "obj", "keypoints": kp_names, "skeleton": []}],
+            "annotations": [
+                {
+                    "id": 1,
+                    "image_id": 1,
+                    "category_id": 1,
+                    "keypoints": [50.0, 50.0, 2],
+                    "num_keypoints": 1,
+                    "area": 2500.0,
+                    "bbox": [0.0, 0.0, 100.0, 100.0],
+                    "iscrowd": 0,
+                },
+                {
+                    "id": 2,
+                    "image_id": 2,
+                    "category_id": 1,
+                    "keypoints": [25.0, 75.0, 2],
+                    "num_keypoints": 1,
+                    "area": 2500.0,
+                    "bbox": [0.0, 0.0, 100.0, 100.0],
+                    "iscrowd": 0,
+                },
+            ],
+        }
+        coco_gt = COCO()
+        coco_gt.dataset = dataset
+        coco_gt.createIndex()
+        sigma = [self._SIGMA]
+
+        # Single batch
+        metric_single = MetricKeypointOKS(coco_gt, keypoint_oks_sigmas=sigma)
+        combined = {}
+        combined.update(_make_keypoint_prediction(1, [(50.0, 50.0)]))
+        combined.update(_make_keypoint_prediction(2, [(25.0, 75.0)]))
+        metric_single.update(combined)
+        result_single = metric_single.compute()
+
+        # Two batches (one image each)
+        metric_split = MetricKeypointOKS(coco_gt, keypoint_oks_sigmas=sigma)
+        metric_split.update(_make_keypoint_prediction(1, [(50.0, 50.0)]))
+        metric_split.update(_make_keypoint_prediction(2, [(25.0, 75.0)]))
+        result_split = metric_split.compute()
+
+        assert result_single["map_50"] == pytest.approx(result_split["map_50"], abs=1e-6)
+        assert result_single["map_50"] == pytest.approx(1.0, abs=1e-3)

tests/export/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/export/test_export.py

@@ -0,0 +1,884 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for model export functionality.
+
+Use cases covered:
+- Segmentation outputs must be present in both train/eval modes to avoid export crashes.
+- Export should not change the original model's training state.
+- CLI export path (deploy.export.main) must include 'masks' in output_names for
+  segmentation models, call make_infer_image with the correct individual args, and call export_onnx with args.output_dir
+  as the first argument.
+"""
+
+import importlib.util
+import inspect
+import types
+import warnings
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Literal
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from torch.jit import TracerWarning
+
+from rfdetr import RFDETRSegNano
+from rfdetr import detr as _detr_module
+from rfdetr.export import main as _cli_export_module
+
+_IS_ONNX_INSTALLED = importlib.util.find_spec("onnx") is not None
+
+
+@contextmanager
+def ignore_tracer_warnings() -> Iterator[None]:
+    """Suppress torch.jit.TracerWarning during export tests to reduce log spam."""
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=TracerWarning)
+        yield
+
+
+class _DummyCoreModel:
+    """Minimal torch.nn.Module stub shared across export tests.
+
+    Avoids real forward passes; returns synthetic detection (and optionally segmentation) outputs matching the shapes
+    expected by RFDETR.export().
+    """
+
+    def __init__(self, *, segmentation_head: bool = False) -> None:
+        self._segmentation_head = segmentation_head
+
+    def to(self, *_args, **_kwargs):
+        return self
+
+    def eval(self):
+        return self
+
+    def cpu(self):
+        return self
+
+    def __call__(self, *_args, **_kwargs):
+        out = {"pred_boxes": torch.zeros(1, 1, 4), "pred_logits": torch.zeros(1, 1, 2)}
+        if self._segmentation_head:
+            out["pred_masks"] = torch.zeros(1, 1, 2, 2)
+        return out
+
+
+def test_export_onnx_uses_legacy_exporter_when_dynamo_flag_exists(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    """`export_onnx` should pass `dynamo=False` when supported by torch.onnx.export."""
+    captured_kwargs: dict = {}
+
+    class _ToyModel(torch.nn.Module):
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return x
+
+    def _fake_onnx_export(*args, **kwargs) -> None:
+        captured_kwargs.update(kwargs)
+
+    monkeypatch.setattr(_cli_export_module.torch.onnx, "export", _fake_onnx_export)
+
+    _cli_export_module.export_onnx(
+        output_dir=str(tmp_path),
+        model=_ToyModel(),
+        input_names=["images"],
+        input_tensors=torch.randn(1, 3, 8, 8),
+        output_names=["dets"],
+        dynamic_axes={},
+        verbose=False,
+    )
+
+    has_dynamo_arg = "dynamo" in inspect.signature(torch.onnx.export).parameters
+    assert ("dynamo" in captured_kwargs) == has_dynamo_arg
+    if has_dynamo_arg:
+        assert captured_kwargs["dynamo"] is False
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for export test")
+@pytest.mark.skipif(not _IS_ONNX_INSTALLED, reason="onnx not installed, run: pip install rfdetr[onnx]")
+def test_segmentation_model_export_no_crash(tmp_path: Path) -> None:
+    """Integration test: exporting a segmentation model should not crash.
+
+    This exercises the full export path to ensure no AttributeError occurs.
+    """
+    model = RFDETRSegNano()
+
+    # This should not crash with "AttributeError: 'dict' object has no attribute 'shape'"
+    with ignore_tracer_warnings():
+        model.export(output_dir=str(tmp_path), verbose=False)
+
+    # Verify export produced output files
+    onnx_files = list(tmp_path.glob("*.onnx"))
+    assert len(onnx_files) > 0, "Export should produce ONNX file(s)"
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for export test")
+@pytest.mark.skipif(not _IS_ONNX_INSTALLED, reason="onnx not installed, run: pip install rfdetr[onnx]")
+def test_export_does_not_change_original_training_state(tmp_path: Path) -> None:
+    """Verify that calling export() does not change the original model's train/eval state.
+
+    This ensures that export() puts a deepcopy of the model in eval mode without mutating the underlying training model
+    used by RF-DETR.
+    """
+    model = RFDETRSegNano()
+
+    # Access the underlying torch module (model.model.model), as in other tests
+    torch_model = model.model.model.to("cuda")
+
+    # Ensure the original model is in training mode
+    torch_model.train()
+    assert torch_model.training is True, "Precondition: original model should start in training mode"
+
+    # Call export() on the high-level model; this should not change the original model's mode
+    with ignore_tracer_warnings():
+        model.export(output_dir=str(tmp_path))
+
+    # After export, the original underlying model should still be in training mode
+    assert torch_model.training is True, "export() should not change the original model's training state"
+
+
+@pytest.mark.parametrize(
+    "dynamic_batch, segmentation_head",
+    [
+        pytest.param(True, False, id="detection_dynamic"),
+        pytest.param(True, True, id="segmentation_dynamic"),
+        pytest.param(False, False, id="detection_static"),
+    ],
+)
+def test_rfdetr_export_dynamic_batch_forwards_dynamic_axes(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+    dynamic_batch: bool,
+    segmentation_head: bool,
+) -> None:
+    """`RFDETR.export(..., dynamic_batch=True)` must pass a non-None `dynamic_axes` dict to `export_onnx`;
+    `dynamic_batch=False` must pass `None`."""
+    model = types.SimpleNamespace(
+        model=types.SimpleNamespace(
+            model=_DummyCoreModel(segmentation_head=segmentation_head), device="cpu", resolution=14
+        ),
+        model_config=types.SimpleNamespace(
+            segmentation_head=segmentation_head,
+            use_grouppose_keypoints=False,
+            num_channels=3,
+        ),
+        size=None,
+    )
+
+    captured: dict = {}
+
+    def _fake_make_infer_image(*_args, **_kwargs):
+        return torch.zeros(1, 3, 14, 14)
+
+    def _fake_export_onnx(*_args, dynamic_axes=None, **_kw):
+        captured["dynamic_axes"] = dynamic_axes
+        return str(tmp_path / "inference_model.onnx")
+
+    monkeypatch.setattr("rfdetr.export.main.make_infer_image", _fake_make_infer_image)
+    monkeypatch.setattr("rfdetr.export.main.export_onnx", _fake_export_onnx)
+    monkeypatch.setattr("rfdetr.detr.deepcopy", lambda x: x)
+
+    _detr_module.RFDETR.export(model, output_dir=str(tmp_path), dynamic_batch=dynamic_batch, shape=(14, 14))
+
+    dynamic_axes = captured.get("dynamic_axes")
+    if not dynamic_batch:
+        assert dynamic_axes is None, f"expected None for static export, got {dynamic_axes!r}"
+        return
+
+    assert isinstance(dynamic_axes, dict), f"expected dict, got {dynamic_axes!r}"
+    for name, axes in dynamic_axes.items():
+        assert axes == {0: "batch"}, f"axis spec for {name!r} should be {{0: 'batch'}}, got {axes!r}"
+
+    expected_names = {"input", "dets", "labels", "masks"} if segmentation_head else {"input", "dets", "labels"}
+    assert set(dynamic_axes.keys()) == expected_names, f"expected keys {expected_names}, got {set(dynamic_axes.keys())}"
+
+
+@pytest.mark.gpu
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("mode", [pytest.param("train", id="train_mode"), pytest.param("eval", id="eval_mode")])
+def test_segmentation_outputs_present_in_train_and_eval(mode: Literal["train", "eval"]) -> None:
+    """Use case: segmentation outputs are present in both train and eval modes."""
+    model = RFDETRSegNano()
+
+    # Access the underlying torch module (model.model.model)
+    torch_model = model.model.model.to("cuda")
+
+    # Use resolution compatible with model's patch size (312 for seg-nano)
+    resolution = model.model.resolution
+    dummy_input = torch.randn(1, 3, resolution, resolution, device="cuda")
+
+    if mode == "train":
+        torch_model.train()
+    else:
+        torch_model.eval()
+
+    with torch.no_grad():
+        output = torch_model(dummy_input)
+
+    assert "pred_boxes" in output
+    assert "pred_logits" in output
+    assert "pred_masks" in output
+
+
+# --------------------------------------------------------------------------
+# Tests for the CLI export path: rfdetr.export.main.main()
+# --------------------------------------------------------------------------
+
+
+class TestCliExportMain:
+    """Unit tests for deploy.export.main() (CLI export path).
+
+    Three bugs were present before the fix:
+    1. output_names omitted 'masks' for segmentation models.
+    2. make_infer_image received the whole args Namespace instead of individual fields.
+    3. export_onnx received model/args in the wrong positions (output_dir was missing).
+    """
+
+    @pytest.fixture
+    def output_dir(self, tmp_path: Path) -> str:
+        return str(tmp_path)
+
+    @staticmethod
+    def _make_args(
+        *,
+        backbone_only: bool = False,
+        segmentation_head: bool = False,
+        output_dir: str,
+        infer_dir: str | None = None,
+        shape: tuple[int, int] = (640, 640),
+        batch_size: int = 1,
+        verbose: bool = False,
+        opset_version: int = 17,
+        tensorrt: bool = False,
+        dynamic_batch: bool = False,
+    ) -> types.SimpleNamespace:
+        return types.SimpleNamespace(
+            device="cpu",
+            seed=42,
+            layer_norm=False,
+            resume=None,
+            backbone_only=backbone_only,
+            segmentation_head=segmentation_head,
+            output_dir=output_dir,
+            infer_dir=infer_dir,
+            shape=shape,
+            batch_size=batch_size,
+            verbose=verbose,
+            opset_version=opset_version,
+            tensorrt=tensorrt,
+            dynamic_batch=dynamic_batch,
+        )
+
+    @staticmethod
+    def _run(args: types.SimpleNamespace) -> tuple[dict, dict]:
+        """Run deploy.export.main(args) with all heavy dependencies mocked.
+
+        Stubs out build_model, make_infer_image, and export_onnx, and injects mock onnx/onnxsim modules so the export
+        module can be imported even when those optional packages are not installed.
+
+        Returns (make_infer_image_captured, export_onnx_captured).
+        """
+        make_infer_image_captured: dict = {}
+        export_onnx_captured: dict = {}
+
+        mock_model = MagicMock()
+        # parameters() must return an iterable of real objects so sum(p.numel()) works
+        mock_model.parameters.return_value = []
+        mock_model.backbone.parameters.return_value = []
+        mock_model.backbone.__getitem__.return_value.projector.parameters.return_value = []
+        mock_model.backbone.__getitem__.return_value.encoder.parameters.return_value = []
+        mock_model.transformer.parameters.return_value = []
+        mock_model.to.return_value = mock_model
+        mock_model.cpu.return_value = mock_model
+        mock_model.eval.return_value = mock_model
+
+        if args.backbone_only:
+            mock_model.return_value = torch.zeros(1, 512, 20, 20)
+        elif args.segmentation_head:
+            mock_model.return_value = {
+                "pred_boxes": torch.zeros(1, 100, 4),
+                "pred_logits": torch.zeros(1, 100, 90),
+                "pred_masks": torch.zeros(1, 100, 27, 27),
+            }
+        else:
+            mock_model.return_value = {
+                "pred_boxes": torch.zeros(1, 300, 4),
+                "pred_logits": torch.zeros(1, 300, 90),
+            }
+
+        mock_tensor = MagicMock()
+        mock_tensor.to.return_value = mock_tensor
+        mock_tensor.cpu.return_value = mock_tensor
+
+        def fake_make_infer_image(*pos_args, **kw_args):
+            make_infer_image_captured["positional"] = pos_args
+            make_infer_image_captured["keyword"] = kw_args
+            return mock_tensor
+
+        def fake_export_onnx(output_dir, model, input_names, input_tensors, output_names, dynamic_axes, **kwargs):
+            export_onnx_captured["output_dir"] = output_dir
+            export_onnx_captured["model"] = model
+            export_onnx_captured["output_names"] = output_names
+            export_onnx_captured["dynamic_axes"] = dynamic_axes
+            export_onnx_captured["kwargs"] = kwargs
+            return str(args.output_dir) + "/inference_model.onnx"
+
+        with (
+            patch.object(_cli_export_module, "build_model", return_value=(mock_model, MagicMock(), MagicMock())),
+            patch.object(_cli_export_module, "make_infer_image", side_effect=fake_make_infer_image),
+            patch.object(_cli_export_module, "export_onnx", side_effect=fake_export_onnx),
+            patch.object(_cli_export_module, "get_rank", return_value=0),
+        ):
+            _cli_export_module.main(args)
+
+        return make_infer_image_captured, export_onnx_captured
+
+    @pytest.mark.parametrize(
+        "segmentation_head, backbone_only, expected_output_names",
+        [
+            pytest.param(True, False, ["dets", "labels", "masks"], id="segmentation"),
+            pytest.param(False, False, ["dets", "labels"], id="detection"),
+            pytest.param(False, True, ["features"], id="backbone_only"),
+        ],
+    )
+    def test_output_names(
+        self,
+        output_dir: str,
+        segmentation_head: bool,
+        backbone_only: bool,
+        expected_output_names: list[str],
+    ) -> None:
+        """export_onnx must receive the correct output_names for every model type.
+
+        Before the fix, deploy/export.py line 253 used:
+
+        output_names = ['features'] if args.backbone_only else ['dets', 'labels']
+
+        which always omitted 'masks' for segmentation models.
+        """
+        args = self._make_args(
+            backbone_only=backbone_only,
+            segmentation_head=segmentation_head,
+            output_dir=output_dir,
+        )
+        _, export_onnx_captured = self._run(args)
+
+        actual = export_onnx_captured.get("output_names")
+        assert actual == expected_output_names, f"expected output_names={expected_output_names}, got {actual!r}"
+
+    def test_make_infer_image_receives_individual_fields(self, output_dir: str) -> None:
+        """make_infer_image must be called with (infer_dir, shape, batch_size, device), not with the whole args
+        Namespace.
+
+        Before the fix, deploy/export.py line 251 used:
+
+        input_tensors = make_infer_image(args, device)
+        """
+        shape = (640, 640)
+        batch_size = 2
+        infer_dir = None
+        args = self._make_args(
+            output_dir=output_dir,
+            infer_dir=infer_dir,
+            shape=shape,
+            batch_size=batch_size,
+        )
+        make_infer_image_captured, _ = self._run(args)
+
+        pos = make_infer_image_captured.get("positional", ())
+        assert pos[:3] == (infer_dir, shape, batch_size), f"expected (infer_dir, shape, batch_size), got {pos[:3]!r}"
+
+    def test_export_onnx_receives_output_dir_and_kwargs(self, output_dir: str) -> None:
+        """export_onnx must be called as export_onnx(output_dir, model, ...) with backbone_only, verbose, and
+        opset_version forwarded as keyword args.
+
+        Before the fix, deploy/export.py line 294 used:
+
+        export_onnx(model, args, input_names, input_tensors, output_names, dynamic_axes)
+
+        which swapped output_dir/model and dropped all keyword args.
+        """
+        args = self._make_args(
+            output_dir=output_dir,
+            verbose=True,
+            opset_version=11,
+        )
+        _, export_onnx_captured = self._run(args)
+
+        assert "output_dir" in export_onnx_captured, "export_onnx was not called"
+        assert export_onnx_captured["output_dir"] == output_dir, (
+            f"expected output_dir={output_dir!r}, got {export_onnx_captured['output_dir']!r}"
+        )
+        kwargs = export_onnx_captured.get("kwargs", {})
+        assert kwargs.get("verbose") == args.verbose, (
+            f"expected verbose={args.verbose!r}, got {kwargs.get('verbose')!r}"
+        )
+        assert kwargs.get("opset_version") == args.opset_version, (
+            f"expected opset_version={args.opset_version!r}, got {kwargs.get('opset_version')!r}"
+        )
+        assert "backbone_only" in kwargs, "backbone_only kwarg missing from export_onnx call"
+
+    @pytest.mark.parametrize(
+        "dynamic_batch, segmentation_head, backbone_only",
+        [
+            pytest.param(True, False, False, id="detection_dynamic"),
+            pytest.param(True, True, False, id="segmentation_dynamic"),
+            pytest.param(True, False, True, id="backbone_only_dynamic"),
+            pytest.param(False, False, False, id="detection_static"),
+        ],
+    )
+    def test_dynamic_batch_forwards_dynamic_axes(
+        self,
+        output_dir: str,
+        dynamic_batch: bool,
+        segmentation_head: bool,
+        backbone_only: bool,
+    ) -> None:
+        """CLI --dynamic_batch=True must pass {name: {0: 'batch'}} for every I/O name.
+
+        When dynamic_batch=False, dynamic_axes must be None (static export).
+        """
+        args = self._make_args(
+            output_dir=output_dir,
+            dynamic_batch=dynamic_batch,
+            segmentation_head=segmentation_head,
+            backbone_only=backbone_only,
+        )
+        _, captured = self._run(args)
+
+        dynamic_axes = captured.get("dynamic_axes")
+        if not dynamic_batch:
+            assert dynamic_axes is None, f"expected None for static export, got {dynamic_axes!r}"
+            return
+
+        assert isinstance(dynamic_axes, dict), f"expected dict, got {dynamic_axes!r}"
+        for name, axes in dynamic_axes.items():
+            assert axes == {0: "batch"}, f"axis spec for {name!r} should be {{0: 'batch'}}, got {axes!r}"
+
+        # Every input/output name must have an entry
+        if backbone_only:
+            expected_names = {"input", "features"}
+        elif segmentation_head:
+            expected_names = {"input", "dets", "labels", "masks"}
+        else:
+            expected_names = {"input", "dets", "labels"}
+        assert set(dynamic_axes.keys()) == expected_names, (
+            f"expected keys {expected_names}, got {set(dynamic_axes.keys())}"
+        )
+
+    @pytest.mark.parametrize(
+        "device",
+        [
+            pytest.param("cpu", id="cpu"),
+            pytest.param("cuda", id="cuda"),
+        ],
+    )
+    def test_model_moved_to_correct_device(self, output_dir: str, device: str, monkeypatch) -> None:
+        """model.to() and input_tensors.to() must use args.device, not a hard-coded 'cuda'.
+
+        Before the fix, export/main.py line 145 called model.eval().to('cuda') unconditionally, which crashed when
+        CUDA_VISIBLE_DEVICES was blank (CPU export).
+        """
+        import torch
+
+        to_calls = []
+
+        mock_model = MagicMock()
+        mock_model.parameters.return_value = []
+        mock_model.backbone.parameters.return_value = []
+        mock_model.backbone.__getitem__.return_value.projector.parameters.return_value = []
+        mock_model.backbone.__getitem__.return_value.encoder.parameters.return_value = []
+        mock_model.transformer.parameters.return_value = []
+
+        # Capture .to() calls
+        def _record_to(dev):
+            to_calls.append(str(dev))
+            return mock_model
+
+        mock_model.to.side_effect = _record_to
+        mock_model.cpu.return_value = mock_model
+        mock_model.eval.return_value = mock_model
+        mock_model.return_value = {"pred_boxes": torch.zeros(1, 300, 4), "pred_logits": torch.zeros(1, 300, 90)}
+
+        mock_tensor = MagicMock()
+        tensor_to_calls = []
+
+        def _tensor_to(dev):
+            tensor_to_calls.append(str(dev))
+            return mock_tensor
+
+        mock_tensor.to.side_effect = _tensor_to
+        mock_tensor.cpu.return_value = mock_tensor
+
+        # When testing "cuda" path, pretend CUDA is not available so the
+        # fallback-to-cpu branch fires and we can verify the warning without
+        # needing a GPU in CI.
+        monkeypatch.setattr(torch, "cuda", MagicMock(is_available=MagicMock(return_value=False)))
+
+        args = self._make_args(output_dir=output_dir)
+        args.device = device
+
+        with (
+            patch.object(_cli_export_module, "build_model", return_value=(mock_model, MagicMock(), MagicMock())),
+            patch.object(_cli_export_module, "make_infer_image", return_value=mock_tensor),
+            patch.object(_cli_export_module, "export_onnx", return_value=str(output_dir) + "/m.onnx"),
+            patch.object(_cli_export_module, "get_rank", return_value=0),
+        ):
+            _cli_export_module.main(args)
+
+        # The model must NOT have been moved to a hard-coded "cuda" device when
+        # device="cpu" — verify the fallback CPU path was taken.
+        assert "cuda" not in to_calls, f"model was moved to 'cuda' unexpectedly: {to_calls}"
+
+
+class TestExportPatchSize:
+    """RFDETR.export() patch_size validation and shape-divisibility tests."""
+
+    @staticmethod
+    def _scaffold(
+        monkeypatch: pytest.MonkeyPatch, tmp_path: Path, patch_size: int, num_windows: int
+    ) -> types.SimpleNamespace:
+        """Build a minimal RFDETR-like namespace with controllable patch_size/num_windows."""
+        model = types.SimpleNamespace(
+            model=types.SimpleNamespace(
+                model=_DummyCoreModel(),
+                device="cpu",
+                resolution=patch_size * num_windows * 2,  # always valid
+            ),
+            model_config=types.SimpleNamespace(
+                segmentation_head=False,
+                use_grouppose_keypoints=False,
+                patch_size=patch_size,
+                num_windows=num_windows,
+                num_channels=3,
+            ),
+            size=None,
+        )
+
+        def _fake_make_infer_image(*_a, **_kw):
+            return torch.zeros(1, 3, 8, 8)
+
+        def _fake_export_onnx(*_a, **_kw):
+            return str(tmp_path / "inference_model.onnx")
+
+        monkeypatch.setattr("rfdetr.export.main.make_infer_image", _fake_make_infer_image)
+        monkeypatch.setattr("rfdetr.export.main.export_onnx", _fake_export_onnx)
+        monkeypatch.setattr("rfdetr.detr.deepcopy", lambda x: x)
+        return model
+
+    def test_export_patch_size_mismatch_raises(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+        """export(patch_size=X) must raise ValueError when X != model_config.patch_size."""
+        model = self._scaffold(monkeypatch, tmp_path, patch_size=14, num_windows=4)
+        with pytest.raises(ValueError, match="patch_size"):
+            _detr_module.RFDETR.export(model, output_dir=str(tmp_path), patch_size=16)
+
+    @pytest.mark.parametrize("bad_patch_size", [0, -1])
+    def test_export_invalid_patch_size_raises(
+        self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, bad_patch_size: int
+    ) -> None:
+        """Export() must raise ValueError when patch_size is not a positive integer."""
+        model = self._scaffold(monkeypatch, tmp_path, patch_size=14, num_windows=4)
+        # Keep model_config.patch_size consistent with the patch_size argument for this test
+        model.model_config.patch_size = bad_patch_size
+        with pytest.raises(ValueError, match="patch_size must be a positive integer"):
+            _detr_module.RFDETR.export(model, output_dir=str(tmp_path), patch_size=bad_patch_size)
+
+    def test_export_shape_must_be_divisible_by_block_size(
+        self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+    ) -> None:
+        """Export() must reject shapes not divisible by patch_size * num_windows."""
+        # patch_size=16, num_windows=2 → block_size=32; shape (48, 64): 48 % 32 != 0
+        model = self._scaffold(monkeypatch, tmp_path, patch_size=16, num_windows=2)
+        with pytest.raises(ValueError, match="divisible by 32"):
+            _detr_module.RFDETR.export(model, output_dir=str(tmp_path), shape=(48, 64))
+
+    @pytest.mark.parametrize(
+        "bad_shape",
+        [
+            pytest.param((-64, 64), id="negative_height"),
+            pytest.param((64, -64), id="negative_width"),
+            pytest.param((0, 64), id="zero_height"),
+            pytest.param((64, 0), id="zero_width"),
+        ],
+    )
+    def test_export_negative_or_zero_shape_raises(
+        self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, bad_shape: tuple[int, int]
+    ) -> None:
+        """Export() must reject non-positive shape dimensions (Python -N % M == 0 wraps silently)."""
+        model = self._scaffold(monkeypatch, tmp_path, patch_size=16, num_windows=2)
+        with pytest.raises(ValueError, match="positive integers"):
+            _detr_module.RFDETR.export(model, output_dir=str(tmp_path), shape=bad_shape)
+
+    def test_export_shape_valid_for_block_size(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+        """Export() accepts shape divisible by patch_size * num_windows without error."""
+        # patch_size=16, num_windows=2 → block_size=32; shape (64, 64) is valid
+        model = self._scaffold(monkeypatch, tmp_path, patch_size=16, num_windows=2)
+        # Should not raise
+        _detr_module.RFDETR.export(model, output_dir=str(tmp_path), shape=(64, 64))
+
+    @pytest.mark.parametrize("bad_patch_size", [True, False])
+    def test_export_bool_patch_size_raises(
+        self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, bad_patch_size: bool
+    ) -> None:
+        """Export() must reject bool values for patch_size (isinstance(True, int) is True)."""
+        model = self._scaffold(monkeypatch, tmp_path, patch_size=14, num_windows=1)
+        with pytest.raises(ValueError, match="patch_size must be a positive integer"):
+            _detr_module.RFDETR.export(model, output_dir=str(tmp_path), patch_size=bad_patch_size)
+
+    def test_export_explicit_patch_size_matching_config_succeeds(
+        self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+    ) -> None:
+        """export(patch_size=X) must succeed when X matches model_config.patch_size."""
+        model = self._scaffold(monkeypatch, tmp_path, patch_size=14, num_windows=4)
+        # patch_size=14 matches model_config.patch_size=14; block_size=56; resolution=112 (56*2)
+        _detr_module.RFDETR.export(model, output_dir=str(tmp_path), patch_size=14)
+
+    @pytest.mark.parametrize(
+        "bad_shape",
+        [
+            pytest.param((14.0, 14.0), id="float_dims"),
+            pytest.param((14,), id="wrong_arity_one_element"),
+            pytest.param((14, 14, 3), id="wrong_arity_three_elements"),
+            pytest.param((True, 14), id="bool_height"),
+            pytest.param((14, False), id="bool_width"),
+        ],
+    )
+    def test_export_invalid_shape_type_raises(
+        self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, bad_shape: tuple
+    ) -> None:
+        """Export() must raise ValueError for float, bool, or wrong-arity shape tuples."""
+        model = self._scaffold(monkeypatch, tmp_path, patch_size=14, num_windows=1)
+        with pytest.raises(ValueError, match="shape"):
+            _detr_module.RFDETR.export(model, output_dir=str(tmp_path), shape=bad_shape)
+
+    @pytest.mark.parametrize("bad_num_windows", [0, -1, True])
+    def test_export_invalid_num_windows_raises(
+        self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path, bad_num_windows: int
+    ) -> None:
+        """Export() must raise ValueError when model_config.num_windows is not a positive integer."""
+        model = self._scaffold(monkeypatch, tmp_path, patch_size=14, num_windows=1)
+        model.model_config.num_windows = bad_num_windows
+        with pytest.raises(ValueError, match="num_windows must be a positive integer"):
+            _detr_module.RFDETR.export(model, output_dir=str(tmp_path))
+
+    def test_export_default_resolution_not_divisible_by_block_size_raises(
+        self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+    ) -> None:
+        """Export() with shape=None must raise ValueError when model.resolution % block_size != 0."""
+        # patch_size=14, num_windows=3 → block_size=42; scaffold sets resolution=84 (42*2) which is valid
+        # Override resolution to 50 (not divisible by 42) to trigger the check
+        model = self._scaffold(monkeypatch, tmp_path, patch_size=14, num_windows=3)
+        model.model.resolution = 50
+        with pytest.raises(ValueError, match="default resolution"):
+            _detr_module.RFDETR.export(model, output_dir=str(tmp_path))
+
+
+def test_make_infer_image_produces_correct_rectangular_shape() -> None:
+    """make_infer_image must produce a (B, C, H, W) tensor for non-square shapes.
+
+    Regression test for the square-resize bug where ``Resize((shape[0], shape[0]))`` was used instead of
+    ``Resize((shape[0], shape[1]))``, causing the output width to silently equal the height.
+    """
+    from rfdetr.export.main import make_infer_image
+
+    h, w, b = 112, 224, 2
+    tensor = make_infer_image(infer_dir=None, shape=(h, w), batch_size=b, device="cpu")
+    assert tensor.shape == (b, 3, h, w), f"Expected shape ({b}, 3, {h}, {w}), got {tensor.shape}"
+
+
+# ---------------------------------------------------------------------------
+# ONNX export variant naming
+# ---------------------------------------------------------------------------
+
+
+class TestExportOnnxVariantNaming:
+    """Verify that export_onnx uses variant_name in the output filename."""
+
+    def test_variant_name_in_filename(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+        """When variant_name is provided, the ONNX file is named after the variant."""
+        captured: dict = {}
+
+        def _fake_onnx_export(*args, **kwargs) -> None:
+            captured["output_file"] = args[2]  # 3rd positional arg is output_file
+
+        monkeypatch.setattr(_cli_export_module.torch.onnx, "export", _fake_onnx_export)
+
+        _cli_export_module.export_onnx(
+            output_dir=str(tmp_path),
+            model=torch.nn.Identity(),
+            input_names=["input"],
+            input_tensors=torch.randn(1, 3, 8, 8),
+            output_names=["dets"],
+            dynamic_axes=None,
+            verbose=False,
+            variant_name="rfdetr-medium",
+        )
+
+        assert captured["output_file"].endswith("rfdetr-medium.onnx")
+
+    def test_variant_name_with_backbone(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+        """backbone_only + variant_name produces '{variant}-backbone.onnx'."""
+        captured: dict = {}
+
+        def _fake_onnx_export(*args, **kwargs) -> None:
+            captured["output_file"] = args[2]
+
+        monkeypatch.setattr(_cli_export_module.torch.onnx, "export", _fake_onnx_export)
+
+        _cli_export_module.export_onnx(
+            output_dir=str(tmp_path),
+            model=torch.nn.Identity(),
+            input_names=["input"],
+            input_tensors=torch.randn(1, 3, 8, 8),
+            output_names=["features"],
+            dynamic_axes=None,
+            backbone_only=True,
+            verbose=False,
+            variant_name="rfdetr-nano",
+        )
+
+        assert captured["output_file"].endswith("rfdetr-nano-backbone.onnx")
+
+    def test_default_name_without_variant(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+        """Without variant_name, falls back to 'inference_model.onnx'."""
+        captured: dict = {}
+
+        def _fake_onnx_export(*args, **kwargs) -> None:
+            captured["output_file"] = args[2]
+
+        monkeypatch.setattr(_cli_export_module.torch.onnx, "export", _fake_onnx_export)
+
+        _cli_export_module.export_onnx(
+            output_dir=str(tmp_path),
+            model=torch.nn.Identity(),
+            input_names=["input"],
+            input_tensors=torch.randn(1, 3, 8, 8),
+            output_names=["dets"],
+            dynamic_axes=None,
+            verbose=False,
+        )
+
+        assert captured["output_file"].endswith("inference_model.onnx")
+
+    def test_default_backbone_name_without_variant(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+        """Without variant_name + backbone_only, falls back to 'backbone_model.onnx'."""
+        captured: dict = {}
+
+        def _fake_onnx_export(*args, **kwargs) -> None:
+            captured["output_file"] = args[2]
+
+        monkeypatch.setattr(_cli_export_module.torch.onnx, "export", _fake_onnx_export)
+
+        _cli_export_module.export_onnx(
+            output_dir=str(tmp_path),
+            model=torch.nn.Identity(),
+            input_names=["input"],
+            input_tensors=torch.randn(1, 3, 8, 8),
+            output_names=["features"],
+            dynamic_axes=None,
+            backbone_only=True,
+            verbose=False,
+        )
+
+        assert captured["output_file"].endswith("backbone_model.onnx")
+
+    def test_rfdetr_export_passes_variant_name(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+        """RFDETR.export() passes self.size as variant_name to export_onnx."""
+        captured: dict = {}
+
+        model = types.SimpleNamespace(
+            model=types.SimpleNamespace(model=_DummyCoreModel(), device="cpu", resolution=14),
+            model_config=types.SimpleNamespace(
+                segmentation_head=False,
+                use_grouppose_keypoints=False,
+                num_channels=3,
+            ),
+            size="rfdetr-medium",
+        )
+
+        def _fake_make_infer_image(*_args, **_kwargs):
+            return torch.zeros(1, 3, 14, 14)
+
+        def _fake_export_onnx(*_args, variant_name=None, **_kw):
+            captured["variant_name"] = variant_name
+            return str(tmp_path / "rfdetr-medium.onnx")
+
+        monkeypatch.setattr("rfdetr.export.main.make_infer_image", _fake_make_infer_image)
+        monkeypatch.setattr("rfdetr.export.main.export_onnx", _fake_export_onnx)
+        monkeypatch.setattr("rfdetr.detr.deepcopy", lambda x: x)
+
+        _detr_module.RFDETR.export(model, output_dir=str(tmp_path), shape=(14, 14))
+
+        assert captured["variant_name"] == "rfdetr-medium"
+
+    def test_rfdetr_export_passes_none_when_size_not_set(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+        """Base RFDETR (size=None) passes None as variant_name."""
+        captured: dict = {}
+
+        model = types.SimpleNamespace(
+            model=types.SimpleNamespace(model=_DummyCoreModel(), device="cpu", resolution=14),
+            model_config=types.SimpleNamespace(
+                segmentation_head=False,
+                use_grouppose_keypoints=False,
+                num_channels=3,
+            ),
+            size=None,
+        )
+
+        def _fake_make_infer_image(*_args, **_kwargs):
+            return torch.zeros(1, 3, 14, 14)
+
+        def _fake_export_onnx(*_args, variant_name=None, **_kw):
+            captured["variant_name"] = variant_name
+            return str(tmp_path / "inference_model.onnx")
+
+        monkeypatch.setattr("rfdetr.export.main.make_infer_image", _fake_make_infer_image)
+        monkeypatch.setattr("rfdetr.export.main.export_onnx", _fake_export_onnx)
+        monkeypatch.setattr("rfdetr.detr.deepcopy", lambda x: x)
+
+        _detr_module.RFDETR.export(model, output_dir=str(tmp_path), shape=(14, 14))
+
+        assert captured["variant_name"] is None
+
+    @pytest.mark.parametrize(
+        "variant_name, expected_suffix",
+        [
+            pytest.param("", "inference_model.onnx", id="empty_string_falls_back_to_default"),
+            pytest.param("foo/bar", "bar.onnx", id="path_separator_stripped_to_basename"),
+            pytest.param("/tmp/x", "x.onnx", id="absolute_path_stripped_to_basename"),
+        ],
+    )
+    def test_variant_name_sanitization(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        tmp_path: Path,
+        variant_name: str,
+        expected_suffix: str,
+    ) -> None:
+        """variant_name edge cases: empty string falls back to default; path separators are stripped."""
+        captured: dict = {}
+
+        def _fake_onnx_export(*args, **kwargs) -> None:
+            captured["output_file"] = args[2]
+
+        monkeypatch.setattr(_cli_export_module.torch.onnx, "export", _fake_onnx_export)
+
+        _cli_export_module.export_onnx(
+            output_dir=str(tmp_path),
+            model=torch.nn.Identity(),
+            input_names=["input"],
+            input_tensors=torch.randn(1, 3, 8, 8),
+            output_names=["dets"],
+            dynamic_axes=None,
+            verbose=False,
+            variant_name=variant_name or None,
+        )
+
+        assert captured["output_file"].endswith(expected_suffix)

tests/export/test_export_for_roboflow.py

@@ -0,0 +1,92 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for ``RFDETR.export_for_roboflow``.
+
+``export_for_roboflow`` is the extracted, network-free core of ``deploy_to_roboflow``: it writes ``weights.pt`` (a dict
+with ``"model"`` and ``"args"`` keys) and ``class_names.txt`` into a target directory.  A lightweight stub stands in for
+``self.model`` so the file-writing contract is exercised without building a real model or downloading weights.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+
+from rfdetr.detr import RFDETR
+
+
+def _make_stub_model(class_names: list[str]) -> RFDETR:
+    """Build an RFDETR instance whose model/state are stubbed for export_for_roboflow.
+
+    ``RFDETR.__init__`` is bypassed; only the attributes ``export_for_roboflow`` reads are populated.
+    """
+    instance = RFDETR.__new__(RFDETR)
+    args = SimpleNamespace(resolution=560)
+    inner_module = SimpleNamespace(state_dict=lambda: {"weight": torch.zeros(2, 2)})
+    instance.model = SimpleNamespace(model=inner_module, args=args, class_names=class_names)
+    return instance
+
+
+class TestExportForRoboflow:
+    """export_for_roboflow writes a deploy-ready bundle into a directory."""
+
+    def test_writes_weights_pt_with_model_and_args(self, tmp_path: Path) -> None:
+        """weights.pt is a dict with 'model' and 'args', and args carries resolution."""
+        model = _make_stub_model(["cat", "dog"])
+
+        model.export_for_roboflow(str(tmp_path))
+
+        bundle = torch.load(tmp_path / "weights.pt", map_location="cpu", weights_only=False)
+        assert set(bundle) == {"model", "args"}
+        assert "weight" in bundle["model"]
+        assert bundle["args"].resolution == 560
+
+    def test_writes_class_names_txt(self, tmp_path: Path) -> None:
+        """class_names.txt lists one class name per line."""
+        model = _make_stub_model(["cat", "dog"])
+
+        model.export_for_roboflow(str(tmp_path))
+
+        assert (tmp_path / "class_names.txt").read_text(encoding="utf-8") == "cat\ndog"
+
+    def test_embeds_class_names_in_args(self, tmp_path: Path) -> None:
+        """class_names are embedded in the saved args namespace when absent."""
+        model = _make_stub_model(["cat", "dog"])
+
+        model.export_for_roboflow(str(tmp_path))
+
+        bundle = torch.load(tmp_path / "weights.pt", map_location="cpu", weights_only=False)
+        assert bundle["args"].class_names == ["cat", "dog"]
+
+    def test_does_not_overwrite_existing_args_class_names(self, tmp_path: Path) -> None:
+        """args.class_names already set on the model is preserved in the saved bundle."""
+        model = _make_stub_model(["cat", "dog"])
+        model.model.args.class_names = ["pre_existing"]
+
+        model.export_for_roboflow(str(tmp_path))
+
+        bundle = torch.load(tmp_path / "weights.pt", map_location="cpu", weights_only=False)
+        assert bundle["args"].class_names == ["pre_existing"]
+
+    def test_empty_class_names_writes_empty_file(self, tmp_path: Path) -> None:
+        """Empty class_names list produces an empty class_names.txt (no trailing newline)."""
+        model = _make_stub_model([])
+
+        model.export_for_roboflow(str(tmp_path))
+
+        assert (tmp_path / "class_names.txt").read_text(encoding="utf-8") == ""
+
+    def test_creates_output_dir_when_missing(self, tmp_path: Path) -> None:
+        """output_dir is created if it does not already exist."""
+        model = _make_stub_model(["cat", "dog"])
+        target = tmp_path / "nested" / "bundle"
+
+        model.export_for_roboflow(str(target))
+
+        assert (target / "weights.pt").exists()
+        assert (target / "class_names.txt").exists()

tests/export/test_onnx_notes.py

@@ -0,0 +1,160 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for the ``notes`` parameter in :func:`~rfdetr.export._onnx.exporter.export_onnx`."""
+
+import json
+from pathlib import Path
+
+import pytest
+import torch
+import torch.nn as nn
+
+onnx = pytest.importorskip("onnx", reason="onnx not installed; skip ONNX notes tests")
+
+
+from rfdetr.export._onnx.exporter import export_onnx  # noqa: E402
+
+
+class _TinyModel(nn.Module):
+    """Minimal model that can be exported to ONNX."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Run a trivial identity-like forward pass.
+
+        Args:
+            x: Input tensor.
+
+        Returns:
+            Input tensor unchanged.
+        """
+        return x
+
+
+def _export_tiny_model(tmp_path: Path, notes: object = None) -> str:
+    """Export a tiny model to ONNX and return the output file path.
+
+    Args:
+        tmp_path: Temporary directory provided by pytest.
+        notes: Optional notes to embed in the ONNX file.
+
+    Returns:
+        Path to the exported ONNX file.
+    """
+    model = _TinyModel().eval()
+    input_tensor = torch.randn(1, 3, 32, 32)
+    return export_onnx(
+        output_dir=str(tmp_path),
+        model=model,
+        input_names=["input"],
+        input_tensors=input_tensor,
+        output_names=["output"],
+        dynamic_axes=None,
+        verbose=False,
+        notes=notes,
+    )
+
+
+class TestExportOnnxNotes:
+    """Verify ``notes`` metadata round-trips through the ONNX export."""
+
+    @pytest.mark.parametrize(
+        "notes, expected_value",
+        [
+            pytest.param("simple string", "simple string", id="string"),
+            pytest.param(
+                {"date": "2026-01-01", "labeller": "Alice"},
+                '{"date": "2026-01-01", "labeller": "Alice"}',
+                id="dict",
+            ),
+            pytest.param(["class_a", "class_b"], '["class_a", "class_b"]', id="list"),
+            pytest.param(42, "42", id="int"),
+        ],
+    )
+    def test_notes_embedded_in_onnx_metadata(self, tmp_path: Path, notes: object, expected_value: str) -> None:
+        """Notes are stored as the 'notes' metadata_props entry in the ONNX model."""
+        output_file = _export_tiny_model(tmp_path, notes=notes)
+
+        model = onnx.load(output_file)
+        meta = {prop.key: prop.value for prop in model.metadata_props}
+        assert "rfdetr_notes" in meta
+        assert meta["rfdetr_notes"] == expected_value
+
+    def test_string_notes_stored_verbatim_without_json_wrapping(self, tmp_path: Path) -> None:
+        """Plain string notes must be stored as-is, not double-encoded as JSON."""
+        notes = "my run description"
+        output_file = _export_tiny_model(tmp_path, notes=notes)
+
+        model = onnx.load(output_file)
+        meta = {prop.key: prop.value for prop in model.metadata_props}
+        assert meta["rfdetr_notes"] == "my run description"
+
+    def test_dict_notes_round_trip_via_json(self, tmp_path: Path) -> None:
+        """Dict notes deserialise back to the original dict via json.loads."""
+        notes = {"project": "ceramics", "batch": 7}
+        output_file = _export_tiny_model(tmp_path, notes=notes)
+
+        model = onnx.load(output_file)
+        meta = {prop.key: prop.value for prop in model.metadata_props}
+        assert json.loads(meta["rfdetr_notes"]) == notes
+
+    def test_no_notes_metadata_when_notes_is_none(self, tmp_path: Path) -> None:
+        """When notes=None (default), no 'rfdetr_notes' metadata entry is written."""
+        output_file = _export_tiny_model(tmp_path, notes=None)
+
+        model = onnx.load(output_file)
+        meta = {prop.key: prop.value for prop in model.metadata_props}
+        assert "rfdetr_notes" not in meta
+
+    @pytest.mark.parametrize(
+        "notes",
+        [
+            pytest.param("", id="empty_string"),
+            pytest.param({}, id="empty_dict"),
+            pytest.param([], id="empty_list"),
+            pytest.param(0, id="zero"),
+            pytest.param(False, id="false"),
+        ],
+    )
+    def test_falsy_notes_still_embedded(self, tmp_path: Path, notes: object) -> None:
+        """Falsy but non-None notes are embedded; guard is 'is not None', not truthiness."""
+        output_file = _export_tiny_model(tmp_path, notes=notes)
+
+        model = onnx.load(output_file)
+        meta = {prop.key: prop.value for prop in model.metadata_props}
+        assert "rfdetr_notes" in meta
+
+    def test_unicode_notes_stored_verbatim(self, tmp_path: Path) -> None:
+        """Unicode string notes survive the ONNX metadata round-trip unchanged."""
+        notes = "Reviewer: Łukasz · 2026-Q2 · ✅"
+        output_file = _export_tiny_model(tmp_path, notes=notes)
+
+        model = onnx.load(output_file)
+        meta = {prop.key: prop.value for prop in model.metadata_props}
+        assert meta["rfdetr_notes"] == notes
+
+    def test_nan_notes_raises_value_error(self, tmp_path: Path) -> None:
+        """Non-finite float notes raise ValueError (allow_nan=False)."""
+        with pytest.raises(ValueError):
+            _export_tiny_model(tmp_path, notes=float("nan"))
+
+    def test_notes_is_keyword_only(self, tmp_path: Path) -> None:
+        """Notes must be passed as a keyword argument; positional use raises TypeError."""
+        model = _TinyModel().eval()
+        input_tensor = torch.randn(1, 3, 32, 32)
+        with pytest.raises(TypeError):
+            export_onnx(  # type: ignore[call-arg]
+                str(tmp_path),
+                model,
+                ["input"],
+                input_tensor,
+                ["output"],
+                None,
+                False,
+                False,
+                17,
+                None,
+                "positional_notes_value",
+            )

tests/export/test_tensorrt_export.py

@@ -0,0 +1,229 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for TensorRT export helpers."""
+
+import argparse
+import subprocess
+
+import pytest
+
+from rfdetr.export import _tensorrt as tensorrt_export
+
+
+def test_run_command_shell_dry_run_handles_missing_cuda_visible_devices(monkeypatch) -> None:
+    """Dry-run logging should not crash when CUDA_VISIBLE_DEVICES is unset."""
+    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+
+    logged_messages = []
+    monkeypatch.setattr(tensorrt_export.logger, "info", logged_messages.append)
+
+    result = tensorrt_export.run_command_shell(["trtexec", "--help"], dry_run=True)
+
+    assert result.returncode == 0
+    assert any("CUDA_VISIBLE_DEVICES=" in message for message in logged_messages)
+
+
+def test_run_command_shell_uses_list_not_string(monkeypatch) -> None:
+    """subprocess.run must be called with a list (shell=False) to prevent injection."""
+    captured = {}
+
+    def _fake_run(command, shell, capture_output, text, check):
+        captured["command"] = command
+        captured["shell"] = shell
+        return subprocess.CompletedProcess(command, 0, stdout="", stderr="")
+
+    monkeypatch.setattr(tensorrt_export.subprocess, "run", _fake_run)
+
+    tensorrt_export.run_command_shell(["trtexec", "--onnx=/some/model.onnx"], dry_run=False)
+
+    assert isinstance(captured["command"], list), "command must be a list, not a string"
+    assert captured["shell"] is False, "shell=False is required to prevent injection"
+
+
+def test_run_command_shell_dry_run_does_not_invoke_subprocess(monkeypatch) -> None:
+    """Dry-run must return early without calling subprocess.run."""
+    was_called = []
+
+    def _should_not_run(*args, **kwargs):
+        was_called.append(True)
+        return subprocess.CompletedProcess([], 0)
+
+    monkeypatch.setattr(tensorrt_export.subprocess, "run", _should_not_run)
+    monkeypatch.setattr(tensorrt_export.logger, "info", lambda _: None)
+
+    result = tensorrt_export.run_command_shell(["trtexec", "--help"], dry_run=True)
+
+    assert not was_called, "subprocess.run must not be called during dry_run"
+    assert result.returncode == 0
+
+
+def test_trtexec_returns_engine_path(monkeypatch) -> None:
+    """Trtexec() must return the .engine path, not None."""
+    captured_argv = []
+
+    def _fake_run(command, **kwargs):
+        captured_argv.extend(command)
+        return subprocess.CompletedProcess(command, 0, stdout="", stderr="")
+
+    monkeypatch.setattr(tensorrt_export.subprocess, "run", _fake_run)
+    monkeypatch.setattr(tensorrt_export, "parse_trtexec_output", lambda _: {})
+
+    args = argparse.Namespace(profile=False, verbose=False, dry_run=False)
+    result = tensorrt_export.trtexec("/tmp/model.onnx", args)
+
+    assert result == "/tmp/model.engine"
+
+
+def test_trtexec_dry_run_returns_engine_path(monkeypatch) -> None:
+    """Trtexec() with dry_run=True must still return the engine path."""
+    monkeypatch.setattr(tensorrt_export.logger, "info", lambda _: None)
+    monkeypatch.setattr(tensorrt_export, "parse_trtexec_output", lambda _: {})
+
+    args = argparse.Namespace(profile=False, verbose=False, dry_run=True)
+    result = tensorrt_export.trtexec("/tmp/model.onnx", args)
+
+    assert result == "/tmp/model.engine"
+
+
+@pytest.mark.parametrize(
+    ("onnx_path", "expected_engine"),
+    [
+        pytest.param("/output/rfdetr.onnx", "/output/rfdetr.engine", id="plain-path"),
+        pytest.param("/path with spaces/model.onnx", "/path with spaces/model.engine", id="path-with-spaces"),
+        pytest.param("/model;rm -rf /.onnx", "/model;rm -rf /.engine", id="shell-metachar"),
+    ],
+)
+def test_trtexec_argv_contains_no_shell_string(monkeypatch, onnx_path: str, expected_engine: str) -> None:
+    """Trtexec builds an argv list; no shell string concatenation of user paths."""
+    captured = {}
+
+    def _fake_run(command, shell, **kwargs):
+        captured["command"] = command
+        captured["shell"] = shell
+        return subprocess.CompletedProcess(command, 0, stdout="", stderr="")
+
+    monkeypatch.setattr(tensorrt_export.subprocess, "run", _fake_run)
+    monkeypatch.setattr(tensorrt_export, "parse_trtexec_output", lambda _: {})
+
+    args = argparse.Namespace(profile=False, verbose=False, dry_run=False)
+    result = tensorrt_export.trtexec(onnx_path, args)
+
+    assert result == expected_engine
+    assert isinstance(captured["command"], list), "argv must be a list"
+    assert captured["shell"] is False, "shell=False required"
+    # Verify the ONNX path appears as a standalone argument element (not shell-expanded)
+    assert any(onnx_path in arg for arg in captured["command"])
+
+
+# ---------------------------------------------------------------------------
+# parse_trtexec_output (#1)
+# ---------------------------------------------------------------------------
+
+_FULL_TRTEXEC_STDOUT = """\
+[I] GPU Compute Time: min = 1.23 ms, max = 4.56 ms, mean = 2.34 ms, median = 2.10 ms
+[I] Host to Device Transfer Time: min = 0.10 ms, max = 0.20 ms, mean = 0.15 ms
+[I] Device to Host Transfer Time: min = 0.05 ms, max = 0.08 ms, mean = 0.06 ms
+[I] Latency: min = 1.40 ms, max = 4.80 ms, mean = 2.55 ms
+[I] Throughput: 391.22 qps
+"""
+
+_PARTIAL_TRTEXEC_STDOUT = """\
+[I] GPU Compute Time: min = 1.00 ms, max = 2.00 ms, mean = 1.50 ms, median = 1.45 ms
+[I] Throughput: 100.00 qps
+"""
+
+
+@pytest.mark.parametrize(
+    ("output_text", "expected"),
+    [
+        pytest.param(
+            _FULL_TRTEXEC_STDOUT,
+            {
+                "compute_min_ms": 1.23,
+                "compute_max_ms": 4.56,
+                "compute_mean_ms": 2.34,
+                "compute_median_ms": 2.10,
+                "h2d_min_ms": 0.10,
+                "h2d_max_ms": 0.20,
+                "h2d_mean_ms": 0.15,
+                "d2h_min_ms": 0.05,
+                "d2h_max_ms": 0.08,
+                "d2h_mean_ms": 0.06,
+                "latency_min_ms": 1.40,
+                "latency_max_ms": 4.80,
+                "latency_mean_ms": 2.55,
+                "throughput_qps": 391.22,
+            },
+            id="all-5-patterns",
+        ),
+        pytest.param(
+            "",
+            {},
+            id="empty-stdout",
+        ),
+        pytest.param(
+            _PARTIAL_TRTEXEC_STDOUT,
+            {
+                "compute_min_ms": 1.00,
+                "compute_max_ms": 2.00,
+                "compute_mean_ms": 1.50,
+                "compute_median_ms": 1.45,
+                "throughput_qps": 100.00,
+            },
+            id="partial-stdout",
+        ),
+    ],
+)
+def test_parse_trtexec_output(output_text: str, expected: dict) -> None:
+    """parse_trtexec_output extracts timing statistics from trtexec stdout."""
+    result = tensorrt_export.parse_trtexec_output(output_text)
+    assert result == pytest.approx(expected, abs=1e-6)
+
+
+# ---------------------------------------------------------------------------
+# CalledProcessError logging path (#15)
+# ---------------------------------------------------------------------------
+
+
+def test_run_command_shell_called_process_error_is_reraised(monkeypatch) -> None:
+    """CalledProcessError from subprocess.run is re-raised after logging."""
+    error_messages = []
+
+    def _fake_run(command, **kwargs):
+        raise subprocess.CalledProcessError(returncode=1, cmd=["trtexec"], stderr="engine build failed")
+
+    monkeypatch.setattr(tensorrt_export.subprocess, "run", _fake_run)
+    monkeypatch.setattr(tensorrt_export.logger, "error", error_messages.append)
+
+    with pytest.raises(subprocess.CalledProcessError):
+        tensorrt_export.run_command_shell(["trtexec", "--onnx=/tmp/model.onnx"], dry_run=False)
+
+    assert error_messages, "logger.error must be called when CalledProcessError is raised"
+
+
+# ---------------------------------------------------------------------------
+# profile=True argv path (#17)
+# ---------------------------------------------------------------------------
+
+
+def test_trtexec_profile_true_wraps_with_nsys(monkeypatch) -> None:
+    """Profile=True wraps trtexec with 'nsys profile …' and the output flag is present."""
+    captured_argv: list[str] = []
+
+    def _fake_run(command, **kwargs):
+        captured_argv.extend(command)
+        return subprocess.CompletedProcess(command, 0, stdout="", stderr="")
+
+    monkeypatch.setattr(tensorrt_export.subprocess, "run", _fake_run)
+    monkeypatch.setattr(tensorrt_export, "parse_trtexec_output", lambda _: {})
+    monkeypatch.setattr(tensorrt_export.logger, "info", lambda _: None)
+
+    args = argparse.Namespace(profile=True, verbose=False, dry_run=False)
+    tensorrt_export.trtexec("/tmp/model.onnx", args)
+
+    assert captured_argv[0] == "nsys", "profile=True must wrap with nsys as argv[0]"
+    argv_str = " ".join(captured_argv)
+    assert "--output=" in argv_str, "nsys profile must include --output= flag"

tests/export/test_tflite_inference.py

@@ -0,0 +1,784 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for TFLite inference helpers.
+
+Covers:
+* ``_create_interpreter()`` — interpreter loading with tflite_runtime / tensorflow fallback
+* ``_run_inference()`` — image preprocessing, invocation, and detection decoding
+* ``_decode_masks()`` — segmentation mask upsampling and thresholding
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from unittest import mock
+
+import numpy as np
+import pytest
+import supervision as sv
+from PIL import Image as PILImage
+
+from rfdetr.export._tflite.inference import (
+    _bilinear_resize_half_pixel,
+    _create_interpreter,
+    _decode_masks,
+    _preprocess_image,
+    _run_inference,
+)
+
+# ---------------------------------------------------------------------------
+# Shared helpers / factories
+# ---------------------------------------------------------------------------
+
+_INPUT_SHAPE = [1, 224, 224, 3]
+_DET_OUTPUT = {"shape": [1, 10, 4], "name": "serving_default_dets:0", "index": 1}
+_LABEL_OUTPUT = {"shape": [1, 10, 82], "name": "serving_default_labels:0", "index": 2}
+
+
+def _make_boxes() -> np.ndarray:
+    """Return (1, 10, 4) array of normalised cxcywh boxes all centred at 0.5."""
+    return np.array([[[0.5, 0.5, 0.1, 0.1]] * 10], dtype=np.float32)
+
+
+def _make_logits(high_conf_idx: int | None = 0) -> np.ndarray:
+    """Return (1, 10, 82) logits with one high-confidence entry when requested.
+
+    Background fill is -10.0 so sigmoid scores are near zero (~0.0001) for all entries except the explicitly boosted one
+    (logit=+10.0, sigmoid≈0.9999). This ensures the helper works correctly under per-class sigmoid scoring.
+    """
+    logits = np.full((1, 10, 82), -10.0, dtype=np.float32)
+    if high_conf_idx is not None:
+        logits[0, high_conf_idx, 0] = 10.0
+    return logits
+
+
+def _make_interp(
+    input_shape: list[int] | None = None,
+    out_dets: list[dict] | None = None,
+    boxes: np.ndarray | None = None,
+    logits: np.ndarray | None = None,
+) -> mock.MagicMock:
+    """Build a mock TFLite interpreter with configurable I/O details."""
+    if input_shape is None:
+        input_shape = _INPUT_SHAPE
+    out_dets = out_dets if out_dets is not None else [_DET_OUTPUT, _LABEL_OUTPUT]
+    if boxes is None:
+        boxes = _make_boxes()
+    if logits is None:
+        logits = _make_logits()
+
+    def _get_tensor(index: int) -> np.ndarray:
+        if index == _DET_OUTPUT["index"]:
+            return boxes
+        if index == _LABEL_OUTPUT["index"]:
+            return logits
+        raise ValueError(f"Unknown tensor index: {index}")
+
+    interp = mock.MagicMock()
+    interp.get_input_details.return_value = [{"shape": input_shape, "index": 0, "dtype": np.float32}]
+    interp.get_output_details.return_value = out_dets
+    interp.get_tensor.side_effect = _get_tensor
+    return interp
+
+
+def _save_rgb_image(path: Path, size: tuple[int, int] = (64, 64)) -> None:
+    """Write a small solid-colour RGB JPEG to *path*."""
+    PILImage.new("RGB", size, color=(100, 150, 200)).save(path)
+
+
+def _save_grayscale_image(path: Path, size: tuple[int, int] = (64, 64)) -> None:
+    """Write a small solid-colour grayscale PNG to *path*."""
+    PILImage.new("L", size, color=128).save(path)
+
+
+# ---------------------------------------------------------------------------
+# TestCreateInterpreter
+# ---------------------------------------------------------------------------
+
+# Shared masking entries for mock.patch.dict(sys.modules, ...) that force
+# ``_create_interpreter`` to skip the ai_edge_litert backend probe.
+_AI_EDGE_LITERT_MASK: dict[str, None] = {
+    "ai_edge_litert": None,
+    "ai_edge_litert.interpreter": None,
+}
+
+
+class TestCreateInterpreter:
+    """Tests for ``_create_interpreter()``."""
+
+    @pytest.fixture()
+    def _mock_tflite_runtime(self):
+        """Inject a fake tflite_runtime.interpreter into sys.modules and mask ai_edge_litert.
+
+        ``_create_interpreter`` probes backends in priority order: ``ai_edge_litert`` first, then ``tflite_runtime``,
+        then ``tensorflow``. Masking ``ai_edge_litert`` and ``ai_edge_litert.interpreter`` to ``None`` forces the import
+        loop to fall through to the ``tflite_runtime`` path so tests exercise that branch regardless of what is
+        installed.
+
+        Python's import machinery resolves ``import tflite_runtime.interpreter`` by looking up
+        ``sys.modules["tflite_runtime.interpreter"]`` directly. We also set the ``interpreter`` attribute on the parent
+        package mock so attribute-path resolution is consistent regardless of Python version.
+        """
+        interp_instance = mock.MagicMock()
+        interp_instance.get_input_details.return_value = [{"shape": [1, 640, 640, 3], "dtype": np.float32}]
+        interp_instance.get_output_details.return_value = [
+            {"shape": [1, 300, 4], "name": "dets"},
+            {"shape": [1, 300, 81], "name": "labels"},
+        ]
+        interp_cls = mock.MagicMock(return_value=interp_instance)
+
+        # Build the submodule with a real Interpreter attribute
+        import types
+
+        mod = types.ModuleType("tflite_runtime.interpreter")
+        mod.Interpreter = interp_cls  # type: ignore[attr-defined]
+
+        # Build parent package that exposes mod as .interpreter
+        parent_mod = types.ModuleType("tflite_runtime")
+        parent_mod.interpreter = mod  # type: ignore[attr-defined]
+
+        with mock.patch.dict(
+            sys.modules,
+            {
+                **_AI_EDGE_LITERT_MASK,
+                "tflite_runtime": parent_mod,
+                "tflite_runtime.interpreter": mod,
+            },
+        ):
+            yield interp_cls, interp_instance
+
+    def test_uses_tflite_runtime_when_ai_edge_litert_absent(self, _mock_tflite_runtime) -> None:
+        """tflite_runtime is used as backend when ai_edge_litert is masked from the environment."""
+        interp_cls, interp_instance = _mock_tflite_runtime
+        _create_interpreter("model.tflite")
+        interp_cls.assert_called_once_with(model_path="model.tflite")
+
+    def test_allocate_tensors_called(self, _mock_tflite_runtime) -> None:
+        """allocate_tensors() is always called after construction."""
+        _, interp_instance = _mock_tflite_runtime
+        _create_interpreter("model.tflite")
+        interp_instance.allocate_tensors.assert_called_once()
+
+    def test_falls_back_to_tensorflow_when_tflite_runtime_missing(self) -> None:
+        """tensorflow.lite.Interpreter is used when tflite_runtime is absent."""
+        interp_instance = mock.MagicMock()
+        interp_instance.get_input_details.return_value = [{"shape": [1, 640, 640, 3], "dtype": np.float32}]
+        interp_instance.get_output_details.return_value = [{"shape": [1, 300, 4], "name": "dets"}]
+        tf_interp_cls = mock.MagicMock(return_value=interp_instance)
+
+        tf_lite_mod = mock.MagicMock()
+        tf_lite_mod.Interpreter = tf_interp_cls
+        tf_mod = mock.MagicMock()
+        tf_mod.lite = tf_lite_mod
+
+        with mock.patch.dict(
+            sys.modules,
+            {
+                **_AI_EDGE_LITERT_MASK,
+                "tflite_runtime": None,
+                "tflite_runtime.interpreter": None,
+                "tensorflow": tf_mod,
+                "tensorflow.lite": tf_lite_mod,
+            },
+        ):
+            _create_interpreter("model.tflite")
+
+        tf_interp_cls.assert_called_once_with(model_path="model.tflite")
+        interp_instance.allocate_tensors.assert_called_once()
+
+    def test_returns_interpreter(self, _mock_tflite_runtime) -> None:
+        """Return value is the interpreter instance (not the class)."""
+        _, interp_instance = _mock_tflite_runtime
+        result = _create_interpreter("model.tflite")
+        assert result is interp_instance
+
+    def test_logs_input_and_output_shapes(self, _mock_tflite_runtime) -> None:
+        """Logger.debug is called with 'Input' and 'Output' shape lines."""
+        with mock.patch("rfdetr.export._tflite.inference.logger") as mock_logger:
+            _create_interpreter("model.tflite")
+        debug_msgs = [call.args[0] for call in mock_logger.debug.call_args_list]
+        assert any("Input" in m for m in debug_msgs)
+        assert any("Output" in m for m in debug_msgs)
+
+    def test_accepts_path_object(self, _mock_tflite_runtime) -> None:
+        """Path objects are converted to strings before passing to Interpreter."""
+        interp_cls, _ = _mock_tflite_runtime
+        _create_interpreter(Path("model.tflite"))
+        call_kwargs = interp_cls.call_args[1]
+        assert call_kwargs["model_path"] == "model.tflite"
+        assert isinstance(call_kwargs["model_path"], str)
+
+    @pytest.fixture()
+    def _mock_ai_edge_litert(self):
+        """Inject a fake ai_edge_litert.interpreter into sys.modules and mask lower-priority backends.
+
+        Mirrors ``_mock_tflite_runtime`` for the first-priority backend so the ``ai_edge_litert.interpreter`` branch of
+        ``_create_interpreter`` can be exercised independently of whether the real package is installed.
+        """
+        interp_instance = mock.MagicMock()
+        interp_instance.get_input_details.return_value = [{"shape": [1, 640, 640, 3], "dtype": np.float32}]
+        interp_instance.get_output_details.return_value = [
+            {"shape": [1, 300, 4], "name": "dets"},
+            {"shape": [1, 300, 81], "name": "labels"},
+        ]
+        interp_cls = mock.MagicMock(return_value=interp_instance)
+
+        import types
+
+        mod = types.ModuleType("ai_edge_litert.interpreter")
+        mod.Interpreter = interp_cls  # type: ignore[attr-defined]
+
+        parent_mod = types.ModuleType("ai_edge_litert")
+        parent_mod.interpreter = mod  # type: ignore[attr-defined]
+
+        with mock.patch.dict(
+            sys.modules,
+            {
+                "ai_edge_litert": parent_mod,
+                "ai_edge_litert.interpreter": mod,
+                "tflite_runtime": None,
+                "tflite_runtime.interpreter": None,
+            },
+        ):
+            yield interp_cls, interp_instance
+
+    def test_uses_ai_edge_litert_when_available(self, _mock_ai_edge_litert) -> None:
+        """ai_edge_litert is used as the first-priority backend when it is importable."""
+        interp_cls, _ = _mock_ai_edge_litert
+        _create_interpreter("model.tflite")
+        interp_cls.assert_called_once_with(model_path="model.tflite")
+
+    def test_ai_edge_litert_allocate_tensors_called(self, _mock_ai_edge_litert) -> None:
+        """allocate_tensors() is called after construction via the ai_edge_litert backend."""
+        _, interp_instance = _mock_ai_edge_litert
+        _create_interpreter("model.tflite")
+        interp_instance.allocate_tensors.assert_called_once()
+
+    def test_ai_edge_litert_returns_interpreter(self, _mock_ai_edge_litert) -> None:
+        """Return value is the ai_edge_litert interpreter instance."""
+        _, interp_instance = _mock_ai_edge_litert
+        result = _create_interpreter("model.tflite")
+        assert result is interp_instance
+
+    def test_raises_when_no_backend_available(self) -> None:
+        """ImportError with a helpful install message is raised when all backends are absent."""
+        with mock.patch.dict(
+            sys.modules,
+            {
+                **_AI_EDGE_LITERT_MASK,
+                "tflite_runtime": None,
+                "tflite_runtime.interpreter": None,
+                "tensorflow": None,
+                "tensorflow.lite": None,
+            },
+        ):
+            with pytest.raises(ImportError, match="TFLite inference requires"):
+                _create_interpreter("model.tflite")
+
+
+# ---------------------------------------------------------------------------
+# TestRunInference
+# ---------------------------------------------------------------------------
+
+
+class TestRunInference:
+    """Tests for ``_run_inference()``."""
+
+    @pytest.fixture()
+    def rgb_image(self, tmp_path: Path) -> Path:
+        """Write a small RGB JPEG to a temp file and return its path."""
+        p = tmp_path / "image.jpg"
+        _save_rgb_image(p)
+        return p
+
+    @pytest.fixture()
+    def grayscale_image(self, tmp_path: Path) -> Path:
+        """Write a small grayscale PNG to a temp file and return its path."""
+        p = tmp_path / "gray.png"
+        _save_grayscale_image(p)
+        return p
+
+    def test_returns_detections_and_image(self, rgb_image: Path) -> None:
+        """Return type is tuple[sv.Detections, PIL.Image.Image]."""
+        interp = _make_interp()
+        result = _run_inference(interp, rgb_image)
+        assert isinstance(result, tuple)
+        dets, img = result
+        assert isinstance(dets, sv.Detections)
+        assert isinstance(img, PILImage.Image)
+
+    def test_detections_above_threshold_kept(self, rgb_image: Path) -> None:
+        """At least one detection is returned when one logit is high-confidence."""
+        interp = _make_interp(logits=_make_logits(high_conf_idx=0))
+        dets, _ = _run_inference(interp, rgb_image, threshold=0.3)
+        assert len(dets) >= 1
+
+    def test_detections_below_threshold_filtered(self, rgb_image: Path) -> None:
+        """No detections survive when all logits are zero (uniform probs < 0.3)."""
+        interp = _make_interp(logits=_make_logits(high_conf_idx=None))
+        dets, _ = _run_inference(interp, rgb_image, threshold=0.3)
+        assert len(dets) == 0
+
+    def test_boxes_in_pixel_space(self, rgb_image: Path) -> None:
+        """Xyxy coordinates are scaled to image pixel dimensions, not 0–1 range."""
+        img_size = (200, 100)  # (width, height) for PIL
+        PILImage.new("RGB", img_size, color=(100, 150, 200)).save(rgb_image)
+
+        # One centred box: cx=0.5, cy=0.5, w=0.2, h=0.2
+        boxes = np.array([[[0.5, 0.5, 0.2, 0.2]] + [[0.0, 0.0, 0.0, 0.0]] * 9], dtype=np.float32)
+        logits = _make_logits(high_conf_idx=0)
+        interp = _make_interp(boxes=boxes, logits=logits)
+
+        dets, _ = _run_inference(interp, rgb_image, threshold=0.3)
+        # With cx=0.5*200=100, cy=0.5*100=50, bw=0.2*200=40, bh=0.2*100=20
+        # xyxy expected: [80, 40, 120, 60]
+        assert dets.xyxy[0, 0] > 1.0  # x1 in pixel coords, not 0–1
+
+    def test_set_tensor_called_with_correct_shape(self, rgb_image: Path) -> None:
+        """set_tensor receives a tensor matching (1, H, W, C)."""
+        _, H, W, C = _INPUT_SHAPE  # noqa: N806
+        interp = _make_interp()
+        _run_inference(interp, rgb_image)
+        call_args = interp.set_tensor.call_args
+        tensor_arg = call_args[0][1]
+        assert tensor_arg.shape == (1, H, W, C)
+
+    def test_invoke_called_exactly_once(self, rgb_image: Path) -> None:
+        """interp.invoke() is called exactly once per inference call."""
+        interp = _make_interp()
+        _run_inference(interp, rgb_image)
+        interp.invoke.assert_called_once()
+
+    def test_grayscale_image_accepted(self, grayscale_image: Path) -> None:
+        """Grayscale (L-mode) input with C=1 is accepted without error."""
+        input_shape = [1, 224, 224, 1]
+        det_out = {"shape": [1, 10, 4], "name": "serving_default_dets:0", "index": 1}
+        label_out = {"shape": [1, 10, 82], "name": "serving_default_labels:0", "index": 2}
+        interp = _make_interp(input_shape=input_shape, out_dets=[det_out, label_out])
+        dets, _ = _run_inference(interp, grayscale_image)
+        assert isinstance(dets, sv.Detections)
+
+    def test_output_lookup_by_name_robust_to_ordering(self, rgb_image: Path) -> None:
+        """Swapping dets/labels order in get_output_details returns same detections."""
+        logits = _make_logits(high_conf_idx=0)
+        boxes = _make_boxes()
+
+        # Canonical order: dets first, labels second
+        interp_normal = _make_interp(boxes=boxes, logits=logits)
+        dets_normal, _ = _run_inference(interp_normal, rgb_image, threshold=0.3)
+
+        # Swapped order: labels first, dets second
+        det_out_swapped = {"shape": [1, 10, 4], "name": "serving_default_dets:0", "index": 1}
+        label_out_swapped = {"shape": [1, 10, 82], "name": "serving_default_labels:0", "index": 2}
+        interp_swapped = _make_interp(
+            out_dets=[label_out_swapped, det_out_swapped],
+            boxes=boxes,
+            logits=logits,
+        )
+        dets_swapped, _ = _run_inference(interp_swapped, rgb_image, threshold=0.3)
+
+        assert len(dets_normal) == len(dets_swapped)
+
+    def test_raises_for_non_float32_input_dtype(self, rgb_image: Path) -> None:
+        """ValueError raised when model input dtype is not float32."""
+        interp = mock.MagicMock()
+        interp.get_input_details.return_value = [{"shape": _INPUT_SHAPE, "index": 0, "dtype": np.uint8}]
+        interp.get_output_details.return_value = [_DET_OUTPUT, _LABEL_OUTPUT]
+        with pytest.raises(ValueError, match="float32"):
+            _run_inference(interp, rgb_image)
+
+
+# ---------------------------------------------------------------------------
+# TestSigmoidScoring
+# ---------------------------------------------------------------------------
+
+
+class TestSigmoidScoring:
+    """Tests for per-class sigmoid scoring introduced in _run_inference."""
+
+    @pytest.fixture()
+    def rgb_image(self, tmp_path: Path) -> Path:
+        """Write a small RGB JPEG to a temp file and return its path."""
+        p = tmp_path / "image.jpg"
+        _save_rgb_image(p)
+        return p
+
+    def test_high_logit_yields_confidence_near_one(self, rgb_image: Path) -> None:
+        """Logit of 10.0 produces sigmoid ≈ 0.9999; confidence[0] > 0.99."""
+        logits = _make_logits(high_conf_idx=0)  # logits[0, 0, 0] = 10.0
+        interp = _make_interp(logits=logits)
+        dets, _ = _run_inference(interp, rgb_image, threshold=0.3)
+        assert dets.confidence[0] > 0.99
+
+    def test_low_logit_filtered_at_threshold(self, rgb_image: Path) -> None:
+        """Logit of -10.0 produces sigmoid ≈ 0.0001; detection filtered at threshold=0.3."""
+        logits = np.full((1, 10, 82), -10.0, dtype=np.float32)
+        interp = _make_interp(logits=logits)
+        dets, _ = _run_inference(interp, rgb_image, threshold=0.3)
+        assert len(dets) == 0
+
+    def test_multiclass_class_id_is_argmax_of_logits(self, rgb_image: Path) -> None:
+        """Argmax of sigmoid equals argmax of logits; query with [5,2,1,...] gets class_id==0."""
+        # Shape (1, 10, 82): first query has logits [5, 2, 1, 0, ...], rest are -100
+        logits = np.full((1, 10, 82), -100.0, dtype=np.float32)
+        logits[0, 0, 0] = 5.0
+        logits[0, 0, 1] = 2.0
+        logits[0, 0, 2] = 1.0
+        interp = _make_interp(logits=logits)
+        dets, _ = _run_inference(interp, rgb_image, threshold=0.3)
+        # argmax of sigmoid == argmax of logits because sigmoid is monotone increasing
+        assert dets.class_id[0] == 0
+
+
+# ---------------------------------------------------------------------------
+# TestShapeBasedOutputFallback
+# ---------------------------------------------------------------------------
+
+# Generic output detail dicts used across shape-based fallback tests.
+# Indices mirror the canonical ones so _make_interp's _get_tensor dispatch works.
+_GENERIC_DET_OUTPUT = {"shape": [1, 10, 4], "name": "Identity_0", "index": 1}
+_GENERIC_LABEL_OUTPUT = {"shape": [1, 10, 82], "name": "Identity_1", "index": 2}
+
+
+class TestShapeBasedOutputFallback:
+    """Tests for the shape-based output matching fallback in _run_inference."""
+
+    @pytest.fixture()
+    def rgb_image(self, tmp_path: Path) -> Path:
+        """Write a small RGB JPEG to a temp file and return its path."""
+        p = tmp_path / "image.jpg"
+        _save_rgb_image(p)
+        return p
+
+    def test_unambiguous_shapes_inferred_correctly(self, rgb_image: Path) -> None:
+        """Generic names with shapes [1,10,4] and [1,10,82] resolve without error."""
+        interp = _make_interp(
+            out_dets=[_GENERIC_DET_OUTPUT, _GENERIC_LABEL_OUTPUT],
+            logits=_make_logits(high_conf_idx=0),
+        )
+        dets, _ = _run_inference(interp, rgb_image, threshold=0.3)
+        assert isinstance(dets, sv.Detections)
+        assert len(dets) >= 1
+
+    def test_ambiguous_shapes_two_outputs_positional_fallback(self, rgb_image: Path) -> None:
+        """When both outputs have last-dim==4 (num_classes==3) and there are exactly 2, positional fallback is used."""
+        # num_classes=3 → logits shape last-dim==4; boxes last-dim==4 → ambiguous
+        # Positional order: index 0 = boxes (Identity_0, tensor index 1), index 1 = logits (Identity_1, tensor index 2)
+        ambiguous_dets = {"shape": [1, 10, 4], "name": "Identity_0", "index": 1}
+        ambiguous_labels = {"shape": [1, 10, 4], "name": "Identity_1", "index": 2}
+        # Build logits of shape (1, 10, 4) so last col is dropped → (10, 3) per-class
+        logits_ambiguous = np.full((1, 10, 4), -10.0, dtype=np.float32)
+        logits_ambiguous[0, 0, 0] = 10.0  # first query, first class → high confidence
+        interp = _make_interp(
+            out_dets=[ambiguous_dets, ambiguous_labels],
+            logits=logits_ambiguous,
+        )
+        dets, _ = _run_inference(interp, rgb_image, threshold=0.3)
+        assert isinstance(dets, sv.Detections)
+        assert len(dets) >= 1
+
+    def test_three_outputs_all_dim4_raises_value_error(self, rgb_image: Path) -> None:
+        """3 outputs all with last-dim==4 and no name match raises ValueError with expected message."""
+        # Need a third tensor index; extend _get_tensor via a custom mock
+        third_output = {"shape": [1, 10, 4], "name": "Identity_2", "index": 3}
+        boxes = _make_boxes()
+        logits = _make_logits()
+
+        def _get_tensor(index: int) -> np.ndarray:
+            if index == 1:
+                return boxes
+            if index in (2, 3):
+                return logits
+            raise ValueError(f"Unknown tensor index: {index}")
+
+        interp = mock.MagicMock()
+        interp.get_input_details.return_value = [{"shape": _INPUT_SHAPE, "index": 0, "dtype": np.float32}]
+        interp.get_output_details.return_value = [
+            {"shape": [1, 10, 4], "name": "Identity_0", "index": 1},
+            {"shape": [1, 10, 4], "name": "Identity_1", "index": 2},
+            third_output,
+        ]
+        interp.get_tensor.side_effect = _get_tensor
+
+        with pytest.raises(ValueError, match="Shape-based TFLite output matching failed"):
+            _run_inference(interp, rgb_image, threshold=0.3)
+
+    def test_three_outputs_with_rank4_masks_resolves_correctly(self, rgb_image: Path) -> None:
+        """3-output segmentation export (boxes/logits/masks) with generic names resolves without error.
+
+        Ensures the shape fallback ignores the rank-4 masks tensor and correctly identifies boxes [1,Q,4] and logits
+        [1,Q,C+1] as rank-3 candidates.
+        """
+        boxes = _make_boxes()
+        logits = _make_logits(high_conf_idx=0)
+        masks = np.zeros((1, 10, 28, 28), dtype=np.float32)
+
+        def _get_tensor(index: int) -> np.ndarray:
+            if index == 1:
+                return boxes
+            if index == 2:
+                return logits
+            if index == 3:
+                return masks
+            raise ValueError(f"Unknown tensor index: {index}")
+
+        interp = mock.MagicMock()
+        interp.get_input_details.return_value = [{"shape": _INPUT_SHAPE, "index": 0, "dtype": np.float32}]
+        interp.get_output_details.return_value = [
+            {"shape": [1, 10, 4], "name": "Identity_0", "index": 1},
+            {"shape": [1, 10, 82], "name": "Identity_1", "index": 2},
+            {"shape": [1, 10, 28, 28], "name": "Identity_2", "index": 3},
+        ]
+        interp.get_tensor.side_effect = _get_tensor
+
+        dets, _ = _run_inference(interp, rgb_image, threshold=0.3)
+        assert isinstance(dets, sv.Detections)
+        assert len(dets) >= 1
+
+
+# ---------------------------------------------------------------------------
+# TestMaskDecoding
+# ---------------------------------------------------------------------------
+
+
+class TestMaskDecoding:
+    """Tests for ``_decode_masks()`` and mask decoding in ``_run_inference()``."""
+
+    @pytest.fixture()
+    def rgb_image(self, tmp_path: Path) -> Path:
+        """Write a small RGB JPEG to a temp file and return its path."""
+        p = tmp_path / "image.jpg"
+        _save_rgb_image(p)
+        return p
+
+    def test_decode_masks_shape_and_dtype(self) -> None:
+        """Output shape is (K, height, width) from out_size=(width, height); dtype is bool."""
+        out = _decode_masks(np.zeros((3, 10, 10), dtype=np.float32), (40, 20))
+        assert out.shape == (3, 20, 40)
+        assert out.dtype == bool
+
+    def test_decode_masks_thresholds_at_zero(self) -> None:
+        """Positive logits decode to True, negative logits to False."""
+        logits = np.stack(
+            [
+                np.full((8, 8), 5.0, dtype=np.float32),
+                np.full((8, 8), -5.0, dtype=np.float32),
+            ]
+        )
+        out = _decode_masks(logits, (16, 16))
+        assert out[0].all()
+        assert not out[1].any()
+
+    def test_decode_masks_empty_input(self) -> None:
+        """Zero masks in yields a (0, height, width) array, not an error."""
+        out = _decode_masks(np.zeros((0, 10, 10), dtype=np.float32), (32, 32))
+        assert out.shape == (0, 32, 32)
+
+    def test_run_inference_decodes_masks_for_seg_model(self, rgb_image: Path) -> None:
+        """A 3-output segmentation export populates Detections.mask at image size."""
+        boxes = _make_boxes()
+        logits = _make_logits(high_conf_idx=0)
+        masks = np.full((1, 10, 28, 28), -10.0, dtype=np.float32)
+        masks[0, 0] = 10.0  # query 0 (the kept detection) gets an all-positive mask
+
+        def _get_tensor(index: int) -> np.ndarray:
+            return {1: boxes, 2: logits, 3: masks}[index]
+
+        interp = mock.MagicMock()
+        interp.get_input_details.return_value = [{"shape": _INPUT_SHAPE, "index": 0, "dtype": np.float32}]
+        interp.get_output_details.return_value = [
+            {"shape": [1, 10, 4], "name": "Identity_0", "index": 1},
+            {"shape": [1, 10, 82], "name": "Identity_1", "index": 2},
+            {"shape": [1, 10, 28, 28], "name": "Identity_2", "index": 3},
+        ]
+        interp.get_tensor.side_effect = _get_tensor
+
+        dets, img = _run_inference(interp, rgb_image, threshold=0.3)
+        assert dets.mask is not None
+        assert dets.mask.shape == (len(dets), img.height, img.width)
+        assert dets.mask.dtype == bool
+        assert dets.mask[0].all()  # query 0's all-positive logits decode to a full mask
+
+    def test_run_inference_no_mask_for_detection_model(self, rgb_image: Path) -> None:
+        """A 2-output detection export leaves Detections.mask as None."""
+        interp = _make_interp(logits=_make_logits(high_conf_idx=0))
+        dets, _ = _run_inference(interp, rgb_image, threshold=0.3)
+        assert dets.mask is None
+
+    def test_run_inference_name_based_mask_detection(self, rgb_image: Path) -> None:
+        """Output named 'masks:0' exercises the name-based path and sets Detections.mask."""
+        boxes = _make_boxes()
+        logits = _make_logits(high_conf_idx=0)
+        masks = np.full((1, 10, 28, 28), 10.0, dtype=np.float32)
+
+        def _get_tensor(index: int) -> np.ndarray:
+            return {1: boxes, 2: logits, 3: masks}[index]
+
+        interp = mock.MagicMock()
+        interp.get_input_details.return_value = [{"shape": _INPUT_SHAPE, "index": 0, "dtype": np.float32}]
+        interp.get_output_details.return_value = [
+            {"shape": [1, 10, 4], "name": "serving_default_dets:0", "index": 1},
+            {"shape": [1, 10, 82], "name": "serving_default_labels:0", "index": 2},
+            {"shape": [1, 10, 28, 28], "name": "serving_default_masks:0", "index": 3},
+        ]
+        interp.get_tensor.side_effect = _get_tensor
+
+        dets, _ = _run_inference(interp, rgb_image, threshold=0.3)
+        assert dets.mask is not None
+
+    def test_run_inference_seg_model_no_detections_returns_none_mask(self, rgb_image: Path) -> None:
+        """Seg model with all scores below threshold returns mask=None (keep.any() is False)."""
+        boxes = _make_boxes()
+        logits = _make_logits(high_conf_idx=None)  # all scores near zero, below threshold
+        masks = np.full((1, 10, 28, 28), 10.0, dtype=np.float32)
+
+        def _get_tensor(index: int) -> np.ndarray:
+            return {1: boxes, 2: logits, 3: masks}[index]
+
+        interp = mock.MagicMock()
+        interp.get_input_details.return_value = [{"shape": _INPUT_SHAPE, "index": 0, "dtype": np.float32}]
+        interp.get_output_details.return_value = [
+            {"shape": [1, 10, 4], "name": "Identity_0", "index": 1},
+            {"shape": [1, 10, 82], "name": "Identity_1", "index": 2},
+            {"shape": [1, 10, 28, 28], "name": "Identity_2", "index": 3},
+        ]
+        interp.get_tensor.side_effect = _get_tensor
+
+        dets, _ = _run_inference(interp, rgb_image, threshold=0.3)
+        assert len(dets) == 0
+        assert dets.mask is None
+
+    def test_decode_masks_raises_on_wrong_rank(self) -> None:
+        """_decode_masks raises ValueError when input is not rank-3."""
+        with pytest.raises(ValueError, match="rank-3"):
+            _decode_masks(np.zeros((10, 28, 28, 1), dtype=np.float32), (56, 56))
+
+    def test_decode_masks_exact_zero_logit_decodes_to_false(self) -> None:
+        """Logit exactly 0.0 is not > 0.0 and decodes to False (strict threshold)."""
+        zero_logits = np.zeros((1, 8, 8), dtype=np.float32)
+        out = _decode_masks(zero_logits, (16, 16))
+        assert not out.any()
+
+    def test_decode_masks_non_square_logit_input(self) -> None:
+        """Non-square logit map (K, Hm, Wm) with Hm != Wm resizes to the correct output shape."""
+        logits = np.full((3, 7, 14), 5.0, dtype=np.float32)
+        out = _decode_masks(logits, (56, 28))  # out_size=(width=56, height=28)
+        assert out.shape == (3, 28, 56)
+        assert out.all()  # all-positive logits → all True
+
+    def test_decode_masks_parity_positive_negative_regions(self) -> None:
+        """Positive/negative logit regions map correctly after bilinear upsample + threshold.
+
+        Uses high-magnitude logits (±10) so no ambiguity near the boundary; verifies the core _decode_masks contract
+        matches the >0 PostProcess.forward equivalent.
+        """
+        logits = np.full((1, 14, 14), -10.0, dtype=np.float32)
+        logits[0, :7, :] = 10.0  # top half strongly positive, bottom half strongly negative
+        out = _decode_masks(logits, (28, 28))
+        # Interior rows well away from the half-way boundary
+        assert out[0, 1:6, :].all()  # top rows → all True
+        assert not out[0, 15:27, :].any()  # bottom rows → all False
+
+
+# ---------------------------------------------------------------------------
+# TestBilinearResizeHalfPixel
+# ---------------------------------------------------------------------------
+
+
+class TestBilinearResizeHalfPixel:
+    """Tests for ``_bilinear_resize_half_pixel()``."""
+
+    def test_output_shape(self) -> None:
+        """Output shape is (K, out_h, out_w)."""
+        src = np.ones((3, 8, 8), dtype=np.float32)
+        out = _bilinear_resize_half_pixel(src, 16, 16)
+        assert out.shape == (3, 16, 16)
+
+    def test_output_dtype_is_float32(self) -> None:
+        """Output dtype is float32 regardless of input magnitude."""
+        src = np.ones((1, 4, 4), dtype=np.float32)
+        out = _bilinear_resize_half_pixel(src, 8, 8)
+        assert out.dtype == np.float32
+
+    def test_identity_when_no_resize(self) -> None:
+        """Output equals input when target dimensions match source dimensions."""
+        rng = np.random.default_rng(0)
+        src = rng.random((2, 8, 8)).astype(np.float32)
+        out = _bilinear_resize_half_pixel(src, 8, 8)
+        np.testing.assert_allclose(out, src, atol=1e-6)
+
+    @pytest.mark.parametrize(
+        ("src_shape", "out_h", "out_w"),
+        [
+            pytest.param((1, 4, 4), 8, 8, id="upsample_square"),
+            pytest.param((3, 7, 5), 14, 10, id="upsample_nonsquare"),
+            pytest.param((2, 8, 8), 4, 4, id="downsample"),
+            pytest.param((1, 1, 1), 3, 3, id="degenerate_1x1"),
+        ],
+    )
+    def test_parity_with_torch_interpolate(self, src_shape: tuple[int, int, int], out_h: int, out_w: int) -> None:
+        """Output matches F.interpolate(mode='bilinear', align_corners=False) to within 1e-5."""
+        torch = pytest.importorskip("torch")
+        import torch.nn.functional as F  # noqa: N812
+
+        rng = np.random.default_rng(42)
+        src = rng.random(src_shape).astype(np.float32)
+
+        result = _bilinear_resize_half_pixel(src, out_h, out_w)
+
+        t = torch.from_numpy(src).unsqueeze(0)
+        with torch.no_grad():
+            ref = F.interpolate(t, size=(out_h, out_w), mode="bilinear", align_corners=False)
+        ref_np = ref.squeeze(0).numpy()
+
+        np.testing.assert_allclose(result, ref_np, atol=1e-5)
+
+
+# ---------------------------------------------------------------------------
+# TestPreprocessImage
+# ---------------------------------------------------------------------------
+
+
+class TestPreprocessImage:
+    """Tests for ``_preprocess_image()``."""
+
+    def test_output_shape_rgb(self) -> None:
+        """RGB image returns float32 array of shape (1, H, W, 3)."""
+        pil_img = PILImage.new("RGB", (100, 80))
+        out = _preprocess_image(pil_img, (64, 64))
+        assert out.shape == (1, 64, 64, 3)
+        assert out.dtype == np.float32
+
+    def test_output_shape_grayscale(self) -> None:
+        """Grayscale image with channels=1 returns float32 array of shape (1, H, W, 1)."""
+        pil_img = PILImage.new("L", (100, 80))
+        out = _preprocess_image(pil_img, (64, 64), channels=1)
+        assert out.shape == (1, 64, 64, 1)
+        assert out.dtype == np.float32
+
+    def test_output_values_are_normalized(self) -> None:
+        """ImageNet normalization shifts black-pixel output below -1.0."""
+        pil_img = PILImage.new("RGB", (32, 32), color=(0, 0, 0))
+        out = _preprocess_image(pil_img, (32, 32))
+        # pixel 0 → 0.0 → (0.0 - 0.485) / 0.229 ≈ -2.12
+        assert out.min() < -1.0
+
+    def test_pil_fallback_when_torch_unavailable(self) -> None:
+        """PIL path is used when torch is masked from sys.modules."""
+        pil_img = PILImage.new("RGB", (100, 80))
+        with mock.patch.dict(
+            sys.modules,
+            {
+                "torch": None,
+                "torchvision": None,
+                "torchvision.transforms": None,
+                "torchvision.transforms.functional": None,
+            },
+        ):
+            out = _preprocess_image(pil_img, (64, 64))
+        assert out.shape == (1, 64, 64, 3)
+        assert out.dtype == np.float32

tests/export/test_tflite_inference_parity.py

@@ -0,0 +1,261 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Preprocessing parity tests: ``_run_inference`` must produce essentially the same input tensor as ``RFDETR.predict``
+for the same source image, otherwise the TFLite-exported model is fed inputs the PyTorch graph never saw and detections
+drift.
+
+History: an earlier version of ``_run_inference`` called ``PIL.Image.resize`` without a filter argument, picking up
+PIL's default (BICUBIC since Pillow 9.1.0). PyTorch's predict() path uses torchvision ``F.resize`` (BILINEAR). The
+mismatch caused IoU drift up to 0.36 on detail-rich images and a 2-class-mismatch FP16 disaster on the ``dog`` test
+image. This test exists to keep ``_preprocess_image`` locked to BILINEAR -- any regression that re-introduces BICUBIC or
+otherwise shifts the resize filter will surface here.
+"""
+
+from __future__ import annotations
+
+import sys
+
+import numpy as np
+import pytest
+import torchvision.transforms.functional as F  # noqa: N812
+from PIL import Image as PILImage
+
+from rfdetr.export._tflite.inference import _bilinear_resize_half_pixel, _preprocess_image
+
+IMAGENET_MEAN = [0.485, 0.456, 0.406]
+IMAGENET_STD = [0.229, 0.224, 0.225]
+
+# Bound for max abs diff in normalised space between the PyTorch and TFLite preprocessing pipelines.
+# With torchvision available (which is always the case in this repo's CI -- it's a hard rfdetr
+# dependency) _preprocess_image runs the same torchvision call PyTorch's predict() uses, so the
+# tensors are bit-exact. The 0.05 bound is generous so the torch-free fallback path (which uses
+# PIL.BILINEAR and shows ~0.016 max diff) also passes; the BICUBIC regression would push max diff
+# to ~0.5, well above the bound.
+MAX_ABS_DIFF_BOUND = 0.05
+# When torchvision is available the inference path matches torchvision-resize byte-for-byte, so
+# the diff is effectively zero modulo floating-point noise.
+BIT_EXACT_BOUND = 1e-5
+
+
+def _pytorch_preprocess(pil_img: PILImage.Image, hw: tuple[int, int]) -> np.ndarray:
+    """Mirror of the PyTorch predict() preprocessing: to_tensor -> resize -> normalize."""
+    img = F.to_tensor(pil_img)
+    img = F.resize(img, list(hw))
+    img = F.normalize(img, IMAGENET_MEAN, IMAGENET_STD)
+    return img.unsqueeze(0).numpy()
+
+
+def _tflite_preprocess_to_nchw(pil_img: PILImage.Image, hw: tuple[int, int]) -> np.ndarray:
+    """Call ``_preprocess_image`` and convert NHWC -> NCHW for apples-to-apples comparison."""
+    nhwc = _preprocess_image(pil_img, hw, channels=3)
+    return nhwc.transpose(0, 3, 1, 2)
+
+
+def _make_synthetic_rgb(seed: int, size: tuple[int, int]) -> PILImage.Image:
+    """Deterministic synthetic RGB image with structure (not pure noise) so resize filtering matters."""
+    rng = np.random.default_rng(seed)
+    height, width = size
+    base = rng.integers(0, 256, size=(height // 8, width // 8, 3), dtype=np.uint8)
+    pil_small = PILImage.fromarray(base, mode="RGB")
+    return pil_small.resize((width, height), getattr(PILImage, "Resampling", PILImage).NEAREST)
+
+
+class TestPreprocessingParity:
+    """``_preprocess_image`` must match PyTorch's predict() preprocessing within MAX_ABS_DIFF_BOUND.
+
+    Three shapes cover the common downscale ratios produced by RFDETR exports:
+      - 1280x720 -> 384x384 (nano default): heavy downscale, the case that surfaced the BICUBIC bug
+      - 800x600  -> 384x384: moderate downscale, mixed aspect ratio
+      - 384x384  -> 384x384: identity resize -- only normalisation differs (rounding noise only)
+    """
+
+    @pytest.mark.parametrize(
+        ("src_size", "tgt_size", "seed"),
+        [
+            pytest.param((1280, 720), (384, 384), 0, id="1280x720_to_384x384"),
+            pytest.param((800, 600), (384, 384), 1, id="800x600_to_384x384"),
+            pytest.param((384, 384), (384, 384), 2, id="identity_384x384"),
+        ],
+    )
+    def test_matches_pytorch_predict_preprocessing(
+        self, src_size: tuple[int, int], tgt_size: tuple[int, int], seed: int
+    ) -> None:
+        pil = _make_synthetic_rgb(seed, (src_size[1], src_size[0]))  # _make takes (H, W)
+        pt = _pytorch_preprocess(pil, tgt_size)
+        tf = _tflite_preprocess_to_nchw(pil, tgt_size)
+
+        assert pt.shape == tf.shape, f"shape mismatch: PT {pt.shape} vs TF {tf.shape}"
+        max_diff = float(np.abs(pt - tf).max())
+        # torchvision is a hard rfdetr dependency, so in this test environment _preprocess_image
+        # uses the torchvision path and matches PyTorch byte-for-byte. The torch-free fallback is
+        # exercised separately by test_torch_free_fallback_still_close.
+        assert max_diff < BIT_EXACT_BOUND, (
+            f"PyTorch vs TFLite preprocessing diverged: max|diff|={max_diff:.6f} exceeds "
+            f"{BIT_EXACT_BOUND}. With torchvision available, _preprocess_image should be using "
+            f"torchvision.transforms.functional.resize and the diff should be effectively zero. "
+            f"If this fires, check that the torch/torchvision import path inside _preprocess_image "
+            f"hasn't been broken."
+        )
+
+    def test_grayscale_channel_handling(self) -> None:
+        """Grayscale (channels=1) path must produce shape (1, H, W, 1)."""
+        rng = np.random.default_rng(3)
+        height, width = 256, 256
+        pil = PILImage.fromarray(rng.integers(0, 256, size=(height, width), dtype=np.uint8), mode="L")
+        tf = _preprocess_image(pil, (128, 128), channels=1)
+        assert tf.shape == (1, 128, 128, 1), f"unexpected shape: {tf.shape}"
+        assert tf.dtype == np.float32
+
+    def test_returns_nhwc_float32(self) -> None:
+        """``_preprocess_image`` returns NHWC float32 with a leading batch dim."""
+        pil = _make_synthetic_rgb(seed=7, size=(64, 64))
+        tf = _preprocess_image(pil, (32, 32), channels=3)
+        assert tf.shape == (1, 32, 32, 3)
+        assert tf.dtype == np.float32
+
+    def test_normalisation_uses_imagenet_stats(self) -> None:
+        """A mid-gray (128) image should land near zero on all channels after normalisation."""
+        gray = np.full((64, 64, 3), 128, dtype=np.uint8)
+        pil = PILImage.fromarray(gray, mode="RGB")
+        tf = _preprocess_image(pil, (32, 32), channels=3)
+        # 128/255 ~= 0.502; expected normalised values per channel:
+        #   (0.502 - 0.485) / 0.229 ~=  0.074
+        #   (0.502 - 0.456) / 0.224 ~=  0.205
+        #   (0.502 - 0.406) / 0.225 ~=  0.426
+        expected = np.array([(128 / 255.0 - IMAGENET_MEAN[c]) / IMAGENET_STD[c] for c in range(3)], dtype=np.float32)
+        per_channel_mean = tf[0].mean(axis=(0, 1))
+        np.testing.assert_allclose(per_channel_mean, expected, atol=1e-3)
+
+    def test_torch_free_fallback_still_close(self) -> None:
+        """Simulate the torch-free environment by masking torch imports; assert the PIL fallback still stays within the
+        looser MAX_ABS_DIFF_BOUND.
+
+        This documents the gap users on edge deployments without torch installed will see (versus the bit-exact
+        torchvision path).
+        """
+        from unittest import mock
+
+        pil = _make_synthetic_rgb(seed=11, size=(720, 1280))
+        tgt = (384, 384)
+        pt = _pytorch_preprocess(pil, tgt)
+
+        # Hide torch from _preprocess_image's lazy import, forcing the PIL fallback.
+        with mock.patch.dict(sys.modules, {"torch": None}):
+            tf = _tflite_preprocess_to_nchw(pil, tgt)
+
+        max_diff = float(np.abs(pt - tf).max())
+        assert max_diff < MAX_ABS_DIFF_BOUND, (
+            f"Torch-free PIL fallback diverged: max|diff|={max_diff:.4f} > {MAX_ABS_DIFF_BOUND}. "
+            "The fallback uses PIL.BILINEAR which should keep diff ~0.016; a regression to "
+            "BICUBIC would push it ~30x larger."
+        )
+
+
+class TestPreprocessingFilterRegression:
+    """Direct comparison of BILINEAR vs BICUBIC.
+
+    Asserts the current code stays on BILINEAR by showing BICUBIC would produce a much larger divergence.
+    """
+
+    def test_bicubic_would_be_much_worse(self) -> None:
+        """If a future change reverts to BICUBIC default, this confirms how much worse it gets."""
+        pil = _make_synthetic_rgb(seed=42, size=(720, 1280))
+        tgt = (384, 384)
+
+        pt = _pytorch_preprocess(pil, tgt)
+        tf_current = _tflite_preprocess_to_nchw(pil, tgt)
+
+        # Simulate the regression: PIL default (BICUBIC since 9.1.0).
+        mean = np.array(IMAGENET_MEAN, dtype=np.float32)
+        std = np.array(IMAGENET_STD, dtype=np.float32)
+        height, width = tgt
+        arr_bicubic = np.array(pil.convert("RGB").resize((width, height)), dtype=np.float32) / 255.0
+        tf_bicubic = ((arr_bicubic - mean) / std)[np.newaxis].transpose(0, 3, 1, 2)
+
+        max_diff_current = float(np.abs(pt - tf_current).max())
+        max_diff_bicubic = float(np.abs(pt - tf_bicubic).max())
+
+        # The BILINEAR fix must be at least 5x closer to PyTorch than the BICUBIC regression.
+        # In practice it's ~30x closer; the 5x floor is forgiving of pillow / numpy drift.
+        assert max_diff_current * 5 < max_diff_bicubic, (
+            f"_preprocess_image is too close to BICUBIC behaviour: "
+            f"current max|diff|={max_diff_current:.4f}, BICUBIC max|diff|={max_diff_bicubic:.4f}. "
+            f"Check that _PIL_BILINEAR is being passed to .resize()."
+        )
+
+
+class TestBilinearResizeHalfPixelParity:
+    """``_bilinear_resize_half_pixel`` is the torch-free fallback used by ``_decode_masks``.
+
+    It must match ``torch.nn.functional.interpolate(..., mode="bilinear", align_corners=False)``
+    -- the same call ``PostProcess.forward`` uses -- byte-for-byte modulo float noise. Sharp-edge
+    inputs are the worst case: even a sub-pixel shift in the half-pixel convention flips boundary
+    pixels and tanks mask IoU.
+    """
+
+    @staticmethod
+    def _torch_interpolate(src: np.ndarray, out_hw: tuple[int, int]) -> np.ndarray:
+        """Reference implementation: ``F.interpolate`` with ``align_corners=False``."""
+        import torch
+        import torch.nn.functional as TF  # noqa: N812
+
+        with torch.no_grad():
+            t = torch.from_numpy(src.astype(np.float32)).unsqueeze(0)
+            out = TF.interpolate(t, size=out_hw, mode="bilinear", align_corners=False)
+        return out.squeeze(0).cpu().numpy()
+
+    @pytest.mark.parametrize(
+        ("src_hw", "out_hw"),
+        [
+            pytest.param((28, 28), (384, 384), id="upsample_28_to_384"),
+            pytest.param((56, 56), (256, 256), id="upsample_56_to_256"),
+            pytest.param((100, 100), (100, 100), id="identity_100"),
+            pytest.param((100, 100), (50, 50), id="downsample_100_to_50"),
+            pytest.param((40, 60), (200, 400), id="non_square_upsample"),
+        ],
+    )
+    def test_matches_torch_interpolate_on_random_logits(self, src_hw: tuple[int, int], out_hw: tuple[int, int]) -> None:
+        """Random logits over a small batch must resize identically to ``F.interpolate``."""
+        rng = np.random.default_rng(0)
+        src = rng.standard_normal((3, *src_hw)).astype(np.float32) * 4.0
+
+        ours = _bilinear_resize_half_pixel(src, out_hw[0], out_hw[1])
+        ref = self._torch_interpolate(src, out_hw)
+
+        max_diff = float(np.abs(ours - ref).max())
+        # 1e-4 absorbs the float32 op-order noise that accumulates on large upsample ratios
+        # (mine: split bilinear sums in pure numpy; torch: fused kernel). Half-pixel-convention
+        # drift would push this several orders of magnitude higher.
+        assert max_diff < 1e-4, (
+            f"_bilinear_resize_half_pixel diverged from F.interpolate(align_corners=False): "
+            f"max|diff|={max_diff:.2e} on shape {src_hw} -> {out_hw}. "
+            "Half-pixel convention drift would surface here."
+        )
+
+    def test_sharp_edge_mask_matches_torch(self) -> None:
+        """A mask with a sharp left/right boundary is the regression-prone case.
+
+        This is the shape ``_decode_masks`` actually consumes (logits with a zero-crossing). A half-pixel shift would
+        flip the boundary column and is exactly what the original PIL.BILINEAR path got wrong.
+        """
+        src = np.full((1, 28, 28), -10.0, dtype=np.float32)
+        src[0, :, 14:] = 10.0  # sharp vertical edge at column 14
+
+        out_hw = (224, 224)
+        ours = _bilinear_resize_half_pixel(src, out_hw[0], out_hw[1])
+        ref = self._torch_interpolate(src, out_hw)
+
+        max_diff = float(np.abs(ours - ref).max())
+        assert max_diff < 1e-4, (
+            f"Sharp-edge resize diverged from F.interpolate: max|diff|={max_diff:.2e}. "
+            "This is the case that previously dropped mask IoU below 0.6 with PIL.BILINEAR."
+        )
+
+        # Also assert the thresholded output matches: this is what _decode_masks actually returns.
+        assert np.array_equal(ours > 0, ref > 0), (
+            "Boolean mask after thresholding diverged from F.interpolate. Even a single column of "
+            "flipped pixels would show up here -- the exact failure mode the original PR fixes."
+        )

tests/inference/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/inference/helpers.py

@@ -0,0 +1,106 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Shared test helpers for the inference test suite.
+
+Plain classes and functions (not pytest fixtures) shared across multiple test modules to avoid verbatim duplication.
+Import with a relative import::
+
+    from .helpers import _BaseFakeRFDETR, _DummyModel, _DummyRFDETR
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+
+from rfdetr.detr import RFDETR
+
+
+class _BaseFakeRFDETR(RFDETR):
+    """RFDETR test double that skips weight downloads and returns a minimal model config.
+
+    Subclasses must override ``get_model`` to supply the model context appropriate for
+    the scenario under test.
+
+    Examples:
+        This class is imported directly by test modules that need a weight-free RFDETR.
+    """
+
+    def maybe_download_pretrain_weights(self) -> None:
+        """Skip weight download in tests."""
+        return None
+
+    def get_model_config(self, **kwargs: object) -> SimpleNamespace:
+        """Return a minimal config sufficient for most test scenarios."""
+        return SimpleNamespace(num_channels=3)
+
+
+class _DummyModel:
+    """Minimal model stub that returns deterministic postprocessed results.
+
+    Examples:
+        >>> m = _DummyModel(labels=[0, 1])
+        >>> len(m._labels)
+        2
+    """
+
+    def __init__(
+        self,
+        class_names: list[str] | None = None,
+        labels: list[int] | None = None,
+        include_keypoints: bool = False,
+        num_keypoints: int = 17,
+    ) -> None:
+        """Initialise stub with optional class names, label list, and keypoint flag."""
+        self.device = torch.device("cpu")
+        self.resolution = 28
+        self.model = torch.nn.Identity()
+        self.class_names = class_names
+        self._labels = labels if labels is not None else [1]
+        self._include_keypoints = include_keypoints
+        self._num_keypoints = num_keypoints
+
+    def postprocess(self, predictions: Any, target_sizes: torch.Tensor) -> list[dict[str, torch.Tensor]]:
+        """Return fixed scores/boxes (and optional keypoints) for every image in the batch."""
+        batch = target_sizes.shape[0]
+        results = []
+        for _ in range(batch):
+            result: dict[str, torch.Tensor] = {
+                "scores": torch.tensor([0.9] * len(self._labels)),
+                "labels": torch.tensor(self._labels),
+                "boxes": torch.tensor([[0.0, 0.0, 1.0, 1.0]] * len(self._labels)),
+            }
+            if self._include_keypoints:
+                result["keypoints"] = torch.full((len(self._labels), self._num_keypoints, 3), 0.5, dtype=torch.float32)
+                result["keypoint_precision_cholesky"] = torch.full(
+                    (len(self._labels), self._num_keypoints, 3), 0.25, dtype=torch.float32
+                )
+            results.append(result)
+        return results
+
+
+class _DummyRFDETR(RFDETR):
+    """Weight-free RFDETR that delegates to ``_DummyModel`` for all inference.
+
+    Examples:
+        >>> m = _DummyRFDETR()
+        >>> isinstance(m.model, _DummyModel)
+        True
+    """
+
+    def maybe_download_pretrain_weights(self) -> None:
+        """Skip weight download in tests."""
+        return None
+
+    def get_model_config(self, **kwargs: object) -> SimpleNamespace:
+        """Return a minimal namespace with just ``num_channels``."""
+        return SimpleNamespace(num_channels=3)
+
+    def get_model(self, config: SimpleNamespace) -> _DummyModel:
+        """Return a fresh ``_DummyModel`` instance."""
+        return _DummyModel()

tests/inference/test_from_checkpoint.py

@@ -0,0 +1,857 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for RFDETR.from_checkpoint classmethod.
+
+The inference logic is isolated by patching ``torch.load`` and the target model class inside ``rfdetr.variants`` (or
+``rfdetr.platform.models`` for plus models).  No model weights are downloaded or GPU memory allocated.
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import warnings
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from rfdetr.config import PretrainWeightsCompatibilityWarning
+from rfdetr.detr import RFDETR
+from rfdetr.detr import logger as detr_logger
+from rfdetr.platform import _IS_RFDETR_PLUS_AVAILABLE
+from rfdetr.variants import RFDETRSmall
+
+
+class _CustomObj:
+    """Module-level class that weights_only=True rejects (not in safe globals).
+
+    Must be at module scope so pickle can resolve the fully-qualified name during torch.save.  Local/nested classes
+    cannot be pickled.
+    """
+
+
+def _ns(pretrain_weights: str, num_classes: int = 80) -> dict:
+    """Fake legacy checkpoint with argparse.Namespace args."""
+    return {"args": argparse.Namespace(pretrain_weights=pretrain_weights, num_classes=num_classes)}
+
+
+def _dict(pretrain_weights: str, num_classes: int = 80) -> dict:
+    """Fake PTL-style checkpoint with dict args."""
+    return {"args": {"pretrain_weights": pretrain_weights, "num_classes": num_classes}}
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _call_from_checkpoint(ckpt: dict, path: Path, cls_patch_target: str, **kwargs):
+    """Invoke RFDETR.from_checkpoint with torch.load mocked to return *ckpt* and the model class at *cls_patch_target*
+    replaced by a MagicMock.
+
+    Returns:
+        Tuple of (result, mock_class).
+    """
+    mock_instance = MagicMock()
+    with (
+        patch("rfdetr.detr.torch.load", return_value=ckpt),
+        patch(cls_patch_target) as mock_cls,
+    ):
+        mock_cls.return_value = mock_instance
+        result = RFDETR.from_checkpoint(path, **kwargs)
+    return result, mock_cls
+
+
+# ---------------------------------------------------------------------------
+# Namespace args (legacy .pth checkpoints)
+# ---------------------------------------------------------------------------
+
+
+class TestFromCheckpointNamespaceArgs:
+    """from_checkpoint with argparse.Namespace args (legacy engine.py format)."""
+
+    @pytest.mark.parametrize(
+        ("pretrain_weights, patch_target"),
+        [
+            ("rf-detr-nano.pth", "RFDETRNano"),
+            ("rf-detr-small.pth", "RFDETRSmall"),
+            ("rf-detr-medium.pth", "RFDETRMedium"),
+            ("rf-detr-large.pth", "RFDETRLarge"),
+            ("rf-detr-keypoint-preview-xlarge.pth", "RFDETRKeypointPreview"),
+            ("rf-detr-base.pth", "RFDETRBase"),
+            ("rf-detr-seg-nano.pt", "RFDETRSegNano"),
+            ("rf-detr-seg-small.pt", "RFDETRSegSmall"),
+            ("rf-detr-seg-medium.pt", "RFDETRSegMedium"),
+            ("rf-detr-seg-large.pt", "RFDETRSegLarge"),
+            ("rf-detr-seg-xlarge.pt", "RFDETRSegXLarge"),
+            ("rf-detr-seg-xxlarge.pt", "RFDETRSeg2XLarge"),
+            ("rf-detr-seg-preview.pt", "RFDETRSegPreview"),
+        ],
+    )
+    def test_characterization_infers_correct_class_namespace(
+        self,
+        tmp_path: Path,
+        pretrain_weights: str,
+        patch_target: str,
+    ) -> None:
+        """Namespace-style args: correct subclass is called for each model size."""
+        result, mock_cls = _call_from_checkpoint(
+            _ns(pretrain_weights), tmp_path / "ckpt.pth", f"rfdetr.variants.{patch_target}"
+        )
+
+        mock_cls.assert_called_once()
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs.get("num_classes") == 80
+        assert call_kwargs.get("pretrain_weights") == str(tmp_path / "ckpt.pth")
+        assert result is mock_cls.return_value
+
+    @pytest.mark.parametrize(
+        "missing_value",
+        [
+            pytest.param("none", id="bare-none"),
+            pytest.param("null", id="bare-null"),
+            pytest.param("", id="empty"),
+            pytest.param("  None  ", id="whitespace-None"),
+            pytest.param("  ", id="whitespace-only"),
+            pytest.param(" null ", id="whitespace-null"),
+            pytest.param(None, id="python-None"),
+        ],
+    )
+    def test_namespace_args_falls_back_to_checkpoint_filename_when_pretrain_weights_missing(
+        self, tmp_path: Path, missing_value: str | None
+    ) -> None:
+        """Namespace args: filename fallback fires when pretrain_weights is unset-like."""
+        ckpt = _ns(missing_value)  # type: ignore[arg-type]
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "rf-detr-small.pth", "rfdetr.variants.RFDETRSmall")
+        mock_cls.assert_called_once()
+        assert mock_cls.call_args.kwargs["num_classes"] == 80
+
+
+# ---------------------------------------------------------------------------
+# Dict args (PTL / converted checkpoints)
+# ---------------------------------------------------------------------------
+
+
+class TestFromCheckpointDictArgs:
+    """from_checkpoint with dict-style args (PTL or convert_legacy_checkpoint output)."""
+
+    @pytest.mark.parametrize(
+        ("pretrain_weights, patch_target"),
+        [
+            ("rf-detr-small.pth", "RFDETRSmall"),
+            ("rf-detr-base.pth", "RFDETRBase"),
+        ],
+    )
+    def test_characterization_infers_correct_class_dict(
+        self,
+        tmp_path: Path,
+        pretrain_weights: str,
+        patch_target: str,
+    ) -> None:
+        """Dict-style args: correct subclass is called without AttributeError."""
+        _, mock_cls = _call_from_checkpoint(
+            _dict(pretrain_weights), tmp_path / "ckpt.pth", f"rfdetr.variants.{patch_target}"
+        )
+
+        mock_cls.assert_called_once()
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs.get("num_classes") == 80
+
+    def test_characterization_dict_args_missing_num_classes_uses_default(self, tmp_path: Path) -> None:
+        """Dict args without num_classes: constructor is called without num_classes kwarg."""
+        ckpt = {"args": {"pretrain_weights": "rf-detr-small.pth"}}
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "ckpt.pth", "rfdetr.variants.RFDETRSmall")
+
+        call_kwargs = mock_cls.call_args.kwargs
+        assert "num_classes" not in call_kwargs
+
+
+# ---------------------------------------------------------------------------
+# Edge cases
+# ---------------------------------------------------------------------------
+
+
+class TestFromCheckpointEdgeCases:
+    """Edge-case handling in from_checkpoint."""
+
+    def test_nonexistent_path_raises_file_not_found(self, tmp_path: Path) -> None:
+        """from_checkpoint raises FileNotFoundError when path does not exist."""
+        with pytest.raises(FileNotFoundError):
+            RFDETR.from_checkpoint(tmp_path / "nope.pth")
+
+    def test_directory_path_raises_os_error(self, tmp_path: Path) -> None:
+        """from_checkpoint raises OSError when path is a directory, not a file."""
+        with pytest.raises((OSError, IsADirectoryError)):
+            RFDETR.from_checkpoint(tmp_path)
+
+    def test_characterization_unknown_pretrain_weights_raises_value_error(self, tmp_path: Path) -> None:
+        """Unrecognised pretrain_weights name raises a descriptive ValueError."""
+        ckpt = _ns("/my/custom/finetuned.pth")
+        with patch("rfdetr.detr.torch.load", return_value=ckpt):
+            with pytest.raises(ValueError, match="Could not infer model class"):
+                RFDETR.from_checkpoint(tmp_path / "ckpt.pth")
+
+    def test_filename_fallback_unrecognized_name_raises_value_error(self, tmp_path: Path) -> None:
+        """ValueError fires via filename-fallback path when filename has no known model token."""
+        ckpt = {"args": {"pretrain_weights": "none", "num_classes": 80}}
+        with patch("rfdetr.detr.torch.load", return_value=ckpt):
+            with pytest.raises(ValueError, match="Could not infer model class"):
+                RFDETR.from_checkpoint(tmp_path / "finetuned.pth")
+
+    @pytest.mark.skipif(_IS_RFDETR_PLUS_AVAILABLE, reason="rfdetr_plus is installed — guard not active")
+    def test_filename_fallback_xlarge_without_plus_raises_import_error(self, tmp_path: Path) -> None:
+        """ImportError fires via filename-fallback path when rfdetr_plus is absent."""
+        ckpt = {"args": {"pretrain_weights": "none", "num_classes": 80}}
+        with patch("rfdetr.detr.torch.load", return_value=ckpt):
+            with pytest.raises(ImportError):
+                RFDETR.from_checkpoint(tmp_path / "rf-detr-xlarge-starter.pth")
+
+    def test_characterization_missing_args_key_raises_key_error(self, tmp_path: Path) -> None:
+        """Checkpoint without 'args' key raises KeyError."""
+        ckpt = {"model": {}}
+        with patch("rfdetr.detr.torch.load", return_value=ckpt):
+            with pytest.raises(KeyError):
+                RFDETR.from_checkpoint(tmp_path / "ckpt.pth")
+
+    def test_characterization_callable_on_subclass(self, tmp_path: Path) -> None:
+        """from_checkpoint can be called on a concrete subclass (RFDETRSmall)."""
+        mock_instance = MagicMock()
+        with (
+            patch("rfdetr.detr.torch.load", return_value=_ns("rf-detr-small.pth")),
+            patch("rfdetr.variants.RFDETRSmall") as mock_cls,
+        ):
+            mock_cls.return_value = mock_instance
+            result = RFDETRSmall.from_checkpoint(tmp_path / "ckpt.pth")
+
+        assert result is mock_instance
+        mock_cls.assert_called_once()
+
+    def test_characterization_extra_kwargs_forwarded(self, tmp_path: Path) -> None:
+        """Extra **kwargs are forwarded to the model constructor."""
+        _, mock_cls = _call_from_checkpoint(
+            _ns("rf-detr-small.pth"),
+            tmp_path / "ckpt.pth",
+            "rfdetr.variants.RFDETRSmall",
+            resolution=640,
+        )
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs.get("resolution") == 640
+
+    def test_characterization_pretrain_weights_in_kwargs_is_overridden(self, tmp_path: Path) -> None:
+        """pretrain_weights passed in **kwargs is silently overridden by the checkpoint path."""
+        _, mock_cls = _call_from_checkpoint(
+            _ns("rf-detr-small.pth"),
+            tmp_path / "ckpt.pth",
+            "rfdetr.variants.RFDETRSmall",
+            pretrain_weights="/should/be/overridden.pth",
+        )
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs["pretrain_weights"] == str(tmp_path / "ckpt.pth")
+
+    def test_characterization_caller_num_classes_overrides_checkpoint(self, tmp_path: Path) -> None:
+        """Caller-supplied num_classes takes precedence over the checkpoint's stored value."""
+        _, mock_cls = _call_from_checkpoint(
+            _ns("rf-detr-small.pth", num_classes=80),
+            tmp_path / "ckpt.pth",
+            "rfdetr.variants.RFDETRSmall",
+            num_classes=5,
+        )
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs["num_classes"] == 5
+
+    def test_checkpoint_model_config_forwarded_to_constructor(self, tmp_path: Path) -> None:
+        """Reload should preserve schema-dependent model config from PTL ``.pth`` checkpoints."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-keypoint-preview-xlarge.pth", "num_classes": 1},
+            "model_name": "RFDETRKeypointPreview",
+            "model_config": {
+                "num_keypoints_per_class": [0, 17],
+                "use_grouppose_keypoints": True,
+                "dual_projector": True,
+                "pretrain_weights": "/old/path.pth",
+            },
+        }
+        _, mock_cls = _call_from_checkpoint(
+            ckpt,
+            tmp_path / "checkpoint_best_total.pth",
+            "rfdetr.variants.RFDETRKeypointPreview",
+        )
+
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs["num_keypoints_per_class"] == [0, 17]
+        assert call_kwargs["use_grouppose_keypoints"] is True
+        assert call_kwargs["dual_projector"] is True
+        assert call_kwargs["num_classes"] == 1
+        assert call_kwargs["pretrain_weights"] == str(tmp_path / "checkpoint_best_total.pth")
+
+    @pytest.mark.skipif(_IS_RFDETR_PLUS_AVAILABLE, reason="rfdetr_plus is installed — guard not active")
+    def test_characterization_xlarge_without_plus_raises_import_error(self, tmp_path: Path) -> None:
+        """Xlarge checkpoint without rfdetr_plus raises ImportError instead of wrong class."""
+        for weights in ("rf-detr-xlarge.pth", "rf-detr-xxlarge.pth"):
+            ckpt = _ns(weights)
+            with patch("rfdetr.detr.torch.load", return_value=ckpt):
+                with pytest.raises(ImportError):
+                    RFDETR.from_checkpoint(tmp_path / "ckpt.pth")
+
+    def test_trust_gate_rejects_custom_class_by_default(self, tmp_path: Path) -> None:
+        """from_checkpoint raises RuntimeError for custom-class checkpoints without trust_checkpoint=True.
+
+        Scenario: user calls from_checkpoint on a file containing an unrecognised Python class.
+        The default trust_checkpoint=False must reject it to prevent arbitrary code execution.
+        """
+        ckpt_path = tmp_path / "custom_obj.pth"
+        torch.save({"model": {}, "args": _CustomObj(), "model_name": "RFDETRSmall"}, ckpt_path)
+
+        with pytest.raises(RuntimeError, match="trust_checkpoint=True"):
+            RFDETR.from_checkpoint(ckpt_path)
+
+
+# ---------------------------------------------------------------------------
+# Deprecated class instantiation
+# ---------------------------------------------------------------------------
+
+
+class TestDeprecatedClassInstantiation:
+    """Deprecated model classes emit deprecation warnings on instantiation."""
+
+    @pytest.mark.parametrize(
+        ("cls_name, import_path"),
+        [
+            ("RFDETRBase", "rfdetr.variants.RFDETRBase"),
+            ("RFDETRLargeDeprecated", "rfdetr.variants.RFDETRLargeDeprecated"),
+            ("RFDETRSegPreview", "rfdetr.variants.RFDETRSegPreview"),
+        ],
+    )
+    def test_direct_instantiation_is_allowed(self, cls_name: str, import_path: str) -> None:
+        """Direct instantiation of a deprecated class does not raise RuntimeError."""
+        import importlib
+
+        module_path, attr = import_path.rsplit(".", 1)
+        module = importlib.import_module(module_path)
+        cls = getattr(module, attr)
+        with patch("rfdetr.detr.RFDETR.__init__", return_value=None):
+            model = cls()
+        assert model.__class__.__name__ == cls_name
+
+    @pytest.mark.parametrize("pretrain_weights", ["rf-detr-base.pth", "rf-detr-seg-preview.pt"])
+    def test_from_checkpoint_resolves_deprecated_class(
+        self,
+        tmp_path: Path,
+        pretrain_weights: str,
+    ) -> None:
+        """from_checkpoint still resolves deprecated classes without KeyError on minimal mocked checkpoints."""
+        ckpt = _ns(pretrain_weights)
+        with (
+            patch("rfdetr.detr.torch.load", return_value=ckpt),
+            patch("rfdetr.detr.RFDETR.__init__", return_value=None),
+        ):
+            model = RFDETR.from_checkpoint(tmp_path / "ckpt.pth")
+        assert model.__class__.__name__ in {"RFDETRBase", "RFDETRSegPreview"}
+
+
+# ---------------------------------------------------------------------------
+# model_name in checkpoint (#887)
+# ---------------------------------------------------------------------------
+
+
+def _ckpt_with_model_name(model_name: str, num_classes: int = 80) -> dict:
+    """Fake checkpoint with model_name key (new format)."""
+    return {
+        "args": {"pretrain_weights": "rf-detr-small.pth", "num_classes": num_classes},
+        "model_name": model_name,
+    }
+
+
+class TestFromCheckpointModelName:
+    """from_checkpoint uses model_name when present in checkpoint."""
+
+    @pytest.mark.parametrize(
+        ("model_name, patch_target"),
+        [
+            ("RFDETRNano", "RFDETRNano"),
+            ("RFDETRSmall", "RFDETRSmall"),
+            ("RFDETRMedium", "RFDETRMedium"),
+            ("RFDETRLarge", "RFDETRLarge"),
+            ("RFDETRKeypointPreview", "RFDETRKeypointPreview"),
+            ("RFDETRBase", "RFDETRBase"),
+            ("RFDETRSegNano", "RFDETRSegNano"),
+            ("RFDETRSegPreview", "RFDETRSegPreview"),
+            ("RFDETRSegSmall", "RFDETRSegSmall"),
+            ("RFDETRSegMedium", "RFDETRSegMedium"),
+            ("RFDETRSegLarge", "RFDETRSegLarge"),
+            ("RFDETRSegXLarge", "RFDETRSegXLarge"),
+            ("RFDETRSeg2XLarge", "RFDETRSeg2XLarge"),
+        ],
+    )
+    def test_model_name_resolves_correct_class(self, tmp_path: Path, model_name: str, patch_target: str) -> None:
+        """model_name in checkpoint maps directly to the correct subclass."""
+        result, mock_cls = _call_from_checkpoint(
+            _ckpt_with_model_name(model_name), tmp_path / "ckpt.pth", f"rfdetr.variants.{patch_target}"
+        )
+        mock_cls.assert_called_once()
+        assert result is mock_cls.return_value
+
+    def test_model_name_takes_priority_over_pretrain_weights(self, tmp_path: Path) -> None:
+        """model_name is used even when pretrain_weights points to a different size."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-nano.pth", "num_classes": 80},
+            "model_name": "RFDETRLarge",
+        }
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "ckpt.pth", "rfdetr.variants.RFDETRLarge")
+        mock_cls.assert_called_once()
+
+    def test_falls_back_to_pretrain_weights_without_model_name(self, tmp_path: Path) -> None:
+        """Old checkpoints without model_name still work via pretrain_weights parsing."""
+        ckpt = _dict("rf-detr-small.pth")
+        assert "model_name" not in ckpt
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "ckpt.pth", "rfdetr.variants.RFDETRSmall")
+        mock_cls.assert_called_once()
+
+    @pytest.mark.parametrize(
+        "missing_value",
+        [
+            pytest.param("none", id="bare-none"),
+            pytest.param("null", id="bare-null"),
+            pytest.param("", id="empty"),
+            pytest.param("  None  ", id="whitespace-None"),
+            pytest.param("  ", id="whitespace-only"),
+            pytest.param(" null ", id="whitespace-null"),
+            pytest.param(None, id="python-None"),
+        ],
+    )
+    def test_falls_back_to_checkpoint_filename_when_pretrain_weights_missing(
+        self, tmp_path: Path, missing_value: str | None
+    ) -> None:
+        """When pretrain_weights is missing-like, from_checkpoint infers class from checkpoint filename."""
+        ckpt = {"args": {"pretrain_weights": missing_value, "num_classes": 80}}
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "rf-detr-small.pth", "rfdetr.variants.RFDETRSmall")
+        mock_cls.assert_called_once()
+        assert mock_cls.call_args.kwargs["num_classes"] == 80
+
+    def test_unknown_model_name_falls_back_to_pretrain_weights(self, tmp_path: Path) -> None:
+        """Unrecognised model_name falls back to pretrain_weights parsing."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-small.pth", "num_classes": 80},
+            "model_name": "UnknownModel",
+        }
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "ckpt.pth", "rfdetr.variants.RFDETRSmall")
+        mock_cls.assert_called_once()
+
+    def test_model_name_with_whitespace_is_stripped(self, tmp_path: Path) -> None:
+        """Leading/trailing whitespace in model_name is stripped before class resolution."""
+        ckpt = _ckpt_with_model_name("  RFDETRSmall  ")
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "ckpt.pth", "rfdetr.variants.RFDETRSmall")
+        mock_cls.assert_called_once()
+
+    @pytest.mark.parametrize(
+        "model_name, expected_class",
+        [
+            ("RFDETRBase", "RFDETRBase"),
+            ("RFDETRSegPreview", "RFDETRSegPreview"),
+        ],
+    )
+    def test_model_name_deprecated_class_resolves_and_instantiates(
+        self, tmp_path: Path, model_name: str, expected_class: str
+    ) -> None:
+        """from_checkpoint resolves deprecated model_name values and instantiates the resolved class."""
+        ckpt = _ckpt_with_model_name(model_name)
+        with (
+            patch("rfdetr.detr.torch.load", return_value=ckpt),
+            patch("rfdetr.detr.RFDETR.__init__", return_value=None),
+        ):
+            model = RFDETR.from_checkpoint(tmp_path / "ckpt.pth")
+        assert model.__class__.__name__ == expected_class
+
+    def test_large_deprecated_model_name_resolves_to_deprecated_class(self, tmp_path: Path) -> None:
+        """Checkpoints saved with model_name='RFDETRLargeDeprecated' must load as RFDETRLargeDeprecated.
+
+        Before the fix, RFDETRLargeDeprecated was absent from _name_map; the substring matcher would pick RFDETRLarge,
+        which fails with a pydantic literal_error when the saved model_config carries encoder='dinov2_windowed_base'
+        (only valid for the deprecated Large configuration).
+        """
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-large.pth", "num_classes": 80},
+            "model_name": "RFDETRLargeDeprecated",
+            "model_config": {
+                "encoder": "dinov2_windowed_base",
+                "projector_scale": "P4",
+            },
+        }
+        result, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "ckpt.pth", "rfdetr.variants.RFDETRLargeDeprecated")
+        mock_cls.assert_called_once()
+        assert result is mock_cls.return_value
+
+    @pytest.mark.skipif(_IS_RFDETR_PLUS_AVAILABLE, reason="rfdetr_plus is installed — guard not active")
+    @pytest.mark.parametrize("model_name", ["RFDETRXLarge", "RFDETR2XLarge"])
+    def test_plus_model_name_without_plus_raises_import_error(self, tmp_path: Path, model_name: str) -> None:
+        """Plus checkpoints using model_name raise install guidance without rfdetr_plus."""
+        ckpt = {
+            "args": {"pretrain_weights": "", "num_classes": 80},
+            "model_name": model_name,
+        }
+        with patch("rfdetr.detr.torch.load", return_value=ckpt):
+            with pytest.raises(ImportError, match="rfdetr_plus package"):
+                RFDETR.from_checkpoint(tmp_path / "ckpt.pth")
+
+
+# ---------------------------------------------------------------------------
+# num_classes provenance (fine-tuning a from_checkpoint model on a new dataset)
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def args_only_checkpoint(tmp_path: Path) -> Path:
+    """Minimal checkpoint with num_classes in args only; model_config carries no num_classes key.
+
+    Covers the legacy checkpoint format where num_classes is embedded in the args dict rather
+    than in model_config.  from_checkpoint extracts it via the args path (detr.py:454-457) and
+    injects it into constructor_kwargs — this fixture verifies that path also clears the
+    Pydantic provenance marker.  Only exercises the args-injection path; model_config path
+    covered by ``two_class_checkpoint``.
+
+    Args:
+        tmp_path: Pytest temporary directory.
+
+    Returns:
+        Path to the saved checkpoint file.
+    """
+    path = tmp_path / "small_two_class_args_only.pth"
+    torch.save(
+        {
+            "model": {"class_embed.bias": torch.zeros(3)},
+            "model_name": "RFDETRSmall",
+            "model_config": {},
+            "args": {"class_names": ["cat", "dog"], "num_classes": 2},
+        },
+        path,
+    )
+    return path
+
+
+@pytest.fixture
+def two_class_checkpoint(tmp_path: Path) -> Path:
+    """Save a minimal synthetic 2-class checkpoint to disk (no downloads, no real weights).
+
+    Follows the lightweight checkpoint pattern used elsewhere in the suite (``test_detr_shim``,
+    ``test_load_pretrain_weights``): write only what ``from_checkpoint``/``load_pretrain_weights`` actually inspect —
+    the ``class_embed.bias`` tensor sized for 2 classes + background, plus the metadata used to resolve the model
+    (``model_name``) and the class count (``model_config`` carrying ``num_classes=2``).  A *non-default*
+    ``num_classes`` is what trips the user-override guards, so it is written explicitly rather than relying on a
+    published checkpoint (whose default 90 would not trip them).  ``from_checkpoint`` still builds a real model from
+    this, which is what the provenance and head-shape assertions exercise.
+
+    Args:
+        tmp_path: Pytest temporary directory.
+
+    Returns:
+        Path to the saved checkpoint file.
+    """
+    path = tmp_path / "small_two_class.pth"
+    torch.save(
+        {
+            "model": {"class_embed.bias": torch.zeros(3)},
+            "model_name": "RFDETRSmall",
+            "model_config": {"num_classes": 2},
+            "args": {"class_names": ["cat", "dog"]},
+        },
+        path,
+    )
+    return path
+
+
+class TestFromCheckpointNumClassesProvenance:
+    """Checkpoint-derived num_classes must not be treated as a user override.
+
+    Regression tests for https://github.com/roboflow/rf-detr/issues/1092: ``from_checkpoint`` copies ``num_classes``
+    out of the checkpoint into the constructor kwargs, which used to mark the field as explicitly user-set.  Both
+    provenance guards (``RFDETR._align_num_classes_from_dataset`` and the head re-init logic in
+    ``rfdetr.models.weights.load_pretrain_weights``) then refused to adapt the detection head to a new dataset's
+    class count, breaking fine-tuning from a checkpoint.
+    """
+
+    def test_checkpoint_num_classes_is_not_marked_user_set(self, two_class_checkpoint: Path) -> None:
+        """from_checkpoint adopts the checkpoint class count without warning about pretrained weights."""
+        with warnings.catch_warnings():
+            warnings.filterwarnings("error", category=PretrainWeightsCompatibilityWarning)
+            model = RFDETR.from_checkpoint(two_class_checkpoint)
+
+        assert model.model_config.num_classes == 2
+        assert model.model.model.class_embed.bias.shape[0] == 3, "Head must match checkpoint (2 classes + background)."
+        assert "num_classes" not in model.model_config.model_fields_set, (
+            "Checkpoint-derived num_classes must not be recorded as explicitly user-set; "
+            "otherwise train() refuses to align the head to a new dataset's class count."
+        )
+
+    def test_train_alignment_adapts_head_to_new_dataset(
+        self, two_class_checkpoint: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Fine-tuning a from_checkpoint model on a dataset with a different class count adapts the head."""
+        model = RFDETR.from_checkpoint(two_class_checkpoint)
+        monkeypatch.setattr(RFDETR, "_detect_num_classes_for_training", staticmethod(lambda *a, **k: 5))
+
+        model._align_num_classes_from_dataset("<five-class-dataset>")
+
+        assert model.model_config.num_classes == 5
+        assert model.model.args.num_classes == 5
+        # train() rebuilds the model from model_config (inside RFDETRModelModule), reloading the checkpoint
+        # weights with the aligned class count; the rebuilt head must adopt the dataset class count.
+        rebuilt = model.get_model(model.model_config)
+        assert rebuilt.model.class_embed.bias.shape[0] == 6, "Rebuilt head must have 5 classes + background."
+
+    def test_explicit_num_classes_kwarg_still_wins(
+        self,
+        two_class_checkpoint: Path,
+        monkeypatch: pytest.MonkeyPatch,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        """An explicit num_classes kwarg to from_checkpoint stays authoritative over the dataset."""
+        model = RFDETR.from_checkpoint(two_class_checkpoint, num_classes=7)
+
+        assert model.model_config.num_classes == 7
+        assert "num_classes" in model.model_config.model_fields_set
+        assert model.model.model.class_embed.bias.shape[0] == 8, "Head must expand to 7 classes + background."
+
+        monkeypatch.setattr(RFDETR, "_detect_num_classes_for_training", staticmethod(lambda *a, **k: 5))
+        monkeypatch.setattr(detr_logger, "propagate", True)
+        with caplog.at_level(logging.WARNING, logger="rf-detr"):
+            model._align_num_classes_from_dataset("<five-class-dataset>")
+
+        assert model.model_config.num_classes == 7, "Explicit user num_classes must be preserved."
+        assert any("Using the model's configured value" in record.message for record in caplog.records)
+
+    def test_checkpoint_num_classes_from_args_not_marked_user_set(self, args_only_checkpoint: Path) -> None:
+        """num_classes injected from checkpoint args (not model_config) is cleared from model_fields_set."""
+        model = RFDETR.from_checkpoint(args_only_checkpoint)
+
+        assert model.model_config.num_classes == 2
+        assert "num_classes" not in model.model_config.model_fields_set, (
+            "num_classes from checkpoint args must not be recorded as explicitly user-set; "
+            "otherwise train() refuses to adapt the head to a new dataset's class count."
+        )
+
+    def test_explicit_default_num_classes_pins_head(
+        self,
+        two_class_checkpoint: Path,
+        monkeypatch: pytest.MonkeyPatch,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        """Passing num_classes equal to the ModelConfig default still pins the detection head.
+
+        An explicit num_classes is honored regardless of whether it equals the class default:
+        ``_align_num_classes_from_dataset`` keys off whether the field was set, not whether the
+        value differs from the default, so the dataset count cannot silently override it.  This
+        guards against re-introducing the ``value != default`` clause, whose asymmetric behavior
+        (default silently aligned, non-default preserved) was the bug this test now pins.
+        """
+        model = RFDETR.from_checkpoint(two_class_checkpoint)
+        default_nc = type(model.model_config).model_fields["num_classes"].default
+        # Simulate calling from_checkpoint(path, num_classes=<default>):
+        # assigning the field adds "num_classes" to model_fields_set automatically (Pydantic v2).
+        model.model_config.num_classes = default_nc
+
+        assert "num_classes" in model.model_config.model_fields_set
+        monkeypatch.setattr(RFDETR, "_detect_num_classes_for_training", staticmethod(lambda *a, **k: 5))
+        monkeypatch.setattr(detr_logger, "propagate", True)
+        with caplog.at_level(logging.WARNING, logger="rf-detr"):
+            model._align_num_classes_from_dataset("<five-class-dataset>")
+
+        assert model.model_config.num_classes == default_nc, (
+            "Explicitly passing the ModelConfig default for num_classes must pin the head; "
+            "the dataset class count must not silently override an explicit user setting."
+        )
+        assert any("Using the model's configured value" in record.message for record in caplog.records)
+
+    def test_explicit_default_num_classes_via_from_checkpoint_integrated(
+        self,
+        two_class_checkpoint: Path,
+        monkeypatch: pytest.MonkeyPatch,
+        caplog: pytest.LogCaptureFixture,
+    ) -> None:
+        """from_checkpoint(path, num_classes=<default>) pins head via the integrated code path.
+
+        Unlike test_explicit_default_num_classes_pins_head which simulates the explicit-default scenario via post-
+        construction assignment, this test calls from_checkpoint directly with num_classes=default_nc.  A regression in
+        how from_checkpoint passes num_classes into the constructor would be caught here but not by the proxy-based
+        test.
+        """
+        default_nc = RFDETRSmall._model_config_class.model_fields["num_classes"].default
+        model = RFDETR.from_checkpoint(two_class_checkpoint, num_classes=default_nc)
+
+        assert model.model_config.num_classes == default_nc
+        assert "num_classes" in model.model_config.model_fields_set, (
+            "from_checkpoint with explicit num_classes must keep it in model_fields_set; "
+            "only checkpoint-derived num_classes should be cleared."
+        )
+
+        monkeypatch.setattr(RFDETR, "_detect_num_classes_for_training", staticmethod(lambda *a, **k: 5))
+        monkeypatch.setattr(detr_logger, "propagate", True)
+        with caplog.at_level(logging.WARNING, logger="rf-detr"):
+            model._align_num_classes_from_dataset("<five-class-dataset>")
+
+        assert model.model_config.num_classes == default_nc, (
+            "Head must remain pinned at default_nc after alignment; "
+            "from_checkpoint-supplied num_classes must not be silently overridden."
+        )
+        assert any("Using the model's configured value" in record.message for record in caplog.records)
+
+    def test_equal_class_count_does_not_rebuild_head(
+        self, two_class_checkpoint: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Checkpoint and dataset sharing the same class count leaves the head unchanged."""
+        model = RFDETR.from_checkpoint(two_class_checkpoint)
+        original_bias_shape = model.model.model.class_embed.bias.shape
+        monkeypatch.setattr(RFDETR, "_detect_num_classes_for_training", staticmethod(lambda *a, **k: 2))
+
+        model._align_num_classes_from_dataset("<two-class-dataset>")
+
+        assert model.model_config.num_classes == 2
+        assert model.model.model.class_embed.bias.shape == original_bias_shape, (
+            "Head must not be rebuilt when dataset class count matches checkpoint class count."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Weight-based schema inference
+# ---------------------------------------------------------------------------
+
+
+def _make_kp_active_mask(schema: list[int]) -> torch.Tensor:
+    """Build a bool _kp_active_mask tensor encoding *schema* (mirrors LwDetr._create_kp_active_mask).
+
+    Args:
+        schema: Keypoints-per-class list, e.g. ``[0, 33]`` for background + 33-kp class.
+
+    Returns:
+        Bool tensor of shape ``[len(schema), max(schema)]`` with True in active keypoint slots.
+    """
+    if not schema or max(schema) == 0:
+        return torch.zeros(0, 0, dtype=torch.bool)
+    max_kp = max(schema)
+    mask = torch.zeros(len(schema), max_kp, dtype=torch.bool)
+    for idx, n_kp in enumerate(schema):
+        mask[idx, :n_kp] = True
+    return mask
+
+
+class TestFromCheckpointWeightInference:
+    """from_checkpoint infers schema from checkpoint weights when model_config is absent or stale.
+
+    Regression tests for the bug where a fine-tuned 33-kp keypoint model loaded with the COCO default [0, 17] schema
+    because model_config["num_keypoints_per_class"] was never updated from the default before the checkpoint was saved.
+    The authoritative schema is embedded in the checkpoint weights via the _kp_active_mask buffer; from_checkpoint now
+    reads it directly.
+    """
+
+    def test_infers_keypoint_schema_from_kp_active_mask(self, tmp_path: Path) -> None:
+        """Stale model_config kp schema [0, 17] is overridden by weight-inferred [0, 33]."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-keypoint-preview-xlarge.pth"},
+            "model_name": "RFDETRKeypointPreview",
+            "model_config": {"num_keypoints_per_class": [0, 17]},
+            "model": {"_kp_active_mask": _make_kp_active_mask([0, 33])},
+        }
+        _, mock_cls = _call_from_checkpoint(
+            ckpt, tmp_path / "checkpoint_best_total.pth", "rfdetr.variants.RFDETRKeypointPreview"
+        )
+
+        assert mock_cls.call_args.kwargs["num_keypoints_per_class"] == [0, 33]
+
+    def test_infers_keypoint_schema_when_model_config_absent(self, tmp_path: Path) -> None:
+        """num_keypoints_per_class is inferred from _kp_active_mask when model_config is missing."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-keypoint-preview-xlarge.pth"},
+            "model_name": "RFDETRKeypointPreview",
+            "model": {"_kp_active_mask": _make_kp_active_mask([0, 33])},
+        }
+        _, mock_cls = _call_from_checkpoint(
+            ckpt, tmp_path / "checkpoint_best_total.pth", "rfdetr.variants.RFDETRKeypointPreview"
+        )
+
+        assert mock_cls.call_args.kwargs["num_keypoints_per_class"] == [0, 33]
+
+    def test_user_kwarg_wins_over_weight_inferred_keypoint_schema(self, tmp_path: Path) -> None:
+        """Explicit num_keypoints_per_class kwarg overrides weight-inferred [0, 33] schema."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-keypoint-preview-xlarge.pth"},
+            "model_name": "RFDETRKeypointPreview",
+            "model": {"_kp_active_mask": _make_kp_active_mask([0, 33])},
+        }
+        _, mock_cls = _call_from_checkpoint(
+            ckpt,
+            tmp_path / "checkpoint_best_total.pth",
+            "rfdetr.variants.RFDETRKeypointPreview",
+            num_keypoints_per_class=[0, 17],
+        )
+
+        assert mock_cls.call_args.kwargs["num_keypoints_per_class"] == [0, 17]
+
+    def test_infers_num_classes_from_class_embed_weight(self, tmp_path: Path) -> None:
+        """Stale model_config num_classes=90 is overridden by class_embed.weight shape inference."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-small.pth"},
+            "model_name": "RFDETRSmall",
+            "model_config": {"num_classes": 90},
+            "model": {"class_embed.weight": torch.zeros(3, 256)},
+        }
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "checkpoint_best_total.pth", "rfdetr.variants.RFDETRSmall")
+
+        assert mock_cls.call_args.kwargs["num_classes"] == 2
+
+    def test_user_kwarg_wins_over_weight_inferred_num_classes(self, tmp_path: Path) -> None:
+        """Explicit num_classes kwarg overrides weight-inferred value from class_embed.weight."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-small.pth"},
+            "model_name": "RFDETRSmall",
+            "model": {"class_embed.weight": torch.zeros(3, 256)},
+        }
+        _, mock_cls = _call_from_checkpoint(
+            ckpt,
+            tmp_path / "checkpoint_best_total.pth",
+            "rfdetr.variants.RFDETRSmall",
+            num_classes=90,
+        )
+
+        assert mock_cls.call_args.kwargs["num_classes"] == 90
+
+    def test_infers_schema_from_ptl_ckpt_state_dict_format(self, tmp_path: Path) -> None:
+        """Weight inference works for PTL-native .ckpt format (state_dict with model.
+
+        prefix).
+        """
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-keypoint-preview-xlarge.pth"},
+            "model_name": "RFDETRKeypointPreview",
+            "state_dict": {
+                "model._kp_active_mask": _make_kp_active_mask([0, 33]),
+                "model.class_embed.weight": torch.zeros(3, 256),
+            },
+        }
+        _, mock_cls = _call_from_checkpoint(ckpt, tmp_path / "checkpoint.ckpt", "rfdetr.variants.RFDETRKeypointPreview")
+
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs["num_keypoints_per_class"] == [0, 33]
+        assert call_kwargs["num_classes"] == 2
+
+    def test_consistent_checkpoint_produces_no_override(self, tmp_path: Path) -> None:
+        """When model_config and weights agree, weight inference leaves constructor_kwargs unchanged."""
+        ckpt = {
+            "args": {"pretrain_weights": "rf-detr-keypoint-preview-xlarge.pth"},
+            "model_name": "RFDETRKeypointPreview",
+            "model_config": {"num_keypoints_per_class": [0, 33], "num_classes": 2},
+            "model": {
+                "_kp_active_mask": _make_kp_active_mask([0, 33]),
+                "class_embed.weight": torch.zeros(3, 256),
+            },
+        }
+        _, mock_cls = _call_from_checkpoint(
+            ckpt, tmp_path / "checkpoint_best_total.pth", "rfdetr.variants.RFDETRKeypointPreview"
+        )
+
+        call_kwargs = mock_cls.call_args.kwargs
+        assert call_kwargs["num_keypoints_per_class"] == [0, 33]
+        assert call_kwargs["num_classes"] == 2

tests/inference/test_inference.py

@@ -0,0 +1,54 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for rfdetr.inference weight-adaptation helpers."""
+
+import pytest
+import torch
+
+from rfdetr.inference import _adapt_input_conv
+
+
+@pytest.fixture(autouse=True)
+def reset_random_seeds():
+    """Ensure reproducible random state for every test in this module."""
+    torch.manual_seed(0)
+
+
+class TestAdaptInputConv:
+    @pytest.mark.parametrize(
+        ("num_channels", "expected_shape", "expected_builder"),
+        [
+            pytest.param(3, (8, 3, 3, 3), lambda weight: weight, id="identity_3ch"),
+            pytest.param(1, (8, 1, 3, 3), lambda weight: weight.mean(dim=1, keepdim=True), id="mean_1ch"),
+            pytest.param(
+                4,
+                (8, 4, 3, 3),
+                lambda weight: torch.cat([weight, weight], dim=1)[:, :4] * (3.0 / 4.0),
+                id="tile_4ch",
+            ),
+            pytest.param(
+                6,
+                (8, 6, 3, 3),
+                lambda weight: torch.cat([weight, weight], dim=1)[:, :6] * (3.0 / 6.0),
+                id="tile_6ch",
+            ),
+            pytest.param(
+                2,
+                (8, 2, 3, 3),
+                lambda weight: weight[:, :2] * (3.0 / 2.0),
+                id="tile_2ch",
+            ),
+        ],
+    )
+    def test_adapt_input_conv(self, num_channels, expected_shape, expected_builder):
+        """Verify shape and values for each _adapt_input_conv branch."""
+        conv_weight = torch.randn(8, 3, 3, 3)
+
+        adapted_weight = _adapt_input_conv(num_channels, conv_weight)
+        expected_weight = expected_builder(conv_weight)
+
+        assert adapted_weight.shape == expected_shape
+        torch.testing.assert_close(adapted_weight, expected_weight)

tests/inference/test_inference_keypoints.py

@@ -0,0 +1,50 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Focused predict() contract tests for keypoint and non-keypoint outputs."""
+
+import numpy as np
+import PIL.Image
+import supervision as sv
+
+from .helpers import _DummyModel, _DummyRFDETR
+
+
+def test_predict_returns_supervision_keypoints() -> None:
+    """Keypoint model predictions return ``sv.KeyPoints`` with detection details."""
+    image = PIL.Image.new("RGB", (64, 48), color=(128, 128, 128))
+    model = _DummyRFDETR()
+    model.model = _DummyModel(labels=[0, 1], include_keypoints=True)
+
+    key_points = model.predict(image)
+
+    assert isinstance(key_points, sv.KeyPoints)
+    assert key_points.xy.shape == (2, 17, 2)
+    assert key_points.keypoint_confidence.shape == (2, 17)
+    assert key_points.data["xyxy"].shape == (2, 4)
+    assert key_points.detection_confidence.shape == (2,)
+    assert np.isfinite(key_points.xy).all()
+    assert np.isfinite(key_points.keypoint_confidence).all()
+    assert "keypoint_precision_cholesky" in key_points.data
+    keypoint_precision = key_points.data["keypoint_precision_cholesky"]
+    assert isinstance(keypoint_precision, np.ndarray)
+    assert keypoint_precision.shape == (2, 17, 3)
+    assert np.isfinite(keypoint_precision).all()
+    assert "source_image" in key_points.data
+    assert len(key_points.data["source_image"]) == 2
+
+
+def test_predict_default_detection_without_keypoints_unchanged() -> None:
+    """Default detection prediction keeps legacy output structure."""
+    image = PIL.Image.new("RGB", (64, 48), color=(128, 128, 128))
+    model = _DummyRFDETR()
+
+    detections = model.predict(image)
+
+    assert "keypoints" not in detections.data
+    assert not hasattr(detections, "keypoints")
+    assert "class_name" in detections.data
+    assert "source_shape" in detections.data
+    assert detections.data["source_shape"].shape[1] == 2

tests/inference/test_model_device_move.py

@@ -0,0 +1,68 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for the lazy device move running under ``torch.inference_mode()``.
+
+``predict()`` stacks ``@torch.inference_mode()`` on top of ``@_ensure_model_on_device``, so the deferred CPU-to-
+accelerator move happens while inference mode is active.  Tensors materialised under inference mode are *inference
+tensors*: they can never require gradients, so a later ``train()`` / auto-batch probe silently produces no gradients.
+The move itself must therefore always run with inference mode disabled.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import Any
+
+import torch
+from torch import nn
+
+from rfdetr.detr import _move_model_context_to_device
+
+
+class _RecordingModule(nn.Module):
+    """Module whose ``to()`` records whether inference mode was active at move time."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = nn.Linear(2, 2)
+        self.inference_mode_at_move: bool | None = None
+
+    def to(self, *args: Any, **kwargs: Any) -> "_RecordingModule":
+        """Record the inference-mode state instead of performing a real device move."""
+        self.inference_mode_at_move = torch.is_inference_mode_enabled()
+        return self
+
+
+class TestMoveModelContextUnderInferenceMode:
+    """The deferred device move must never materialise parameters as inference tensors."""
+
+    def test_moved_params_are_not_inference_tensors(self) -> None:
+        """A real ``.to()`` move inside ``torch.inference_mode()`` must not create inference-tensor parameters."""
+        ctx = SimpleNamespace(device=torch.device("meta"), model=nn.Linear(2, 2))
+
+        with torch.inference_mode():
+            _move_model_context_to_device(ctx)
+
+        assert not any(p.is_inference() for p in ctx.model.parameters())
+
+    def test_move_still_materializes_on_target_device(self) -> None:
+        """The inference-mode guard must not suppress the device move itself."""
+        ctx = SimpleNamespace(device=torch.device("meta"), model=nn.Linear(2, 2))
+
+        with torch.inference_mode():
+            _move_model_context_to_device(ctx)
+
+        assert all(p.device.type == "meta" for p in ctx.model.parameters())
+
+    def test_move_runs_with_inference_mode_disabled(self) -> None:
+        """The ``.to()`` call itself must observe inference mode as disabled."""
+        module = _RecordingModule()
+        ctx = SimpleNamespace(device=torch.device("meta"), model=module)
+
+        with torch.inference_mode():
+            _move_model_context_to_device(ctx)
+
+        assert module.inference_mode_at_move is False

tests/inference/test_optimize_for_inference.py

@@ -0,0 +1,532 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for RFDETR.optimize_for_inference()."""
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from rfdetr.detr import RFDETR
+
+
+class _FakeModel(torch.nn.Module):
+    """Minimal nn.Module that satisfies the optimize_for_inference contract."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = torch.nn.Linear(1, 1)
+
+    def forward(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
+        return {"pred_boxes": self.linear(x[:, :1, :1, :1].squeeze(-1).squeeze(-1))}
+
+    def export(self) -> None:
+        pass
+
+
+class _FakeModelContext:
+    def __init__(self, device: torch.device | str = torch.device("cpu"), resolution: int = 28) -> None:
+        self.device = torch.device(device) if not isinstance(device, torch.device) else device
+        self.resolution = resolution
+        self.model = _FakeModel()
+        self.inference_model = None
+
+
+class _FakeRFDETR(RFDETR):
+    def maybe_download_pretrain_weights(self) -> None:
+        return None
+
+    def get_model_config(self, **kwargs) -> SimpleNamespace:
+        return SimpleNamespace(num_channels=3)
+
+    def get_model(self, config: SimpleNamespace) -> _FakeModelContext:
+        return _FakeModelContext()
+
+
+class TestOptimizeForInferenceDtype:
+    """Dtype coercion and validation tests."""
+
+    def test_string_dtype_float32_is_accepted(self) -> None:
+        """Passing dtype='float32' (str) should be coerced to torch.float32."""
+        rfdetr = _FakeRFDETR()
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False, dtype="float32")
+
+        assert rfdetr._optimized_dtype == torch.float32
+
+    def test_string_dtype_float16_is_accepted(self) -> None:
+        """Passing dtype='float16' (str) should be coerced to torch.float16."""
+        rfdetr = _FakeRFDETR()
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False, dtype="float16")
+
+        assert rfdetr._optimized_dtype == torch.float16
+
+    def test_torch_dtype_is_passed_through(self) -> None:
+        """Passing dtype=torch.float32 directly should work as before."""
+        rfdetr = _FakeRFDETR()
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False, dtype=torch.float32)
+
+        assert rfdetr._optimized_dtype == torch.float32
+
+    def test_invalid_dtype_type_raises_type_error(self) -> None:
+        """Passing an invalid dtype type (e.g. int) should raise TypeError."""
+        rfdetr = _FakeRFDETR()
+
+        with pytest.raises(TypeError, match="dtype must be a torch.dtype or a string name of a dtype"):
+            rfdetr.optimize_for_inference(compile=False, dtype=42)  # type: ignore[arg-type]
+
+    def test_invalid_dtype_string_raises_type_error(self) -> None:
+        """Passing a non-existent dtype string should raise TypeError with a descriptive message."""
+        rfdetr = _FakeRFDETR()
+
+        with pytest.raises(TypeError, match="dtype must be a torch.dtype or a string name of a dtype"):
+            rfdetr.optimize_for_inference(compile=False, dtype="not_a_dtype")
+
+    def test_valid_torch_attr_that_is_not_dtype_raises_type_error(self) -> None:
+        """'Tensor' is a valid torch attribute but not a torch.dtype — should raise TypeError."""
+        rfdetr = _FakeRFDETR()
+
+        with pytest.raises(TypeError, match="dtype must be a torch.dtype or a string name of a dtype"):
+            rfdetr.optimize_for_inference(compile=False, dtype="Tensor")  # type: ignore[arg-type]
+
+    @pytest.mark.parametrize("dtype_str", ["float32", "float16", "bfloat16"])
+    def test_string_dtype_variants_are_accepted(self, dtype_str: str) -> None:
+        """Common dtype string names should be accepted and coerced to the matching torch.dtype."""
+        rfdetr = _FakeRFDETR()
+        expected = getattr(torch, dtype_str)
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False, dtype=dtype_str)
+
+        assert rfdetr._optimized_dtype == expected
+
+
+class TestOptimizeForInferenceCudaDeviceContext:
+    """Verify that optimize_for_inference wraps operations in the correct device context."""
+
+    @pytest.mark.gpu
+    @patch("rfdetr.detr._move_model_context_to_device")
+    @patch("rfdetr.detr.deepcopy")
+    @patch("torch.cuda.device")
+    def test_cuda_device_context_manager_is_used_for_cuda_device(
+        self,
+        mock_cuda_device,
+        mock_deepcopy,
+        _mock_move_model_context_to_device,
+    ) -> None:
+        """torch.cuda.device() context should be entered when model is on CUDA."""
+        rfdetr = _FakeRFDETR()
+        # Simulate a CUDA device without actually requiring CUDA hardware
+        rfdetr.model.device = torch.device("cuda", 0)
+        mock_deepcopy.return_value = rfdetr.model.model
+
+        entered_devices: list[torch.device] = []
+
+        class _CapturingDeviceCtx:
+            def __init__(self, captured_device):
+                entered_devices.append(captured_device)
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, *args):
+                pass
+
+        mock_cuda_device.side_effect = _CapturingDeviceCtx
+        rfdetr.optimize_for_inference(compile=False, dtype=torch.float32)
+
+        assert len(entered_devices) == 1
+        assert entered_devices[0] == torch.device("cuda", 0)
+
+    def test_nullcontext_used_for_cpu_device(self) -> None:
+        """contextlib.nullcontext() should be used when model is on CPU (no CUDA init)."""
+        rfdetr = _FakeRFDETR()
+        rfdetr.model.device = torch.device("cpu")
+
+        # torch.cuda.device should NOT be called for CPU devices
+        with (
+            patch("torch.cuda.device") as mock_cuda_device,
+            patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model),
+        ):
+            rfdetr.optimize_for_inference(compile=False, dtype=torch.float32)
+
+        mock_cuda_device.assert_not_called()
+
+    @pytest.mark.gpu
+    @patch("rfdetr.detr._move_model_context_to_device")
+    @patch("rfdetr.detr.deepcopy")
+    @patch("torch.cuda.device")
+    def test_cuda_device_context_uses_model_device(
+        self,
+        mock_cuda_device,
+        mock_deepcopy,
+        _mock_move_model_context_to_device,
+    ) -> None:
+        """The device passed to torch.cuda.device() should match self.model.device."""
+        rfdetr = _FakeRFDETR()
+        expected_device = torch.device("cuda", 2)
+        rfdetr.model.device = expected_device
+        mock_deepcopy.return_value = rfdetr.model.model
+
+        captured: dict[str, torch.device] = {}
+
+        class _CapturingCtx:
+            def __init__(self, captured_device):
+                captured["device"] = captured_device
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, *args):
+                pass
+
+        mock_cuda_device.side_effect = _CapturingCtx
+        rfdetr.optimize_for_inference(compile=False)
+
+        assert captured.get("device") == expected_device
+
+
+class TestOptimizeForInferenceCompile:
+    """Tests for the compile=True path (JIT trace)."""
+
+    def test_compile_true_calls_jit_trace(self) -> None:
+        """torch.jit.trace should be called with the model and a correctly-shaped dummy input."""
+        rfdetr = _FakeRFDETR()
+        mock_traced = rfdetr.model.model
+
+        with (
+            patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model),
+            patch("torch.jit.trace", return_value=mock_traced) as mock_trace,
+        ):
+            rfdetr.optimize_for_inference(compile=True, batch_size=2)
+
+        assert mock_trace.called
+        dummy_input: torch.Tensor = mock_trace.call_args.args[1]
+        resolution = rfdetr.model.resolution
+        assert dummy_input.shape == (2, 3, resolution, resolution)
+
+    def test_compile_true_sets_compiled_flags(self) -> None:
+        """_optimized_has_been_compiled=True and _optimized_batch_size should be set after compile=True."""
+        rfdetr = _FakeRFDETR()
+
+        with (
+            patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model),
+            patch("torch.jit.trace", return_value=rfdetr.model.model),
+        ):
+            rfdetr.optimize_for_inference(compile=True, batch_size=4)
+
+        assert rfdetr._optimized_has_been_compiled is True
+        assert rfdetr._optimized_batch_size == 4
+
+    def test_compile_false_skips_jit_trace(self) -> None:
+        """torch.jit.trace should NOT be called when compile=False."""
+        rfdetr = _FakeRFDETR()
+
+        with (
+            patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model),
+            patch("torch.jit.trace") as mock_trace,
+        ):
+            rfdetr.optimize_for_inference(compile=False)
+
+        mock_trace.assert_not_called()
+        assert rfdetr._optimized_has_been_compiled is False
+        assert rfdetr._optimized_batch_size is None
+
+
+class TestOptimizeForInferenceState:
+    """Verify that optimize_for_inference correctly sets internal state flags."""
+
+    def test_is_optimized_inplace_false_before_optimization(self) -> None:
+        """is_optimized_inplace is False before any optimization is applied."""
+        rfdetr = _FakeRFDETR()
+        assert rfdetr.is_optimized_inplace is False
+
+    def test_is_optimized_flag_set(self) -> None:
+        """_is_optimized_for_inference should be True after optimization."""
+        rfdetr = _FakeRFDETR()
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False)
+
+        assert rfdetr._is_optimized_for_inference is True
+
+    def test_inference_model_set(self) -> None:
+        """model.inference_model should be set after optimization."""
+        rfdetr = _FakeRFDETR()
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False)
+
+        assert rfdetr.model.inference_model is not None
+
+    def test_remove_optimized_model_clears_state(self) -> None:
+        """remove_optimized_model() should clear all optimization flags."""
+        rfdetr = _FakeRFDETR()
+
+        with patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model):
+            rfdetr.optimize_for_inference(compile=False)
+
+        rfdetr.remove_optimized_model()
+
+        assert rfdetr._is_optimized_for_inference is False
+        assert rfdetr.model.inference_model is None
+        assert rfdetr._optimized_dtype is None
+        assert rfdetr._optimized_resolution is None
+        assert rfdetr._optimized_has_been_compiled is False
+        assert rfdetr._optimized_batch_size is None
+        assert rfdetr.is_optimized_inplace is False
+
+
+class TestOptimizeForInferenceInplace:
+    """Tests for the low-memory in-place optimization path."""
+
+    def test_inplace_false_keeps_deepcopy_behavior(self) -> None:
+        """The default path should still deep-copy the loaded module."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+        copied_model = _FakeModel()
+
+        with patch("rfdetr.detr.deepcopy", return_value=copied_model) as mock_deepcopy:
+            rfdetr.optimize_for_inference(compile=False)
+
+        mock_deepcopy.assert_called_once_with(original_model)
+        assert rfdetr.model.model is original_model
+        assert rfdetr.model.inference_model is copied_model
+        assert rfdetr._is_optimized_for_inference is True
+        assert rfdetr.is_optimized_inplace is False
+
+    def test_inplace_true_compile_false_does_not_deepcopy(self) -> None:
+        """Inplace=True with compile=False should use the loaded module directly."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+
+        with patch("rfdetr.detr.deepcopy") as mock_deepcopy:
+            rfdetr.optimize_for_inference(compile=False, inplace=True)
+
+        mock_deepcopy.assert_not_called()
+        assert rfdetr.model.model is None
+        assert rfdetr.model.inference_model is original_model
+        assert rfdetr._is_optimized_for_inference is True
+        assert rfdetr.is_optimized_inplace is True
+
+    def test_remove_optimized_model_after_inplace_warns_and_preserves_state(self) -> None:
+        """remove_optimized_model() after inplace optimization issues UserWarning and no-ops."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+
+        rfdetr.optimize_for_inference(compile=False, inplace=True)
+
+        with pytest.warns(UserWarning, match="no effect after inplace optimization"):
+            rfdetr.remove_optimized_model()
+
+        assert rfdetr.model.model is None
+        assert rfdetr.model.inference_model is original_model
+        assert rfdetr._is_optimized_for_inference is True
+        assert rfdetr.is_optimized_inplace is True
+
+    def test_second_optimize_after_inplace_raises_runtime_error(self) -> None:
+        """Calling optimize_for_inference() again after inplace=True raises RuntimeError."""
+        rfdetr = _FakeRFDETR()
+
+        rfdetr.optimize_for_inference(compile=False, inplace=True)
+
+        with pytest.raises(RuntimeError, match="base model has been cleared"):
+            rfdetr.optimize_for_inference(compile=False)
+
+    def test_inplace_true_default_dtype_float32_does_not_cast(self) -> None:
+        """Inplace=True with default dtype (float32) leaves weights unchanged — no casting occurs."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+        original_dtype = original_model.linear.weight.dtype
+
+        rfdetr.optimize_for_inference(compile=False, inplace=True)
+
+        assert rfdetr.model.inference_model is original_model
+        assert original_model.linear.weight.dtype == original_dtype
+        assert rfdetr._optimized_dtype == torch.float32
+        assert rfdetr._optimized_has_been_compiled is False
+        assert rfdetr._optimized_batch_size is None
+
+    @pytest.mark.parametrize(
+        "dtype",
+        [
+            pytest.param(torch.float16, id="float16"),
+            pytest.param(torch.bfloat16, id="bfloat16"),
+        ],
+    )
+    def test_inplace_true_allows_destructive_dtype_casting(self, dtype: torch.dtype) -> None:
+        """In-place optimization may cast the original module to the target dtype."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+
+        rfdetr.optimize_for_inference(compile=False, dtype=dtype, inplace=True)
+
+        assert rfdetr.model.model is None
+        assert rfdetr.model.inference_model is original_model
+        assert original_model.linear.weight.dtype == dtype
+        assert rfdetr._optimized_dtype == dtype
+
+    def test_inplace_export_failure_keeps_base_model(self) -> None:
+        """Export failure in the in-place path should not clear model.model."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+
+        with (
+            patch("rfdetr.detr.deepcopy") as mock_deepcopy,
+            patch.object(original_model, "export", side_effect=RuntimeError("export failed")),
+            pytest.raises(RuntimeError, match="export failed"),
+        ):
+            rfdetr.optimize_for_inference(compile=False, inplace=True)
+
+        mock_deepcopy.assert_not_called()
+        assert rfdetr.model.model is original_model
+        assert rfdetr.model.inference_model is None
+        assert rfdetr._is_optimized_for_inference is False
+        assert rfdetr.is_optimized_inplace is False
+
+    @pytest.mark.parametrize(
+        "dtype",
+        [
+            pytest.param(torch.int8, id="torch-int8"),
+            pytest.param("int8", id="string-int8"),
+        ],
+    )
+    def test_inplace_non_floating_dtype_raises_before_export(self, dtype: torch.dtype | str) -> None:
+        """In-place optimization rejects non-floating dtypes before mutating the base model."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+
+        with (
+            patch("rfdetr.detr.deepcopy") as mock_deepcopy,
+            patch.object(original_model, "export") as mock_export,
+            pytest.raises(ValueError, match="floating-point torch.dtype"),
+        ):
+            rfdetr.optimize_for_inference(compile=False, dtype=dtype, inplace=True)
+
+        mock_deepcopy.assert_not_called()
+        mock_export.assert_not_called()
+        assert rfdetr.model.model is original_model
+        assert rfdetr.model.inference_model is None
+        assert rfdetr._is_optimized_for_inference is False
+        assert rfdetr.is_optimized_inplace is False
+
+    def test_inplace_compile_true_raises_before_export_or_trace(self) -> None:
+        """In-place optimization rejects compile=True before mutating the base model."""
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+
+        with (
+            patch("rfdetr.detr.deepcopy") as mock_deepcopy,
+            patch.object(original_model, "export") as mock_export,
+            patch("torch.jit.trace") as mock_trace,
+            pytest.raises(ValueError, match="inplace=True.*compile=False"),
+        ):
+            rfdetr.optimize_for_inference(compile=True, inplace=True)
+
+        mock_deepcopy.assert_not_called()
+        mock_export.assert_not_called()
+        mock_trace.assert_not_called()
+        assert rfdetr.model.model is original_model
+        assert rfdetr.model.inference_model is None
+        assert rfdetr._is_optimized_for_inference is False
+        assert rfdetr._optimized_has_been_compiled is False
+        assert rfdetr._optimized_batch_size is None
+        assert rfdetr.is_optimized_inplace is False
+
+
+class TestOptimizeForInferenceExceptionRecovery:
+    """Verify state consistency when optimization fails mid-execution."""
+
+    def test_deepcopy_failure_leaves_clean_state(self) -> None:
+        """If deepcopy raises, inference_model should be None and _is_optimized_for_inference False."""
+        rfdetr = _FakeRFDETR()
+        # Simulate a previously-optimized state to confirm remove_optimized_model ran
+        rfdetr._is_optimized_for_inference = True
+        rfdetr.model.inference_model = rfdetr.model.model
+
+        with (
+            patch("rfdetr.detr.deepcopy", side_effect=RuntimeError("deepcopy failed")),
+            pytest.raises(RuntimeError, match="deepcopy failed"),
+        ):
+            rfdetr.optimize_for_inference(compile=False)
+
+        assert rfdetr.model.inference_model is None
+        assert rfdetr._is_optimized_for_inference is False
+
+    def test_export_failure_leaves_is_optimized_false(self) -> None:
+        """If export() raises after deepcopy succeeds, _is_optimized_for_inference stays False."""
+        rfdetr = _FakeRFDETR()
+        fake_copy = _FakeModel()
+
+        with (
+            patch("rfdetr.detr.deepcopy", return_value=fake_copy),
+            patch.object(fake_copy, "export", side_effect=RuntimeError("export failed")),
+            pytest.raises(RuntimeError, match="export failed"),
+        ):
+            rfdetr.optimize_for_inference(compile=False)
+
+        assert rfdetr._is_optimized_for_inference is False
+
+    def test_jit_trace_failure_leaves_compiled_flags_false(self) -> None:
+        """If jit.trace raises, _optimized_has_been_compiled and _optimized_batch_size stay unset."""
+        rfdetr = _FakeRFDETR()
+
+        with (
+            patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model),
+            patch("torch.jit.trace", side_effect=RuntimeError("trace failed")),
+            pytest.raises(RuntimeError, match="trace failed"),
+        ):
+            rfdetr.optimize_for_inference(compile=True, batch_size=2)
+
+        assert rfdetr._optimized_has_been_compiled is False
+        assert rfdetr._optimized_batch_size is None
+
+    def test_jit_trace_failure_leaves_model_fully_unoptimized(self) -> None:
+        """jit.trace failure leaves both _is_optimized_for_inference=False and inference_model=None."""
+        rfdetr = _FakeRFDETR()
+
+        with (
+            patch("rfdetr.detr.deepcopy", return_value=rfdetr.model.model),
+            patch("torch.jit.trace", side_effect=RuntimeError("trace failed")),
+            pytest.raises(RuntimeError, match="trace failed"),
+        ):
+            rfdetr.optimize_for_inference(compile=True)
+
+        assert rfdetr._is_optimized_for_inference is False
+        assert rfdetr.model.inference_model is None
+
+    def test_inplace_export_failure_module_mutations_are_not_undone(self) -> None:
+        """RFDETR resets flags on export failure but cannot undo module-level mutations.
+
+        Production export() may mutate the module (e.g. forward->forward_export) before raising; those changes are not
+        reversed by the exception-recovery path.
+        """
+        rfdetr = _FakeRFDETR()
+        original_model = rfdetr.model.model
+        mutated: dict[str, bool] = {"happened": False}
+
+        def _mutating_export() -> None:
+            mutated["happened"] = True
+            raise RuntimeError("export failed mid-mutation")
+
+        with (
+            patch("rfdetr.detr.deepcopy"),
+            patch.object(original_model, "export", side_effect=_mutating_export),
+            pytest.raises(RuntimeError, match="export failed mid-mutation"),
+        ):
+            rfdetr.optimize_for_inference(compile=False, inplace=True)
+
+        assert rfdetr._is_optimized_for_inference is False
+        assert rfdetr.is_optimized_inplace is False
+        assert rfdetr.model.model is original_model
+        # The mutation happened and cannot be undone by RFDETR's recovery path
+        assert mutated["happened"] is True

tests/inference/test_predict_eval_mode.py

@@ -0,0 +1,136 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests that unoptimized inference always runs the module in eval mode."""
+
+from types import SimpleNamespace
+
+import PIL.Image
+import pytest
+import torch
+
+from rfdetr import detr as detr_module
+
+from .helpers import _BaseFakeRFDETR
+
+
+class _FakeModelWithDropout(torch.nn.Module):
+    """Minimal module whose behavior differs between train and eval mode."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.dropout = torch.nn.Dropout(p=0.5)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Pass input through dropout, active only in train mode."""
+        return self.dropout(x)
+
+
+class _FakeModelContext:
+    """Minimal model context supplying the attributes predict() and train() need."""
+
+    def __init__(self) -> None:
+        self.device = torch.device("cpu")
+        self.resolution = 28
+        self.model = _FakeModelWithDropout()
+        self.inference_model = None
+
+
+class _FakeRFDETR(_BaseFakeRFDETR):
+    """Concrete test double: provides a dropout-bearing model for eval-mode tests."""
+
+    def get_model(self, config: SimpleNamespace) -> _FakeModelContext:
+        """Return a minimal model context with a dropout-bearing module."""
+        return _FakeModelContext()
+
+
+class TestUnoptimizedInferenceEvalMode:
+    """`_ensure_eval_mode_for_unoptimized_inference` must keep the module in eval mode."""
+
+    def test_eval_mode_reasserted_after_train_round_trip(self) -> None:
+        """Eval mode must be applied to whatever self.model.model currently points to.
+
+        ``train()`` rebinds ``self.model.model`` to a brand-new module left in training mode, so eval must be re-applied
+        to the *current* object on every call — not to a cached reference captured at init.
+        """
+        rfdetr = _FakeRFDETR()
+
+        # First inference call: warns once and switches to eval mode.
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+        assert rfdetr.model.model.training is False
+
+        # Simulate train() rebinding self.model.model to a fresh training-mode module.
+        rfdetr.model.model = _FakeModelWithDropout()
+        assert rfdetr.model.model.training is True  # new object starts in train mode
+
+        # Every subsequent inference call must re-assert eval on the *new* object.
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+        assert rfdetr.model.model.training is False
+
+    def test_optimized_model_skips_eval_assertion(self) -> None:
+        """When _is_optimized_for_inference is True, the method must be a no-op.
+
+        The compiled inference_model snapshot is already in eval mode; calling eval() on the stale self.model.model
+        would target the wrong object.
+        """
+        rfdetr = _FakeRFDETR()
+        rfdetr._is_optimized_for_inference = True
+        rfdetr.model.model.train()
+        assert rfdetr.model.model.training is True  # confirm starting state
+
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+
+        assert rfdetr.model.model.training is True  # must remain unchanged
+
+    def test_not_optimized_warning_emitted_only_once(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """The not-optimized warning is logged at most once across repeated calls."""
+        warnings: list[str] = []
+        monkeypatch.setattr(detr_module.logger, "warning", lambda msg, *a, **k: warnings.append(msg))
+
+        rfdetr = _FakeRFDETR()
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+        rfdetr.model.model.train()
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+
+        assert len(warnings) == 1
+
+    def test_eval_mode_applied_on_every_call(self) -> None:
+        """Eval() must run on every call, not just when the warning fires.
+
+        Simulate the code path where the warning has already been emitted
+        (``_has_warned_about_not_being_optimized_for_inference=True``) and verify
+        that ``eval()`` is still applied to the current module.
+        """
+        rfdetr = _FakeRFDETR()
+        rfdetr._has_warned_about_not_being_optimized_for_inference = True
+        rfdetr.model.model.train()
+
+        rfdetr._ensure_eval_mode_for_unoptimized_inference()
+
+        assert rfdetr.model.model.training is False
+
+    def test_predict_puts_module_in_eval_mode(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """Predict() must delegate to _ensure_eval_mode_for_unoptimized_inference, leaving module in eval mode."""
+        rfdetr = _FakeRFDETR()
+        img = PIL.Image.new("RGB", (640, 640), color=(128, 128, 128))
+
+        monkeypatch.setattr(
+            rfdetr.model.model,
+            "forward",
+            lambda batch: {"pred_logits": torch.zeros(1, 10, 81), "pred_boxes": torch.zeros(1, 10, 4)},
+        )
+        monkeypatch.setattr(
+            rfdetr.model,
+            "postprocess",
+            lambda preds, target_sizes: [
+                {"scores": torch.zeros(0), "labels": torch.zeros(0, dtype=torch.long), "boxes": torch.zeros(0, 4)}
+            ],
+            raising=False,
+        )
+
+        rfdetr.predict(img)
+
+        assert rfdetr.model.model.training is False

tests/inference/test_public_api_keypoints.py

@@ -0,0 +1,34 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Public API tests for the keypoint preview variant."""
+
+from rfdetr import RFDETRKeypointPreview
+from rfdetr.config import KeypointTrainConfig, RFDETRKeypointPreviewConfig
+from rfdetr.detr import RFDETRKeypointPreview as RFDETRKeypointPreviewFromDetr
+from rfdetr.variants import RFDETRKeypointPreview as RFDETRKeypointPreviewFromVariants
+
+
+def test_keypoint_preview_top_level_import() -> None:
+    """RFDETRKeypointPreview must be importable from top-level package and keep shared identity."""
+    assert RFDETRKeypointPreview is RFDETRKeypointPreviewFromVariants
+    assert RFDETRKeypointPreview is RFDETRKeypointPreviewFromDetr
+
+
+def test_keypoint_preview_variant_metadata() -> None:
+    """RFDETRKeypointPreview exposes the expected variant metadata and config class."""
+    assert RFDETRKeypointPreview.size == "rfdetr-keypoint-preview"
+    assert RFDETRKeypointPreview._model_config_class is RFDETRKeypointPreviewConfig
+    assert RFDETRKeypointPreview._train_config_class is KeypointTrainConfig
+    assert RFDETRKeypointPreviewConfig.model_fields["pretrain_weights"].default == "rf-detr-keypoint-preview-xlarge.pth"
+
+
+def test_predict_docstring_mentions_one_class_keypoint_mapping() -> None:
+    """The public predict() contract must document active-first keypoint class-id mapping."""
+    doc = RFDETRKeypointPreview.predict.__doc__ or ""
+    assert "one-class preview keypoint setup" in doc
+    assert "class_id=0" in doc
+    assert "class_id=1" in doc
+    assert "__background__" in doc

tests/inference/test_supervision_editable_keypoints.py

@@ -0,0 +1,24 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Smoke test for RF-DETR keypoints carried by Supervision KeyPoints."""
+
+import numpy as np
+import supervision as sv
+
+
+def test_rfdetr_keypoints_include_detection_details() -> None:
+    """RF-DETR-style keypoints preserve detection boxes and scores."""
+    key_points = sv.KeyPoints(
+        xy=np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32),
+        keypoint_confidence=np.array([[0.9, 0.8]], dtype=np.float32),
+        detection_confidence=np.array([0.95], dtype=np.float32),
+        class_id=np.array([1], dtype=int),
+        data={"xyxy": np.array([[0, 0, 10, 10]], dtype=np.float32)},
+    )
+
+    assert key_points.xy.shape == (1, 2, 2)
+    np.testing.assert_array_equal(key_points.data["xyxy"], np.array([[0, 0, 10, 10]], dtype=np.float32))
+    np.testing.assert_array_equal(key_points.detection_confidence, np.array([0.95], dtype=np.float32))

tests/inference/test_trt_inference.py

@@ -0,0 +1,110 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+from PIL import Image
+
+from rfdetr.export.benchmark import TRTInference, infer_transforms
+
+
+class TestTRTInference:
+    def test_synchronize_sync_mode_does_not_require_stream(self, monkeypatch) -> None:
+        """`synchronize()` should not access stream in sync mode."""
+        inference = TRTInference.__new__(TRTInference)
+        inference.sync_mode = True
+
+        mock_is_available = Mock(return_value=True)
+        mock_cuda_sync = Mock()
+        monkeypatch.setattr("torch.cuda.is_available", mock_is_available)
+        monkeypatch.setattr("torch.cuda.synchronize", mock_cuda_sync)
+
+        inference.synchronize()
+
+        mock_is_available.assert_called_once()
+        mock_cuda_sync.assert_called_once()
+
+    def test_synchronize_async_mode_uses_stream_sync(self, monkeypatch) -> None:
+        """`synchronize()` should use stream synchronization in async mode."""
+        inference = TRTInference.__new__(TRTInference)
+        inference.sync_mode = False
+        inference.stream = Mock()
+
+        mock_cuda_sync = Mock()
+        monkeypatch.setattr("torch.cuda.synchronize", mock_cuda_sync)
+
+        inference.synchronize()
+
+        inference.stream.synchronize.assert_called_once()
+        mock_cuda_sync.assert_not_called()
+
+    def test_infer_transforms_accepts_none_target(self) -> None:
+        """Benchmark inference preprocessing should support image-only input."""
+        image = Image.new("RGB", (320, 240))
+
+        image_tensor, target = infer_transforms()(image, None)
+
+        assert isinstance(image_tensor, torch.Tensor)
+        assert image_tensor.shape == (3, 640, 640)
+        assert image_tensor.dtype == torch.float32
+        assert target is None
+
+
+class TestBenchmarkShapeParameterization:
+    """Benchmark preprocessing/postprocessing read input size and query count instead of hardcoding 640/300."""
+
+    def test_infer_transforms_uses_requested_size(self) -> None:
+        """infer_transforms resizes to the caller-supplied (height, width)."""
+        image = Image.new("RGB", (320, 240))
+
+        image_tensor, _ = infer_transforms((512, 384))(image, None)
+
+        assert image_tensor.shape == (3, 512, 384)
+
+    def test_infer_transforms_defaults_to_640(self) -> None:
+        """The default input size stays 640x640 for callers that do not pass a size."""
+        image = Image.new("RGB", (320, 240))
+
+        image_tensor, _ = infer_transforms()(image, None)
+
+        assert image_tensor.shape == (3, 640, 640)
+
+    def test_static_dim_returns_concrete_int(self) -> None:
+        """A concrete positive dimension is returned unchanged."""
+        from rfdetr.export.benchmark import _static_dim
+
+        assert _static_dim(384, 640) == 384
+
+    @pytest.mark.parametrize(
+        "value",
+        [
+            pytest.param("height", id="dynamic-string"),
+            pytest.param(None, id="none"),
+            pytest.param(-1, id="negative"),
+        ],
+    )
+    def test_static_dim_falls_back_for_dynamic_axis(self, value) -> None:
+        """Dynamic/unknown axes fall back to the provided default."""
+        from rfdetr.export.benchmark import _static_dim
+
+        assert _static_dim(value, 640) == 640
+
+    def test_post_process_respects_num_queries(self) -> None:
+        """post_process selects exactly num_queries detections per image."""
+        from rfdetr.export.benchmark import post_process
+
+        num_queries = 5
+        outputs = {
+            "labels": torch.rand(1, 20, 3),
+            "dets": torch.rand(1, 20, 4),
+        }
+        target_sizes = torch.tensor([[480, 640]])
+
+        results = post_process(outputs, target_sizes, num_queries=num_queries)
+
+        assert results[0]["scores"].shape == (num_queries,)

tests/models/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/models/backbone/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/models/backbone/test_backbone.py

@@ -0,0 +1,119 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for dual-projector backbone joiner routing."""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from rfdetr.models.backbone import Joiner
+from rfdetr.utilities.tensors import NestedTensor
+
+
+class _FakeBackbone(nn.Module):
+    """Backbone shim used to validate Joiner contract changes."""
+
+    def __init__(
+        self,
+        features: list[NestedTensor],
+        cross_attention_features: list[object] | None,
+    ) -> None:
+        super().__init__()
+        self._features = features
+        self._cross_attention_features = cross_attention_features
+
+    def forward(self, tensor: torch.Tensor | NestedTensor):
+        if isinstance(tensor, torch.Tensor):
+            feats = [f.tensors for f in self._features]
+            masks = [f.mask for f in self._features]
+            return feats, masks, self._cross_attention_features
+        return self._features, self._cross_attention_features
+
+
+class _FakePositionEncoding(nn.Module):
+    """Tiny callable that behaves like a position encoder."""
+
+    def forward(self, nested_tensor: NestedTensor | torch.Tensor, align_dim_orders: bool = False) -> torch.Tensor:
+        if isinstance(nested_tensor, NestedTensor):
+            base = nested_tensor.tensors
+        else:
+            base = nested_tensor
+        if base.dim() == 3:
+            base = base[:, None]
+        return torch.zeros((base.shape[0], 1, base.shape[-2], base.shape[-1]), dtype=base.dtype, device=base.device)
+
+
+def _feature(shape: tuple[int, ...], batch_size: int = 2) -> NestedTensor:
+    channels, height, width = shape
+    return NestedTensor(
+        tensors=torch.ones((batch_size, channels, height, width), dtype=torch.float32),
+        mask=torch.zeros((batch_size, height, width), dtype=torch.bool),
+    )
+
+
+def _input_tensor(batch_size: int = 2) -> tuple[NestedTensor, torch.Tensor]:
+    return (
+        NestedTensor(
+            tensors=torch.ones((batch_size, 3, 16, 16), dtype=torch.float32),
+            mask=torch.zeros((batch_size, 16, 16), dtype=torch.bool),
+        ),
+        torch.ones((batch_size, 3, 16, 16), dtype=torch.float32),
+    )
+
+
+def test_joiner_dual_projector_disabled_contract() -> None:
+    """Joiner should forward one feature stream and a ``None`` cross-attention stream when disabled."""
+    features = [_feature((256, 16, 16))]
+    joiner = Joiner(_FakeBackbone(features, None), _FakePositionEncoding())
+
+    input_tensor, image = _input_tensor()
+
+    _, _, cross_attention = joiner(input_tensor)
+    assert cross_attention is None
+    assert len(joiner(input_tensor)[0]) == 1
+
+    exported = joiner.forward_export(image)
+    assert exported[3] is None
+    assert len(exported[0]) == 1
+    assert exported[2][0].shape == (2, 16, 16)
+
+
+def test_joiner_dual_projector_enabled_contract() -> None:
+    """Joiner should forward cross-attention features in parallel with feature features when enabled."""
+    features = [_feature((256, 16, 16)), _feature((256, 8, 8))]
+    cross_attention_features = [_feature((256, 16, 16)), _feature((256, 8, 8))]
+    joiner = Joiner(_FakeBackbone(features, cross_attention_features), _FakePositionEncoding())
+
+    input_tensor, _ = _input_tensor()
+
+    feature_tensors, _, cross_attention = joiner(input_tensor)
+    assert len(feature_tensors) == len(cross_attention)
+    assert all(f.tensors.shape == c.tensors.shape for f, c in zip(feature_tensors, cross_attention))
+    assert all(f.mask is not None for f in cross_attention)
+
+
+def test_joiner_forward_export_contract() -> None:
+    """Exported joiner contracts should remain 4-tuples and preserve cross-attention stream arity."""
+    exported_features = [torch.ones(2, 256, 16, 16), torch.ones(2, 256, 8, 8)]
+    exported_masks = [torch.zeros(2, 16, 16, dtype=torch.bool), torch.zeros(2, 8, 8, dtype=torch.bool)]
+    export_backbone = _FakeBackbone(
+        [NestedTensor(t, mask) for t, mask in zip(exported_features, exported_masks)],
+        [torch.ones(2, 256, 16, 16), torch.ones(2, 256, 8, 8)],
+    )
+    joiner = Joiner(export_backbone, _FakePositionEncoding())
+
+    outputs = joiner.forward_export(torch.ones(2, 3, 16, 16))
+    feats_out, masks_out, poss, cross_attention = outputs
+
+    assert len(feats_out) == len(exported_features)
+    assert len(masks_out) == len(exported_masks)
+    assert feats_out[0].shape == exported_features[0].shape
+    assert masks_out[0].shape == exported_masks[0].shape
+    assert len(outputs) == 4
+    assert poss[0].shape == exported_features[0][:, :1, :, :].shape
+    assert isinstance(cross_attention, list)
+    assert all(isinstance(feature, torch.Tensor) for feature in cross_attention)

tests/models/backbone/test_backbone_export.py

@@ -0,0 +1,52 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for backbone export behavior."""
+
+import sys
+from types import ModuleType
+
+from rfdetr.models.backbone.backbone import Backbone
+
+
+class TestBackboneExport:
+    """Tests for ``Backbone.export``."""
+
+    def test_export_without_lora_encoder_skips_peft_import_and_warning(self, monkeypatch) -> None:
+        """Non-LoRA exports should not warn just because peft is unavailable."""
+        backbone = object.__new__(Backbone)
+        backbone.encoder = object()
+        warning_messages: list[str] = []
+
+        monkeypatch.delitem(sys.modules, "peft", raising=False)
+        monkeypatch.setattr("rfdetr.models.backbone.backbone.logger.warning", warning_messages.append)
+
+        backbone.export()
+
+        assert warning_messages == []
+
+    def test_export_replaces_peft_encoder_with_merged_encoder(self, monkeypatch) -> None:
+        """Export should replace PEFT wrapper with merged base encoder."""
+
+        class _MergedEncoder:
+            pass
+
+        class _FakePeftModel:
+            def __init__(self) -> None:
+                self._merged = _MergedEncoder()
+
+            def merge_and_unload(self):
+                return self._merged
+
+        peft_module = ModuleType("peft")
+        peft_module.PeftModel = _FakePeftModel
+        monkeypatch.setitem(sys.modules, "peft", peft_module)
+
+        backbone = object.__new__(Backbone)
+        backbone.encoder = _FakePeftModel()
+
+        backbone.export()
+
+        assert isinstance(backbone.encoder, _MergedEncoder)

tests/models/backbone/test_windowed_dino.py

@@ -0,0 +1,544 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+import pytest
+import torch
+
+from rfdetr.models.backbone.dinov2_with_windowed_attn import (
+    Dinov2WithRegistersAttention,
+    Dinov2WithRegistersSdpaAttention,
+    WindowedDinov2WithRegistersBackbone,
+    WindowedDinov2WithRegistersConfig,
+    WindowedDinov2WithRegistersEmbeddings,
+    WindowedDinov2WithRegistersModel,
+    _find_pruneable_heads_and_indices,
+    _get_aligned_output_features_output_indices,
+)
+
+
+def test_window_partition_forward_rectangular_preserves_shapes():
+    """Regression test for WindowedDinov2WithRegistersEmbeddings.forward with rectangular input.
+
+    Ensures window partitioning logic correctly handles H != W.
+    """
+    # Params: H_patches=6, W_patches=4, num_windows=2 -> 3x2 patches per window
+    batch_size, hidden_size, patch_size, num_windows = 1, 64, 16, 2
+    hp, wp, nr = 6, 4, 4
+    h, w = hp * patch_size, wp * patch_size
+
+    config = WindowedDinov2WithRegistersConfig(
+        hidden_size=hidden_size,
+        patch_size=patch_size,
+        num_windows=num_windows,
+        image_size=h,  # square image_size for positional embeddings
+        num_register_tokens=nr,
+    )
+    model = WindowedDinov2WithRegistersEmbeddings(config)
+
+    # Input is rectangular
+    pixel_values = torch.randn(batch_size, 3, h, w)
+    result = model(pixel_values)
+
+    expected_batch = batch_size * (num_windows**2)
+    expected_seq_len = 1 + nr + (hp // num_windows) * (wp // num_windows)
+
+    assert result.shape == (expected_batch, expected_seq_len, hidden_size)
+
+
+# Before fix in PR #448 the reshape used num_h_patches_per_window in both the height
+# AND width dimension. This only fails when height and width produce different patch
+# counts, so all tests below use non-square images (hp != wp).
+
+
+@pytest.mark.parametrize(
+    "hp, wp, num_windows",
+    [
+        (4, 6, 2),  # wider than tall
+        (6, 4, 2),  # taller than wide
+        (6, 9, 3),  # 3-window grid, non-square
+        (8, 4, 2),  # 2:1 aspect ratio
+    ],
+)
+def test_window_partition_nonsquare_does_not_raise(hp, wp, num_windows):
+    """Before the fix, the reshape used num_h_patches_per_window for the width dimension, so the total element count
+    mismatched and PyTorch raised a RuntimeError for any non-square image.
+
+    The fix replaces that variable with num_w_patches_per_window, making the operation valid for all shapes.
+    """
+    hidden_size, patch_size, nr = 32, 16, 0
+    h, w = hp * patch_size, wp * patch_size
+
+    config = WindowedDinov2WithRegistersConfig(
+        hidden_size=hidden_size,
+        patch_size=patch_size,
+        num_windows=num_windows,
+        image_size=max(h, w),
+        num_register_tokens=nr,
+    )
+    model = WindowedDinov2WithRegistersEmbeddings(config)
+    pixel_values = torch.randn(1, 3, h, w)
+
+    # This line would raise RuntimeError before the fix
+    result = model(pixel_values)
+
+    expected_batch = num_windows**2
+    expected_seq_len = 1 + (hp // num_windows) * (wp // num_windows)
+    assert result.shape == (expected_batch, expected_seq_len, hidden_size)
+
+
+def test_window_partition_correct_window_content():
+    """Verifies that after windowing each window contains the spatially correct patch tokens — not just that the shape
+    is right.
+
+    Layout with hp=4, wp=6, num_windows=2 (2x2 grid of windows):   Window (0,0): rows 0-1, cols 0-2   Window (0,1): rows
+    0-1, cols 3-5   Window (1,0): rows 2-3, cols 0-2   Window (1,1): rows 2-3, cols 3-5
+
+    Before the fix the reshape used num_h_patches_per_window for the width dim so it raised an error and never produced
+    window content at all.
+    """
+    hidden_size, patch_size, num_windows, nr = 1, 16, 2, 0
+    hp, wp = 4, 6
+    h, w = hp * patch_size, wp * patch_size
+    batch_size = 1
+
+    config = WindowedDinov2WithRegistersConfig(
+        hidden_size=hidden_size,
+        patch_size=patch_size,
+        num_windows=num_windows,
+        image_size=max(h, w),
+        num_register_tokens=nr,
+        num_hidden_layers=1,
+        num_attention_heads=1,
+    )
+    model = WindowedDinov2WithRegistersEmbeddings(config)
+
+    # Disable position embeddings and cls token so we can track patch identity.
+    # Each patch gets a unique value equal to its flat index (row * wp + col).
+    with torch.no_grad():
+        model.position_embeddings.zero_()
+        model.cls_token.zero_()
+
+    # Build a synthetic patch embedding: patch at (row, col) has value row*wp+col.
+    # Shape after patch projection: (1, hp*wp, 1) — hidden_size=1 for simplicity.
+    patch_ids = torch.arange(hp * wp, dtype=torch.float).view(1, hp * wp, 1)
+
+    # Bypass the full forward pass and exercise the windowing logic directly.
+    pixel_tokens = patch_ids  # (1, 24, 1)
+    pixel_tokens_2d = pixel_tokens.view(batch_size, hp, wp, hidden_size)  # (1,4,6,1)
+
+    num_h_patches_per_window = hp // num_windows  # 2
+    num_w_patches_per_window = wp // num_windows  # 3
+
+    # --- correct reshape (the fix) ---
+    windowed = pixel_tokens_2d.reshape(
+        batch_size * num_windows,
+        num_h_patches_per_window,
+        num_windows,
+        num_w_patches_per_window,
+        hidden_size,
+    )
+    windowed = windowed.permute(0, 2, 1, 3, 4)
+    windowed = windowed.reshape(
+        batch_size * num_windows**2,
+        num_h_patches_per_window * num_w_patches_per_window,
+        hidden_size,
+    )
+
+    # Expected content for each of the 4 windows (6 patches each):
+    expected = torch.tensor(
+        [
+            # Window 0 (rows 0-1, cols 0-2): ids 0,1,2, 6,7,8
+            [[0.0], [1.0], [2.0], [6.0], [7.0], [8.0]],
+            # Window 1 (rows 0-1, cols 3-5): ids 3,4,5, 9,10,11
+            [[3.0], [4.0], [5.0], [9.0], [10.0], [11.0]],
+            # Window 2 (rows 2-3, cols 0-2): ids 12,13,14, 18,19,20
+            [[12.0], [13.0], [14.0], [18.0], [19.0], [20.0]],
+            # Window 3 (rows 2-3, cols 3-5): ids 15,16,17, 21,22,23
+            [[15.0], [16.0], [17.0], [21.0], [22.0], [23.0]],
+        ]
+    )
+    assert torch.equal(windowed, expected), f"Window content mismatch:\n{windowed}\n!=\n{expected}"
+
+
+def test_buggy_reshape_raises_for_nonsquare():
+    """Directly demonstrates what the pre-fix code did: using num_h_patches_per_window in the width position of the
+    reshape causes a RuntimeError when the element count is not divisible by the (wrong) shape.
+
+    With hidden_size=1 and hp=4, wp=6, num_windows=2 the total elements are 24 but the buggy target dims (2,2,2,2,-1)
+    require a non-integer last dimension, so PyTorch raises RuntimeError.
+    """
+    hp, wp = 4, 6  # non-square: width > height
+    num_windows = 2
+    hidden_size = 1  # chosen so total / buggy-fixed-dims is non-integer
+
+    num_h_patches_per_window = hp // num_windows  # 2
+    num_w_patches_per_window = wp // num_windows  # 3
+    batch_size = 1
+
+    # Simulate pixel_tokens_with_pos_embed after the .view() call
+    pixel_tokens_2d = torch.randn(batch_size, hp, wp, hidden_size)
+
+    # The correct reshape (post-fix) must succeed
+    pixel_tokens_2d.reshape(
+        batch_size * num_windows,
+        num_h_patches_per_window,
+        num_windows,
+        num_w_patches_per_window,  # correct
+        hidden_size,
+    )
+
+    # The buggy reshape (pre-fix) must raise RuntimeError:
+    # total elements = 1*4*6*1 = 24,  fixed-dims product = 2*2*2*2 = 16,  16 ∤ 24.
+    with pytest.raises(RuntimeError):
+        pixel_tokens_2d.reshape(
+            batch_size * num_windows,
+            num_h_patches_per_window,
+            num_windows,
+            num_h_patches_per_window,  # bug: height used for width
+            -1,
+        )
+
+
+def test_buggy_reshape_silent_corruption_for_nonsquare():
+    """When hidden_size happens to make the total element count divisible by the buggy target shape, PyTorch does NOT
+    raise — instead the last dimension is inflated, which silently corrupts the tensor layout.
+
+    Pre-fix with hp=4, wp=6, hidden_size=8, num_windows=2:   total elements = 1*4*6*8 = 192 buggy fixed dims = 2*2*2*2 =
+    16  →  last dim inferred as 192/16 = 12 (not 8)
+
+    The fix ensures the correct reshape always yields a last dim equal to hidden_size.
+    """
+    hp, wp = 4, 6
+    num_windows = 2
+    hidden_size = 8
+
+    num_h_patches_per_window = hp // num_windows  # 2
+    num_w_patches_per_window = wp // num_windows  # 3
+    batch_size = 1
+
+    pixel_tokens_2d = torch.randn(batch_size, hp, wp, hidden_size)
+
+    # Buggy reshape silently infers last dim = 12 (not 8)
+    buggy_out = pixel_tokens_2d.reshape(
+        batch_size * num_windows,
+        num_h_patches_per_window,
+        num_windows,
+        num_h_patches_per_window,  # bug
+        -1,
+    )
+    assert buggy_out.shape[-1] != hidden_size, "Buggy reshape should produce wrong last dim"
+
+    # Correct reshape always yields last dim == hidden_size
+    correct_out = pixel_tokens_2d.reshape(
+        batch_size * num_windows,
+        num_h_patches_per_window,
+        num_windows,
+        num_w_patches_per_window,  # fix
+        hidden_size,
+    )
+    assert correct_out.shape[-1] == hidden_size
+
+
+# ---------------------------------------------------------------------------
+# Tests for locally-copied utility functions (removed from transformers v5 public API)
+# ---------------------------------------------------------------------------
+
+
+class TestGetAlignedOutputFeaturesOutputIndices:
+    """Tests for the local copy of get_aligned_output_features_output_indices."""
+
+    def test_both_none_returns_last_stage(self):
+        stage_names = ["stage1", "stage2", "stage3"]
+        features, indices = _get_aligned_output_features_output_indices(None, None, stage_names)
+        assert features == ["stage3"]
+        assert indices == [2]
+
+    def test_only_out_features_derives_indices(self):
+        stage_names = ["stem", "layer1", "layer2", "layer3"]
+        features, indices = _get_aligned_output_features_output_indices(["layer1", "layer3"], None, stage_names)
+        assert features == ["layer1", "layer3"]
+        assert indices == [1, 3]
+
+    def test_only_out_indices_derives_features(self):
+        stage_names = ["stem", "layer1", "layer2", "layer3"]
+        features, indices = _get_aligned_output_features_output_indices(None, [0, 2], stage_names)
+        assert features == ["stem", "layer2"]
+        assert indices == [0, 2]
+
+    def test_both_provided_returns_as_is(self):
+        stage_names = ["stem", "layer1", "layer2"]
+        features, indices = _get_aligned_output_features_output_indices(["layer1"], [1], stage_names)
+        assert features == ["layer1"]
+        assert indices == [1]
+
+    def test_out_indices_converted_to_list(self):
+        """out_indices supplied as a tuple must be returned as a list."""
+        stage_names = ["stem", "layer1", "layer2"]
+        _, indices = _get_aligned_output_features_output_indices(None, (1, 2), stage_names)
+        assert isinstance(indices, list)
+        assert indices == [1, 2]
+
+
+class TestFindPruneableHeadsAndIndices:
+    """Tests for the local copy of find_pruneable_heads_and_indices."""
+
+    def test_no_pruning_returns_full_index(self):
+        heads, index = _find_pruneable_heads_and_indices(set(), n_heads=4, head_size=3, already_pruned_heads=set())
+        assert len(heads) == 0
+        assert len(index) == 12  # 4 * 3, nothing masked
+
+    @pytest.mark.parametrize(
+        "head_to_prune, expected_index",
+        [
+            pytest.param({0}, list(range(3, 12)), id="prune-first-head"),
+            pytest.param({3}, list(range(9)), id="prune-last-head"),
+        ],
+    )
+    def test_prune_single_head_removes_correct_rows(self, head_to_prune, expected_index):
+        # Head N masked → N*head_size indices removed; remaining = n_heads*head_size - head_size = 9
+        heads, index = _find_pruneable_heads_and_indices(
+            head_to_prune, n_heads=4, head_size=3, already_pruned_heads=set()
+        )
+        assert heads == head_to_prune
+        assert len(index) == 9
+        assert index.tolist() == expected_index
+
+    def test_already_pruned_head_adjusts_offset(self):
+        # Head 0 was already pruned. Now pruning head 1 (which is now effective head 0
+        # after offset adjustment) should remove 3 more indices from the effective mask.
+        heads, index = _find_pruneable_heads_and_indices({1}, n_heads=4, head_size=3, already_pruned_heads={0})
+        assert 1 in heads
+        assert len(index) == 9  # 4*3 - 3 pruned
+
+
+# ---------------------------------------------------------------------------
+# Smoke tests for WindowedDinov2WithRegistersBackbone
+# ---------------------------------------------------------------------------
+
+
+def _minimal_backbone_config(**kwargs) -> WindowedDinov2WithRegistersConfig:
+    """Return the smallest valid config for backbone instantiation tests."""
+    defaults = dict(
+        hidden_size=32,
+        num_hidden_layers=1,
+        num_attention_heads=2,
+        intermediate_size=64,
+        patch_size=16,
+        image_size=64,
+        num_register_tokens=0,
+        num_windows=1,
+    )
+    defaults.update(kwargs)
+    return WindowedDinov2WithRegistersConfig(**defaults)
+
+
+class TestWindowedDinov2WithRegistersBackbone:
+    """Smoke tests that guard against _init_transformers_backbone() API regressions."""
+
+    @pytest.mark.parametrize(
+        "attr",
+        [
+            pytest.param("stage_names", id="stage_names"),
+            pytest.param("out_features", id="out_features"),
+        ],
+    )
+    def test_instantiation_sets_list_attribute(self, attr):
+        config = _minimal_backbone_config()
+        backbone = WindowedDinov2WithRegistersBackbone(config)
+        assert hasattr(backbone, attr)
+        assert isinstance(getattr(backbone, attr), list)
+        assert len(getattr(backbone, attr)) > 0
+
+    def test_forward_returns_backbone_output(self):
+        config = _minimal_backbone_config()
+        backbone = WindowedDinov2WithRegistersBackbone(config)
+        backbone.eval()
+        pixel_values = torch.randn(1, 3, 64, 64)
+        with torch.no_grad():
+            output = backbone(pixel_values)
+        assert hasattr(output, "feature_maps")
+        assert len(output.feature_maps) == len(backbone.out_features)
+
+
+# ---------------------------------------------------------------------------
+# Test for output_attentions=True SDPA fallback path
+# ---------------------------------------------------------------------------
+
+
+class TestSdpaFallbackWithOutputAttentions:
+    """Guards the output_attentions behaviour in windowed attention."""
+
+    def test_output_attentions_true_raises(self):
+        """Windowed attention explicitly does not support output_attentions=True."""
+        config = _minimal_backbone_config()
+        model = WindowedDinov2WithRegistersModel(config)
+        model.eval()
+        pixel_values = torch.randn(1, 3, 64, 64)
+        with torch.no_grad():
+            with pytest.raises(AssertionError, match="output_attentions is not supported for windowed attention"):
+                model(pixel_values, output_attentions=True)
+
+
+class TestSetAttnImplementation:
+    """Tests for WindowedDinov2WithRegistersModel.set_attn_implementation."""
+
+    @pytest.mark.parametrize(
+        "switches, expected_impl, expected_cls",
+        [
+            pytest.param(["eager"], "eager", Dinov2WithRegistersAttention, id="sdpa-to-eager"),
+            pytest.param(["eager", "sdpa"], "sdpa", Dinov2WithRegistersSdpaAttention, id="roundtrip-back-to-sdpa"),
+        ],
+    )
+    def test_switch_updates_config_and_layers(self, switches, expected_impl, expected_cls):
+        """After each call in *switches*, config and all layer attention modules reflect the final impl."""
+        config = _minimal_backbone_config()
+        model = WindowedDinov2WithRegistersModel(config)
+
+        for impl in switches:
+            model.set_attn_implementation(impl)
+
+        assert model.config._attn_implementation == expected_impl
+        for layer in model.encoder.layer:
+            assert type(layer.attention) is expected_cls
+
+    def test_invalid_implementation_raises(self):
+        """Passing an unknown key raises ValueError with a clear message."""
+        config = _minimal_backbone_config()
+        model = WindowedDinov2WithRegistersModel(config)
+
+        with pytest.raises(ValueError, match="Unknown attn_implementation"):
+            model.set_attn_implementation("flash_attention_2")
+
+
+@pytest.mark.parametrize(
+    "h, w, num_windows, should_raise",
+    [
+        pytest.param(64, 64, 2, False, id="valid-square"),
+        pytest.param(64, 96, 2, False, id="valid-rectangular"),
+        pytest.param(32, 32, 1, False, id="num_windows-1-valid"),
+        pytest.param(33, 64, 2, True, id="h-not-divisible"),
+        pytest.param(64, 33, 2, True, id="w-not-divisible"),
+        pytest.param(33, 33, 2, True, id="both-not-divisible"),
+    ],
+)
+def test_forward_validates_spatial_dims(h: int, w: int, num_windows: int, should_raise: bool) -> None:
+    """WindowedDinov2WithRegistersEmbeddings raises ValueError for incompatible dims.
+
+    Both H and W must be divisible by patch_size * num_windows.  The check must survive Python's -O flag (assert would
+    be silently stripped).
+    """
+    patch_size = 16
+    config = WindowedDinov2WithRegistersConfig(
+        hidden_size=32,
+        patch_size=patch_size,
+        num_windows=num_windows,
+        image_size=max(h, w),
+        num_register_tokens=0,
+    )
+    model = WindowedDinov2WithRegistersEmbeddings(config)
+    pixel_values = torch.randn(1, 3, h, w)
+    if should_raise:
+        with pytest.raises(ValueError, match="divisible"):
+            model(pixel_values)
+    else:
+        model(pixel_values)  # must not raise
+
+
+def _make_small_model() -> WindowedDinov2WithRegistersModel:
+    """Return the smallest valid WindowedDinov2WithRegistersModel for unit tests."""
+    config = WindowedDinov2WithRegistersConfig(
+        hidden_size=32,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        image_size=32,
+        patch_size=16,
+        num_register_tokens=2,
+    )
+    return WindowedDinov2WithRegistersModel(config)
+
+
+class TestSetAttnImplementationPreservesWeights:
+    """set_attn_implementation must transfer trained weights to the new attention module.
+
+    Before the fix the method replaced each layer's attention with a freshly constructed (randomly initialised) module,
+    silently discarding all trained q/k/v/output weights.
+    """
+
+    @pytest.mark.parametrize(
+        "from_impl, to_impl",
+        [
+            pytest.param("sdpa", "eager", id="sdpa_to_eager"),
+            pytest.param("eager", "sdpa", id="eager_to_sdpa"),
+        ],
+    )
+    def test_query_weight_preserved_after_switch(self, from_impl: str, to_impl: str) -> None:
+        """After switching implementation the query weight tensor must be unchanged."""
+        model = _make_small_model()
+        model.set_attn_implementation(from_impl)
+
+        # Record the query weights of every layer before switching.
+        before = [layer.attention.attention.query.weight.clone() for layer in model.encoder.layer]
+
+        model.set_attn_implementation(to_impl)
+
+        after = [layer.attention.attention.query.weight for layer in model.encoder.layer]
+        for layer_idx, (w_before, w_after) in enumerate(zip(before, after)):
+            assert torch.equal(w_before, w_after), (
+                f"Layer {layer_idx}: query weight changed after set_attn_implementation({from_impl!r} → {to_impl!r})"
+            )
+
+    def test_key_and_value_weights_preserved(self) -> None:
+        """Key and value weights must also survive the implementation switch."""
+        model = _make_small_model()
+        model.set_attn_implementation("sdpa")
+
+        key_before = [layer.attention.attention.key.weight.clone() for layer in model.encoder.layer]
+        val_before = [layer.attention.attention.value.weight.clone() for layer in model.encoder.layer]
+
+        model.set_attn_implementation("eager")
+
+        for layer_idx, layer in enumerate(model.encoder.layer):
+            assert torch.equal(key_before[layer_idx], layer.attention.attention.key.weight), (
+                f"Layer {layer_idx}: key weight changed after implementation switch"
+            )
+            assert torch.equal(val_before[layer_idx], layer.attention.attention.value.weight), (
+                f"Layer {layer_idx}: value weight changed after implementation switch"
+            )
+
+    def test_output_dense_weight_preserved(self) -> None:
+        """The output projection (dense) weight must survive the implementation switch."""
+        model = _make_small_model()
+        model.set_attn_implementation("sdpa")
+
+        dense_before = [layer.attention.output.dense.weight.clone() for layer in model.encoder.layer]
+
+        model.set_attn_implementation("eager")
+
+        for layer_idx, layer in enumerate(model.encoder.layer):
+            assert torch.equal(dense_before[layer_idx], layer.attention.output.dense.weight), (
+                f"Layer {layer_idx}: output dense weight changed after implementation switch"
+            )
+
+    def test_config_updated_after_switch(self) -> None:
+        """config._attn_implementation must reflect the new implementation after switching."""
+        model = _make_small_model()
+        model.set_attn_implementation("eager")
+        assert model.config._attn_implementation == "eager"
+        model.set_attn_implementation("sdpa")
+        assert model.config._attn_implementation == "sdpa"
+
+    def test_attention_module_type_after_switch(self) -> None:
+        """After switching to eager, every layer must hold a non-SDPA attention class."""
+        model = _make_small_model()
+        model.set_attn_implementation("eager")
+        for layer in model.encoder.layer:
+            assert isinstance(layer.attention, Dinov2WithRegistersAttention)
+            assert not isinstance(layer.attention, Dinov2WithRegistersSdpaAttention)
+
+    def test_invalid_implementation_raises_value_error(self) -> None:
+        """An unknown implementation name must raise ValueError before touching any layer."""
+        model = _make_small_model()
+        with pytest.raises(ValueError, match="Unknown attn_implementation"):
+            model.set_attn_implementation("flash_attn")

tests/models/conftest.py

@@ -0,0 +1,24 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Shared fixtures for the models test suite."""
+
+import pytest
+import torch
+
+
+@pytest.fixture(autouse=True)
+def reset_torch_safe_globals():
+    """Reset torch serialization safe globals after each test.
+
+    Prevents cross-test state contamination caused by ``_safe_torch_load``'s Attempt 2 path, which calls
+    ``torch.serialization.add_safe_globals``. Without this reset, globals registered by one test bleed into subsequent
+    tests and can mask trust-gate failures.
+    """
+    yield
+    try:
+        torch.serialization.clear_safe_globals()
+    except AttributeError:
+        pass  # torch <2.4 does not have clear_safe_globals

tests/models/heads/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/models/heads/test_detection_head.py

@@ -0,0 +1,274 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests for _resize_linear(), LWDETR.reinitialize_detection_head(), and _aggregate_keypoint_class_logits().
+
+These tests guard against the out_features staleness bug where in-place .data mutation did not update
+nn.Linear.out_features, causing ONNX export to emit stale (pre-fine-tuning) class counts.
+
+Also covers the spurious "Keypoint class-logit boost has N classes but detection head has M" warning that fired when
+num_keypoints_per_class exactly covered all foreground classes (correct configuration) but the comparison was against
+class_embed.out_features which includes the background slot (+1).
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from torch import nn
+
+from rfdetr.models.lwdetr import LWDETR, _resize_linear
+
+
+def _make_minimal_lwdetr(num_classes: int = 91, two_stage: bool = False) -> LWDETR:
+    """Construct the smallest viable LWDETR without loading pretrained weights.
+
+    Uses a MagicMock backbone and transformer with hidden_dim=4 so the model can be constructed in milliseconds without
+    any network I/O.
+
+    Args:
+        num_classes: Initial number of output classes passed to LWDETR.
+        two_stage: Whether to enable two-stage mode (creates enc_out_class_embed).
+
+    Returns:
+        An LWDETR instance with hidden_dim=4, num_queries=2, group_detr=1.
+
+    Examples:
+        >>> model = _make_minimal_lwdetr(num_classes=91)
+        >>> isinstance(model, LWDETR)
+        True
+    """
+    hidden_dim = 4
+    backbone = MagicMock()
+    transformer = MagicMock()
+    transformer.d_model = hidden_dim
+    transformer.decoder = MagicMock()
+    transformer.decoder.bbox_embed = None
+    return LWDETR(
+        backbone=backbone,
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=num_classes,
+        num_queries=2,
+        group_detr=1,
+        two_stage=two_stage,
+    )
+
+
+def _make_keypoint_lwdetr(num_classes: int, num_keypoints_per_class: list[int]) -> LWDETR:
+    """Construct a minimal keypoint-capable LWDETR with detection head resized to num_classes+1.
+
+    Mirrors what happens after loading a pretrained checkpoint and fine-tuning to num_classes
+    foreground categories: reinitialize_detection_head is called with num_classes+1 (includes
+    background), so class_embed.out_features == num_classes+1 in the returned model.
+
+    Args:
+        num_classes: Number of foreground detection classes.
+        num_keypoints_per_class: Keypoint count per foreground class.
+
+    Returns:
+        An LWDETR with use_grouppose_keypoints=True and class_embed.out_features==num_classes+1.
+
+    Examples:
+        >>> model = _make_keypoint_lwdetr(num_classes=2, num_keypoints_per_class=[17, 4])
+        >>> model.class_embed.out_features
+        3
+    """
+    hidden_dim = 4
+    backbone = MagicMock()
+    transformer = MagicMock()
+    transformer.d_model = hidden_dim
+    transformer.decoder = MagicMock()
+    transformer.decoder.bbox_embed = None
+    transformer.decoder.num_keypoints_per_class = num_keypoints_per_class
+    transformer.decoder.keypoint_class_mask = torch.zeros(1, 1, dtype=torch.bool)
+    transformer.num_keypoints_per_class = num_keypoints_per_class
+    model = LWDETR(
+        backbone=backbone,
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=num_classes,
+        num_queries=2,
+        group_detr=1,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=num_keypoints_per_class,
+    )
+    # Simulate post-checkpoint-load state: detection head includes background slot.
+    model.reinitialize_detection_head(num_classes + 1)
+    return model
+
+
+def _keypoint_tensor(num_keypoints_per_class: list[int], batch: int = 1, seq: int = 1) -> torch.Tensor:
+    """Build a zero keypoint prediction tensor with the shape expected by _aggregate_keypoint_class_logits.
+
+    The second-to-last dimension must equal num_kp_classes * max_kp (padded layout).
+
+    Args:
+        num_keypoints_per_class: Keypoint schema for the model.
+        batch: Batch size dimension.
+        seq: Sequence (query) dimension.
+
+    Returns:
+        Zero tensor of shape (batch, seq, num_kp_classes * max_kp, 8).
+
+    Examples:
+        >>> t = _keypoint_tensor([17, 4])
+        >>> t.shape
+        torch.Size([1, 1, 34, 8])
+    """
+    num_kp_classes = len(num_keypoints_per_class)
+    max_kp = max(num_keypoints_per_class) if any(num_keypoints_per_class) else 1
+    total_padded = num_kp_classes * max_kp
+    return torch.zeros(batch, seq, total_padded, 8)
+
+
+class TestResizeLinear:
+    """Unit tests for _resize_linear() — verifies out_features, weight shape, and bias shape."""
+
+    def test_shrink_out_features(self) -> None:
+        """Shrink: out_features equals the requested smaller class count."""
+        result = _resize_linear(nn.Linear(256, 91), 8)
+        assert result.out_features == 8, f"Expected out_features=8, got {result.out_features}"
+        assert result.weight.shape == (8, 256), f"Expected weight (8, 256), got {result.weight.shape}"
+        assert result.bias is not None
+        assert result.bias.shape == (8,), f"Expected bias (8,), got {result.bias.shape}"
+
+    def test_expand_out_features(self) -> None:
+        """Expand: out_features equals the requested larger class count via tiling."""
+        result = _resize_linear(nn.Linear(256, 10), 25)
+        assert result.out_features == 25, f"Expected out_features=25, got {result.out_features}"
+        assert result.weight.shape == (25, 256), f"Expected weight (25, 256), got {result.weight.shape}"
+        assert result.bias is not None
+        assert result.bias.shape == (25,), f"Expected bias (25,), got {result.bias.shape}"
+
+    def test_same_size_preserves_values(self) -> None:
+        """Same size: shapes and weight/bias values are preserved exactly."""
+        linear = nn.Linear(256, 91)
+        result = _resize_linear(linear, 91)
+        assert result.out_features == 91
+        assert result.weight.shape == (91, 256)
+        assert result.bias is not None
+        assert result.bias.shape == (91,)
+        assert torch.allclose(result.weight.data, linear.weight.data)
+        assert torch.allclose(result.bias.data, linear.bias.data)
+
+    def test_no_bias_returns_no_bias(self) -> None:
+        """Bias=False input: returned module has bias=None and out_features is correct."""
+        linear = nn.Linear(256, 91, bias=False)
+        result = _resize_linear(linear, 8)
+        assert result.out_features == 8, f"Expected out_features=8, got {result.out_features}"
+        assert result.bias is None, "Expected bias=None for bias=False input"
+
+
+class TestReinitializeDetectionHead:
+    """Integration tests for LWDETR.reinitialize_detection_head().
+
+    Uses a minimal LWDETR (hidden_dim=4, no real backbone) to verify that out_features is updated on the replaced
+    nn.Linear modules — the core invariant required for correct ONNX export.
+    """
+
+    def test_updates_class_embed_out_features(self) -> None:
+        """class_embed.out_features must reflect num_classes after reinitialize.
+
+        The `num_outputs_including_background` argument represents the total number of classifier outputs (foreground
+        classes plus background).
+        """
+        num_outputs_including_background = 8
+        model = _make_minimal_lwdetr(num_classes=91)
+        model.reinitialize_detection_head(num_outputs_including_background)
+        assert model.class_embed.out_features == num_outputs_including_background, (
+            f"Expected class_embed.out_features={num_outputs_including_background}, "
+            f"got {model.class_embed.out_features}"
+        )
+        assert model.class_embed.weight.shape == (num_outputs_including_background, 4), (
+            f"Expected weight ({num_outputs_including_background}, 4), got {model.class_embed.weight.shape}"
+        )
+
+    def test_two_stage_updates_enc_out_class_embed(self) -> None:
+        """enc_out_class_embed entries must also have updated out_features in two-stage mode.
+
+        The `num_outputs_including_background` argument represents the total number of classifier outputs (foreground
+        classes plus background).
+        """
+        num_outputs_including_background = 8
+        model = _make_minimal_lwdetr(num_classes=91, two_stage=True)
+        model.reinitialize_detection_head(num_outputs_including_background)
+        enc_embeds = model.transformer.enc_out_class_embed
+        assert len(enc_embeds) > 0, "enc_out_class_embed should be non-empty in two-stage mode"
+        for i, embed in enumerate(enc_embeds):
+            assert embed.out_features == num_outputs_including_background, (
+                f"enc_out_class_embed[{i}].out_features={embed.out_features}, "
+                f"expected {num_outputs_including_background}"
+            )
+            assert embed.weight.shape == (num_outputs_including_background, 4), (
+                f"enc_out_class_embed[{i}].weight.shape={embed.weight.shape}, "
+                f"expected ({num_outputs_including_background}, 4)"
+            )
+
+
+class TestAggregateKeypointClassLogits:
+    """Regression tests for LWDETR._aggregate_keypoint_class_logits().
+
+    Guards against a spurious warning that fired when num_keypoints_per_class exactly covered all
+    foreground classes: class_embed.out_features includes background (+1), so len(schema)==num_classes
+    always satisfied schema_len < detection_num_classes, triggering the warning incorrectly.
+
+    Uses _kp_zero_pad_warned as a proxy for whether the warning fired — the rf-detr logger uses
+    propagate=False which prevents standard caplog capture.
+    """
+
+    @pytest.mark.parametrize(
+        "num_classes,num_keypoints_per_class",
+        [
+            pytest.param(1, [17], id="coco-person-1class"),
+            pytest.param(2, [17, 4], id="basketball-2class"),
+            pytest.param(3, [17, 4, 0], id="3class-schema-covers-all"),
+        ],
+    )
+    def test_no_warning_when_schema_covers_all_foreground_classes(
+        self,
+        num_classes: int,
+        num_keypoints_per_class: list[int],
+    ) -> None:
+        """No warning when num_keypoints_per_class covers exactly all foreground detection classes.
+
+        Regression: the comparison used class_embed.out_features (num_classes+1) instead of
+        num_classes, so a fully correct schema always triggered the spurious mismatch warning.
+        """
+        model = _make_keypoint_lwdetr(num_classes=num_classes, num_keypoints_per_class=num_keypoints_per_class)
+        fake_kp = _keypoint_tensor(num_keypoints_per_class)
+
+        model._aggregate_keypoint_class_logits(fake_kp)
+
+        assert not model._kp_zero_pad_warned, (
+            f"Spurious warning fired for schema={num_keypoints_per_class} with num_classes={num_classes}"
+        )
+
+    def test_warning_fires_when_schema_shorter_than_foreground_classes(self) -> None:
+        """Warning fires when schema covers fewer classes than foreground detection classes.
+
+        Scenario: 3 foreground classes but schema only covers 1 (e.g. only person has keypoints).
+        The two uncovered foreground classes receive zero boost — a real mismatch worth warning about.
+        """
+        model = _make_keypoint_lwdetr(num_classes=3, num_keypoints_per_class=[17])
+        fake_kp = _keypoint_tensor([17])
+
+        model._aggregate_keypoint_class_logits(fake_kp)
+
+        assert model._kp_zero_pad_warned, "Expected warning flag set for schema shorter than foreground class count"
+
+    def test_output_shape_matches_detection_head(self) -> None:
+        """Output shape is (batch, seq, detection_num_classes) regardless of schema length."""
+        num_classes = 2
+        schema = [17, 4]
+        model = _make_keypoint_lwdetr(num_classes=num_classes, num_keypoints_per_class=schema)
+        batch, seq = 2, 10
+        fake_kp = _keypoint_tensor(schema, batch=batch, seq=seq)
+
+        out = model._aggregate_keypoint_class_logits(fake_kp)
+
+        assert out.shape == (batch, seq, num_classes + 1), (
+            f"Expected shape {(batch, seq, num_classes + 1)}, got {out.shape}"
+        )

tests/models/heads/test_keypoint_heads.py

@@ -0,0 +1,278 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+import pytest
+import torch
+
+from rfdetr.models.heads import ConditionalQueryInitializer
+from rfdetr.models.heads.keypoints import (
+    compute_keypoint_matching_cost,
+    compute_l1_keypoint_loss,
+)
+
+
+def test_conditional_query_initializer_shape() -> None:
+    """Initializer output should have expected batch/query/out dimensions."""
+    initializer = ConditionalQueryInitializer(dim=32, num_queries=11, out_dim=16)
+    query_features = torch.randn(3, 32)
+    queries = initializer(query_features)
+
+    assert queries.shape == (3, 11, 16)
+
+
+def test_conditional_query_initializer_zero_adaln_identity() -> None:
+    """A zeroed AdaLN gate should make initializer return the unmodified learned queries."""
+    initializer = ConditionalQueryInitializer(dim=16, num_queries=5, out_dim=16)
+    query_features = torch.randn(4, 16)
+    output = initializer(query_features)
+    expected = initializer.queries.unsqueeze(0).expand_as(output)
+
+    assert torch.equal(output, expected)
+
+
+def test_compute_l1_keypoint_loss_smoke() -> None:
+    """Loss helper should emit four finite vectors with matching target batch shape."""
+    pred_keypoints = torch.randn(3, 17, 7)
+    target_keypoints = torch.rand(3, 17, 3)
+    target_keypoints[:, :, 2] = 2.0
+    target_classes = torch.tensor([0, 0, 0], dtype=torch.int64)
+    target_areas = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
+    losses = compute_l1_keypoint_loss(
+        all_pred_keypoints=pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=target_classes,
+        target_areas=target_areas,
+        num_keypoints_per_class=[17],
+    )
+
+    assert len(losses) == 4
+    for loss in losses:
+        assert loss.shape == (3,)
+        assert torch.isfinite(loss).all()
+
+
+def test_compute_l1_keypoint_loss_skips_visible_zero_area_nll_residuals() -> None:
+    """Visible keypoints on zero-area targets should not produce non-finite Gaussian NLL."""
+    pred_keypoints = torch.zeros(1, 17, 7)
+    target_keypoints = torch.rand(1, 17, 3)
+    target_keypoints[:, :, 2] = 2.0
+    losses = compute_l1_keypoint_loss(
+        all_pred_keypoints=pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=torch.tensor([0], dtype=torch.int64),
+        target_areas=torch.tensor([0.0], dtype=torch.float32),
+        num_keypoints_per_class=[17],
+    )
+
+    for loss in losses:
+        assert torch.isfinite(loss).all()
+
+
+def test_compute_l1_keypoint_loss_uses_raw_rflow_gaussian_nll() -> None:
+    """Perfect keypoints should use raw r-flow NLL without a floor shift."""
+    pred_keypoints = torch.zeros(1, 1, 7)
+    pred_keypoints[:, :, 4] = 0.3
+    pred_keypoints[:, :, 6] = -0.2
+    target_keypoints = torch.tensor([[[0.0, 0.0, 2.0]]], dtype=torch.float32)
+
+    _, _, _, nll = compute_l1_keypoint_loss(
+        all_pred_keypoints=pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=torch.tensor([0], dtype=torch.int64),
+        target_areas=torch.tensor([1.0], dtype=torch.float32),
+        num_keypoints_per_class=[1],
+    )
+
+    torch.testing.assert_close(nll, torch.tensor([-0.1]), rtol=1e-4, atol=1e-6)
+
+
+def test_compute_l1_keypoint_loss_does_not_clamp_log_cholesky_nll() -> None:
+    """Large precision log-diagonals should remain raw to match r-flow."""
+    pred_keypoints = torch.zeros(1, 1, 7)
+    pred_keypoints[:, :, 4] = 25.0
+    target_keypoints = torch.tensor([[[0.0, 0.0, 2.0]]], dtype=torch.float32)
+
+    _, _, _, nll = compute_l1_keypoint_loss(
+        all_pred_keypoints=pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=torch.tensor([0], dtype=torch.int64),
+        target_areas=torch.tensor([1.0], dtype=torch.float32),
+        num_keypoints_per_class=[1],
+    )
+
+    torch.testing.assert_close(nll, torch.tensor([-25.0]), rtol=1e-4, atol=1e-6)
+
+
+def test_compute_l1_keypoint_loss_raw_nll_gradients_match_reference_formula() -> None:
+    """The implemented NLL gradients should match the raw r-flow Gaussian formula."""
+    pred_keypoints = torch.tensor([[[0.2, -0.1, 0.0, 0.0, 0.3, 0.1, -0.2]]], requires_grad=True)
+    target_keypoints = torch.tensor([[[0.0, 0.0, 2.0]]], dtype=torch.float32)
+    target_areas = torch.tensor([1.0], dtype=torch.float32)
+    _, _, _, nll = compute_l1_keypoint_loss(
+        all_pred_keypoints=pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=torch.tensor([0], dtype=torch.int64),
+        target_areas=target_areas,
+        num_keypoints_per_class=[1],
+    )
+    nll.sum().backward()
+    grad = pred_keypoints.grad.detach().clone()
+
+    raw_pred_keypoints = pred_keypoints.detach().clone().requires_grad_(True)
+    dx = raw_pred_keypoints[:, :, 0] - target_keypoints[:, :, 0]
+    dy = raw_pred_keypoints[:, :, 1] - target_keypoints[:, :, 1]
+    log_l11 = raw_pred_keypoints[:, :, 4]
+    l21 = raw_pred_keypoints[:, :, 5]
+    log_l22 = raw_pred_keypoints[:, :, 6]
+    u0 = log_l11.exp() * dx + l21 * dy
+    u1 = log_l22.exp() * dy
+    raw_nll = 0.5 * (u0 * u0 + u1 * u1) / target_areas.unsqueeze(1) - (log_l11 + log_l22)
+    raw_nll.sum().backward()
+
+    torch.testing.assert_close(nll.detach(), raw_nll.detach().reshape(-1), rtol=1e-4, atol=1e-6)
+    torch.testing.assert_close(grad, raw_pred_keypoints.grad, rtol=1e-4, atol=1e-6)
+
+
+def test_compute_l1_keypoint_loss_rejects_missing_schema() -> None:
+    """Missing keypoint schema should fail before producing zero supervision."""
+    pred_keypoints = torch.randn(1, 17, 7)
+    target_keypoints = torch.rand(1, 17, 3)
+
+    with pytest.raises(ValueError, match="num_keypoints_per_class must be non-empty"):
+        compute_l1_keypoint_loss(
+            all_pred_keypoints=pred_keypoints,
+            target_keypoints=target_keypoints,
+            target_classes=torch.tensor([0], dtype=torch.int64),
+            target_areas=torch.tensor([1.0], dtype=torch.float32),
+            num_keypoints_per_class=[],
+        )
+
+
+def test_compute_keypoint_matching_cost_smoke() -> None:
+    """Matching-cost helper should return a four-term cost tensor for each target."""
+    all_pred_keypoints = torch.randn(2, 4, 17, 7)
+    target_keypoints = torch.rand(2, 17, 3)
+    target_keypoints[:, :, 2] = 2.0
+    target_classes = torch.tensor([0, 0], dtype=torch.int64)
+    target_areas = torch.tensor([1.0, 2.0], dtype=torch.float32)
+    cost_l1, cost_findable, cost_visible, cost_nll = compute_keypoint_matching_cost(
+        all_pred_keypoints=all_pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=target_classes,
+        target_areas=target_areas,
+        num_keypoints_per_class=[17],
+    )
+
+    assert cost_l1.shape == (2, 4, 2)
+    assert cost_findable.shape == (2, 4, 2)
+    assert cost_visible.shape == (2, 4, 2)
+    assert cost_nll.shape == (2, 4, 2)
+    assert torch.isfinite(cost_l1).all()
+    assert torch.isfinite(cost_findable).all()
+    assert torch.isfinite(cost_visible).all()
+    assert torch.isfinite(cost_nll).all()
+
+
+def test_compute_keypoint_matching_cost_skips_zero_area_nll_residuals() -> None:
+    """Zero-area targets should not produce non-finite keypoint matching costs."""
+    all_pred_keypoints = torch.zeros(1, 2, 17, 7)
+    target_keypoints = torch.rand(1, 17, 3)
+    target_keypoints[:, :, 2] = 2.0
+    costs = compute_keypoint_matching_cost(
+        all_pred_keypoints=all_pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=torch.tensor([0], dtype=torch.int64),
+        target_areas=torch.tensor([0.0], dtype=torch.float32),
+        num_keypoints_per_class=[17],
+    )
+
+    for cost in costs:
+        assert torch.isfinite(cost).all()
+
+
+def test_compute_keypoint_matching_cost_does_not_clamp_log_cholesky_nll() -> None:
+    """Matching NLL should use raw precision log-diagonals to match r-flow."""
+    all_pred_keypoints = torch.zeros(1, 1, 1, 7)
+    all_pred_keypoints[:, :, :, 4] = 25.0
+    target_keypoints = torch.tensor([[[0.0, 0.0, 2.0]]], dtype=torch.float32)
+
+    _, _, _, cost_nll = compute_keypoint_matching_cost(
+        all_pred_keypoints=all_pred_keypoints,
+        target_keypoints=target_keypoints,
+        target_classes=torch.tensor([0], dtype=torch.int64),
+        target_areas=torch.tensor([1.0], dtype=torch.float32),
+        num_keypoints_per_class=[1],
+    )
+
+    torch.testing.assert_close(cost_nll, torch.tensor([[[-25.0]]]), rtol=1e-4, atol=1e-6)
+
+
+def test_compute_keypoint_matching_cost_rejects_missing_schema() -> None:
+    """Missing keypoint schema should fail before matcher costs become keypoint no-ops."""
+    all_pred_keypoints = torch.randn(1, 2, 17, 7)
+    target_keypoints = torch.rand(1, 17, 3)
+
+    with pytest.raises(ValueError, match="num_keypoints_per_class must be non-empty"):
+        compute_keypoint_matching_cost(
+            all_pred_keypoints=all_pred_keypoints,
+            target_keypoints=target_keypoints,
+            target_classes=torch.tensor([0], dtype=torch.int64),
+            target_areas=torch.tensor([1.0], dtype=torch.float32),
+            num_keypoints_per_class=[],
+        )
+
+
+class TestComputeKeypointMatchingCostSmoke:
+    """Group: compute_keypoint_matching_cost — shape and boundary checks."""
+
+    def test_n_targets_zero_returns_four_empty_cost_tensors(self) -> None:
+        """Empty target set should return four finite (B, Q, 0) cost tensors immediately."""
+        b, q = 2, 4
+        all_pred_keypoints = torch.randn(b, q, 17, 7)
+
+        cost_l1, cost_findable, cost_visible, cost_nll = compute_keypoint_matching_cost(
+            all_pred_keypoints=all_pred_keypoints,
+            target_keypoints=torch.empty(0, 17, 3),
+            target_classes=torch.empty(0, dtype=torch.int64),
+            target_areas=torch.empty(0),
+            num_keypoints_per_class=[17],
+        )
+
+        for cost, name in (
+            (cost_l1, "cost_l1"),
+            (cost_findable, "cost_findable"),
+            (cost_visible, "cost_visible"),
+            (cost_nll, "cost_nll"),
+        ):
+            assert cost.shape == (b, q, 0), f"{name}: expected shape ({b}, {q}, 0), got {cost.shape}"
+            assert torch.isfinite(cost).all(), f"{name}: expected all-finite tensor, got non-finite values"
+
+
+class TestComputeL1KeypointLossOobClass:
+    """Group: compute_l1_keypoint_loss — out-of-range class index handling."""
+
+    def test_class_index_out_of_range_returns_zero_losses_without_raising(self) -> None:
+        """Out-of-range class index should emit a warning and return zeros, not raise."""
+        pred_keypoints = torch.randn(1, 17, 7)
+        target_keypoints = torch.rand(1, 17, 3)
+        target_keypoints[:, :, 2] = 2.0
+        # class index 2 is out of range for num_keypoints_per_class=[17] (only class 0 defined)
+        result = compute_l1_keypoint_loss(
+            all_pred_keypoints=pred_keypoints,
+            target_keypoints=target_keypoints,
+            target_classes=torch.tensor([2], dtype=torch.int64),
+            target_areas=torch.tensor([1.0], dtype=torch.float32),
+            num_keypoints_per_class=[17],
+        )
+
+        assert len(result) == 4, f"Expected 4-tuple, got {len(result)} elements"
+        for i, loss in enumerate(result):
+            assert loss.shape == (1,), f"Loss[{i}]: expected shape (1,), got {loss.shape}"
+            torch.testing.assert_close(
+                loss,
+                torch.zeros(1),
+                msg=f"Loss[{i}]: expected all zeros for out-of-range class, got {loss}",
+            )

tests/models/heads/test_segmentation_head.py

@@ -0,0 +1,263 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for DepthwiseConvBlock and _DepthwiseConvWithoutCuDNN (segmentation head)."""
+
+from contextlib import contextmanager
+
+import pytest
+import torch
+
+from rfdetr.models.heads.segmentation import DepthwiseConvBlock
+
+
+@pytest.fixture(autouse=True)
+def _reset_random_seeds() -> None:
+    """Reset random seeds before each test for reproducibility."""
+    torch.manual_seed(42)
+
+
+@pytest.mark.parametrize(
+    "device",
+    [
+        pytest.param("cpu", id="cpu"),
+        pytest.param(
+            "cuda",
+            id="gpu",
+            marks=[
+                pytest.mark.gpu,
+                pytest.mark.skipif(
+                    not torch.cuda.is_available(),
+                    reason="CUDA is not available",
+                ),
+            ],
+        ),
+    ],
+)
+def test_depthwise_conv_block_forward(device: str) -> None:
+    """DepthwiseConvBlock forward pass produces correct output shape without error."""
+    block = DepthwiseConvBlock(dim=8).to(device)
+    x = torch.randn(1, 8, 4, 4, device=device)
+    y = block(x)
+    assert y.shape == x.shape
+
+
+def test_depthwise_conv_forward_disables_cudnn(monkeypatch) -> None:
+    """Depthwise conv should execute with cuDNN disabled during forward."""
+    block = DepthwiseConvBlock(dim=8)
+    enabled_calls: list[bool] = []
+    original_flags = torch.backends.cudnn.flags
+
+    @contextmanager
+    def _tracking_flags(*, enabled: bool):
+        enabled_calls.append(enabled)
+        with original_flags(enabled=enabled):
+            yield
+
+    monkeypatch.setattr(torch.backends.cudnn, "flags", _tracking_flags)
+
+    x = torch.randn(1, 8, 4, 4)
+    y = block(x)
+    assert y.shape == x.shape
+    assert enabled_calls, "torch.backends.cudnn.flags was never called"
+    assert all(not e for e in enabled_calls)
+
+
+def test_depthwise_conv_backward_disables_cudnn(monkeypatch) -> None:
+    """Backward pass must also run with cuDNN disabled (issue #731).
+
+    The previous fix (PR #728) only wrapped the forward pass in a context manager.  The backward kernels ran with cuDNN
+    re-enabled, causing RuntimeError on T4/P100 GPUs.
+    """
+    block = DepthwiseConvBlock(dim=8)
+    enabled_calls: list[bool] = []
+    original_flags = torch.backends.cudnn.flags
+
+    @contextmanager
+    def _tracking_flags(*, enabled: bool):
+        enabled_calls.append(enabled)
+        with original_flags(enabled=enabled):
+            yield
+
+    monkeypatch.setattr(torch.backends.cudnn, "flags", _tracking_flags)
+
+    x = torch.randn(1, 8, 4, 4, requires_grad=True)
+    y = block(x)
+    y.sum().backward()
+
+    assert x.grad is not None
+    assert x.grad.shape == x.shape
+    # cuDNN must be disabled for both forward and backward
+    assert len(enabled_calls) >= 2
+    assert all(not e for e in enabled_calls)
+
+
+@pytest.mark.parametrize(
+    "device",
+    [
+        pytest.param("cpu", id="cpu"),
+        pytest.param(
+            "cuda",
+            id="gpu",
+            marks=[
+                pytest.mark.gpu,
+                pytest.mark.skipif(
+                    not torch.cuda.is_available(),
+                    reason="CUDA is not available",
+                ),
+            ],
+        ),
+    ],
+)
+def test_depthwise_conv_backward_produces_correct_gradients(device: str) -> None:
+    """Backward pass through DepthwiseConvBlock produces valid gradients."""
+    block = DepthwiseConvBlock(dim=8).to(device)
+    x = torch.randn(1, 8, 4, 4, device=device, requires_grad=True)
+    y = block(x)
+    y.sum().backward()
+    assert x.grad is not None
+    assert x.grad.shape == x.shape
+    assert torch.isfinite(x.grad).all()
+
+
+def test_depthwise_conv_gradients_match_reference() -> None:
+    """Custom autograd Function gradients match nn.Conv2d gradients.
+
+    Verifies that _DepthwiseConvWithoutCuDNN produces the same gradients as a standard nn.Conv2d forward+backward (run
+    with cuDNN disabled globally).
+    """
+    torch.manual_seed(42)
+    dim = 8
+    block = DepthwiseConvBlock(dim=dim)
+
+    # Reference: run standard nn.Conv2d with cuDNN globally disabled
+    x_ref = torch.randn(1, dim, 4, 4, requires_grad=True)
+    with torch.backends.cudnn.flags(enabled=False):
+        y_ref = block.dwconv(x_ref)
+    y_ref.sum().backward()
+
+    x_ref_grad = x_ref.grad.clone()
+    weight_ref_grad = block.dwconv.weight.grad.clone()
+    bias_ref_grad = block.dwconv.bias.grad.clone()
+
+    # Our implementation via _depthwise_conv.  zero_grad() so that the second
+    # backward does not accumulate into weight.grad from the first run.
+    block.zero_grad()
+    x_test = x_ref.detach().clone().requires_grad_(True)
+    y_test = block._depthwise_conv(x_test)
+    y_test.sum().backward()
+
+    assert torch.allclose(y_ref, y_test, atol=1e-6)
+    assert torch.allclose(x_ref_grad, x_test.grad, atol=1e-6)
+    assert torch.allclose(weight_ref_grad, block.dwconv.weight.grad, atol=1e-6)
+    assert torch.allclose(bias_ref_grad, block.dwconv.bias.grad, atol=1e-6)
+
+
+def test_depthwise_conv_backward_fp16_grad_output() -> None:
+    """Backward must not crash when grad_output is fp16 (AMP 16-mixed on T4/P100).
+
+    On T4/P100, trainer resolves amp=True to '16-mixed'.  In that mode the backward receives fp16 grad_output while the
+    saved weight stays fp32. Without explicit dtype casting, conv2d_input raises:
+        RuntimeError: expected scalar type Half but found Float
+    """
+    dim = 8
+    block = DepthwiseConvBlock(dim=dim)
+    x = torch.randn(1, dim, 4, 4, requires_grad=True)
+
+    # Simulate 16-mixed backward: forward in fp32, grad_output arrives as fp16
+    y = block._depthwise_conv(x)
+    grad_output = torch.ones_like(y, dtype=torch.float16)
+    y.backward(grad_output)
+
+    assert x.grad is not None
+    assert x.grad.dtype == torch.float32
+    assert torch.isfinite(x.grad).all()
+
+
+def test_depthwise_conv_backward_bf16_activation_keeps_grads_fp32() -> None:
+    """grad_input and weight.grad must be fp32 when saved activation x is bf16 (issue #959).
+
+    Under bf16-mixed AMP, the activation x entering _DepthwiseConvWithoutCuDNN is bf16 while weight stays fp32.  The old
+    code cast grad_input back to x.dtype (bf16), propagating bf16 gradients to fp32 backbone parameters so that
+    param.grad.dtype became bf16.  Fused AdamW then crashed with 'params, grads, exp_avgs, and exp_avg_sqs must have
+    same dtype, device, and layout' (see issue #959).  The fix keeps grad_input in weight.dtype (fp32).
+
+    This test drives the backward directly with a bf16 saved activation to reproduce the dtype that is present at
+    training time without requiring a GPU.
+    """
+    import types
+
+    from rfdetr.models.heads.segmentation import _DepthwiseConvWithoutCuDNN
+
+    dim = 8
+    weight = torch.randn(dim, 1, 3, 3, requires_grad=True)  # fp32 parameter (never cast by AMP)
+    x_bf16 = torch.randn(1, dim, 4, 4, dtype=torch.bfloat16)  # bf16 activation (cast by AMP forward)
+    grad_output = torch.ones(1, dim, 4, 4, dtype=torch.bfloat16)  # bf16 grad (from bf16 backward)
+
+    # Build a minimal context mirroring what ctx would contain after the AMP forward pass.
+    ctx = types.SimpleNamespace()
+    ctx.saved_tensors = (x_bf16, weight)
+    ctx.has_bias = False
+    ctx.stride = (1, 1)
+    ctx.padding = (1, 1)
+    ctx.dilation = (1, 1)
+    ctx.groups = dim
+    ctx.needs_input_grad = [True, True, False, False, False, False, False]
+
+    grad_input, grad_weight, *_ = _DepthwiseConvWithoutCuDNN.backward(ctx, grad_output)
+
+    assert grad_input is not None, "grad_input should not be None when needs_input_grad[0] is True"
+    assert grad_input.dtype == torch.float32, (
+        f"grad_input is {grad_input.dtype} — bf16 grad_input propagates to fp32 backbone params "
+        "and crashes fused AdamW (issue #959)"
+    )
+    assert grad_weight is not None, "grad_weight should not be None when needs_input_grad[1] is True"
+    assert grad_weight.dtype == torch.float32, (
+        f"grad_weight is {grad_weight.dtype} — weight grad must stay fp32 to match param dtype"
+    )
+
+
+def test_depthwise_conv_no_cudnn_bias_none() -> None:
+    """_DepthwiseConvWithoutCuDNN forward and backward work correctly with bias=None.
+
+    Exercises the ctx.has_bias=False branch in forward and the grad_bias=None return in backward — never reached via
+    DepthwiseConvBlock (always has bias).
+    """
+    from rfdetr.models.heads.segmentation import _DepthwiseConvWithoutCuDNN
+
+    dim = 8
+    weight = torch.randn(dim, 1, 3, 3, requires_grad=True)
+    x = torch.randn(1, dim, 4, 4, requires_grad=True)
+    y = _DepthwiseConvWithoutCuDNN.apply(x, weight, None, (1, 1), (1, 1), (1, 1), dim)
+    y_ref = torch.nn.functional.conv2d(x.detach(), weight.detach(), None, stride=1, padding=1, dilation=1, groups=dim)
+    assert torch.allclose(y.detach(), y_ref, atol=1e-6)
+    y.sum().backward()
+    assert x.grad is not None
+    assert x.grad.shape == x.shape
+    assert torch.isfinite(x.grad).all()
+    assert weight.grad is not None
+    assert weight.grad.shape == weight.shape
+    assert torch.isfinite(weight.grad).all()
+
+
+@pytest.mark.parametrize(
+    "layer_scale_init_value", [pytest.param(0, id="no_gamma"), pytest.param(1e-6, id="with_gamma")]
+)
+def test_depthwise_conv_block_layer_scale(layer_scale_init_value: float) -> None:
+    """DepthwiseConvBlock with and without layer scaling produces valid output and gradients.
+
+    Exercises the gamma=None (layer_scale_init_value=0) and gamma!=None (layer_scale_init_value>0) branches in
+    DepthwiseConvBlock.forward().
+    """
+    block = DepthwiseConvBlock(dim=8, layer_scale_init_value=layer_scale_init_value)
+    x = torch.randn(1, 8, 4, 4, requires_grad=True)
+    y = block(x)
+    assert y.shape == x.shape
+    y.sum().backward()
+    assert x.grad is not None
+    assert torch.isfinite(x.grad).all()
+    if layer_scale_init_value > 0:
+        assert block.gamma is not None
+        assert block.gamma.grad is not None

tests/models/test_builder_characterization.py

@@ -0,0 +1,433 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Characterization tests for build_model() and build_criterion_and_postprocessors().
+
+These tests pin the current behavior of the legacy namespace-based builder functions. They serve as a safety net during
+the config-native builder refactoring: any change that alters these outputs is a regression.
+
+All tests in this file must pass against the CURRENT codebase.
+"""
+
+import pytest
+import torch
+
+from rfdetr._namespace import _namespace_from_configs
+from rfdetr.config import (
+    RFDETRBaseConfig,
+    RFDETRKeypointPreviewConfig,
+    RFDETRNanoConfig,
+    RFDETRSegNanoConfig,
+    SegmentationTrainConfig,
+    TrainConfig,
+)
+from rfdetr.models.criterion import SetCriterion
+from rfdetr.models.lwdetr import LWDETR, build_criterion_and_postprocessors, build_model
+from rfdetr.models.postprocess import PostProcess
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_ns(mc=None, tc=None):
+    """Build a namespace suitable for builder functions."""
+    mc = mc or RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+    tc = tc or TrainConfig(dataset_dir="/tmp")
+    return _namespace_from_configs(mc, tc)
+
+
+# ---------------------------------------------------------------------------
+# build_model characterization
+# ---------------------------------------------------------------------------
+
+
+class TestBuildModelCharacterization:
+    """Pin current build_model() behaviour for the standard code path."""
+
+    def test_returns_lwdetr_instance(self) -> None:
+        ns = _make_ns()
+        model = build_model(ns)
+        assert isinstance(model, LWDETR)
+
+    def test_num_classes_plus_one(self) -> None:
+        """build_model applies the +1 background class convention."""
+        mc = RFDETRBaseConfig(num_classes=5, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        model = build_model(ns)
+        assert model.class_embed.out_features == 6
+
+    def test_num_queries_forwarded(self) -> None:
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        model = build_model(ns)
+        assert model.num_queries == mc.num_queries
+
+    @pytest.mark.parametrize(
+        "config_class, expected_queries",
+        [
+            pytest.param(RFDETRBaseConfig, 300, id="base"),
+            pytest.param(RFDETRNanoConfig, 300, id="nano"),
+            pytest.param(RFDETRSegNanoConfig, 100, id="seg_nano"),
+        ],
+    )
+    def test_num_queries_per_config_variant(self, config_class, expected_queries) -> None:
+        mc = config_class(pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        model = build_model(ns)
+        assert model.num_queries == expected_queries
+
+    def test_segmentation_head_none_for_detection(self) -> None:
+        ns = _make_ns()
+        model = build_model(ns)
+        assert model.segmentation_head is None
+
+    def test_segmentation_head_present_for_seg_config(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        model = build_model(ns)
+        assert model.segmentation_head is not None
+
+    def test_aux_loss_enabled_by_default(self) -> None:
+        ns = _make_ns()
+        model = build_model(ns)
+        assert model.aux_loss is True
+
+    def test_group_detr_forwarded(self) -> None:
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        model = build_model(ns)
+        assert model.group_detr == mc.group_detr
+
+    def test_num_feature_levels_set_on_args(self) -> None:
+        """build_model mutates args.num_feature_levels = len(projector_scale)."""
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        build_model(ns)
+        assert ns.num_feature_levels == len(mc.projector_scale)
+
+    @pytest.mark.parametrize(
+        "config_class, expected_param_count_range",
+        [
+            pytest.param(RFDETRBaseConfig, (25_000_000, 40_000_000), id="base"),
+            pytest.param(RFDETRNanoConfig, (25_000_000, 40_000_000), id="nano"),
+        ],
+    )
+    def test_param_count_in_expected_range(self, config_class, expected_param_count_range) -> None:
+        """Sanity check that the model has a plausible number of parameters."""
+        mc = config_class(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        model = build_model(ns)
+        total = sum(p.numel() for p in model.parameters())
+        low, high = expected_param_count_range
+        assert low <= total <= high, f"Expected param count in [{low}, {high}], got {total}"
+
+    def test_encoder_only_returns_triple(self) -> None:
+        """When encoder_only=True, build_model returns (encoder, None, None)."""
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        ns.encoder_only = True
+        result = build_model(ns)
+        assert isinstance(result, tuple), f"Expected tuple, got {type(result)}"
+        assert len(result) == 3
+        encoder, second, third = result
+        assert second is None
+        assert third is None
+        assert encoder is not None
+
+    def test_backbone_only_returns_triple(self) -> None:
+        """When backbone_only=True, build_model returns (backbone, None, None)."""
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        ns.backbone_only = True
+        result = build_model(ns)
+        assert isinstance(result, tuple), f"Expected tuple, got {type(result)}"
+        assert len(result) == 3
+        backbone, second, third = result
+        assert second is None
+        assert third is None
+        assert backbone is not None
+
+
+# ---------------------------------------------------------------------------
+# build_criterion_and_postprocessors characterization
+# ---------------------------------------------------------------------------
+
+
+class TestBuildCriterionCharacterization:
+    """Pin current build_criterion_and_postprocessors() behaviour."""
+
+    def test_returns_criterion_and_postprocess(self) -> None:
+        ns = _make_ns()
+        criterion, postprocess = build_criterion_and_postprocessors(ns)
+        assert isinstance(criterion, SetCriterion)
+        assert isinstance(postprocess, PostProcess)
+
+    def test_detection_losses_list(self) -> None:
+        """Detection-only config has exactly ['labels', 'boxes', 'cardinality']."""
+        ns = _make_ns()
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.losses == ["labels", "boxes", "cardinality"]
+
+    def test_segmentation_losses_include_masks(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        ns = _make_ns(mc=mc, tc=tc)
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert "masks" in criterion.losses
+
+    def test_num_select_forwarded_to_postprocess(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        _, postprocess = build_criterion_and_postprocessors(ns)
+        assert postprocess.num_select == 100
+
+    def test_num_select_default_for_base(self) -> None:
+        ns = _make_ns()
+        _, postprocess = build_criterion_and_postprocessors(ns)
+        assert postprocess.num_select == 300
+
+    def test_weight_dict_contains_base_losses(self) -> None:
+        ns = _make_ns()
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert "loss_ce" in criterion.weight_dict
+        assert "loss_bbox" in criterion.weight_dict
+        assert "loss_giou" in criterion.weight_dict
+
+    def test_weight_dict_values_match_namespace(self) -> None:
+        ns = _make_ns()
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.weight_dict["loss_ce"] == ns.cls_loss_coef
+        assert criterion.weight_dict["loss_bbox"] == ns.bbox_loss_coef
+        assert criterion.weight_dict["loss_giou"] == ns.giou_loss_coef
+
+    def test_segmentation_weight_dict_contains_mask_losses(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        ns = _make_ns(mc=mc, tc=tc)
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert "loss_mask_ce" in criterion.weight_dict
+        assert "loss_mask_dice" in criterion.weight_dict
+
+    def test_aux_loss_expands_weight_dict(self) -> None:
+        """With aux_loss=True and 3 dec_layers, weight_dict has aux entries _0 and _1."""
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        assert ns.aux_loss is True
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        # dec_layers=3 -> 2 aux layers (0 and 1)
+        assert "loss_ce_0" in criterion.weight_dict
+        assert "loss_ce_1" in criterion.weight_dict
+
+    def test_two_stage_adds_enc_losses(self) -> None:
+        """With two_stage=True, weight_dict has '_enc' suffix entries."""
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        assert ns.two_stage is True
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert "loss_ce_enc" in criterion.weight_dict
+        assert "loss_bbox_enc" in criterion.weight_dict
+        assert "loss_giou_enc" in criterion.weight_dict
+
+    def test_criterion_num_classes_plus_one(self) -> None:
+        mc = RFDETRBaseConfig(num_classes=5, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.num_classes == 6
+
+    def test_focal_alpha_forwarded(self) -> None:
+        ns = _make_ns()
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.focal_alpha == pytest.approx(0.25)
+
+    def test_group_detr_forwarded_to_criterion(self) -> None:
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.group_detr == mc.group_detr
+
+    def test_segmentation_criterion_has_mask_point_sample_ratio(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        ns = _make_ns(mc=mc, tc=tc)
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.mask_point_sample_ratio == 16
+
+    def test_ia_bce_loss_forwarded(self) -> None:
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ns = _make_ns(mc=mc)
+        criterion, _ = build_criterion_and_postprocessors(ns)
+        assert criterion.ia_bce_loss == mc.ia_bce_loss
+
+
+# ---------------------------------------------------------------------------
+# _build_model_context characterization
+# ---------------------------------------------------------------------------
+
+
+class TestBuildModelContextCharacterization:
+    """Pin current _build_model_context() behaviour.
+
+    _build_model_context is the inference-path factory used by RFDETR.get_model(). It has zero test coverage today.
+    """
+
+    def test_returns_model_context(self) -> None:
+        from rfdetr.detr import ModelContext, _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert isinstance(ctx, ModelContext)
+
+    def test_model_is_lwdetr(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert isinstance(ctx.model, LWDETR)
+
+    def test_postprocess_is_postprocess(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert isinstance(ctx.postprocess, PostProcess)
+
+    def test_resolution_from_config(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.resolution == mc.resolution
+
+    def test_device_from_config(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.device == torch.device("cpu")
+
+    def test_torch_device_cpu_from_config(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device=torch.device("cpu"))
+        ctx = _build_model_context(mc)
+        assert ctx.device == torch.device("cpu")
+
+    def test_class_names_none_without_pretrain(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.class_names is None
+
+    def test_num_select_on_postprocess(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.postprocess.num_select == 100
+
+    def test_keypoint_preview_postprocess_has_keypoint_schema(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRKeypointPreviewConfig(pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.postprocess.num_keypoints_per_class == [17]
+
+    def test_args_namespace_attached(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert hasattr(ctx.args, "num_classes")
+        assert hasattr(ctx.args, "num_select")
+
+    def test_inference_model_initially_none(self) -> None:
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.inference_model is None
+
+    def test_args_dataset_dir_does_not_leak_cwd(self) -> None:
+        """The serialized namespace must not embed the caller's realpathed CWD as dataset_dir."""
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.args.dataset_dir is None
+
+    def test_args_output_dir_does_not_leak_cwd(self) -> None:
+        """The serialized namespace must keep a relative output_dir, not the caller's realpathed CWD."""
+        from rfdetr.detr import _build_model_context
+
+        mc = RFDETRBaseConfig(num_classes=80, pretrain_weights=None, device="cpu")
+        ctx = _build_model_context(mc)
+        assert ctx.args.output_dir == "output"
+
+
+# ---------------------------------------------------------------------------
+# RFDETRModelModule.__init__ characterization
+# ---------------------------------------------------------------------------
+
+
+class TestRFDETRModelModuleInitCharacterization:
+    """Pin RFDETRModelModule.__init__() structural outputs.
+
+    The existing test_module_model.py tests the init via mocked build_model and build_namespace. These tests exercise
+    the REAL init path (no mocks) to characterize what a freshly built module looks like.
+    """
+
+    def _make_module(self, mc=None, tc=None):
+        from rfdetr.training.module_model import RFDETRModelModule
+
+        mc = mc or RFDETRBaseConfig(num_classes=5, pretrain_weights=None, device="cpu")
+        tc = tc or TrainConfig(dataset_dir="/tmp")
+        return RFDETRModelModule(mc, tc)
+
+    def test_model_attribute_is_lwdetr(self) -> None:
+        module = self._make_module()
+        # model could be wrapped by torch.compile, so check the underlying type
+        underlying = getattr(module.model, "_orig_mod", module.model)
+        assert isinstance(underlying, LWDETR)
+
+    def test_criterion_is_set_criterion(self) -> None:
+        module = self._make_module()
+        assert isinstance(module.criterion, SetCriterion)
+
+    def test_postprocess_is_postprocess(self) -> None:
+        module = self._make_module()
+        assert isinstance(module.postprocess, PostProcess)
+
+    def test_strict_loading_false(self) -> None:
+        """strict_loading=False allows partial state-dict loading."""
+        module = self._make_module()
+        assert module.strict_loading is False
+
+    def test_configs_stored(self) -> None:
+        mc = RFDETRBaseConfig(num_classes=5, pretrain_weights=None, device="cpu")
+        tc = TrainConfig(dataset_dir="/tmp")
+        module = self._make_module(mc=mc, tc=tc)
+        assert module.model_config is mc
+        assert module.train_config is tc
+
+    def test_criterion_num_classes_matches_model(self) -> None:
+        """Criterion and model must agree on num_classes (both use +1 convention)."""
+        mc = RFDETRBaseConfig(num_classes=5, pretrain_weights=None, device="cpu")
+        module = self._make_module(mc=mc)
+        underlying = getattr(module.model, "_orig_mod", module.model)
+        assert module.criterion.num_classes == underlying.class_embed.out_features
+
+    def test_postprocess_num_select_matches_config(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        module = self._make_module(mc=mc, tc=tc)
+        assert module.postprocess.num_select == mc.num_select
+
+    def test_segmentation_criterion_with_seg_config(self) -> None:
+        mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cpu")
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        module = self._make_module(mc=mc, tc=tc)
+        assert "masks" in module.criterion.losses

tests/models/test_builders.py

@@ -0,0 +1,162 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Characterization tests for config-native builder functions.
+
+These tests validate build_model_from_config() and build_criterion_from_config() which accept Pydantic config objects
+directly instead of requiring a pre-built SimpleNamespace. If these functions cannot be imported, all tests skip via the
+module-level pytestmark.
+"""
+
+import pytest
+
+from rfdetr.config import (
+    RFDETRBaseConfig,
+    RFDETRSegNanoConfig,
+    SegmentationTrainConfig,
+    TrainConfig,
+)
+
+try:
+    from rfdetr.models import build_criterion_from_config, build_model_from_config
+
+    HAS_CONFIG_BUILDERS = True
+except ImportError:
+    HAS_CONFIG_BUILDERS = False
+
+pytestmark = pytest.mark.skipif(
+    not HAS_CONFIG_BUILDERS,
+    reason="config-native builder functions are not importable",
+)
+
+
+class TestBuildModelFromConfig:
+    """Tests for build_model_from_config(model_config, defaults=MODEL_DEFAULTS)."""
+
+    def test_returns_lwdetr_for_base_config(self) -> None:
+        """build_model_from_config with RFDETRBaseConfig returns an LWDETR instance."""
+        from rfdetr.models.lwdetr import LWDETR
+
+        mc = RFDETRBaseConfig(num_classes=80)
+        model = build_model_from_config(mc)
+        assert isinstance(model, LWDETR), f"Expected LWDETR instance, got {type(model).__name__}"
+
+    def test_num_classes_correct(self) -> None:
+        """num_classes=5 in config should produce class_embed with out_features=6.
+
+        build_model adds +1 to num_classes (background class convention).
+        """
+        mc = RFDETRBaseConfig(num_classes=5)
+        model = build_model_from_config(mc)
+        assert model.class_embed.out_features == 6, (
+            f"Expected class_embed.out_features=6 (num_classes+1), got {model.class_embed.out_features}"
+        )
+
+    def test_parity_with_build_model_via_namespace(self) -> None:
+        """Parameter count must match between config-native and namespace paths."""
+        from rfdetr._namespace import _namespace_from_configs
+        from rfdetr.models.lwdetr import build_model
+
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc = TrainConfig(dataset_dir="/tmp")
+
+        model_config_native = build_model_from_config(mc, tc)
+        ns = _namespace_from_configs(mc, tc)
+        model_namespace = build_model(ns)
+
+        params_native = sum(p.numel() for p in model_config_native.parameters())
+        params_namespace = sum(p.numel() for p in model_namespace.parameters())
+        assert params_native == params_namespace, (
+            f"Parameter count mismatch: config-native={params_native}, namespace={params_namespace}"
+        )
+
+    def test_segmentation_head_created_when_true(self) -> None:
+        """RFDETRSegNanoConfig has segmentation_head=True; model must have it."""
+        mc = RFDETRSegNanoConfig()
+        model = build_model_from_config(mc)
+        assert model.segmentation_head is not None, "Expected segmentation_head to be created for RFDETRSegNanoConfig"
+
+    def test_drop_path_uses_train_config_value(self) -> None:
+        """Non-default TrainConfig.drop_path must reach the model builder path."""
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc = TrainConfig(dataset_dir="/tmp", drop_path=0.2)
+        model = build_model_from_config(mc, tc)
+
+        layers = model._get_backbone_encoder_layers()
+        assert layers is not None
+        assert hasattr(layers[-1], "drop_path")
+        assert layers[-1].drop_path.drop_prob == pytest.approx(0.2)
+
+    def test_rejects_encoder_only_defaults(self) -> None:
+        """The config-native builder guarantees an LWDETR return value."""
+        from dataclasses import replace
+
+        from rfdetr.models import MODEL_DEFAULTS
+
+        mc = RFDETRBaseConfig(num_classes=80)
+
+        with pytest.raises(ValueError, match="encoder_only=False"):
+            build_model_from_config(mc, defaults=replace(MODEL_DEFAULTS, encoder_only=True))
+
+    def test_rejects_backbone_only_defaults(self) -> None:
+        """backbone_only=True in defaults must also raise ValueError."""
+        from dataclasses import replace
+
+        from rfdetr.models import MODEL_DEFAULTS
+
+        mc = RFDETRBaseConfig(num_classes=80)
+
+        with pytest.raises(ValueError, match="backbone_only=False"):
+            build_model_from_config(mc, defaults=replace(MODEL_DEFAULTS, backbone_only=True))
+
+    def test_none_train_config_uses_dummy(self) -> None:
+        """build_model_from_config with train_config=None must not raise."""
+        mc = RFDETRBaseConfig(num_classes=80)
+        model = build_model_from_config(mc, train_config=None)
+        assert model is not None, "Expected a model, got None"
+
+
+class TestBuildCriterionFromConfig:
+    """Tests for build_criterion_from_config(model_config, train_config, defaults)."""
+
+    def test_returns_tuple(self) -> None:
+        """build_criterion_from_config must return a 2-tuple (SetCriterion, PostProcess)."""
+        from rfdetr.models.criterion import SetCriterion
+        from rfdetr.models.postprocess import PostProcess
+
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc = TrainConfig(dataset_dir="/tmp")
+        result = build_criterion_from_config(mc, tc)
+        assert isinstance(result, tuple), f"Expected tuple, got {type(result).__name__}"
+        assert len(result) == 2, f"Expected 2-tuple, got {len(result)}-tuple"
+        criterion, postprocess = result
+        assert isinstance(criterion, SetCriterion), f"Expected SetCriterion, got {type(criterion).__name__}"
+        assert isinstance(postprocess, PostProcess), f"Expected PostProcess, got {type(postprocess).__name__}"
+
+    def test_num_select_postprocess(self) -> None:
+        """RFDETRSegNanoConfig has num_select=100; PostProcess must reflect it."""
+        mc = RFDETRSegNanoConfig()
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        _, postprocess = build_criterion_from_config(mc, tc)
+        assert postprocess.num_select == 100, f"Expected PostProcess.num_select=100, got {postprocess.num_select}"
+
+    def test_segmentation_losses_included(self) -> None:
+        """With segmentation config, 'masks' must be in criterion.losses."""
+        mc = RFDETRSegNanoConfig()
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        criterion, _ = build_criterion_from_config(mc, tc)
+        assert "masks" in criterion.losses, f"Expected 'masks' in criterion.losses, got {criterion.losses}"
+
+    def test_custom_defaults_focal_alpha_applied(self) -> None:
+        """Custom focal_alpha in ModelDefaults must reach SetCriterion."""
+        from dataclasses import replace
+
+        from rfdetr.models import MODEL_DEFAULTS
+
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc = TrainConfig(dataset_dir="/tmp")
+        custom_defaults = replace(MODEL_DEFAULTS, focal_alpha=0.5)
+        criterion, _ = build_criterion_from_config(mc, tc, defaults=custom_defaults)
+        assert criterion.focal_alpha == pytest.approx(0.5), f"Expected focal_alpha=0.5, got {criterion.focal_alpha}"

tests/models/test_config.py

@@ -0,0 +1,925 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+import os
+import warnings
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from pydantic import ValidationError
+
+from rfdetr.config import (
+    KeypointTrainConfig,
+    ModelConfig,
+    PretrainWeightsCompatibilityWarning,
+    RFDETRBaseConfig,
+    RFDETRLargeConfig,
+    RFDETRMediumConfig,
+    RFDETRNanoConfig,
+    RFDETRSeg2XLargeConfig,
+    RFDETRSegLargeConfig,
+    RFDETRSegMediumConfig,
+    RFDETRSegNanoConfig,
+    RFDETRSegSmallConfig,
+    RFDETRSegXLargeConfig,
+    RFDETRSmallConfig,
+    SegmentationTrainConfig,
+    TrainConfig,
+    _detect_device,
+)
+
+
+@pytest.fixture
+def sample_model_config() -> dict[str, object]:
+    return {
+        "encoder": "dinov2_windowed_small",
+        "out_feature_indexes": [1, 2, 3],
+        "dec_layers": 3,
+        "projector_scale": ["P3"],
+        "hidden_dim": 256,
+        "patch_size": 14,
+        "num_windows": 2,
+        "sa_nheads": 8,
+        "ca_nheads": 8,
+        "dec_n_points": 4,
+        "resolution": 384,
+        "positional_encoding_size": 256,
+    }
+
+
+class TestModelConfigValidation:
+    def test_rejects_unknown_fields(self, sample_model_config) -> None:
+        sample_model_config["unknown"] = "value"
+
+        with pytest.raises(ValidationError, match=r"Unknown parameter\(s\): 'unknown'"):
+            ModelConfig(**sample_model_config)
+
+    def test_rejects_unknown_attribute_assignment(self, sample_model_config) -> None:
+        config = ModelConfig(**sample_model_config)
+
+        with pytest.raises(ValueError, match=r"Unknown attribute: 'unknown'\."):
+            setattr(config, "unknown", "value")
+
+    def test_accepts_indexed_cuda_device_string(self, sample_model_config) -> None:
+        config = ModelConfig(**sample_model_config, device="cuda:1")
+        assert config.device == "cuda:1"
+
+    def test_accepts_torch_device(self, sample_model_config) -> None:
+        config = ModelConfig(**sample_model_config, device=torch.device("cuda:2"))
+        assert config.device == "cuda:2"
+
+    def test_rejects_non_string_non_torch_device_with_validation_error(self, sample_model_config) -> None:
+        with pytest.raises(ValidationError, match="device must be a string or torch\\.device\\."):
+            ModelConfig(**sample_model_config, device=123)
+
+    def test_rejects_invalid_device_string(self, sample_model_config) -> None:
+        with pytest.raises(ValidationError, match="Invalid device specifier: 'notadevice'\\."):
+            ModelConfig(**sample_model_config, device="notadevice")
+
+    @pytest.mark.parametrize(
+        "encoder",
+        [
+            pytest.param("dinov2_windowed_small", id="windowed_small"),
+            pytest.param("dinov2_windowed_base", id="windowed_base"),
+            pytest.param("dinov2_registers_windowed_small", id="registers_windowed_small"),
+        ],
+    )
+    def test_accepts_valid_encoder(self, sample_model_config, encoder: str) -> None:
+        """ModelConfig accepts every value in the EncoderName Literal."""
+        config = ModelConfig(**{**sample_model_config, "encoder": encoder})
+        assert config.encoder == encoder
+
+    def test_rejects_invalid_encoder(self, sample_model_config) -> None:
+        """ModelConfig raises ValidationError for encoder strings outside the Literal."""
+        with pytest.raises(ValidationError):
+            ModelConfig(**{**sample_model_config, "encoder": "dinov2_invalid_backbone"})
+
+    def test_rejects_negative_postprocess_trace_alpha(self, sample_model_config) -> None:
+        """ModelConfig rejects negative uncertainty score-fusion exponents."""
+        with pytest.raises(ValidationError):
+            ModelConfig(**sample_model_config, postprocess_trace_alpha=-0.1)
+
+    def test_postprocess_trace_alpha_defaults_to_keypoint_fusion_value(self, sample_model_config) -> None:
+        """ModelConfig defaults to the keypoint uncertainty score-fusion exponent."""
+        config = ModelConfig(**sample_model_config)
+
+        assert config.postprocess_trace_alpha == 0.2
+
+    def test_pretrain_weights_absolute_path_realpath_normalised(self, tmp_path) -> None:
+        """An absolute pathlib.Path for pretrain_weights is stored as the realpath-normalised string."""
+        weights_path = tmp_path / "weights.pth"
+
+        config = RFDETRBaseConfig(pretrain_weights=weights_path)
+
+        assert config.pretrain_weights == os.path.realpath(os.fspath(weights_path))
+
+    @pytest.mark.parametrize(
+        "field",
+        [
+            pytest.param("dataset_dir", id="dataset_dir"),
+            pytest.param("output_dir", id="output_dir"),
+        ],
+    )
+    def test_train_dir_fields_accept_path(self, tmp_path, field: str) -> None:
+        """TrainConfig dataset/output dir fields accept pathlib.Path and store the realpath-normalised string."""
+        path = tmp_path / "artifact"
+        kwargs = {"dataset_dir": str(tmp_path)}
+        kwargs[field] = path
+
+        config = TrainConfig(**kwargs)
+
+        assert getattr(config, field) == os.path.realpath(os.fspath(path))
+
+    def test_accepts_bare_path_object_for_pretrain_weights(self) -> None:
+        """Bare pretrained weight Path values resolve the same as bare strings."""
+        path_config = RFDETRBaseConfig(pretrain_weights=Path("rf-detr-base.pth"))
+        string_config = RFDETRBaseConfig(pretrain_weights="rf-detr-base.pth")
+
+        assert path_config.pretrain_weights == string_config.pretrain_weights
+
+    @pytest.mark.parametrize(
+        "value",
+        [
+            # PTL trainer.fit(ckpt_path=...) sentinels — must pass through verbatim.
+            pytest.param("best", id="sentinel_best"),
+            pytest.param("last", id="sentinel_last"),
+            pytest.param("hpc", id="sentinel_hpc"),
+            pytest.param("registry:model-name", id="sentinel_registry"),
+            # A relative Path proves os.fspath coercion without realpath resolution.
+            pytest.param(Path("checkpoints/last.ckpt"), id="path_object"),
+        ],
+    )
+    def test_resume_coerced_via_fspath_without_realpath(self, value) -> None:
+        """``resume`` accepts pathlib.Path and is coerced to ``str`` via ``os.fspath`` only.
+
+        Unlike ``dataset_dir``/``output_dir``, ``resume`` is forwarded verbatim to PyTorch Lightning's
+        ``trainer.fit(ckpt_path=...)``, which also accepts sentinels such as ``"last"``. Realpath-normalising it would
+        rewrite those sentinels (and relative paths) into spurious absolute paths, so the value must be preserved.
+        """
+        config = TrainConfig(dataset_dir="/tmp", resume=value)
+
+        assert config.resume == os.fspath(value)
+
+
+class TestRFDETRBaseConfigEncoder:
+    """Encoder field validation on RFDETRBaseConfig (no fixture needed — has defaults)."""
+
+    def test_accepts_registers_windowed_small(self) -> None:
+        """RFDETRBaseConfig accepts the new dinov2_registers_windowed_small encoder."""
+        config = RFDETRBaseConfig(encoder="dinov2_registers_windowed_small", pretrain_weights=None)
+        assert config.encoder == "dinov2_registers_windowed_small"
+
+    def test_rejects_invalid_encoder(self) -> None:
+        """RFDETRBaseConfig raises ValidationError for unknown encoder strings."""
+        with pytest.raises(ValidationError):
+            RFDETRBaseConfig(encoder="not_a_real_encoder", pretrain_weights=None)
+
+
+class TestSegmentationTrainConfigNumSelect:
+    """Unit tests for SegmentationTrainConfig.num_select default and per-model values."""
+
+    def test_defaults_to_none(self) -> None:
+        config = SegmentationTrainConfig(dataset_dir="/tmp")
+        assert config.num_select is None
+
+    def test_explicit_value_is_accepted(self) -> None:
+        # Explicitly setting num_select on SegmentationTrainConfig is deprecated (Item #3).
+        with pytest.warns(DeprecationWarning, match="TrainConfig.num_select is deprecated"):
+            config = SegmentationTrainConfig(dataset_dir="/tmp", num_select=42)
+        assert config.num_select == 42
+
+    @pytest.mark.parametrize(
+        "config_class, expected_num_select",
+        [
+            (RFDETRSegNanoConfig, 100),
+            (RFDETRSegSmallConfig, 100),
+            (RFDETRSegMediumConfig, 200),
+            (RFDETRSegLargeConfig, 200),
+            (RFDETRSegXLargeConfig, 300),
+            (RFDETRSeg2XLargeConfig, 300),
+        ],
+    )
+    def test_model_config_has_variant_specific_num_select(self, config_class, expected_num_select) -> None:
+        assert config_class().num_select == expected_num_select
+
+
+class TestTrainConfigRejectsUnknownKwargs:
+    """TrainConfig must raise on unknown/typo'd kwargs instead of silently ignoring them (extra='forbid')."""
+
+    def test_typo_kwarg_raises_with_helpful_message(self, tmp_path) -> None:
+        """A typo'd kwarg (epoch instead of epochs) raises listing the unknown and available parameters."""
+        with pytest.raises(ValidationError, match=r"Unknown parameter\(s\): 'epoch'"):
+            TrainConfig(dataset_dir=str(tmp_path), output_dir=str(tmp_path), epoch=5)
+
+    def test_typo_error_lists_available_parameters(self, tmp_path) -> None:
+        """The rejection message includes the available parameter list so the typo is easy to fix."""
+        with pytest.raises(ValidationError, match=r"Available parameter\(s\):.*epochs"):
+            TrainConfig(dataset_dir=str(tmp_path), output_dir=str(tmp_path), epoch=5)
+
+    @pytest.mark.parametrize(
+        "config_class",
+        [
+            pytest.param(SegmentationTrainConfig, id="segmentation"),
+            pytest.param(KeypointTrainConfig, id="keypoint"),
+        ],
+    )
+    def test_subclasses_reject_unknown_kwargs(self, tmp_path, config_class) -> None:
+        """TrainConfig subclasses inherit the forbid behaviour."""
+        with pytest.raises(ValidationError, match=r"Unknown parameter\(s\): 'epoch'"):
+            config_class(dataset_dir=str(tmp_path), output_dir=str(tmp_path), epoch=5)
+
+    def test_get_train_config_raises_for_typo_kwarg(self, tmp_path) -> None:
+        """The public RFDETR.get_train_config path surfaces the typo instead of swallowing it."""
+        from types import SimpleNamespace
+
+        from rfdetr.detr import RFDETR
+
+        stub = SimpleNamespace(_train_config_class=TrainConfig)
+        with pytest.raises(ValidationError, match=r"Unknown parameter\(s\): 'epoch'"):
+            RFDETR.get_train_config(stub, dataset_dir=str(tmp_path), output_dir=str(tmp_path), epoch=5)
+
+
+class TestTrainConfigT42PromotedFields:
+    """T4-2: Promoted fields exist with correct defaults; device field is absent."""
+
+    def _tc(self, tmp_path, **kwargs):
+        defaults = dict(dataset_dir=str(tmp_path), output_dir=str(tmp_path), tensorboard=False)
+        defaults.update(kwargs)
+        return TrainConfig(**defaults)
+
+    # --- device field removed ---
+
+    def test_device_not_in_model_fields(self):
+        """Device must not appear in TrainConfig.model_fields (PTL auto-detects accelerator)."""
+        assert "device" not in TrainConfig.model_fields
+
+    def test_device_kwarg_rejected(self, tmp_path):
+        """Passing device= directly to TrainConfig raises (extra='forbid'); RFDETR.train() pops it beforehand."""
+        with pytest.raises(ValidationError, match=r"Unknown parameter\(s\): 'device'"):
+            self._tc(tmp_path, device="cpu")
+
+    # --- promoted fields: defaults ---
+
+    def test_clip_max_norm_default(self, tmp_path):
+        """clip_max_norm defaults to 0.1."""
+        assert self._tc(tmp_path).clip_max_norm == pytest.approx(0.1)
+
+    def test_seed_default_is_none(self, tmp_path):
+        """Seed defaults to None (no seeding)."""
+        assert self._tc(tmp_path).seed is None
+
+    def test_sync_bn_default_is_false(self, tmp_path):
+        """sync_bn defaults to False."""
+        assert self._tc(tmp_path).sync_bn is False
+
+    def test_fp16_eval_default_is_false(self, tmp_path):
+        """fp16_eval defaults to False."""
+        assert self._tc(tmp_path).fp16_eval is False
+
+    def test_lr_scheduler_default_is_step(self, tmp_path):
+        """lr_scheduler defaults to 'step'."""
+        assert self._tc(tmp_path).lr_scheduler == "step"
+
+    def test_lr_min_factor_default(self, tmp_path):
+        """lr_min_factor defaults to 0.0."""
+        assert self._tc(tmp_path).lr_min_factor == pytest.approx(0.0)
+
+    def test_dont_save_weights_default_is_false(self, tmp_path):
+        """dont_save_weights defaults to False."""
+        assert self._tc(tmp_path).dont_save_weights is False
+
+    def test_run_test_default_is_false(self, tmp_path):
+        """run_test defaults to False to avoid extra full-dataset test passes."""
+        assert self._tc(tmp_path).run_test is False
+
+    def test_eval_interval_default_is_one(self, tmp_path):
+        """eval_interval defaults to 1 (evaluate each epoch)."""
+        assert self._tc(tmp_path).eval_interval == 1
+
+    def test_skip_best_epochs_default_is_zero(self, tmp_path):
+        """skip_best_epochs defaults to 0 for backward compatibility."""
+        assert self._tc(tmp_path).skip_best_epochs == 0
+
+    def test_ema_update_interval_default_is_one(self, tmp_path):
+        """ema_update_interval defaults to 1 (update every step)."""
+        assert self._tc(tmp_path).ema_update_interval == 1
+
+    def test_compute_val_loss_default_is_true(self, tmp_path):
+        """compute_val_loss defaults to True."""
+        assert self._tc(tmp_path).compute_val_loss is True
+
+    def test_compute_test_loss_default_is_true(self, tmp_path):
+        """compute_test_loss defaults to True."""
+        assert self._tc(tmp_path).compute_test_loss is True
+
+    # --- promoted fields: accept explicit values ---
+
+    @pytest.mark.parametrize(
+        "field, value",
+        [
+            pytest.param("clip_max_norm", 0.5, id="clip_max_norm"),
+            pytest.param("seed", 42, id="seed"),
+            pytest.param("sync_bn", True, id="sync_bn"),
+            pytest.param("fp16_eval", True, id="fp16_eval"),
+            pytest.param("lr_scheduler", "cosine", id="lr_scheduler_cosine"),
+            pytest.param("lr_min_factor", 0.01, id="lr_min_factor"),
+            pytest.param("dont_save_weights", True, id="dont_save_weights"),
+            pytest.param("run_test", True, id="run_test"),
+            pytest.param("eval_interval", 3, id="eval_interval"),
+            pytest.param("skip_best_epochs", 3, id="skip_best_epochs"),
+            pytest.param("ema_update_interval", 4, id="ema_update_interval"),
+            pytest.param("compute_val_loss", False, id="compute_val_loss"),
+            pytest.param("compute_test_loss", False, id="compute_test_loss"),
+            pytest.param("train_log_sync_dist", True, id="train_log_sync_dist"),
+            pytest.param("train_log_on_step", True, id="train_log_on_step"),
+            pytest.param("log_per_class_metrics", False, id="log_per_class_metrics"),
+            pytest.param("prefetch_factor", 4, id="prefetch_factor"),
+            pytest.param("pin_memory", False, id="pin_memory"),
+            pytest.param("persistent_workers", False, id="persistent_workers"),
+        ],
+    )
+    def test_promoted_field_accepts_explicit_value(self, tmp_path, field, value):
+        """Each promoted field accepts an explicit value."""
+        tc = self._tc(tmp_path, **{field: value})
+        assert getattr(tc, field) == value
+
+    def test_lr_scheduler_rejects_invalid_value(self, tmp_path):
+        """lr_scheduler must reject values other than 'step' and 'cosine'."""
+        with pytest.raises((ValueError, ValidationError)):
+            self._tc(tmp_path, lr_scheduler="cyclic")
+
+    @pytest.mark.parametrize(
+        ("field", "value"),
+        [
+            pytest.param("eval_interval", 0, id="eval_interval_zero"),
+            pytest.param("skip_best_epochs", -1, id="skip_best_epochs_negative"),
+            pytest.param("ema_update_interval", 0, id="ema_update_interval_zero"),
+            pytest.param("prefetch_factor", 0, id="prefetch_factor_zero"),
+        ],
+    )
+    def test_interval_and_prefetch_reject_non_positive_values(self, tmp_path, field, value):
+        """Eval/EMA intervals and prefetch_factor must be >= 1 when provided."""
+        with pytest.raises((ValueError, ValidationError)):
+            self._tc(tmp_path, **{field: value})
+
+    def test_batch_size_auto_is_accepted(self, tmp_path):
+        """batch_size accepts the special 'auto' value."""
+        tc = self._tc(tmp_path, batch_size="auto")
+        assert tc.batch_size == "auto"
+
+    @pytest.mark.parametrize(
+        "field,value",
+        [
+            ("batch_size", 0),
+            ("grad_accum_steps", 0),
+            ("auto_batch_target_effective", 0),
+            ("auto_batch_max_targets_per_image", 0),
+        ],
+    )
+    def test_auto_batch_related_fields_reject_non_positive_values(self, tmp_path, field, value):
+        """batch/accum/target-effective/max_targets fields must be >= 1 (except batch_size='auto')."""
+        with pytest.raises((ValueError, ValidationError)):
+            self._tc(tmp_path, **{field: value})
+
+    @pytest.mark.parametrize("ema_headroom", [0.0, 1.5])
+    def test_auto_batch_ema_headroom_must_be_in_open_one(self, tmp_path, ema_headroom):
+        """auto_batch_ema_headroom must be in (0, 1]."""
+        with pytest.raises((ValueError, ValidationError)):
+            self._tc(tmp_path, auto_batch_ema_headroom=ema_headroom)
+
+
+class TestBuildTrainerUsesRealFields:
+    """build_trainer() must read clip_max_norm, seed, sync_bn from real TrainConfig fields."""
+
+    def _tc(self, tmp_path, **kwargs):
+        defaults = dict(
+            dataset_dir=str(tmp_path),
+            output_dir=str(tmp_path),
+            tensorboard=False,
+            wandb=False,
+            mlflow=False,
+            clearml=False,
+            use_ema=False,
+        )
+        defaults.update(kwargs)
+        return TrainConfig(**defaults)
+
+    def _kp_tc(self, tmp_path, **kwargs):
+        defaults = dict(
+            dataset_dir=str(tmp_path),
+            output_dir=str(tmp_path),
+            tensorboard=False,
+            wandb=False,
+            mlflow=False,
+            clearml=False,
+            use_ema=False,
+        )
+        defaults.update(kwargs)
+        return KeypointTrainConfig(**defaults)
+
+    def _mc(self, **kwargs):
+        from rfdetr.config import RFDETRBaseConfig
+
+        defaults = dict(pretrain_weights=None, device="cpu", num_classes=3)
+        defaults.update(kwargs)
+        return RFDETRBaseConfig(**defaults)
+
+    def test_clip_max_norm_forwarded_to_trainer_for_detection(self, tmp_path):
+        """Detection models use Lightning's automatic optimization, so ``gradient_clip_val`` flows through to the
+        Trainer from ``TrainConfig.clip_max_norm`` unchanged."""
+        from rfdetr.training import build_trainer
+
+        trainer = build_trainer(self._tc(tmp_path, clip_max_norm=0.25), self._mc())
+        assert trainer.gradient_clip_val == pytest.approx(0.25)
+
+    def test_clip_max_norm_owned_by_model_module_for_keypoints(self, tmp_path):
+        """Keypoint models use manual optimization; trainer-owned clipping is disabled and ``clip_max_norm`` is applied
+        inside ``RFDETRModelModule._step_optimizer`` instead."""
+        from rfdetr.training import build_trainer
+
+        trainer = build_trainer(
+            self._kp_tc(tmp_path, clip_max_norm=0.25),
+            self._mc(use_grouppose_keypoints=True),
+        )
+        assert trainer.gradient_clip_val is None
+
+    def test_seed_not_applied_in_build_trainer_factory(self, tmp_path):
+        """Seeding is deferred to RFDETRModule.on_fit_start, not build_trainer()."""
+        import unittest.mock as mock
+
+        from rfdetr.training import build_trainer
+
+        with mock.patch("pytorch_lightning.seed_everything") as mock_seed:
+            build_trainer(self._tc(tmp_path, seed=99), self._mc())
+        mock_seed.assert_not_called()
+
+    def test_sync_bn_forwarded_to_trainer(self, tmp_path):
+        """sync_batchnorm=True is passed to Trainer when TrainConfig.sync_bn is True."""
+        import unittest.mock as mock
+
+        from rfdetr.training import build_trainer
+
+        captured_kwargs = {}
+
+        real_trainer_init = __import__("pytorch_lightning").Trainer.__init__
+
+        def _capture_init(self_t, **kwargs):
+            captured_kwargs.update(kwargs)
+            real_trainer_init(self_t, **kwargs)
+
+        with mock.patch("rfdetr.training.trainer.Trainer.__init__", _capture_init):
+            build_trainer(self._tc(tmp_path, sync_bn=True), self._mc())
+
+        assert captured_kwargs.get("sync_batchnorm") is True
+
+
+class TestDeprecatedTrainConfigFields:
+    """Item #3 Phase A: TrainConfig fields deprecated in favour of ModelConfig ownership."""
+
+    def _tc(self, **kwargs):
+        defaults = dict(dataset_dir="/tmp")
+        defaults.update(kwargs)
+        return TrainConfig(**defaults)
+
+    @pytest.mark.parametrize(
+        "field,value",
+        [
+            pytest.param("group_detr", 5, id="group_detr"),
+            pytest.param("ia_bce_loss", False, id="ia_bce_loss"),
+            pytest.param("segmentation_head", True, id="segmentation_head"),
+            pytest.param("num_select", 100, id="num_select"),
+        ],
+    )
+    def test_explicitly_set_deprecated_field_emits_warning(self, field, value) -> None:
+        """Setting a deprecated TrainConfig field explicitly must emit DeprecationWarning."""
+        with pytest.warns(DeprecationWarning, match=f"TrainConfig\\.{field} is deprecated"):
+            self._tc(**{field: value})
+
+    def test_default_group_detr_no_warning(self, recwarn) -> None:
+        """TrainConfig() without explicit group_detr must NOT warn."""
+        self._tc()
+        depr_warnings = [w for w in recwarn.list if issubclass(w.category, DeprecationWarning)]
+        assert not depr_warnings, f"Unexpected DeprecationWarning: {depr_warnings}"
+
+    def test_segmentation_train_config_no_warning_on_default_fields(self, recwarn) -> None:
+        """SegmentationTrainConfig() must NOT warn for its class-level defaults.
+
+        segmentation_head=True and num_select=None are SegmentationTrainConfig defaults, not explicitly set by the user
+        — they must not trigger DeprecationWarning.
+        """
+        SegmentationTrainConfig(dataset_dir="/tmp")
+        depr_warnings = [w for w in recwarn.list if issubclass(w.category, DeprecationWarning)]
+        assert not depr_warnings, f"Unexpected DeprecationWarning: {depr_warnings}"
+
+
+class TestDeprecatedModelConfigClsLossCoef:
+    """Item #3 Phase A: ModelConfig.cls_loss_coef deprecated in favour of TrainConfig ownership."""
+
+    def test_explicit_cls_loss_coef_emits_warning(self) -> None:
+        """Setting cls_loss_coef on ModelConfig explicitly must emit DeprecationWarning."""
+        sample = dict(
+            encoder="dinov2_windowed_small",
+            out_feature_indexes=[1, 2, 3],
+            dec_layers=3,
+            projector_scale=["P3"],
+            hidden_dim=256,
+            patch_size=14,
+            num_windows=2,
+            sa_nheads=8,
+            ca_nheads=8,
+            dec_n_points=4,
+            resolution=384,
+            positional_encoding_size=256,
+        )
+        with pytest.warns(DeprecationWarning, match="ModelConfig\\.cls_loss_coef is deprecated"):
+            ModelConfig(**sample, cls_loss_coef=2.0)
+
+    def test_default_cls_loss_coef_no_warning(self, recwarn) -> None:
+        """RFDETRBaseConfig() without explicit cls_loss_coef must NOT warn."""
+        RFDETRBaseConfig(pretrain_weights=None, device="cpu")
+        depr_warnings = [w for w in recwarn.list if issubclass(w.category, DeprecationWarning)]
+        assert not depr_warnings, f"Unexpected DeprecationWarning: {depr_warnings}"
+
+
+class TestSyncPEWithResolutionAtConstruction:
+    """Tests for the _sync_pe_with_resolution model_validator.
+
+    When a user provides a custom resolution at construction time (e.g., ``RFDETRLarge(resolution=640)``),
+    positional_encoding_size must be updated proportionally for configs where the default PE is formula-derived
+    (``default_pe == default_resolution // patch_size``).
+    """
+
+    @pytest.mark.parametrize(
+        "config_cls, new_resolution, expected_pe",
+        [
+            pytest.param(RFDETRLargeConfig, 640, 640 // 16, id="large_640"),
+            pytest.param(RFDETRLargeConfig, 576, 576 // 16, id="large_576"),
+            pytest.param(RFDETRSmallConfig, 640, 640 // 16, id="small_640"),
+            pytest.param(RFDETRMediumConfig, 640, 640 // 16, id="medium_640"),
+            pytest.param(RFDETRNanoConfig, 416, 416 // 16, id="nano_416"),
+            pytest.param(RFDETRSegNanoConfig, 360, 360 // 12, id="seg_nano_360"),
+            pytest.param(RFDETRSegSmallConfig, 480, 480 // 12, id="seg_small_480"),
+            pytest.param(RFDETRSegMediumConfig, 480, 480 // 12, id="seg_medium_480"),
+            pytest.param(RFDETRSegLargeConfig, 576, 576 // 12, id="seg_large_576"),
+            pytest.param(RFDETRSegXLargeConfig, 576, 576 // 12, id="seg_xlarge_576"),
+            pytest.param(RFDETRSeg2XLargeConfig, 720, 720 // 12, id="seg_2xlarge_720"),
+        ],
+    )
+    def test_positional_encoding_size_updated_for_formula_derived_configs(
+        self,
+        config_cls: type,
+        new_resolution: int,
+        expected_pe: int,
+    ) -> None:
+        """PE is auto-derived from the custom resolution for formula-derived model configs."""
+        cfg = config_cls(resolution=new_resolution, pretrain_weights=None)
+        assert cfg.positional_encoding_size == expected_pe
+
+    def test_explicit_positional_encoding_size_is_not_overridden(self) -> None:
+        """When positional_encoding_size is explicitly provided, the validator must not override it."""
+        cfg = RFDETRLargeConfig(resolution=640, positional_encoding_size=50, pretrain_weights=None)
+        assert cfg.positional_encoding_size == 50
+
+    def test_default_resolution_preserves_default_pe(self) -> None:
+        """Constructing with default resolution (no explicit resolution) must not change PE."""
+        cfg = RFDETRLargeConfig(pretrain_weights=None)
+        assert cfg.resolution == 704
+        assert cfg.positional_encoding_size == 44  # 704 // 16
+
+
+class TestDetectDevice:
+    """Tests for _detect_device() covering PyTorch accelerator detection paths."""
+
+    @patch("rfdetr.config.torch")
+    def test_falls_back_to_cuda_when_accelerator_module_absent(self, mock_torch: MagicMock) -> None:
+        """Returns 'cuda' via legacy fallback when torch.accelerator lacks current_accelerator (PyTorch < 2.4)."""
+        mock_torch.accelerator = MagicMock(spec=[])  # no current_accelerator → hasattr returns False → fallback
+        mock_torch.cuda.is_available.return_value = True
+        mock_torch.backends.mps.is_available.return_value = False
+        assert _detect_device() == "cuda"
+
+    @patch("rfdetr.config.torch")
+    def test_returns_cpu_when_current_accelerator_raises(self, mock_torch: MagicMock) -> None:
+        """Returns 'cpu' directly from the except handler when current_accelerator() raises RuntimeError."""
+        mock_torch.accelerator.current_accelerator.side_effect = RuntimeError("no device")
+        assert _detect_device() == "cpu"
+
+    @patch("rfdetr.config.torch")
+    def test_returns_cpu_when_no_gpu_available(self, mock_torch: MagicMock) -> None:
+        """Returns 'cpu' when accelerator is absent and neither CUDA nor MPS is available."""
+        mock_torch.accelerator = MagicMock(spec=[])  # no current_accelerator → fallback branch
+        mock_torch.cuda.is_available.return_value = False
+        mock_torch.backends.mps.is_available.return_value = False
+        assert _detect_device() == "cpu"
+
+    @patch("rfdetr.config.torch")
+    def test_returns_cpu_when_accelerator_compiled_in_but_unavailable(self, mock_torch: MagicMock) -> None:
+        """Returns 'cpu' when torch was compiled with CUDA but no driver is present at runtime.
+
+        Without ``check_available=True``, ``current_accelerator()`` reports the compile-time accelerator, so the default
+        CUDA wheel on a driverless machine yields ``device("cuda")`` and every model build crashes with "Found no NVIDIA
+        driver". The runtime availability check must win.
+        """
+
+        def fake_current_accelerator(check_available: bool = False) -> "torch.device | None":
+            return None if check_available else torch.device("cuda")
+
+        mock_torch.accelerator.current_accelerator = fake_current_accelerator
+        assert _detect_device() == "cpu"
+
+    @patch("rfdetr.config.torch")
+    def test_returns_accelerator_when_runtime_available(self, mock_torch: MagicMock) -> None:
+        """Returns the accelerator when it passes the runtime availability check."""
+
+        def fake_current_accelerator(check_available: bool = False) -> "torch.device | None":
+            return torch.device("cuda") if check_available else None
+
+        mock_torch.accelerator.current_accelerator = fake_current_accelerator
+        assert _detect_device() == "cuda"
+
+    @patch("rfdetr.config.torch")
+    def test_legacy_signature_unavailable_accelerator_returns_cpu(self, mock_torch: MagicMock) -> None:
+        """Falls back to ``is_available()`` when ``current_accelerator`` lacks ``check_available`` (PyTorch < 2.7)."""
+
+        def legacy_current_accelerator() -> "torch.device":
+            return torch.device("cuda")
+
+        mock_torch.accelerator.current_accelerator = legacy_current_accelerator
+        mock_torch.accelerator.is_available.return_value = False
+        assert _detect_device() == "cpu"
+
+    @patch("rfdetr.config.torch")
+    def test_legacy_signature_available_accelerator_is_kept(self, mock_torch: MagicMock) -> None:
+        """Keeps the accelerator on pre-``check_available`` builds when ``is_available()`` confirms it."""
+
+        def legacy_current_accelerator() -> "torch.device":
+            return torch.device("cuda")
+
+        mock_torch.accelerator.current_accelerator = legacy_current_accelerator
+        mock_torch.accelerator.is_available.return_value = True
+        assert _detect_device() == "cuda"
+
+    @patch("rfdetr.config.torch")
+    def test_legacy_signature_runtime_error_on_fallback_returns_cpu(self, mock_torch: MagicMock) -> None:
+        """Outer RuntimeError handler catches error from legacy fallback call.
+
+        Control-flow: ``current_accelerator(check_available=True)`` raises ``TypeError`` (inner except),
+        then ``current_accelerator()`` raises ``RuntimeError`` (outer except catches) → ``"cpu"``.
+        """
+        call_count = 0
+
+        def raises_on_fallback(**kwargs: object) -> "torch.device":
+            nonlocal call_count
+            call_count += 1
+            if "check_available" in kwargs:
+                raise TypeError("unexpected keyword argument 'check_available'")
+            raise RuntimeError("NVML error on legacy fallback")
+
+        mock_torch.accelerator.current_accelerator = raises_on_fallback
+        assert _detect_device() == "cpu"
+
+
+class TestPretrainWeightsCompatibilityWarning:
+    """Config-time warning for overrides that prevent pretrained weights from loading.
+
+    These tests instantiate the variant *config* directly (not the wrapper class) so they do not touch the network, the
+    cache, or any model construction.
+    """
+
+    def _capture(self, config_cls: type, **kwargs: object) -> list[warnings.WarningMessage]:
+        """Instantiate ``config_cls(**kwargs)`` and return only the pretrain-compat warnings."""
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            config_cls(**kwargs)
+            return [w for w in caught if issubclass(w.category, PretrainWeightsCompatibilityWarning)]
+
+    def test_default_construction_emits_no_warning(self) -> None:
+        """Default variant construction must not warn — defaults match the published checkpoint."""
+        assert self._capture(RFDETRNanoConfig) == []
+
+    def test_encoder_registers_override_warns(self) -> None:
+        """The dinov2-with-registers footgun: switching encoder away from the variant default."""
+        captured = self._capture(RFDETRNanoConfig, encoder="dinov2_registers_windowed_small")
+        assert len(captured) == 1
+        message = str(captured[0].message)
+        assert "encoder" in message
+        assert "dinov2_registers_windowed_small" in message
+        assert "dinov2_windowed_small" in message
+
+    @pytest.mark.parametrize(
+        "field, value",
+        [
+            pytest.param("hidden_dim", 384, id="hidden_dim"),
+            pytest.param("dec_layers", 6, id="dec_layers"),
+            pytest.param("num_windows", 4, id="num_windows"),
+            pytest.param("sa_nheads", 4, id="sa_nheads"),
+            pytest.param("ca_nheads", 8, id="ca_nheads"),
+            pytest.param("dec_n_points", 4, id="dec_n_points"),
+            pytest.param("out_feature_indexes", [2, 5, 8, 11], id="out_feature_indexes"),
+            pytest.param("projector_scale", ["P3", "P4"], id="projector_scale"),
+            pytest.param("bbox_reparam", False, id="bbox_reparam"),
+            pytest.param("lite_refpoint_refine", False, id="lite_refpoint_refine"),
+            pytest.param("layer_norm", False, id="layer_norm"),
+            pytest.param("two_stage", False, id="two_stage"),
+            pytest.param("num_channels", 1, id="num_channels"),
+        ],
+    )
+    def test_load_breaking_override_warns(self, field: str, value: object) -> None:
+        """Each load-breaking architecture override fires the warning."""
+        captured = self._capture(RFDETRNanoConfig, **{field: value})
+        assert len(captured) == 1
+        assert field in str(captured[0].message)
+
+    def test_mask_downsample_ratio_warns_on_seg_variant(self) -> None:
+        """``mask_downsample_ratio`` change is silently miscalibrating; must warn at config time."""
+        captured = self._capture(RFDETRSegNanoConfig, mask_downsample_ratio=2)
+        assert len(captured) == 1
+        assert "mask_downsample_ratio" in str(captured[0].message)
+
+    def test_patch_size_override_warns_defense_in_depth(self) -> None:
+        """patch_size already raises in load_pretrain_weights; the new warning is defense-in-depth.
+
+        We change patch_size to a value that differs from RFDETRNanoConfig's default (16).
+        """
+        captured = self._capture(RFDETRNanoConfig, patch_size=14)
+        assert len(captured) == 1
+        assert "patch_size" in str(captured[0].message)
+
+    def test_segmentation_head_override_warns(self) -> None:
+        """segmentation_head also raises at load time but warning fires first."""
+        # RFDETRNanoConfig has segmentation_head=False; flipping it to True is the override.
+        captured = self._capture(RFDETRNanoConfig, segmentation_head=True)
+        assert len(captured) == 1
+        assert "segmentation_head" in str(captured[0].message)
+
+    @pytest.mark.parametrize(
+        "field, value",
+        [
+            pytest.param("num_queries", 200, id="num_queries_decrease"),
+            pytest.param("num_queries", 300, id="num_queries_equal"),
+            pytest.param("group_detr", 8, id="group_detr_decrease"),
+            pytest.param("num_classes", 5, id="num_classes"),
+            pytest.param("resolution", 448, id="resolution"),
+            pytest.param("positional_encoding_size", 20, id="positional_encoding_size"),
+        ],
+    )
+    def test_silent_field_overrides(self, field: str, value: object) -> None:
+        """Fields that are auto-handled at load time must not emit a warning at config construction."""
+        assert self._capture(RFDETRNanoConfig, **{field: value}) == []
+
+    @pytest.mark.parametrize(
+        "field, value",
+        [
+            pytest.param("num_queries", 400, id="num_queries"),
+            pytest.param("group_detr", 20, id="group_detr"),
+        ],
+    )
+    def test_increase_field_warns(self, field: str, value: object) -> None:
+        """Increasing an integer field above the variant default warns — extra slots are randomly initialised."""
+        captured = self._capture(RFDETRNanoConfig, **{field: value})
+        assert len(captured) == 1
+        assert field in str(captured[0].message)
+
+    def test_pretrain_weights_none_warns(self) -> None:
+        """Explicitly opting out of pretrained weights warns about training from scratch."""
+        captured = self._capture(RFDETRNanoConfig, pretrain_weights=None)
+        assert len(captured) == 1
+        message = str(captured[0].message)
+        assert "from scratch" in message
+        assert "rf-detr-nano.pth" in message
+
+    def test_pretrain_weights_none_only_one_warning(self) -> None:
+        """When pretrain_weights=None, the architecture-overrides warning is suppressed.
+
+        The from-scratch warning is the dominant message; we don't pile on with arch warnings.
+        """
+        captured = self._capture(
+            RFDETRNanoConfig,
+            pretrain_weights=None,
+            encoder="dinov2_registers_windowed_small",
+            hidden_dim=384,
+        )
+        assert len(captured) == 1
+        assert "from scratch" in str(captured[0].message)
+
+    def test_custom_pretrain_weights_path_suppresses_arch_warning(self) -> None:
+        """Custom pretrain_weights path → defer to load-time detector — no config-time arch warning."""
+        captured = self._capture(
+            RFDETRNanoConfig,
+            pretrain_weights="/tmp/my_custom.pth",
+            encoder="dinov2_registers_windowed_small",
+        )
+        assert captured == []
+
+    def test_multiple_overrides_consolidated_into_one_warning(self) -> None:
+        """All overrides are listed in a single warning, not one warning per field."""
+        captured = self._capture(
+            RFDETRNanoConfig,
+            encoder="dinov2_registers_windowed_small",
+            hidden_dim=384,
+            num_queries=400,
+        )
+        assert len(captured) == 1
+        message = str(captured[0].message)
+        for needle in ("encoder", "hidden_dim", "num_queries"):
+            assert needle in message, f"expected {needle!r} in consolidated warning message"
+
+    def test_warning_is_user_warning_subclass(self) -> None:
+        """Confirms downstream filtering via UserWarning works."""
+        assert issubclass(PretrainWeightsCompatibilityWarning, UserWarning)
+
+    def test_modelconfig_with_required_fields_does_not_warn(self, sample_model_config: dict[str, object]) -> None:
+        """Constructing the abstract ModelConfig with required fields cannot compare to defaults — no warning."""
+        assert self._capture(ModelConfig, **sample_model_config) == []
+
+    def test_breaking_field_with_default_factory_skips_comparison(self) -> None:
+        """A subclass whose breaking field uses ``default_factory`` (so ``.default`` is ``PydanticUndefined``) must be
+        silently skipped — we have nothing to compare against."""
+        from pydantic import Field
+
+        class _DefaultFactoryConfig(RFDETRNanoConfig):
+            # Field uses default_factory → FieldInfo.default is PydanticUndefined,
+            # but is_required() is False.  Hits the `continue` on the
+            # PydanticUndefined check.
+            encoder: str = Field(default_factory=lambda: "dinov2_windowed_small")
+
+        assert self._capture(_DefaultFactoryConfig, encoder="dinov2_registers_windowed_small") == []
+
+    def test_increase_field_when_required_skips_comparison(self) -> None:
+        """A subclass where ``num_queries`` becomes required (no default) must be skipped."""
+
+        class _RequiredNumQueriesConfig(RFDETRNanoConfig):
+            num_queries: int  # type: ignore[misc]  # no default → required
+
+        assert self._capture(_RequiredNumQueriesConfig, num_queries=400) == []
+
+    def test_increase_field_with_non_int_default_skips_comparison(self) -> None:
+        """A subclass where ``num_queries`` has a non-int default must be skipped (can't ``>`` compare)."""
+        from typing import Any
+
+        class _NonIntDefaultConfig(RFDETRNanoConfig):
+            num_queries: Any = "300"  # type: ignore[assignment]  # non-int default
+
+        assert self._capture(_NonIntDefaultConfig, num_queries="400") == []
+
+    def test_explicit_variant_default_path_runs_arch_override_check(self) -> None:
+        """Passing the variant's own published-default path string must still check arch overrides.
+
+        Before the case-2 fix, any non-None explicit pretrain_weights bypassed the architecture-override check entirely
+        — including when the user passed the exact variant default string such as "rf-detr-nano.pth".
+        """
+        captured = self._capture(
+            RFDETRNanoConfig,
+            pretrain_weights="rf-detr-nano.pth",
+            encoder="dinov2_registers_windowed_small",
+        )
+        assert len(captured) == 1
+        assert "encoder" in str(captured[0].message)
+
+    def test_product_preserving_group_detr_increase_still_warns(self) -> None:
+        """Increasing group_detr while halving num_queries still warns — check is per-field, not product-aware.
+
+        This documents known current behaviour: the validator compares each field to its variant default independently,
+        not the combined query-slot product.  A product- preserving change (group_detr=26, num_queries=150 vs defaults
+        13, 300) warns for group_detr because 26 > 13, regardless of whether total slots are the same.
+        """
+        captured = self._capture(RFDETRNanoConfig, num_queries=150, group_detr=26)
+        assert len(captured) == 1
+        assert "group_detr" in str(captured[0].message)
+
+
+class TestBreakingListIntegrity:
+    """Guards against stale entries in the pretrain-compatibility breaking-field lists."""
+
+    def test_all_breaking_fields_exist_in_model_config(self) -> None:
+        """Every field guarded by the pretrain-compatibility check must exist in ModelConfig.model_fields.
+
+        Catches typos and fields renamed/removed without updating the breaking lists.
+        """
+        all_breaking = {
+            "encoder",
+            "hidden_dim",
+            "dec_layers",
+            "num_windows",
+            "sa_nheads",
+            "ca_nheads",
+            "dec_n_points",
+            "out_feature_indexes",
+            "projector_scale",
+            "bbox_reparam",
+            "lite_refpoint_refine",
+            "layer_norm",
+            "two_stage",
+            "patch_size",
+            "segmentation_head",
+            "num_channels",
+            "num_queries",
+            "group_detr",
+        }
+        stale = all_breaking - set(ModelConfig.model_fields.keys())
+        assert not stale, f"Fields in breaking lists not in ModelConfig.model_fields: {stale}"

tests/models/test_config_keypoints.py

@@ -0,0 +1,112 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests for keypoint config defaults and namespace forwarding."""
+
+import pytest
+
+from rfdetr._namespace import _namespace_from_configs
+from rfdetr.config import (
+    KeypointTrainConfig,
+    RFDETRBaseConfig,
+    RFDETRKeypointPreviewConfig,
+    SegmentationTrainConfig,
+)
+
+
+def test_keypoint_config_defaults() -> None:
+    """Default model/train keypoint configuration values should match the preview contract."""
+    model = RFDETRKeypointPreviewConfig()
+    train = KeypointTrainConfig(dataset_dir="/tmp")
+
+    assert model.use_grouppose_keypoints is True
+    assert model.dual_projector is True
+    assert model.dual_projector_kp_only is True
+    assert model.num_keypoints_per_class == [17]
+    assert model.positional_encoding_size == 576 // 12
+
+    assert train.keypoint_l1_loss_coef == pytest.approx(1.0)
+    assert train.keypoint_findable_loss_coef == pytest.approx(1.0)
+    assert train.keypoint_visible_loss_coef == pytest.approx(1.0)
+    assert train.keypoint_nll_loss_coef == pytest.approx(1.0)
+    assert train.cls_loss_coef == pytest.approx(2.0)
+
+
+def test_keypoint_preview_config_person_schema() -> None:
+    """Person-keypoint preview config must expose a person-only schema."""
+    model = RFDETRKeypointPreviewConfig()
+
+    assert model.num_keypoints_per_class == [17]
+    assert sum(model.num_keypoints_per_class) == 17
+    assert model.out_feature_indexes == [3, 6, 9, 12]
+    assert model.num_windows == 2
+    assert model.dec_layers == 4
+    assert model.patch_size == 12
+    assert model.resolution == 576
+    assert model.pretrain_weights == "rf-detr-keypoint-preview-xlarge.pth"
+
+
+def test_keypoint_fields_propagate_to_namespace(tmp_path) -> None:
+    """All keypoint config fields are forwarded through _namespace_from_configs."""
+    model = RFDETRKeypointPreviewConfig()
+    train = KeypointTrainConfig(
+        dataset_dir=str(tmp_path),
+        keypoint_flip_pairs=[0, 1, 2, 3],
+        keypoint_l1_loss_coef=1.5,
+        keypoint_findable_loss_coef=2.5,
+        keypoint_visible_loss_coef=3.5,
+        keypoint_nll_loss_coef=4.5,
+    )
+
+    namespace = _namespace_from_configs(model, train)
+
+    assert namespace.use_grouppose_keypoints is True
+    assert namespace.keypoint_cross_attn is True
+    assert namespace.inter_instance_kp_attn is False
+    assert namespace.grouppose_keypoint_dim_downscale == 1
+    assert namespace.dual_projector is True
+    assert namespace.dual_projector_kp_only is True
+    assert namespace.num_keypoints_per_class == [17]
+    assert namespace.keypoint_flip_pairs == [0, 1, 2, 3]
+    assert namespace.keypoint_l1_loss_coef == pytest.approx(1.5)
+    assert namespace.keypoint_findable_loss_coef == pytest.approx(2.5)
+    assert namespace.keypoint_visible_loss_coef == pytest.approx(3.5)
+    assert namespace.keypoint_nll_loss_coef == pytest.approx(4.5)
+
+
+def test_keypoint_nll_loss_coef_default_restored_to_1_0() -> None:
+    """keypoint_nll_loss_coef must default to 1.0 after the 0.5 revert.
+
+    The 0.5 default was introduced to dampen OKS@75 oscillation.  It was later reverted to 1.0 to align with all other
+    keypoint loss terms (l1, findable, visible).  This test guards against silent regressions.
+    """
+    train = KeypointTrainConfig(dataset_dir="/tmp")
+    assert train.keypoint_nll_loss_coef == pytest.approx(1.0)
+
+
+def test_segmentation_train_config_cls_loss_coef_default() -> None:
+    """SegmentationTrainConfig.cls_loss_coef must default to 1.0, not the erroneous 5.0.
+
+    The 5.0 value was always present in SegmentationTrainConfig but was dead code pre-v1.7 (namespace builder read from
+    ModelConfig=1.0).  The v1.7 TrainConfig ownership migration silently activated it.  This test guards against re-
+    introducing that regression.
+    """
+    tc = SegmentationTrainConfig(dataset_dir="/tmp")
+    assert tc.cls_loss_coef == pytest.approx(1.0)
+
+
+def test_unknown_keypoint_fields_are_not_public_config_fields() -> None:
+    """Private keypoint implementation fields are not accepted as public model config."""
+    with pytest.raises(ValueError, match="Unknown parameter"):
+        RFDETRBaseConfig(num_classes=1, keypoint_private_hidden_dim=256)
+
+    # KeypointTrainConfig (a TrainConfig subclass) uses extra="forbid", so unknown
+    # kwargs raise with a helpful message rather than being silently dropped.
+    with pytest.raises(ValueError, match="Unknown parameter"):
+        KeypointTrainConfig(
+            dataset_dir="/tmp",
+            keypoint_private_hidden_dim=256,
+            keypoint_private_loss_coef=1.0,
+        )

tests/models/test_criterion.py

@@ -0,0 +1,134 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for SetCriterion edge paths: _output_device and num_boxes_for_targets."""
+
+import pytest
+import torch
+
+from rfdetr.models.criterion import SetCriterion
+
+
+class _MatcherStub:
+    """Minimal matcher that returns identity indices for every target in the batch."""
+
+    def __call__(self, outputs, targets, group_detr=1):
+        return [(torch.arange(len(t["labels"])), torch.arange(len(t["labels"]))) for t in targets]
+
+
+def _bare_criterion() -> SetCriterion:
+    """Return a SetCriterion with no losses so forward() is a no-op."""
+    criterion = SetCriterion.__new__(SetCriterion)
+    criterion.training = True
+    criterion.group_detr = 1
+    criterion.sum_group_losses = False
+    criterion.losses = []
+    criterion.weight_dict = {}
+    criterion.matcher = _MatcherStub()
+    criterion.num_keypoints_per_class = []
+    return criterion
+
+
+class TestOutputDevice:
+    """Tests for SetCriterion._output_device — probes top-level tensor values only."""
+
+    def test_returns_device_of_first_tensor(self):
+        """Device inferred from the first tensor value in outputs."""
+        outputs = {"pred_logits": torch.zeros(1, 1, 1)}
+
+        device = SetCriterion._output_device(outputs)
+
+        assert device == torch.device("cpu")
+
+    def test_raises_when_no_tensor_present(self):
+        """ValueError raised when no top-level value is a tensor."""
+        outputs = {"meta": "string_value", "count": 42}
+
+        with pytest.raises(ValueError, match="at least one tensor"):
+            SetCriterion._output_device(outputs)
+
+    def test_skips_non_tensor_values(self):
+        """Non-tensor entries at the top level are skipped; first tensor wins."""
+        outputs = {"meta": "ignored", "pred_logits": torch.zeros(1, 1, 1)}
+
+        device = SetCriterion._output_device(outputs)
+
+        assert device == torch.device("cpu")
+
+
+class TestNumBoxesForTargets:
+    """Tests for SetCriterion.num_boxes_for_targets — clamp and empty-target edge cases."""
+
+    def test_returns_tensor_gte_one(self):
+        """Result must be clamped to >= 1.0 to prevent division by zero."""
+        criterion = _bare_criterion()
+        outputs = {"pred_logits": torch.zeros(1, 1, 1)}
+        targets = [{"labels": torch.tensor([0, 1])}]
+
+        result = criterion.num_boxes_for_targets(outputs, targets)
+
+        assert result.item() >= 1.0
+
+    def test_clamps_zero_box_count_to_one(self):
+        """Empty targets (no labels) must clamp to 1.0 to avoid zero denominator."""
+        criterion = _bare_criterion()
+        outputs = {"pred_logits": torch.zeros(1, 1, 1)}
+        targets = [{"labels": torch.zeros(0, dtype=torch.int64)}]
+
+        result = criterion.num_boxes_for_targets(outputs, targets)
+
+        assert result.item() == pytest.approx(1.0)
+
+    def test_clamps_empty_target_list(self):
+        """Empty target list (batch_size=0 edge case) must also clamp to 1.0."""
+        criterion = _bare_criterion()
+        outputs = {"pred_logits": torch.zeros(1, 1, 1)}
+        targets = []
+
+        result = criterion.num_boxes_for_targets(outputs, targets)
+
+        assert result.item() == pytest.approx(1.0)
+
+    def test_counts_labels_correctly(self):
+        """Box count equals total number of labels across all targets in the batch."""
+        criterion = _bare_criterion()
+        outputs = {"pred_logits": torch.zeros(1, 1, 1)}
+        targets = [
+            {"labels": torch.tensor([0, 1])},
+            {"labels": torch.tensor([0])},
+        ]
+
+        result = criterion.num_boxes_for_targets(outputs, targets)
+
+        # 2 + 1 = 3 boxes; single-process so no all-reduce
+        assert result.item() == pytest.approx(3.0)
+
+
+class TestLossMasksEmptyMatch:
+    """Tests for the dict-path zero-GT branch of SetCriterion.loss_masks."""
+
+    def test_dict_path_zero_gt_stays_connected_to_graph(self):
+        """Zero-match dict path returns a loss that back-propagates to every segmentation-head output."""
+        criterion = _bare_criterion()
+        spatial_features = torch.randn(1, 4, 8, 8, requires_grad=True)
+        query_features = torch.randn(1, 5, 4, requires_grad=True)
+        bias = torch.randn(1, requires_grad=True)
+        outputs = {
+            "pred_masks": {
+                "spatial_features": spatial_features,
+                "query_features": query_features,
+                "bias": bias,
+            }
+        }
+        empty = torch.empty(0, dtype=torch.long)
+        indices = [(empty, empty)]
+
+        losses = criterion.loss_masks(outputs, targets=[{}], indices=indices, num_boxes=1)
+
+        assert losses["loss_mask_ce"].requires_grad
+        (losses["loss_mask_ce"] + losses["loss_mask_dice"]).backward()
+        assert spatial_features.grad is not None
+        assert query_features.grad is not None
+        assert bias.grad is not None

tests/models/test_criterion_keypoints.py

@@ -0,0 +1,161 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for keypoint losses in SetCriterion."""
+
+import torch
+
+from rfdetr.models.criterion import SetCriterion
+
+
+class _MatcherStub:
+    """Matcher stub used to avoid depending on Hungarian matching internals."""
+
+    def __call__(self, outputs, targets, group_detr=1):
+        indices = []
+        for target in targets:
+            num_targets = int(target["labels"].shape[0])
+            idx = torch.arange(num_targets, dtype=torch.int64)
+            indices.append((idx, idx))
+        return indices
+
+
+def _make_outputs(
+    batch_size: int,
+    num_queries: int,
+    num_keypoints: int,
+) -> dict[str, torch.Tensor]:
+    return {
+        "pred_logits": torch.zeros(batch_size, num_queries, 2),
+        "pred_boxes": torch.rand(batch_size, num_queries, 4).clamp(0.05, 0.95),
+        "pred_keypoints": torch.randn(batch_size, num_queries, num_keypoints, 8),
+    }
+
+
+def test_loss_keypoints_list_of_dicts_targets() -> None:
+    """Keypoint loss should consume list-of-dicts targets used by public training."""
+    criterion = SetCriterion(
+        num_classes=2,
+        matcher=_MatcherStub(),
+        weight_dict={},
+        focal_alpha=0.25,
+        losses=["keypoints"],
+        num_keypoints_per_class=[17],
+    )
+    outputs = _make_outputs(batch_size=1, num_queries=1, num_keypoints=17)
+    targets = [
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.4, 0.4]], dtype=torch.float32),
+            "keypoints": torch.cat(
+                [
+                    torch.rand(1, 17, 2),
+                    torch.full((1, 17, 1), 2.0),
+                ],
+                dim=-1,
+            ),
+        }
+    ]
+
+    losses = criterion(outputs, targets)
+
+    assert "loss_keypoints_l1" in losses
+    assert "loss_keypoints_findable" in losses
+    assert "loss_keypoints_visible" in losses
+    assert "loss_keypoints_nll" in losses
+    assert all(torch.isfinite(value) for value in losses.values())
+
+
+def test_loss_keypoints_empty_targets() -> None:
+    """Empty target batches should produce finite zero-valued keypoint losses."""
+    criterion = SetCriterion(
+        num_classes=2,
+        matcher=_MatcherStub(),
+        weight_dict={},
+        focal_alpha=0.25,
+        losses=["keypoints"],
+        num_keypoints_per_class=[17],
+    )
+    outputs = _make_outputs(batch_size=1, num_queries=1, num_keypoints=17)
+    targets = [
+        {
+            "labels": torch.zeros((0,), dtype=torch.int64),
+            "boxes": torch.zeros((0, 4), dtype=torch.float32),
+            "keypoints": torch.zeros((0, 17, 3), dtype=torch.float32),
+        }
+    ]
+
+    losses = criterion(outputs, targets)
+
+    assert losses["loss_keypoints_l1"].item() == 0.0
+    assert losses["loss_keypoints_findable"].item() == 0.0
+    assert losses["loss_keypoints_visible"].item() == 0.0
+    assert losses["loss_keypoints_nll"].item() == 0.0
+
+
+def test_loss_keypoints_person_schema_shape() -> None:
+    """Person-only schema `[17]` should be consumed without shape mismatches."""
+    criterion = SetCriterion(
+        num_classes=2,
+        matcher=_MatcherStub(),
+        weight_dict={},
+        focal_alpha=0.25,
+        losses=["keypoints"],
+        num_keypoints_per_class=[17],
+    )
+    outputs = _make_outputs(batch_size=2, num_queries=2, num_keypoints=17)
+    targets = [
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.4, 0.4]], dtype=torch.float32),
+            "keypoints": torch.rand(1, 17, 3),
+        },
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.4, 0.6, 0.3, 0.5]], dtype=torch.float32),
+            "keypoints": torch.rand(1, 17, 3),
+        },
+    ]
+
+    losses = criterion(outputs, targets)
+
+    assert losses["loss_keypoints_l1"].ndim == 0
+    assert losses["loss_keypoints_findable"].ndim == 0
+    assert losses["loss_keypoints_visible"].ndim == 0
+    assert losses["loss_keypoints_nll"].ndim == 0
+
+
+def test_loss_keypoints_multiclass_schema_kmax_targets() -> None:
+    """Heterogeneous keypoint classes should consume Kmax-padded targets."""
+    criterion = SetCriterion(
+        num_classes=3,
+        matcher=_MatcherStub(),
+        weight_dict={},
+        focal_alpha=0.25,
+        losses=["keypoints"],
+        num_keypoints_per_class=[2, 1],
+    )
+    outputs = _make_outputs(batch_size=1, num_queries=2, num_keypoints=4)
+    targets = [
+        {
+            "labels": torch.tensor([0, 1], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.4, 0.4], [0.4, 0.6, 0.3, 0.5]], dtype=torch.float32),
+            "keypoints": torch.tensor(
+                [
+                    [[0.2, 0.3, 2.0], [0.4, 0.5, 2.0]],
+                    [[0.6, 0.7, 2.0], [0.0, 0.0, 0.0]],
+                ],
+                dtype=torch.float32,
+            ),
+        }
+    ]
+
+    losses = criterion(outputs, targets)
+
+    assert losses["loss_keypoints_l1"].ndim == 0
+    assert losses["loss_keypoints_findable"].ndim == 0
+    assert losses["loss_keypoints_visible"].ndim == 0
+    assert losses["loss_keypoints_nll"].ndim == 0
+    assert all(torch.isfinite(value) for value in losses.values())

tests/models/test_lwdetr_backbone_contract.py

@@ -0,0 +1,60 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests proving LWDETR forward contract survives joiner return-shape changes."""
+
+from unittest.mock import MagicMock
+
+import torch
+
+from rfdetr.models.lwdetr import LWDETR
+from rfdetr.utilities.tensors import NestedTensor
+
+
+def test_lwdetr_default_detection_forward_after_backbone_change() -> None:
+    """LWDETR should accept the updated 3-tuple backbone output in non-keypoint mode."""
+    batch_size = 2
+    num_queries = 3
+    hidden_dim = 4
+    num_classes = 7
+
+    features = [
+        NestedTensor(
+            torch.zeros(batch_size, hidden_dim, 4, 4),
+            torch.zeros(batch_size, 4, 4, dtype=torch.bool),
+        )
+    ]
+    poss = [torch.zeros(batch_size, 4, 4, dtype=torch.bool)]
+
+    backbone = MagicMock()
+    backbone.return_value = (features, poss, None)
+
+    transformer = MagicMock()
+    transformer.d_model = hidden_dim
+    transformer_out = (
+        torch.zeros(1, batch_size, num_queries, hidden_dim),
+        torch.zeros(1, batch_size, num_queries, hidden_dim),
+        torch.zeros(batch_size, num_queries, hidden_dim),
+        torch.zeros(batch_size, num_queries, hidden_dim),
+    )
+    transformer.return_value = transformer_out
+
+    model = LWDETR(
+        backbone=backbone,
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=num_classes,
+        num_queries=num_queries,
+        aux_loss=False,
+        group_detr=1,
+        two_stage=False,
+        lite_refpoint_refine=False,
+        bbox_reparam=False,
+    )
+
+    outputs = model(torch.ones(batch_size, 3, 8, 8))
+
+    assert outputs["pred_logits"].shape == (batch_size, num_queries, num_classes)
+    assert outputs["pred_boxes"].shape == (batch_size, num_queries, 4)

tests/models/test_lwdetr_keypoints.py

@@ -0,0 +1,226 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for GroupPose keypoint output wiring in LWDETR."""
+
+from unittest.mock import MagicMock
+
+import torch
+from torch import nn
+
+from rfdetr.models.heads import ConditionalQueryInitializer
+from rfdetr.models.lwdetr import LWDETR
+from rfdetr.utilities.tensors import NestedTensor
+
+
+def _build_feature_batch(batch_size: int, hidden_dim: int) -> list[NestedTensor]:
+    return [
+        NestedTensor(
+            torch.zeros(batch_size, hidden_dim, 4, 4),
+            torch.zeros(batch_size, 4, 4, dtype=torch.bool),
+        )
+    ]
+
+
+class _DummyKeypointDecoder(nn.Module):
+    """Minimal decoder surface needed for keypoint schema resizing."""
+
+    def __init__(self, hidden_dim: int, num_keypoints_per_class: list[int]) -> None:
+        super().__init__()
+        self.num_keypoints_per_class = num_keypoints_per_class
+        self.keypoint_pos_embed = nn.Parameter(torch.randn(sum(num_keypoints_per_class), hidden_dim))
+        self.register_buffer(
+            "keypoint_class_mask",
+            torch.zeros(1 + sum(num_keypoints_per_class), 1 + sum(num_keypoints_per_class), dtype=torch.bool),
+        )
+
+
+class _DummyKeypointTransformer(nn.Module):
+    """Minimal transformer surface needed for LWDETR construction and keypoint schema resizing."""
+
+    def __init__(self, hidden_dim: int, num_keypoints_per_class: list[int]) -> None:
+        super().__init__()
+        self.d_model = hidden_dim
+        self.num_keypoints_per_class = num_keypoints_per_class
+        self.decoder = _DummyKeypointDecoder(hidden_dim, num_keypoints_per_class)
+        self.keypoint_query_initializer = ConditionalQueryInitializer(hidden_dim, sum(num_keypoints_per_class))
+        self.keypoint_query_initializer_enc = ConditionalQueryInitializer(hidden_dim, sum(num_keypoints_per_class))
+
+
+def test_lwdetr_keypoint_forward_outputs() -> None:
+    """GroupPose mode should expose keypoint tensors in model outputs."""
+    batch_size = 2
+    num_queries = 3
+    hidden_dim = 8
+    num_classes = 6
+
+    features = _build_feature_batch(batch_size=batch_size, hidden_dim=hidden_dim)
+    poss = [torch.zeros(batch_size, hidden_dim, 4, 4)]
+
+    backbone = MagicMock()
+    backbone.return_value = (features, poss, None)
+
+    transformer = MagicMock()
+    transformer.d_model = hidden_dim
+    transformer.return_value = (
+        torch.zeros(2, batch_size, num_queries, hidden_dim),  # hs
+        torch.zeros(2, batch_size, num_queries, 4),  # ref_unsigmoid
+        torch.zeros(batch_size, num_queries, hidden_dim),  # hs_enc
+        torch.zeros(batch_size, num_queries, 4),  # ref_enc
+        torch.zeros(2, batch_size, num_queries, 17, hidden_dim),  # keypoint_hs
+        torch.zeros(batch_size, num_queries, 17, 8),  # enc_kp_predictions
+        torch.zeros(batch_size, num_queries, 17, hidden_dim),  # unused keypoint encoder hidden state
+    )
+
+    model = LWDETR(
+        backbone=backbone,
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=num_classes,
+        num_queries=num_queries,
+        aux_loss=True,
+        group_detr=1,
+        two_stage=False,
+        lite_refpoint_refine=False,
+        bbox_reparam=False,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=[17],
+        grouppose_keypoint_dim_downscale=1,
+    )
+
+    outputs = model(torch.ones(batch_size, 3, 8, 8))
+
+    assert outputs["pred_logits"].shape == (batch_size, num_queries, num_classes)
+    assert outputs["pred_boxes"].shape == (batch_size, num_queries, 4)
+    assert outputs["pred_keypoints"].shape == (batch_size, num_queries, 17, 8)
+    assert "keypoint_hidden_states" not in outputs
+    assert "pred_keypoints" in outputs["aux_outputs"][0]
+    assert "keypoint_hidden_states" not in outputs["aux_outputs"][0]
+
+
+def test_lwdetr_reinitialize_keypoint_head_updates_schema_dependent_state() -> None:
+    """Keypoint schema reinit should resize masks and learned keypoint query embeddings."""
+    hidden_dim = 8
+    transformer = _DummyKeypointTransformer(hidden_dim=hidden_dim, num_keypoints_per_class=[17])
+    model = LWDETR(
+        backbone=MagicMock(),
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=3,
+        num_queries=2,
+        aux_loss=False,
+        group_detr=1,
+        two_stage=True,
+        lite_refpoint_refine=True,
+        bbox_reparam=False,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=[17],
+        grouppose_keypoint_dim_downscale=1,
+    )
+
+    model.reinitialize_keypoint_head([2, 1])
+
+    assert model.num_keypoints_per_class == [2, 1]
+    assert model.get_num_keypoints_per_class() == [2, 1]
+    assert model._kp_active_mask.shape == (2, 2)
+    assert model._kp_active_mask.tolist() == [[True, True], [True, False]]
+    assert transformer.num_keypoints_per_class == [2, 1]
+    assert transformer.decoder.num_keypoints_per_class == [2, 1]
+    assert transformer.decoder.keypoint_pos_embed.shape == (3, hidden_dim)
+    assert transformer.decoder.keypoint_class_mask.shape == (4, 4)
+    assert transformer.keypoint_query_initializer.queries.shape == (3, hidden_dim)
+    assert transformer.keypoint_query_initializer_enc.queries.shape == (3, hidden_dim)
+
+
+def test_lwdetr_reset_keypoint_gaussian_parameters_preserves_non_gaussian_rows() -> None:
+    """Gaussian reset should only zero precision-Cholesky output rows on decoder and encoder keypoint heads."""
+    hidden_dim = 8
+    transformer = _DummyKeypointTransformer(hidden_dim=hidden_dim, num_keypoints_per_class=[17])
+    model = LWDETR(
+        backbone=MagicMock(),
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=3,
+        num_queries=2,
+        aux_loss=False,
+        group_detr=1,
+        two_stage=True,
+        lite_refpoint_refine=True,
+        bbox_reparam=False,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=[17],
+        grouppose_keypoint_dim_downscale=1,
+    )
+    with torch.no_grad():
+        model.keypoint_embed.layers[-1].weight.fill_(3.0)
+        model.keypoint_embed.layers[-1].bias.fill_(4.0)
+        model.transformer.enc_out_keypoint_embed[0].layers[-1].weight.fill_(5.0)
+        model.transformer.enc_out_keypoint_embed[0].layers[-1].bias.fill_(6.0)
+
+    model.reset_keypoint_gaussian_parameters()
+
+    torch.testing.assert_close(model.keypoint_embed.layers[-1].weight[:4], torch.full((4, hidden_dim), 3.0))
+    torch.testing.assert_close(model.keypoint_embed.layers[-1].weight[4:7], torch.zeros(3, hidden_dim))
+    torch.testing.assert_close(model.keypoint_embed.layers[-1].weight[7:], torch.full((1, hidden_dim), 3.0))
+    torch.testing.assert_close(model.keypoint_embed.layers[-1].bias[:4], torch.full((4,), 4.0))
+    torch.testing.assert_close(model.keypoint_embed.layers[-1].bias[4:7], torch.zeros(3))
+    torch.testing.assert_close(model.keypoint_embed.layers[-1].bias[7:], torch.full((1,), 4.0))
+    torch.testing.assert_close(
+        model.transformer.enc_out_keypoint_embed[0].layers[-1].weight[4:7], torch.zeros(3, hidden_dim)
+    )
+    torch.testing.assert_close(model.transformer.enc_out_keypoint_embed[0].layers[-1].bias[4:7], torch.zeros(3))
+
+
+def test_lwdetr_get_num_keypoints_per_class_from_checkpoint() -> None:
+    """Checkpoint keypoint schema should be recoverable from `_kp_active_mask`."""
+    state_dict = {"_kp_active_mask": torch.tensor([[True, True], [True, False]])}
+
+    assert LWDETR.get_num_keypoints_per_class_from_checkpoint(state_dict) == [2, 1]
+
+
+def test_lwdetr_default_detection_contract_unchanged() -> None:
+    """Default detection mode should not expose keypoint outputs."""
+    batch_size = 2
+    num_queries = 3
+    hidden_dim = 8
+    num_classes = 6
+
+    features = _build_feature_batch(batch_size=batch_size, hidden_dim=hidden_dim)
+    poss = [torch.zeros(batch_size, hidden_dim, 4, 4)]
+
+    backbone = MagicMock()
+    backbone.return_value = (features, poss, None)
+
+    transformer = MagicMock()
+    transformer.d_model = hidden_dim
+    transformer.return_value = (
+        torch.zeros(1, batch_size, num_queries, hidden_dim),
+        torch.zeros(1, batch_size, num_queries, 4),
+        torch.zeros(batch_size, num_queries, hidden_dim),
+        torch.zeros(batch_size, num_queries, 4),
+    )
+
+    model = LWDETR(
+        backbone=backbone,
+        transformer=transformer,
+        segmentation_head=None,
+        num_classes=num_classes,
+        num_queries=num_queries,
+        aux_loss=False,
+        group_detr=1,
+        two_stage=False,
+        lite_refpoint_refine=False,
+        bbox_reparam=False,
+        use_grouppose_keypoints=False,
+        num_keypoints_per_class=[],
+        grouppose_keypoint_dim_downscale=1,
+    )
+
+    outputs = model(torch.ones(batch_size, 3, 8, 8))
+
+    assert outputs["pred_logits"].shape == (batch_size, num_queries, num_classes)
+    assert outputs["pred_boxes"].shape == (batch_size, num_queries, 4)
+    assert "pred_keypoints" not in outputs
+    assert "keypoint_hidden_states" not in outputs

tests/models/test_matcher.py

@@ -0,0 +1,401 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+import pytest
+import torch
+
+from rfdetr.models import matcher as matcher_module
+from rfdetr.models.matcher import HungarianMatcher
+
+
+@pytest.fixture()
+def matcher() -> HungarianMatcher:
+    """Shared HungarianMatcher instance."""
+    return HungarianMatcher()
+
+
+@pytest.fixture()
+def standard_target() -> dict[str, torch.Tensor]:
+    """Single-class target with one box at (0.5, 0.5, 0.2, 0.2)."""
+    return {
+        "labels": torch.tensor([0], dtype=torch.int64),
+        "boxes": torch.tensor([[0.5, 0.5, 0.2, 0.2]], dtype=torch.float32),
+    }
+
+
+class TestHungarianMatcherNonFiniteCosts:
+    """Tests for non-finite cost matrix sanitization in the Hungarian matcher."""
+
+    @pytest.mark.parametrize(
+        "invalid_value",
+        [
+            pytest.param(float("nan"), id="nan"),
+            pytest.param(float("inf"), id="inf"),
+            pytest.param(float("-inf"), id="-inf"),
+        ],
+    )
+    def test_replaces_non_finite_costs_before_assignment(
+        self,
+        matcher: HungarianMatcher,
+        standard_target: dict[str, torch.Tensor],
+        invalid_value: float,
+    ) -> None:
+        """Matcher should sanitize non-finite costs so assignment still succeeds."""
+        outputs = {
+            "pred_logits": torch.tensor([[[0.0], [10.0]]], dtype=torch.float32),
+            "pred_boxes": torch.tensor(
+                [
+                    [
+                        [invalid_value, 0.5, 0.2, 0.2],
+                        [0.5, 0.5, 0.2, 0.2],
+                    ]
+                ],
+                dtype=torch.float32,
+            ),
+        }
+
+        matched_queries, matched_targets = matcher(outputs, [standard_target])[0]
+
+        assert matched_queries.tolist() == [1]
+        assert matched_targets.tolist() == [0]
+
+    def test_all_nonfinite_produces_valid_assignment(
+        self,
+        matcher: HungarianMatcher,
+        standard_target: dict[str, torch.Tensor],
+    ) -> None:
+        """When ALL costs are non-finite, the fallback sentinel (``dtype_info.max``)
+        should allow ``linear_sum_assignment`` to complete with a valid 1-to-1
+        assignment: exactly one match, query index in [0, num_queries), target index 0.
+
+        This exercises the ``else: replacement_cost = C.new_tensor(dtype_info.max)`` branch.
+        """
+        nan = float("nan")
+        outputs = {
+            "pred_logits": torch.tensor([[[nan], [nan]]], dtype=torch.float32),
+            "pred_boxes": torch.tensor(
+                [
+                    [
+                        [nan, nan, nan, nan],
+                        [nan, nan, nan, nan],
+                    ]
+                ],
+                dtype=torch.float32,
+            ),
+        }
+
+        matched_queries, matched_targets = matcher(outputs, [standard_target])[0]
+
+        assert len(matched_queries) == len(matched_targets) == 1
+        assert 0 <= matched_queries.item() < 2
+        assert matched_targets.item() == 0
+
+    def test_negative_costs_with_nan_selects_valid_query(
+        self,
+        matcher: HungarianMatcher,
+        standard_target: dict[str, torch.Tensor],
+    ) -> None:
+        """Regression test: when all finite costs are negative and one query produces NaN, the matcher must select the
+        valid query, not the NaN one.
+
+        This guards against the bug where ``max_cost * 2`` (the old replacement formula) could be smaller than
+        ``max_cost`` when all costs are negative, causing the NaN query to appear cheaper than valid queries.
+        """
+        nan = float("nan")
+        # Query 0: NaN box coordinates -> produces non-finite costs
+        # Query 1: valid box, low logit -> all-negative but finite costs
+        outputs = {
+            "pred_logits": torch.tensor([[[0.0], [-10.0]]], dtype=torch.float32),
+            "pred_boxes": torch.tensor(
+                [
+                    [
+                        [nan, nan, nan, nan],
+                        [0.5, 0.5, 0.2, 0.2],
+                    ]
+                ],
+                dtype=torch.float32,
+            ),
+        }
+
+        matched_queries, matched_targets = matcher(outputs, [standard_target])[0]
+
+        # The valid query (index 1) must be matched, not the NaN query.
+        assert matched_queries.tolist() == [1]
+        assert matched_targets.tolist() == [0]
+
+    @pytest.mark.parametrize(
+        "image_idx, expected_query_idx",
+        [
+            pytest.param(0, 1, id="image0"),
+            pytest.param(1, 0, id="image1"),
+        ],
+    )
+    def test_batch_size_greater_than_one(
+        self,
+        matcher: HungarianMatcher,
+        image_idx: int,
+        expected_query_idx: int,
+    ) -> None:
+        """Exercises the ``C.split(sizes, -1)`` loop with batch_size > 1.
+
+        Each image has 2 queries and 1 target. One query per image has NaN coordinates; the matcher must select the
+        valid query in each case.
+        """
+        nan = float("nan")
+        outputs = {
+            "pred_logits": torch.tensor(
+                [
+                    [[0.0], [10.0]],  # image 0: query 1 is valid
+                    [[10.0], [0.0]],  # image 1: query 0 is valid
+                ],
+                dtype=torch.float32,
+            ),
+            "pred_boxes": torch.tensor(
+                [
+                    [
+                        [nan, 0.5, 0.2, 0.2],  # image 0, query 0: NaN
+                        [0.5, 0.5, 0.2, 0.2],  # image 0, query 1: valid
+                    ],
+                    [
+                        [0.5, 0.5, 0.2, 0.2],  # image 1, query 0: valid
+                        [nan, 0.5, 0.2, 0.2],  # image 1, query 1: NaN
+                    ],
+                ],
+                dtype=torch.float32,
+            ),
+        }
+        targets = [
+            {
+                "labels": torch.tensor([0], dtype=torch.int64),
+                "boxes": torch.tensor([[0.5, 0.5, 0.2, 0.2]], dtype=torch.float32),
+            },
+            {
+                "labels": torch.tensor([0], dtype=torch.int64),
+                "boxes": torch.tensor([[0.5, 0.5, 0.2, 0.2]], dtype=torch.float32),
+            },
+        ]
+
+        results = matcher(outputs, targets)
+
+        assert len(results) == 2
+
+        matched_queries, matched_targets = results[image_idx]
+        assert matched_queries.tolist() == [expected_query_idx]
+        assert matched_targets.tolist() == [0]
+
+    def test_group_detr_with_nonfinite_costs(
+        self,
+        matcher: HungarianMatcher,
+        standard_target: dict[str, torch.Tensor],
+    ) -> None:
+        """Sanitization runs on the full cost matrix before splitting by group, so non-finite entries must be handled
+        correctly when ``group_detr > 1``.
+
+        4 queries, 2 groups of 2. Query 0 has a NaN box; query 2 (the best valid match in group 1) must be selected
+        across groups.
+        """
+        nan = float("nan")
+        outputs = {
+            "pred_logits": torch.tensor(
+                [[[0.0], [10.0], [0.0], [10.0]]],
+                dtype=torch.float32,
+            ),
+            "pred_boxes": torch.tensor(
+                [
+                    [
+                        [nan, nan, nan, nan],  # group 0, query 0: NaN
+                        [0.5, 0.5, 0.2, 0.2],  # group 0, query 1: valid
+                        [nan, nan, nan, nan],  # group 1, query 0: NaN
+                        [0.5, 0.5, 0.2, 0.2],  # group 1, query 1: valid
+                    ]
+                ],
+                dtype=torch.float32,
+            ),
+        }
+
+        results = matcher(outputs, [standard_target], group_detr=2)
+
+        assert len(results) == 1
+        matched_queries, matched_targets = results[0]
+        # Each group contributes one match; both must map to target 0
+        assert matched_targets.tolist() == [0, 0]
+        # The valid query in each group (indices 1 and 3) must be selected
+        assert set(matched_queries.tolist()) == {1, 3}
+
+    def test_warns_once_per_matcher_instance(
+        self, standard_target: dict[str, torch.Tensor], monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        """Non-finite-cost warning should be emitted once per matcher instance."""
+        expected_warning = (
+            "Non-finite values detected in matcher cost matrix; "
+            "replacing with finite sentinel. "
+            "Check for numerical instability."
+        )
+        warning_messages: list[str] = []
+
+        def record_warning(msg: str, *args: object, **kwargs: object) -> None:
+            warning_messages.append(msg)
+
+        monkeypatch.setattr(matcher_module.logger, "warning", record_warning)
+
+        outputs = {
+            "pred_logits": torch.tensor([[[0.0], [10.0]]], dtype=torch.float32),
+            "pred_boxes": torch.tensor(
+                [
+                    [
+                        [float("nan"), 0.5, 0.2, 0.2],
+                        [0.5, 0.5, 0.2, 0.2],
+                    ]
+                ],
+                dtype=torch.float32,
+            ),
+        }
+
+        first_matcher = HungarianMatcher()
+        second_matcher = HungarianMatcher()
+
+        first_matcher(outputs, [standard_target])
+        first_matcher(outputs, [standard_target])
+        second_matcher(outputs, [standard_target])
+
+        assert warning_messages == [expected_warning, expected_warning]
+
+
+class TestHungarianMatcherSanitization:
+    """Unit tests for the private matcher cost sanitization helper."""
+
+    def test_sanitize_cost_matrix_replaces_non_finite_entries(self) -> None:
+        """Non-finite entries should be replaced with a larger finite sentinel."""
+        cost_matrix = torch.tensor(
+            [
+                [1.0, float("nan")],
+                [float("inf"), -2.0],
+            ],
+            dtype=torch.float32,
+        )
+
+        sanitized = HungarianMatcher._sanitize_cost_matrix(cost_matrix)
+
+        assert torch.isfinite(sanitized).all()
+        assert sanitized[0, 1] == 4.0
+        assert sanitized[1, 0] == 4.0
+        assert sanitized[0, 0] == 1.0
+        assert sanitized[1, 1] == -2.0
+
+    def test_sanitize_cost_matrix_all_non_finite_fallback(self) -> None:
+        """All-non-finite matrices should fall back to the dtype maximum."""
+        cost_matrix = torch.tensor(
+            [
+                [float("nan"), float("inf")],
+                [float("-inf"), float("nan")],
+            ],
+            dtype=torch.float32,
+        )
+
+        sanitized = HungarianMatcher._sanitize_cost_matrix(cost_matrix)
+
+        assert torch.isfinite(sanitized).all()
+        assert torch.all(sanitized == torch.finfo(cost_matrix.dtype).max)
+
+    def test_sanitize_cost_matrix_clamps_overflowing_replacement_cost(self) -> None:
+        """Overflow in the computed replacement cost should clamp to dtype max."""
+        dtype_max = torch.finfo(torch.float32).max
+        cost_matrix = torch.tensor(
+            [
+                [dtype_max, float("nan")],
+                [0.0, 1.0],
+            ],
+            dtype=torch.float32,
+        )
+
+        sanitized = HungarianMatcher._sanitize_cost_matrix(cost_matrix)
+
+        assert torch.isfinite(sanitized).all()
+        assert sanitized[0, 1] == dtype_max
+
+
+class TestHungarianMatcherFocalAlpha:
+    """The configured ``focal_alpha`` must drive the classification matching cost."""
+
+    def test_focal_alpha_changes_assignment(self) -> None:
+        """Two matchers differing only in ``focal_alpha`` must be able to produce different assignments.
+
+        ``focal_alpha`` is accepted, documented as "used in the classification cost", and stored on the matcher, so it
+        must actually influence matching. This input is chosen so the optimal query->target pairing flips between
+        ``focal_alpha=0.25`` and ``focal_alpha=0.90``; if the cost ignores the configured alpha, both assignments
+        collapse to the same result.
+        """
+        outputs = {
+            "pred_logits": torch.tensor(
+                [[[2.3936, -1.4217], [2.3731, -2.1974]]],
+                dtype=torch.float32,
+            ),
+            "pred_boxes": torch.tensor(
+                [[[0.3898, 0.4340, 0.5331, 0.1901], [0.4256, 0.1002, 0.6955, 0.7815]]],
+                dtype=torch.float32,
+            ),
+        }
+        targets = [
+            {
+                "labels": torch.tensor([0, 1], dtype=torch.int64),
+                "boxes": torch.tensor(
+                    [[0.2111, 0.6630, 0.7569, 0.8855], [0.7750, 0.4393, 0.8838, 0.8792]],
+                    dtype=torch.float32,
+                ),
+            }
+        ]
+
+        def assignment(focal_alpha: float) -> list[int]:
+            matcher = HungarianMatcher(cost_class=2.0, cost_bbox=5.0, cost_giou=2.0, focal_alpha=focal_alpha)
+            matched_queries, matched_targets = matcher(outputs, targets)[0]
+            # Queries ordered by the target index they are matched to.
+            return matched_queries[matched_targets.argsort()].tolist()
+
+        assert assignment(0.25) != assignment(0.90)
+        # Pin the exact expected mappings so a misapplied-alpha refactor is caught even when
+        # the two values remain different for unrelated reasons.
+        assert assignment(0.25) == [0, 1]
+        assert assignment(0.90) == [1, 0]
+
+    @pytest.mark.parametrize(
+        "focal_alpha, expected",
+        [
+            pytest.param(0.0, [0, 1], id="alpha_zero_pos_cost_zeroed"),
+            pytest.param(1.0, [1, 0], id="alpha_one_neg_cost_zeroed"),
+        ],
+    )
+    def test_focal_alpha_boundary_values_no_nan(self, focal_alpha: float, expected: list[int]) -> None:
+        """Degenerate focal_alpha values (0.0 and 1.0) must not produce NaN and must yield a valid assignment.
+
+        focal_alpha=0.0 zeroes ``pos_cost_class``; focal_alpha=1.0 zeroes ``neg_cost_class``. Neither path touches
+        ``log(prob)`` directly (formula uses logsigmoid of logits), so no division-by-zero or NaN can occur.
+        """
+        outputs = {
+            "pred_logits": torch.tensor(
+                [[[2.3936, -1.4217], [2.3731, -2.1974]]],
+                dtype=torch.float32,
+            ),
+            "pred_boxes": torch.tensor(
+                [[[0.3898, 0.4340, 0.5331, 0.1901], [0.4256, 0.1002, 0.6955, 0.7815]]],
+                dtype=torch.float32,
+            ),
+        }
+        targets = [
+            {
+                "labels": torch.tensor([0, 1], dtype=torch.int64),
+                "boxes": torch.tensor(
+                    [[0.2111, 0.6630, 0.7569, 0.8855], [0.7750, 0.4393, 0.8838, 0.8792]],
+                    dtype=torch.float32,
+                ),
+            }
+        ]
+
+        matcher = HungarianMatcher(cost_class=2.0, cost_bbox=5.0, cost_giou=2.0, focal_alpha=focal_alpha)
+        matched_queries, matched_targets = matcher(outputs, targets)[0]
+
+        assert not matcher._warned_non_finite_costs, "boundary focal_alpha produced non-finite costs"
+        result = matched_queries[matched_targets.argsort()].tolist()
+        assert result == expected

tests/models/test_matcher_keypoints.py

@@ -0,0 +1,112 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for keypoint matching costs in HungarianMatcher."""
+
+import torch
+
+from rfdetr.models.matcher import HungarianMatcher
+
+
+def _base_outputs(num_queries: int = 2) -> dict[str, torch.Tensor]:
+    """Build minimal detection outputs used across matcher keypoint tests."""
+    pred_logits = torch.full((1, num_queries, 1), 5.0, dtype=torch.float32)
+    pred_boxes = torch.tensor([0.5, 0.5, 0.2, 0.2], dtype=torch.float32).view(1, 1, 4).repeat(1, num_queries, 1)
+    return {
+        "pred_logits": pred_logits,
+        "pred_boxes": pred_boxes,
+    }
+
+
+def test_matcher_keypoint_cost_list_of_dicts_targets() -> None:
+    """Keypoint matching costs should work with public list-of-dicts targets."""
+    matcher = HungarianMatcher(
+        cost_class=0.0,
+        cost_bbox=1.0,
+        cost_giou=0.0,
+        num_keypoints_per_class=[1],
+        keypoint_l1_loss_coef=10.0,
+        keypoint_findable_loss_coef=0.0,
+        keypoint_visible_loss_coef=0.0,
+        keypoint_nll_loss_coef=0.0,
+    )
+    outputs = _base_outputs()
+    outputs["pred_keypoints"] = torch.zeros((1, 2, 1, 8), dtype=torch.float32)
+    outputs["pred_keypoints"][0, 0, 0, :2] = torch.tensor([0.5, 0.5], dtype=torch.float32)
+    outputs["pred_keypoints"][0, 1, 0, :2] = torch.tensor([0.0, 0.0], dtype=torch.float32)
+    targets = [
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.2, 0.2]], dtype=torch.float32),
+            "keypoints": torch.tensor([[[0.5, 0.5, 2.0]]], dtype=torch.float32),
+        }
+    ]
+
+    matched_queries, matched_targets = matcher(outputs, targets)[0]
+
+    assert matched_queries.tolist() == [0]
+    assert matched_targets.tolist() == [0]
+
+
+def test_matcher_keypoint_cost_coefficients_off() -> None:
+    """Zero keypoint coefficients should preserve non-keypoint matching behavior."""
+    base_matcher = HungarianMatcher(cost_class=1.0, cost_bbox=1.0, cost_giou=1.0)
+    keypoint_matcher = HungarianMatcher(
+        cost_class=1.0,
+        cost_bbox=1.0,
+        cost_giou=1.0,
+        num_keypoints_per_class=[1],
+        keypoint_l1_loss_coef=0.0,
+        keypoint_findable_loss_coef=0.0,
+        keypoint_visible_loss_coef=0.0,
+        keypoint_nll_loss_coef=0.0,
+    )
+    outputs = _base_outputs()
+    outputs["pred_logits"][0, 0, 0] = 10.0
+    outputs["pred_logits"][0, 1, 0] = -10.0
+    outputs["pred_boxes"][0, 1, :] = torch.tensor([0.1, 0.1, 0.1, 0.1], dtype=torch.float32)
+    targets = [
+        {
+            "labels": torch.tensor([0], dtype=torch.int64),
+            "boxes": torch.tensor([[0.5, 0.5, 0.2, 0.2]], dtype=torch.float32),
+            "keypoints": torch.tensor([[[0.0, 0.0, 2.0]]], dtype=torch.float32),
+        }
+    ]
+    outputs_with_keypoints = dict(outputs)
+    outputs_with_keypoints["pred_keypoints"] = torch.zeros((1, 2, 1, 8), dtype=torch.float32)
+
+    base_indices = base_matcher(outputs, targets)[0]
+    keypoint_indices = keypoint_matcher(outputs_with_keypoints, targets)[0]
+
+    assert base_indices[0].tolist() == keypoint_indices[0].tolist()
+    assert base_indices[1].tolist() == keypoint_indices[1].tolist()
+
+
+def test_matcher_keypoint_empty_targets() -> None:
+    """Empty keypoint targets should return valid empty match results."""
+    matcher = HungarianMatcher(
+        cost_class=1.0,
+        cost_bbox=1.0,
+        cost_giou=1.0,
+        num_keypoints_per_class=[1],
+        keypoint_l1_loss_coef=1.0,
+        keypoint_findable_loss_coef=1.0,
+        keypoint_visible_loss_coef=1.0,
+        keypoint_nll_loss_coef=1.0,
+    )
+    outputs = _base_outputs(num_queries=3)
+    outputs["pred_keypoints"] = torch.zeros((1, 3, 1, 8), dtype=torch.float32)
+    targets = [
+        {
+            "labels": torch.zeros((0,), dtype=torch.int64),
+            "boxes": torch.zeros((0, 4), dtype=torch.float32),
+            "keypoints": torch.zeros((0, 1, 3), dtype=torch.float32),
+        }
+    ]
+
+    matched_queries, matched_targets = matcher(outputs, targets)[0]
+
+    assert matched_queries.numel() == 0
+    assert matched_targets.numel() == 0

tests/models/test_math.py

@@ -0,0 +1,102 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for rfdetr.models.math utility functions."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from rfdetr.models.math import accuracy, interpolate, inverse_sigmoid
+
+
+class TestInterpolate:
+    """Verify interpolate() delegates to F.interpolate across torchvision versions."""
+
+    def test_resizes_to_target_size(self) -> None:
+        """Interpolate() upsamples a 4-D tensor to the requested spatial size."""
+        x = torch.randn(2, 3, 4, 4)
+
+        out = interpolate(x, size=[8, 8], mode="bilinear", align_corners=False)
+
+        assert out.shape == (2, 3, 8, 8)
+
+    def test_handles_empty_batch(self) -> None:
+        """Interpolate() supports an empty batch dimension without error."""
+        x = torch.randn(0, 3, 4, 4)
+
+        out = interpolate(x, size=[8, 8], mode="nearest")
+
+        assert out.shape == (0, 3, 8, 8)
+
+
+class TestAccuracy:
+    """Verify accuracy() computes precision@k correctly."""
+
+    def test_top1_perfect_batch(self) -> None:
+        """All predictions correct returns top-1 accuracy of 100.0."""
+        output = torch.tensor([[0.0, 10.0], [10.0, 0.0], [0.0, 10.0]])
+        target = torch.tensor([1, 0, 1])
+        result = accuracy(output, target, topk=(1,))
+        assert len(result) == 1
+        assert result[0].item() == pytest.approx(100.0)
+
+    def test_top1_zero_accuracy(self) -> None:
+        """All predictions wrong returns top-1 accuracy of 0.0."""
+        output = torch.tensor([[10.0, 0.0], [0.0, 10.0]])
+        target = torch.tensor([1, 0])
+        result = accuracy(output, target, topk=(1,))
+        assert result[0].item() == pytest.approx(0.0)
+
+    def test_topk_returns_list_of_correct_length(self) -> None:
+        """Topk=(1, 5) returns a list of length 2."""
+        output = torch.randn(10, 10)
+        target = torch.zeros(10, dtype=torch.long)
+        result = accuracy(output, target, topk=(1, 5))
+        assert len(result) == 2
+
+    def test_empty_target_returns_single_zero_regardless_of_topk(self) -> None:
+        """Empty target returns list of length 1 with value 0 regardless of topk length."""
+        output = torch.zeros(0, 5)
+        target = torch.zeros(0, dtype=torch.long)
+        result = accuracy(output, target, topk=(1, 5))
+        assert len(result) == 1
+        assert result[0].item() == pytest.approx(0.0)
+
+
+class TestInverseSigmoid:
+    """Verify inverse_sigmoid() computes the logit function correctly."""
+
+    def test_identity_at_half(self) -> None:
+        """inverse_sigmoid(0.5) equals 0.0 since sigmoid(0.0) = 0.5."""
+        x = torch.tensor([0.5])
+        result = inverse_sigmoid(x)
+        assert result.item() == pytest.approx(0.0, abs=1e-5)
+
+    def test_clamping_at_zero_is_finite(self) -> None:
+        """inverse_sigmoid(0.0) is finite due to eps clamping."""
+        x = torch.tensor([0.0])
+        result = inverse_sigmoid(x)
+        assert torch.isfinite(result).all()
+
+    def test_clamping_at_one_is_finite(self) -> None:
+        """inverse_sigmoid(1.0) is finite due to eps clamping."""
+        x = torch.tensor([1.0])
+        result = inverse_sigmoid(x)
+        assert torch.isfinite(result).all()
+
+    def test_output_shape_matches_input(self) -> None:
+        """Output shape matches input shape for a multi-dimensional tensor."""
+        x = torch.rand(3, 4)
+        result = inverse_sigmoid(x)
+        assert result.shape == x.shape
+
+    def test_gradient_flows_for_non_saturated_input(self) -> None:
+        """Gradients are non-zero for a non-saturated input value."""
+        x = torch.tensor([0.3], requires_grad=True)
+        inverse_sigmoid(x).sum().backward()
+        assert x.grad is not None
+        assert x.grad.abs().item() > 0.0

tests/models/test_model.py

@@ -0,0 +1,54 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+import pytest
+import torch
+
+from rfdetr import RFDETRBase, RFDETRLarge
+
+
+def _get_patch_embed_projection(model) -> torch.nn.Conv2d:
+    """Return the patch-embedding projection layer for an RF-DETR model.
+
+    RFDETR wrappers are not nn.Module; the underlying PyTorch module lives at ``model.model.model``.  Walk
+    named_modules() on that object.
+
+    Args:
+        model: Instantiated RF-DETR wrapper (RFDETRBase / RFDETRLarge).
+
+    Returns:
+        The convolution used to project image channels into patch embeddings.
+
+    Raises:
+        AssertionError: If the patch-embedding projection cannot be located.
+    """
+    # model.model → model context; model.model.model → nn.Module
+    nn_model = model.model.model
+    proj = nn_model.backbone[0].encoder.encoder.embeddings.patch_embeddings.projection
+    if isinstance(proj, torch.nn.Conv2d):
+        return proj
+
+    # Fallback: scan named_modules on the underlying nn.Module
+    for name, module in nn_model.named_modules():
+        if "patch_embeddings" in name and "projection" in name and isinstance(module, torch.nn.Conv2d):
+            return module
+
+    msg = "Could not find patch embedding projection on model"
+    raise AssertionError(msg)
+
+
+@pytest.mark.parametrize("model_class", [RFDETRBase, RFDETRLarge])
+@pytest.mark.parametrize("channels", [1, 4])
+def test_multispectral_support(model_class, channels: int) -> None:
+    model = model_class(
+        num_channels=channels,
+        device="cpu",
+        pretrain_weights=None,
+    )
+
+    patch_embed_projection = _get_patch_embed_projection(model)
+
+    assert patch_embed_projection.in_channels == channels

tests/models/test_namespace.py

@@ -0,0 +1,189 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests for _namespace_from_configs() config forwarding."""
+
+import sys
+from typing import Any
+
+import pytest
+
+from rfdetr._namespace import _namespace_from_configs
+from rfdetr.config import RFDETRBaseConfig, RFDETRSegNanoConfig, SegmentationTrainConfig, TrainConfig
+from rfdetr.models._types import BuilderArgs
+
+
+class TestNamespaceForwarding:
+    """Verify that _namespace_from_configs() forwards TrainConfig fields that were previously hardcoded to wrong
+    defaults."""
+
+    def _make_ns(self: "TestNamespaceForwarding", **tc_kwargs: Any) -> Any:
+        """Build a namespace for tests with minimal default TrainConfig values."""
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc_kwargs.setdefault("dataset_dir", "/tmp")
+        tc = TrainConfig(**tc_kwargs)
+        return _namespace_from_configs(mc, tc)
+
+    def test_aug_config_forwarded_when_set(self: "TestNamespaceForwarding") -> None:
+        aug = {"hsv_h": 0.015, "hsv_s": 0.7}
+        ns = self._make_ns(aug_config=aug)
+        assert ns.aug_config == aug
+
+    def test_aug_config_none_by_default(self: "TestNamespaceForwarding") -> None:
+        ns = self._make_ns()
+        assert ns.aug_config is None
+
+    def test_use_ema_forwarded_true(self: "TestNamespaceForwarding") -> None:
+        ns = self._make_ns(use_ema=True)
+        assert ns.use_ema is True
+
+    def test_use_ema_forwarded_false(self: "TestNamespaceForwarding") -> None:
+        ns = self._make_ns(use_ema=False)
+        assert ns.use_ema is False
+
+    def test_early_stopping_use_ema_forwarded_true(self: "TestNamespaceForwarding") -> None:
+        ns = self._make_ns(early_stopping_use_ema=True)
+        assert ns.early_stopping_use_ema is True
+
+    def test_early_stopping_use_ema_forwarded_false(self: "TestNamespaceForwarding") -> None:
+        ns = self._make_ns(early_stopping_use_ema=False)
+        assert ns.early_stopping_use_ema is False
+
+
+class TestNamespaceProtocol:
+    """_namespace_from_configs() output must satisfy the BuilderArgs Protocol."""
+
+    def _make_ns(self, mc=None, tc=None):
+        mc = mc or RFDETRBaseConfig(num_classes=80)
+        tc = tc or TrainConfig(dataset_dir="/tmp")
+        return _namespace_from_configs(mc, tc)
+
+    @pytest.mark.skipif(
+        sys.version_info < (3, 12),
+        reason="Runtime Protocol attribute checks require Python 3.12+",
+    )
+    def test_namespace_satisfies_builderargs_protocol_py312(self) -> None:
+        """On Python 3.12+, isinstance() verifies data-attribute presence."""
+        ns = self._make_ns()
+        assert isinstance(ns, BuilderArgs)
+
+    def test_namespace_is_builderargs_instance(self) -> None:
+        """Isinstance() check passes on all supported Python versions.
+
+        On Python 3.10/3.11 this is a structural no-op (no method members to check).  On 3.12+ it verifies attribute
+        presence.  The test documents the intent regardless of Python version.
+        """
+        ns = self._make_ns()
+        assert isinstance(ns, BuilderArgs)
+
+
+class TestNamespaceFieldOwnership:
+    """Verify that the namespace reads each field from the authoritative owner."""
+
+    def _make_ns(self, mc=None, tc=None):
+        mc = mc or RFDETRBaseConfig(num_classes=80)
+        tc = tc or TrainConfig(dataset_dir="/tmp")
+        return _namespace_from_configs(mc, tc)
+
+    # --- cls_loss_coef must come from TrainConfig ---
+
+    def test_cls_loss_coef_from_train_config(self) -> None:
+        """ns.cls_loss_coef must reflect TrainConfig.cls_loss_coef, not ModelConfig."""
+        mc = RFDETRBaseConfig(num_classes=80)  # cls_loss_coef=1.0 (ModelConfig default)
+        tc = TrainConfig(dataset_dir="/tmp", cls_loss_coef=2.5)
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.cls_loss_coef == pytest.approx(2.5)
+
+    def test_cls_loss_coef_segmentation_default_matches_pre_1_7_effective_value(self) -> None:
+        """SegmentationTrainConfig default must preserve the pre-1.7 effective loss_ce weight."""
+        mc = RFDETRSegNanoConfig()
+        tc = SegmentationTrainConfig(dataset_dir="/tmp")
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.cls_loss_coef == pytest.approx(1.0)
+
+    def test_cls_loss_coef_segmentation_explicit_train_config_value_wins(self) -> None:
+        """Explicit SegmentationTrainConfig.cls_loss_coef values must propagate to namespace."""
+        mc = RFDETRSegNanoConfig()
+        tc = SegmentationTrainConfig(dataset_dir="/tmp", cls_loss_coef=5.0)
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.cls_loss_coef == pytest.approx(5.0)
+
+    def test_cls_loss_coef_train_config_wins_over_explicit_model_config(self) -> None:
+        """When both are explicitly set, TrainConfig.cls_loss_coef takes precedence."""
+        with pytest.warns(DeprecationWarning, match="ModelConfig\\.cls_loss_coef is deprecated"):
+            mc = RFDETRBaseConfig(num_classes=80, cls_loss_coef=0.5)
+        tc = TrainConfig(dataset_dir="/tmp", cls_loss_coef=3.0)
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.cls_loss_coef == pytest.approx(3.0)
+
+    def test_cls_loss_coef_model_config_explicit_is_preserved_during_deprecation(self) -> None:
+        """Explicit ModelConfig.cls_loss_coef remains effective until removal."""
+        with pytest.warns(DeprecationWarning, match="ModelConfig\\.cls_loss_coef is deprecated"):
+            mc = RFDETRBaseConfig(num_classes=80, cls_loss_coef=2.5)
+        tc = TrainConfig(dataset_dir="/tmp")
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.cls_loss_coef == pytest.approx(2.5)
+
+    # --- num_select must come from ModelConfig unconditionally ---
+
+    def test_num_select_from_model_config(self) -> None:
+        """ns.num_select must equal mc.num_select regardless of tc.num_select."""
+        mc = RFDETRSegNanoConfig()  # num_select=100
+        tc = TrainConfig(dataset_dir="/tmp")  # num_select=300 (default — was the bug)
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.num_select == 100
+
+    @pytest.mark.parametrize(
+        "config_class, expected_num_select",
+        [
+            pytest.param(RFDETRSegNanoConfig, 100, id="seg_nano"),
+            pytest.param(RFDETRBaseConfig, 300, id="base"),
+        ],
+    )
+    def test_num_select_matches_model_config_variant(self, config_class, expected_num_select) -> None:
+        """ns.num_select must equal the model config's num_select for each variant."""
+        mc = config_class()
+        tc = TrainConfig(dataset_dir="/tmp")
+        ns = _namespace_from_configs(mc, tc)
+        assert ns.num_select == expected_num_select
+
+
+class TestBuildNamespaceDeprecated:
+    """build_namespace() is a deprecated shim — verify the warning fires."""
+
+    def test_emits_deprecation_warning(self, reset_build_namespace_warning_state) -> None:
+        """Every call to build_namespace() must emit a DeprecationWarning."""
+        from rfdetr._namespace import build_namespace
+
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc = TrainConfig(dataset_dir="/tmp")
+
+        with pytest.warns(FutureWarning, match="build_namespace"):
+            build_namespace(mc, tc)
+
+    def test_result_identical_to_namespace_from_configs(self, reset_build_namespace_warning_state) -> None:
+        """build_namespace output must equal _namespace_from_configs output."""
+        from rfdetr._namespace import build_namespace
+        from rfdetr.models._defaults import MODEL_DEFAULTS
+
+        mc = RFDETRBaseConfig(num_classes=80)
+        tc = TrainConfig(dataset_dir="/tmp")
+
+        with pytest.warns(FutureWarning):
+            ns_legacy = build_namespace(mc, tc)
+        ns_new = _namespace_from_configs(mc, tc, MODEL_DEFAULTS)
+
+        legacy_attrs = vars(ns_legacy)
+        new_attrs = vars(ns_new)
+
+        assert set(legacy_attrs.keys()) == set(new_attrs.keys()), (
+            f"Key mismatch: "
+            f"legacy_only={set(legacy_attrs) - set(new_attrs)}, "
+            f"new_only={set(new_attrs) - set(legacy_attrs)}"
+        )
+        for key in sorted(legacy_attrs):
+            assert legacy_attrs[key] == new_attrs[key], (
+                f"Value mismatch for '{key}': legacy={legacy_attrs[key]!r}, new={new_attrs[key]!r}"
+            )

tests/models/test_position_encoding.py

@@ -0,0 +1,88 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for rfdetr.models.position_encoding.build_position_encoding.
+
+Covers:
+- Supported aliases ``"sine"`` / ``"v2"`` return a ``PositionEmbeddingSine`` instance.
+- Unsupported but previously accepted aliases ``"learned"`` / ``"v3"`` now raise
+  ``ValueError`` with a message that names supported alternatives.
+- Fully unsupported values raise ``ValueError`` with the same pattern.
+"""
+
+import pytest
+
+from rfdetr.models.position_encoding import (
+    PositionEmbeddingSine,
+    build_position_encoding,
+)
+
+
+class TestBuildPositionEncodingSupportedValues:
+    """build_position_encoding returns valid modules for supported aliases."""
+
+    @pytest.mark.parametrize(
+        "alias",
+        [
+            pytest.param("sine", id="sine"),
+            pytest.param("v2", id="v2"),
+        ],
+    )
+    def test_returns_sine_embedding(self, alias: str) -> None:
+        """Supported aliases produce a PositionEmbeddingSine with normalized=True."""
+        enc = build_position_encoding(hidden_dim=256, position_embedding=alias)
+        assert isinstance(enc, PositionEmbeddingSine)
+        assert enc.normalize is True
+
+    @pytest.mark.parametrize(
+        "hidden_dim, expected_num_pos_feats",
+        [
+            pytest.param(256, 128, id="dim256"),
+            pytest.param(512, 256, id="dim512"),
+        ],
+    )
+    def test_num_pos_feats_is_half_hidden_dim(self, hidden_dim: int, expected_num_pos_feats: int) -> None:
+        """The sine encoding uses hidden_dim // 2 positional feature dimensions."""
+        enc = build_position_encoding(hidden_dim=hidden_dim, position_embedding="sine")
+        assert enc.num_pos_feats == expected_num_pos_feats
+
+
+class TestBuildPositionEncodingUnsupportedValues:
+    """build_position_encoding raises ValueError for broken or unknown aliases."""
+
+    @pytest.mark.parametrize(
+        "alias",
+        [
+            pytest.param("learned", id="learned"),
+            pytest.param("v3", id="v3"),
+        ],
+    )
+    def test_learned_raises_value_error(self, alias: str) -> None:
+        """'learned' and 'v3' are doubly broken and must raise ValueError immediately.
+
+        The PositionEmbeddingLearned class has two bugs:
+        1. forward() signature is incompatible with Joiner.forward() (no align_dim_orders param).
+        2. h, w = x.shape[:2] unpacks batch and channels instead of height and width.
+        Rejecting them at build time is preferable to a silent or confusing runtime failure.
+        """
+        with pytest.raises(ValueError, match="not supported"):
+            build_position_encoding(hidden_dim=256, position_embedding=alias)
+
+    def test_unknown_value_raises_value_error(self) -> None:
+        """A fully unknown alias raises ValueError naming the supported alternatives."""
+        with pytest.raises(ValueError, match="not supported"):
+            build_position_encoding(hidden_dim=256, position_embedding="unknown_variant")
+
+    @pytest.mark.parametrize(
+        "alias",
+        [
+            pytest.param("learned", id="learned"),
+            pytest.param("v3", id="v3"),
+        ],
+    )
+    def test_error_message_mentions_supported_alternatives(self, alias: str) -> None:
+        """Error message for 'learned'/'v3' mentions at least one supported alternative."""
+        with pytest.raises(ValueError, match="sine"):
+            build_position_encoding(hidden_dim=256, position_embedding=alias)

tests/models/test_postprocess.py

@@ -0,0 +1,170 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for PostProcess box clamping behaviour."""
+
+import pytest
+import torch
+
+from rfdetr.models.postprocess import PostProcess
+
+
+class TestGatherAndScaleBoxes:
+    """Tests for :meth:`PostProcess._gather_and_scale_boxes`."""
+
+    def test_clamps_boxes_to_image_bounds(self):
+        """Boxes that extrapolate beyond [0, 1] in normalized space are clamped to pixel-space image dimensions after
+        scaling."""
+        # Three synthetic boxes in cxcywh normalized coords:
+        #   [0] cx=0.01, w=0.10 → x1 = (0.01 - 0.05) * 640 = -25.6  ← negative
+        #   [1] cx=0.99, w=0.10 → x2 = (0.99 + 0.05) * 640 = 665.6  ← overflow
+        #   [2] cx=0.50, w=0.20 → fully in-bounds
+        out_bbox = torch.tensor(
+            [
+                [
+                    [0.01, 0.01, 0.10, 0.10],  # negative x1, y1 after scale
+                    [0.99, 0.99, 0.10, 0.10],  # x2 > img_w, y2 > img_h after scale
+                    [0.50, 0.50, 0.20, 0.20],  # in-bounds control
+                ]
+            ]
+        )  # shape (B=1, Q=3, 4)
+
+        topk_boxes = torch.tensor([[0, 1, 2]])  # select all three
+        target_sizes = torch.tensor([[480, 640]])  # (h, w)
+
+        boxes = PostProcess._gather_and_scale_boxes(out_bbox, topk_boxes, target_sizes)
+
+        img_h, img_w = 480, 640
+
+        # All coords must be >= 0
+        assert (boxes >= 0).all(), f"Negative coords present: {boxes[boxes < 0]}"
+        # x1, x2 must be <= image width
+        assert (boxes[..., 0] <= img_w).all()
+        assert (boxes[..., 2] <= img_w).all()
+        # y1, y2 must be <= image height
+        assert (boxes[..., 1] <= img_h).all()
+        assert (boxes[..., 3] <= img_h).all()
+
+        # Exact clamped values — bounds-only check cannot catch a clamp returning e.g. 1.0 instead of 0.0
+        # box [0]: x1_raw=-25.6, y1_raw=-19.2 → clamped to 0.0
+        assert boxes[0, 0, 0].item() == pytest.approx(0.0), "x1 of underflowing box must clamp to 0"
+        assert boxes[0, 0, 1].item() == pytest.approx(0.0), "y1 of underflowing box must clamp to 0"
+        # box [1]: x2_raw=665.6 → clamped to img_w=640.0; y2_raw=499.2 → clamped to img_h=480.0
+        assert boxes[0, 1, 2].item() == pytest.approx(640.0), "x2 of overflowing box must clamp to img_w"
+        assert boxes[0, 1, 3].item() == pytest.approx(480.0), "y2 of overflowing box must clamp to img_h"
+
+    def test_in_bounds_boxes_unchanged(self):
+        """Boxes already within image bounds are not altered by clamping."""
+        out_bbox = torch.tensor(
+            [
+                [
+                    [0.30, 0.30, 0.20, 0.20],
+                    [0.70, 0.60, 0.30, 0.40],
+                ]
+            ]
+        )
+
+        topk_boxes = torch.tensor([[0, 1]])
+        target_sizes = torch.tensor([[480, 640]])
+
+        boxes = PostProcess._gather_and_scale_boxes(out_bbox, topk_boxes, target_sizes)
+
+        # Manually computed expected values (no clamping needed)
+        expected = torch.tensor(
+            [
+                [
+                    [128.0, 96.0, 256.0, 192.0],  # cx=0.30,cy=0.30,w=0.20,h=0.20
+                    [352.0, 192.0, 544.0, 384.0],  # cx=0.70,cy=0.60,w=0.30,h=0.40
+                ]
+            ]
+        )
+
+        assert torch.allclose(boxes, expected, atol=1e-4), (
+            f"In-bounds boxes were altered.\nExpected:\n{expected}\nGot:\n{boxes}"
+        )
+
+    def test_multiple_images_in_batch(self):
+        """Clamping works correctly across a batch of mixed image sizes."""
+        out_bbox = torch.tensor(
+            [
+                [
+                    [0.01, 0.50, 0.10, 0.20],  # image 0: negative x1
+                ],
+                [
+                    [0.99, 0.50, 0.10, 0.20],  # image 1: x2 overflow
+                ],
+            ]
+        )
+
+        topk_boxes = torch.tensor([[0], [0]])
+        target_sizes = torch.tensor(
+            [
+                [300, 400],  # image 0: 400×300
+                [600, 800],  # image 1: 800×600
+            ]
+        )
+
+        boxes = PostProcess._gather_and_scale_boxes(out_bbox, topk_boxes, target_sizes)
+
+        # Image 0: all coords must be in [0, 400]×[0, 300]
+        assert (boxes[0, :, 0] >= 0).all(), "img0 x1: expected >= 0"
+        assert (boxes[0, :, 0] <= 400).all(), "img0 x1: expected <= img_w (400)"
+        assert (boxes[0, :, 1] >= 0).all(), "img0 y1: expected >= 0"
+        assert (boxes[0, :, 1] <= 300).all(), "img0 y1: expected <= img_h (300)"
+        assert (boxes[0, :, 2] >= 0).all(), "img0 x2: expected >= 0"
+        assert (boxes[0, :, 2] <= 400).all(), "img0 x2: expected <= img_w (400)"
+        assert (boxes[0, :, 3] >= 0).all(), "img0 y2: expected >= 0"
+        assert (boxes[0, :, 3] <= 300).all(), "img0 y2: expected <= img_h (300)"
+
+        # Image 1: all coords must be in [0, 800]×[0, 600]
+        assert (boxes[1, :, 0] >= 0).all(), "img1 x1: expected >= 0"
+        assert (boxes[1, :, 0] <= 800).all(), "img1 x1: expected <= img_w (800)"
+        assert (boxes[1, :, 1] >= 0).all(), "img1 y1: expected >= 0"
+        assert (boxes[1, :, 1] <= 600).all(), "img1 y1: expected <= img_h (600)"
+        assert (boxes[1, :, 2] >= 0).all(), "img1 x2: expected >= 0"
+        assert (boxes[1, :, 2] <= 800).all(), "img1 x2: expected <= img_w (800)"
+        assert (boxes[1, :, 3] >= 0).all(), "img1 y2: expected >= 0"
+        assert (boxes[1, :, 3] <= 600).all(), "img1 y2: expected <= img_h (600)"
+
+
+class TestPostProcessForward:
+    """Integration tests for :meth:`PostProcess.forward`."""
+
+    def test_forward_clamps_edge_boxes_to_bounds(self):
+        """PostProcess.forward returns non-negative in-bounds boxes for edge-hugging predictions."""
+        postprocess = PostProcess(num_select=2)
+        outputs = {
+            "pred_logits": torch.tensor([[[10.0, -10.0], [9.0, -10.0]]]),
+            "pred_boxes": torch.tensor([[[0.01, 0.01, 0.10, 0.10], [0.99, 0.99, 0.10, 0.10]]]),
+        }
+        target_sizes = torch.tensor([[480, 640]])
+        results = postprocess(outputs, target_sizes)
+        boxes = results[0]["boxes"]
+        assert (boxes >= 0).all(), f"Negative coords present: {boxes[boxes < 0]}"
+        assert (boxes[..., 0] <= 640).all(), "x1 exceeds img_w (640)"
+        assert (boxes[..., 2] <= 640).all(), "x2 exceeds img_w (640)"
+        assert (boxes[..., 1] <= 480).all(), "y1 exceeds img_h (480)"
+        assert (boxes[..., 3] <= 480).all(), "y2 exceeds img_h (480)"
+
+
+class TestPostProcessMasks:
+    """Tests for :meth:`PostProcess._postprocess_masks` mask resizing."""
+
+    def test_chunked_upsample_preserves_shape_for_large_k(self):
+        """Chunked upsampling of K=64 masks returns full-resolution boolean masks of shape [K, 1, H, W]."""
+        batch, num_queries, mask_h, mask_w = 1, 64, 16, 16
+        num_select, img_h, img_w = 64, 512, 512
+        out_masks = torch.randn(batch, num_queries, mask_h, mask_w)
+        scores = torch.rand(batch, num_select)
+        labels = torch.zeros(batch, num_select, dtype=torch.long)
+        boxes = torch.zeros(batch, num_select, 4)
+        topk_boxes = torch.arange(num_select).unsqueeze(0)
+        target_sizes = torch.tensor([[img_h, img_w]])
+
+        results = PostProcess._postprocess_masks(out_masks, scores, labels, boxes, topk_boxes, target_sizes)
+
+        masks = results[0]["masks"]
+        assert masks.shape == (num_select, 1, img_h, img_w)
+        assert masks.dtype == torch.bool

tests/models/test_postprocess_keypoints.py

@@ -0,0 +1,174 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for keypoint decoding in PostProcess."""
+
+import pytest
+import torch
+
+from rfdetr.models.postprocess import PostProcess
+
+
+def test_postprocess_keypoints_shape_and_scores() -> None:
+    """PostProcess should emit keypoints and raw precision parameters for top detections."""
+    postprocess = PostProcess(num_select=2, num_keypoints_per_class=[17])
+    outputs = {
+        "pred_logits": torch.tensor([[[10.0, -10.0], [9.0, -10.0]]], dtype=torch.float32),
+        "pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5], [0.4, 0.6, 0.2, 0.3]]], dtype=torch.float32),
+        "pred_keypoints": torch.zeros((1, 2, 17, 8), dtype=torch.float32),
+    }
+    outputs["pred_keypoints"][0, :, :, 0] = 0.5
+    outputs["pred_keypoints"][0, :, :, 1] = 0.25
+    outputs["pred_keypoints"][0, :, :, 2] = 3.0
+    outputs["pred_keypoints"][0, :, :, 4] = 0.25
+    outputs["pred_keypoints"][0, :, :, 5] = 0.5
+    outputs["pred_keypoints"][0, :, :, 6] = -0.25
+
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+    results = postprocess(outputs, target_sizes)
+    keypoints = results[0]["keypoints"]
+    keypoint_precision = results[0]["keypoint_precision_cholesky"]
+
+    assert keypoints.shape == (2, 17, 3)
+    assert torch.allclose(keypoints[:, :, 0], torch.full((2, 17), 100.0))
+    assert torch.allclose(keypoints[:, :, 1], torch.full((2, 17), 25.0))
+    assert torch.all((keypoints[:, :, 2] > 0) & (keypoints[:, :, 2] < 1))
+    assert keypoint_precision.shape == (2, 17, 3)
+    torch.testing.assert_close(keypoint_precision[:, :, 0], torch.full((2, 17), 0.25))
+    torch.testing.assert_close(keypoint_precision[:, :, 1], torch.full((2, 17), 0.5))
+    torch.testing.assert_close(keypoint_precision[:, :, 2], torch.full((2, 17), -0.25))
+
+
+def test_postprocess_keypoints_class_filtering() -> None:
+    """Class-specific keypoint slots should be selected from padded per-class keypoint tensors."""
+    postprocess = PostProcess(num_select=1, num_keypoints_per_class=[2, 1])
+    outputs = {
+        "pred_logits": torch.tensor([[[0.0, 10.0]]], dtype=torch.float32),
+        "pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
+        "pred_keypoints": torch.zeros((1, 1, 4, 8), dtype=torch.float32),
+    }
+    # class 0 slots: [0, 1], class 1 slots: [2, 3]
+    outputs["pred_keypoints"][0, 0, 2, 0] = 0.25
+    outputs["pred_keypoints"][0, 0, 2, 1] = 0.4
+    outputs["pred_keypoints"][0, 0, 2, 2] = 2.0
+
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+    results = postprocess(outputs, target_sizes)
+    keypoints = results[0]["keypoints"]
+    keypoint_precision = results[0]["keypoint_precision_cholesky"]
+
+    assert keypoints.shape == (1, 2, 3)
+    assert torch.allclose(keypoints[0, 0, 0], torch.tensor(50.0))
+    assert torch.allclose(keypoints[0, 0, 1], torch.tensor(40.0))
+    assert 0.0 < keypoints[0, 0, 2].item() < 1.0
+    torch.testing.assert_close(keypoints[0, 1], torch.zeros(3))
+    torch.testing.assert_close(keypoint_precision[0, 1], torch.full((3,), float("nan")), equal_nan=True)
+
+
+def test_postprocess_keypoints_trace_alpha_rescores_active_keypoints_only() -> None:
+    """Trace fusion should use active keypoints for the predicted class and ignore padded slots."""
+    postprocess = PostProcess(num_select=1, num_keypoints_per_class=[2, 1], trace_alpha=1.0)
+    outputs = {
+        "pred_logits": torch.tensor([[[-10.0, 0.0]]], dtype=torch.float32),
+        "pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
+        "pred_keypoints": torch.zeros((1, 1, 4, 8), dtype=torch.float32),
+    }
+    # class 1 has one active slot at flat index 2 and one padded inactive slot at flat index 3.
+    outputs["pred_keypoints"][0, 0, 2, 2] = 10.0
+    outputs["pred_keypoints"][0, 0, 2, 4] = 0.0
+    outputs["pred_keypoints"][0, 0, 2, 5] = 0.0
+    outputs["pred_keypoints"][0, 0, 2, 6] = 0.0
+    outputs["pred_keypoints"][0, 0, 3, 2] = 10.0
+    outputs["pred_keypoints"][0, 0, 3, 4] = -2.0
+    outputs["pred_keypoints"][0, 0, 3, 6] = -2.0
+
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+    results = postprocess(outputs, target_sizes)
+
+    expected_score = torch.tensor([0.2], dtype=torch.float32)
+    torch.testing.assert_close(results[0]["scores"], expected_score, rtol=1e-4, atol=1e-6)
+
+
+def test_postprocess_keypoints_trace_alpha_normalizes_large_fused_scores() -> None:
+    """Trace-fused keypoint scores should be bounded after empirical normalization."""
+    postprocess = PostProcess(num_select=1, num_keypoints_per_class=[1], trace_alpha=1.0)
+    outputs = {
+        "pred_logits": torch.tensor([[[10.0]]], dtype=torch.float32),
+        "pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
+        "pred_keypoints": torch.zeros((1, 1, 1, 8), dtype=torch.float32),
+    }
+    outputs["pred_keypoints"][0, 0, 0, 2] = 10.0
+    outputs["pred_keypoints"][0, 0, 0, 4] = 2.0
+    outputs["pred_keypoints"][0, 0, 0, 5] = 0.0
+    outputs["pred_keypoints"][0, 0, 0, 6] = 2.0
+
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+    results = postprocess(outputs, target_sizes)
+
+    original_score = torch.sigmoid(torch.tensor([10.0], dtype=torch.float32))
+    mean_trace = torch.tensor([2.0], dtype=torch.float32) * torch.exp(torch.tensor([-4.0], dtype=torch.float32))
+    fused_score = original_score * mean_trace.pow(-1.0)
+    expected_score = fused_score / (1.0 + fused_score)
+    assert fused_score.item() > 1.0
+    assert 0.0 < results[0]["scores"].item() < 1.0
+    torch.testing.assert_close(results[0]["scores"], expected_score, rtol=1e-4, atol=1e-6)
+
+
+def test_postprocess_keypoints_trace_alpha_clamps_overflowing_fused_scores() -> None:
+    """Trace fusion should stay finite and strictly below 1.0 when the raw fused score overflows."""
+    postprocess = PostProcess(num_select=1, num_keypoints_per_class=[1], trace_alpha=1.0)
+    outputs = {
+        "pred_logits": torch.tensor([[[0.0]]], dtype=torch.float32),
+        "pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
+        "pred_keypoints": torch.zeros((1, 1, 1, 8), dtype=torch.float32),
+    }
+    outputs["pred_keypoints"][0, 0, 0, 2] = 10.0
+    outputs["pred_keypoints"][0, 0, 0, 4] = 50.0
+    outputs["pred_keypoints"][0, 0, 0, 5] = 0.0
+    outputs["pred_keypoints"][0, 0, 0, 6] = 50.0
+
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+    results = postprocess(outputs, target_sizes)
+
+    score = results[0]["scores"]
+    expected_score = torch.nextafter(torch.ones_like(score), torch.zeros_like(score))
+    assert torch.isfinite(score).all()
+    assert 0.0 < score.item() < 1.0
+    torch.testing.assert_close(score, expected_score, rtol=0.0, atol=0.0)
+
+
+def test_postprocess_keypoints_trace_alpha_uses_log_space_for_extreme_trace() -> None:
+    """Trace fusion should stay finite for extreme covariance terms."""
+    postprocess = PostProcess(num_select=1, num_keypoints_per_class=[1])
+    outputs = {
+        "pred_logits": torch.tensor([[[0.0]]], dtype=torch.float32),
+        "pred_boxes": torch.tensor([[[0.5, 0.5, 0.5, 0.5]]], dtype=torch.float32),
+        "pred_keypoints": torch.zeros((1, 1, 1, 8), dtype=torch.float32),
+    }
+    outputs["pred_keypoints"][0, 0, 0, 2] = 10.0
+    outputs["pred_keypoints"][0, 0, 0, 4] = -50.0
+    outputs["pred_keypoints"][0, 0, 0, 5] = 0.0
+    outputs["pred_keypoints"][0, 0, 0, 6] = 0.0
+
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+    results = postprocess(outputs, target_sizes)
+
+    expected_score = torch.tensor([0.5], dtype=torch.float32) * torch.exp(torch.tensor([-20.0], dtype=torch.float32))
+    torch.testing.assert_close(results[0]["scores"], expected_score, rtol=1e-4, atol=1e-12)
+
+
+def test_postprocess_validate_outputs_raises_when_masks_and_keypoints_both_present() -> None:
+    """PostProcess should raise ValueError when both pred_masks and pred_keypoints are present."""
+    postprocess = PostProcess(num_select=10)
+    outputs = {
+        "pred_logits": torch.zeros((1, 2, 2)),
+        "pred_boxes": torch.zeros((1, 2, 4)),
+        "pred_masks": torch.zeros((1, 2, 4, 4)),
+        "pred_keypoints": torch.zeros((1, 2, 17, 8)),
+    }
+    target_sizes = torch.tensor([[100, 200]], dtype=torch.int64)
+
+    with pytest.raises(ValueError, match="cannot be used together"):
+        postprocess(outputs, target_sizes)

tests/models/test_safe_torch_load.py

@@ -0,0 +1,175 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for the ``_safe_torch_load`` helper in ``rfdetr.util.io``.
+
+Covers the three-stage safe-load strategy: strict weights_only, safe-globals fallback, and opt-in pickle fallback.
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from rfdetr.util.io import _safe_torch_load
+
+# ---------------------------------------------------------------------------
+# Fixtures / helpers
+# ---------------------------------------------------------------------------
+
+
+def _write_tensor_only_checkpoint(path: Path) -> None:
+    """Save a checkpoint containing only tensors and plain dicts to *path*."""
+    ckpt = {"model": {"weight": torch.tensor([1.0, 2.0]), "bias": torch.tensor([0.0])}}
+    torch.save(ckpt, path)
+
+
+def _write_namespace_checkpoint(path: Path) -> None:
+    """Save a checkpoint with an ``argparse.Namespace`` args value to *path*.
+
+    Legacy RF-DETR engine.py checkpoints embed a Namespace; strict ``weights_only=True`` (without safe globals) would
+    reject these.
+    """
+    ckpt = {
+        "model": {"weight": torch.tensor([1.0])},
+        "args": argparse.Namespace(pretrain_weights="rf-detr-small.pth", num_classes=80),
+    }
+    torch.save(ckpt, path)
+
+
+def _write_simple_namespace_checkpoint(path: Path) -> None:
+    """Save a checkpoint with a ``types.SimpleNamespace`` to *path*."""
+    ckpt = {
+        "model": {"weight": torch.tensor([1.0])},
+        "args": SimpleNamespace(pretrain_weights="rf-detr-small.pth"),
+    }
+    torch.save(ckpt, path)
+
+
+class _ArbitraryObject:
+    """Module-level object that torch.save can pickle but weights_only=True rejects.
+
+    Must be defined at module scope so pickle can resolve its fully-qualified name during serialisation (local/nested
+    classes cannot be pickled by torch.save).
+    """
+
+    value = 42
+
+
+def _write_arbitrary_pickle_checkpoint(path: Path) -> None:
+    """Save a checkpoint that embeds an arbitrary class (requires pickle)."""
+    ckpt = {"model": {"weight": torch.tensor([1.0])}, "extra": _ArbitraryObject()}
+    torch.save(ckpt, path)
+
+
+# ---------------------------------------------------------------------------
+# Safe path (weights_only=True)
+# ---------------------------------------------------------------------------
+
+
+class TestSafeTorchLoadSafePath:
+    """Tensor-only checkpoints load without trust=True."""
+
+    def test_tensor_only_checkpoint_loads(self, tmp_path: Path) -> None:
+        """Pure-tensor checkpoint succeeds on the first safe-load attempt."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_tensor_only_checkpoint(ckpt_path)
+
+        result = _safe_torch_load(ckpt_path)
+
+        assert "model" in result
+        assert torch.allclose(result["model"]["weight"], torch.tensor([1.0, 2.0]))
+
+    def test_accepts_pathlib_path(self, tmp_path: Path) -> None:
+        """Helper accepts a :class:`pathlib.Path` argument without error."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_tensor_only_checkpoint(ckpt_path)
+
+        result = _safe_torch_load(ckpt_path)
+
+        assert "model" in result
+
+    def test_accepts_string_path(self, tmp_path: Path) -> None:
+        """Helper accepts a :class:`str` path argument without error."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_tensor_only_checkpoint(ckpt_path)
+
+        result = _safe_torch_load(str(ckpt_path))
+
+        assert "model" in result
+
+
+# ---------------------------------------------------------------------------
+# Safe-globals fallback (legacy Namespace checkpoints)
+# ---------------------------------------------------------------------------
+
+
+class TestSafeTorchLoadSafeGlobals:
+    """Checkpoints with argparse.Namespace / SimpleNamespace load without trust=True."""
+
+    def test_argparse_namespace_loads_without_trust(self, tmp_path: Path) -> None:
+        """argparse.Namespace checkpoint succeeds via the safe-globals retry."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_namespace_checkpoint(ckpt_path)
+
+        result = _safe_torch_load(ckpt_path)
+
+        assert isinstance(result["args"], argparse.Namespace)
+        assert result["args"].num_classes == 80
+
+    def test_simple_namespace_loads_without_trust(self, tmp_path: Path) -> None:
+        """SimpleNamespace checkpoint succeeds via the safe-globals retry."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_simple_namespace_checkpoint(ckpt_path)
+
+        result = _safe_torch_load(ckpt_path)
+
+        assert isinstance(result["args"], SimpleNamespace)
+
+
+# ---------------------------------------------------------------------------
+# Arbitrary pickle — trust=False must raise, trust=True must succeed
+# ---------------------------------------------------------------------------
+
+
+class TestSafeTorchLoadTrustGate:
+    """Arbitrary-pickle checkpoints require explicit trust=True."""
+
+    def test_arbitrary_pickle_raises_without_trust(self, tmp_path: Path) -> None:
+        """Checkpoint with unknown Python object raises RuntimeError when trust=False."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_arbitrary_pickle_checkpoint(ckpt_path)
+
+        with pytest.raises(RuntimeError, match="trust_checkpoint=True"):
+            _safe_torch_load(ckpt_path, trust=False)
+
+    def test_arbitrary_pickle_raises_by_default(self, tmp_path: Path) -> None:
+        """Checkpoint with unknown Python object raises RuntimeError when trust omitted (default=False)."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_arbitrary_pickle_checkpoint(ckpt_path)
+
+        with pytest.raises(RuntimeError, match="trust_checkpoint=True"):
+            _safe_torch_load(ckpt_path)
+
+    def test_arbitrary_pickle_succeeds_with_trust(self, tmp_path: Path) -> None:
+        """Checkpoint with unknown Python object loads when trust=True."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_arbitrary_pickle_checkpoint(ckpt_path)
+
+        result = _safe_torch_load(ckpt_path, trust=True)
+
+        assert "model" in result
+
+    def test_trust_true_emits_warning(self, tmp_path: Path) -> None:
+        """Trust=True triggers a UserWarning about unsafe loading."""
+        ckpt_path = tmp_path / "ckpt.pth"
+        _write_arbitrary_pickle_checkpoint(ckpt_path)
+
+        with pytest.warns(UserWarning, match="weights_only=False"):
+            _safe_torch_load(ckpt_path, trust=True)

tests/models/test_transformer.py

@@ -0,0 +1,501 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for transformer utilities, MS deformable attention core, and MSDeformAttn module."""
+
+import io
+
+import numpy as np
+import pytest
+import torch
+
+from rfdetr.models.ops.functions import ms_deform_attn_core_pytorch
+from rfdetr.models.ops.modules.ms_deform_attn import MSDeformAttn
+from rfdetr.models.transformer import gen_encoder_output_proposals, gen_sineembed_for_position
+
+
+@pytest.fixture(autouse=True)
+def _reset_random_seeds() -> None:
+    """Ensure reproducible random state for every test."""
+    torch.manual_seed(42)
+    torch.cuda.manual_seed_all(42)
+
+
+_MSDeformInputs = tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[tuple[int, int]]]
+
+
+def _build_ms_deform_inputs(
+    bsz: int = 1,
+    n_heads: int = 2,
+    head_dim: int = 4,
+    len_q: int = 3,
+    npts: int = 1,
+    levels: list[tuple[int, int]] | None = None,
+) -> _MSDeformInputs:
+    """Build minimal valid inputs for ms_deform_attn_core_pytorch.
+
+    Args:
+        bsz: Batch size.
+        n_heads: Number of attention heads.
+        head_dim: Dimension per head.
+        len_q: Number of query elements.
+        npts: Number of sampling points per level.
+        levels: List of (H, W) int pairs; defaults to [(4, 4), (2, 2)].
+
+    Returns:
+        Tuple of (value, spatial_shapes_tensor, sampling_locations,
+                  attention_weights, spatial_shapes_hw).
+    """
+    if levels is None:
+        levels = [(4, 4), (2, 2)]
+    nlvl = len(levels)
+
+    total_hw = sum(ht * wd for ht, wd in levels)
+    spatial_shapes_tensor = torch.tensor(levels, dtype=torch.long)
+    value = torch.randn(bsz, n_heads, head_dim, total_hw)
+    # sampling_locations: (bsz, len_q, n_heads, nlvl, npts, 2) in [0, 1]
+    sampling_locations = torch.rand(bsz, len_q, n_heads, nlvl, npts, 2)
+    # attention_weights: (bsz, len_q, n_heads, nlvl * npts)
+    attention_weights = torch.softmax(torch.randn(bsz, len_q, n_heads, nlvl * npts), dim=-1)
+
+    return value, spatial_shapes_tensor, sampling_locations, attention_weights, levels
+
+
+def test_gen_encoder_output_proposals_passes_ij_indexing_to_meshgrid(monkeypatch) -> None:
+    """`gen_encoder_output_proposals` should call `torch.meshgrid` with explicit ij indexing."""
+    original_meshgrid = torch.meshgrid
+    call_count = 0
+
+    def _meshgrid_with_indexing_assertion(*args, **kwargs):
+        nonlocal call_count
+        call_count += 1
+        if kwargs.get("indexing") != "ij":
+            raise AssertionError("torch.meshgrid must be called with indexing='ij'")
+        return original_meshgrid(*args, **kwargs)
+
+    monkeypatch.setattr(torch, "meshgrid", _meshgrid_with_indexing_assertion)
+
+    memory = torch.randn(1, 4, 8)
+    spatial_shapes = torch.tensor([[2, 2]], dtype=torch.long)
+
+    output_memory, output_proposals = gen_encoder_output_proposals(
+        memory,
+        spatial_shapes=spatial_shapes,
+    )
+
+    assert call_count == 1
+
+
+def test_gen_sineembed_for_position_keeps_box_dimensions_in_sin_cos_order() -> None:
+    """4D box positional embeddings must use the pretrained sin/cos order for all dimensions."""
+    pos_tensor = torch.tensor([[[0.125, 0.25, 0.5, 0.75]]], dtype=torch.float32)
+    dim = 4
+    scale = 2 * torch.pi
+    dim_t = torch.arange(dim, dtype=pos_tensor.dtype)
+    dim_t = 10000 ** (2 * (dim_t // 2) / dim)
+
+    expected_parts = []
+    for coord_idx in (1, 0, 2, 3):
+        coord = pos_tensor[:, :, coord_idx] * scale
+        encoded = coord[:, :, None] / dim_t
+        expected_parts.append(torch.stack((encoded[:, :, 0::2].sin(), encoded[:, :, 1::2].cos()), dim=3).flatten(2))
+    expected = torch.cat(expected_parts, dim=2)
+
+    actual = gen_sineembed_for_position(pos_tensor, dim=dim)
+
+    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-6)
+
+
+def test_gen_encoder_output_proposals_rejects_non_square_ij_indexing(monkeypatch) -> None:
+    """Wrong meshgrid indexing (xy vs ij) produces different proposals for non-square spatial shapes."""
+    original_meshgrid = torch.meshgrid
+
+    def _meshgrid_wrong_indexing(*args, **kwargs):
+        kwargs["indexing"] = "xy"
+        return original_meshgrid(*args, **kwargs)
+
+    # Use non-square spatial shapes so that ij vs xy indexing produces observably different outputs.
+    memory = torch.randn(1, 8, 8)
+    spatial_shapes = torch.tensor([[2, 4]], dtype=torch.long)
+
+    correct_memory, correct_proposals = gen_encoder_output_proposals(memory, spatial_shapes=spatial_shapes)
+
+    monkeypatch.setattr(torch, "meshgrid", _meshgrid_wrong_indexing)
+
+    wrong_memory, wrong_proposals = gen_encoder_output_proposals(memory, spatial_shapes=spatial_shapes)
+
+    assert not torch.allclose(correct_proposals, wrong_proposals), (
+        "xy indexing must produce different proposals than ij indexing for non-square spatial shapes"
+    )
+
+
+def test_gen_encoder_output_proposals_accepts_int_tuple_spatial_shapes() -> None:
+    """`gen_encoder_output_proposals` must accept `spatial_shapes` as a tensor of int pairs."""
+    batch = 2
+    ht, wd = 4, 4
+    memory = torch.randn(batch, ht * wd, 8)
+    spatial_shapes = torch.tensor([[ht, wd]], dtype=torch.long)
+
+    output_memory, output_proposals = gen_encoder_output_proposals(memory, spatial_shapes=spatial_shapes)
+
+    assert output_memory.shape == memory.shape
+    assert output_proposals.shape == (batch, ht * wd, 4)
+
+
+def test_gen_encoder_output_proposals_accepts_python_int_pair_spatial_shapes() -> None:
+    """`gen_encoder_output_proposals` must accept `spatial_shapes` as `list[tuple[int, int]]` with no padding mask.
+
+    Regression: `Transformer.forward` passes Python int pairs derived from `src.shape`, so the
+    export-driven call path uses `list[tuple[int, int]]` rather than a tensor.
+    """
+    batch, ht, wd, dim = 2, 4, 4, 8
+    memory = torch.randn(batch, ht * wd, dim)
+    spatial_shapes = [(ht, wd)]  # Python int pairs, as produced by Transformer.forward()
+
+    output_memory, output_proposals = gen_encoder_output_proposals(
+        memory,
+        memory_padding_mask=None,
+        spatial_shapes=spatial_shapes,
+    )
+
+    assert output_memory.shape == memory.shape
+    assert output_proposals.shape == (batch, ht * wd, 4)
+
+
+class TestMSDeformAttnCorePytorch:
+    """Tests for ms_deform_attn_core_pytorch with Python int pair spatial shapes.
+
+    Regression suite for torch.export.export compatibility: iterating over a spatial_shapes tensor yields FakeTensor
+    scalars during FakeTensor tracing, which cannot be used as Python int split/view sizes.  The function now accepts an
+    optional ``value_spatial_shapes_hw`` list of Python int pairs that bypasses tensor iteration.
+    """
+
+    @pytest.fixture
+    def make_inputs(self) -> _MSDeformInputs:
+        """Default two-level inputs: levels=[(4, 4), (2, 2)]."""
+        return _build_ms_deform_inputs()
+
+    @pytest.fixture
+    def single_level_inputs(self) -> _MSDeformInputs:
+        """Single-level inputs: levels=[(8, 8)]."""
+        return _build_ms_deform_inputs(levels=[(8, 8)])
+
+    def test_with_tensor_spatial_shapes(self, make_inputs: _MSDeformInputs) -> None:
+        """Baseline: passing only the tensor spatial_shapes still works."""
+        value, spatial_shapes_tensor, sampling_locations, attention_weights, _ = make_inputs
+
+        output = ms_deform_attn_core_pytorch(value, spatial_shapes_tensor, sampling_locations, attention_weights)
+
+        bsz, n_heads, head_dim, _ = value.shape
+        len_q = sampling_locations.shape[1]
+        assert output.shape == (bsz, len_q, n_heads * head_dim)
+
+    def test_with_python_int_pair_spatial_shapes(self, make_inputs: _MSDeformInputs) -> None:
+        """Regression: value_spatial_shapes_hw list of Python int pairs must be accepted.
+
+        This is the torch.export.export-compatible code path: tensor scalar values (from iterating over a FakeTensor)
+        cannot be used as split/view sizes, so the caller passes explicit Python int pairs via value_spatial_shapes_hw
+        instead.
+        """
+        value, spatial_shapes_tensor, sampling_locations, attention_weights, levels = make_inputs
+
+        output = ms_deform_attn_core_pytorch(
+            value,
+            spatial_shapes_tensor,
+            sampling_locations,
+            attention_weights,
+            value_spatial_shapes_hw=levels,
+        )
+
+        bsz, n_heads, head_dim, _ = value.shape
+        len_q = sampling_locations.shape[1]
+        assert output.shape == (bsz, len_q, n_heads * head_dim)
+
+    def test_tensor_and_hw_paths_produce_identical_outputs(self, make_inputs: _MSDeformInputs) -> None:
+        """Python int pair path and tensor iteration path must produce the same result."""
+        value, spatial_shapes_tensor, sampling_locations, attention_weights, levels = make_inputs
+
+        out_tensor_path = ms_deform_attn_core_pytorch(
+            value, spatial_shapes_tensor, sampling_locations, attention_weights
+        )
+        out_hw_path = ms_deform_attn_core_pytorch(
+            value,
+            spatial_shapes_tensor,
+            sampling_locations,
+            attention_weights,
+            value_spatial_shapes_hw=levels,
+        )
+
+        torch.testing.assert_close(out_tensor_path, out_hw_path)
+
+    def test_single_level(self, single_level_inputs: _MSDeformInputs) -> None:
+        """Single-level case with Python int pair path must not crash."""
+        value, spatial_shapes_tensor, sampling_locations, attention_weights, levels = single_level_inputs
+
+        output = ms_deform_attn_core_pytorch(
+            value,
+            spatial_shapes_tensor,
+            sampling_locations,
+            attention_weights,
+            value_spatial_shapes_hw=levels,
+        )
+
+        assert output.shape[0] == 1
+
+
+class TestMSDeformAttnModule:
+    """Tests for MSDeformAttn.forward covering the export-compatibility changes.
+
+    Validates the module-level parameter threading and export-mode assert guard introduced in the torch.export.export
+    compatibility fix.
+    """
+
+    _d_model = 32
+    _n_heads = 4
+    _n_levels = 2
+    _n_points = 1
+    _hw_pairs: list[tuple[int, int]] = [(4, 4), (2, 2)]
+
+    def _make_module_inputs(
+        self,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        list[tuple[int, int]],
+    ]:
+        """Build minimal valid inputs for MSDeformAttn.forward.
+
+        Returns:
+            Tuple of (query, reference_points, input_flatten,
+                      input_spatial_shapes, input_level_start_index, hw_pairs).
+        """
+        hw_pairs = self._hw_pairs
+        total_len = sum(ht * wd for ht, wd in hw_pairs)
+        bsz, len_q = 1, 3
+
+        query = torch.randn(bsz, len_q, self._d_model)
+        reference_points = torch.rand(bsz, len_q, self._n_levels, 2)
+        input_flatten = torch.randn(bsz, total_len, self._d_model)
+        input_spatial_shapes = torch.tensor(hw_pairs, dtype=torch.long)
+        # Cumulative start index per level: [0, H0*W0]
+        starts = [sum(ht * wd for ht, wd in hw_pairs[:idx]) for idx in range(self._n_levels)]
+        input_level_start_index = torch.tensor(starts, dtype=torch.long)
+
+        return query, reference_points, input_flatten, input_spatial_shapes, input_level_start_index, hw_pairs
+
+    def test_forward_without_hw_param_backward_compat(self) -> None:
+        """MSDeformAttn.forward without hw param produces correct output shape."""
+        module = MSDeformAttn(
+            d_model=self._d_model, n_levels=self._n_levels, n_heads=self._n_heads, n_points=self._n_points
+        )
+        query, ref_pts, input_flatten, spatial_shapes, level_start_index, _ = self._make_module_inputs()
+
+        output = module(query, ref_pts, input_flatten, spatial_shapes, level_start_index)
+
+        bsz, len_q, _ = query.shape
+        assert output.shape == (bsz, len_q, self._d_model)
+
+    def test_forward_with_hw_param_produces_correct_shape(self) -> None:
+        """MSDeformAttn.forward with input_spatial_shapes_hw produces correct output shape."""
+        module = MSDeformAttn(
+            d_model=self._d_model, n_levels=self._n_levels, n_heads=self._n_heads, n_points=self._n_points
+        )
+        query, ref_pts, input_flatten, spatial_shapes, level_start_index, hw_pairs = self._make_module_inputs()
+
+        output = module(
+            query, ref_pts, input_flatten, spatial_shapes, level_start_index, input_spatial_shapes_hw=hw_pairs
+        )
+
+        bsz, len_q, _ = query.shape
+        assert output.shape == (bsz, len_q, self._d_model)
+
+    def test_export_mode_forward_with_hw_param(self) -> None:
+        """MSDeformAttn.forward in export mode with hw param must not raise."""
+        module = MSDeformAttn(
+            d_model=self._d_model, n_levels=self._n_levels, n_heads=self._n_heads, n_points=self._n_points
+        )
+        module.export()
+        query, ref_pts, input_flatten, spatial_shapes, level_start_index, hw_pairs = self._make_module_inputs()
+
+        output = module(
+            query, ref_pts, input_flatten, spatial_shapes, level_start_index, input_spatial_shapes_hw=hw_pairs
+        )
+
+        bsz, len_q, _ = query.shape
+        assert output.shape == (bsz, len_q, self._d_model)
+
+    def test_export_flag_set_after_export_call(self) -> None:
+        """Calling .export() must set _export=True, enabling the torch._assert guard path."""
+        module = MSDeformAttn(
+            d_model=self._d_model, n_levels=self._n_levels, n_heads=self._n_heads, n_points=self._n_points
+        )
+        assert not module._export
+
+        module.export()
+
+        assert module._export
+
+
+class TestGenEncoderOutputProposalsDynamicBatch:
+    """Regression tests for dynamic batch support in gen_encoder_output_proposals.
+
+    Ensures that the ONNX-symbolic refactoring (PR #950 / issue #949) does not bake a fixed batch dimension into
+    proposals and that output shapes are correct for varying batch sizes.
+    """
+
+    @pytest.mark.parametrize("batch_size", [1, 2, 4, 8])
+    def test_output_shape_invariant_across_batch_sizes(self, batch_size: int) -> None:
+        """Output shapes must scale correctly with batch size, with no baked constants.
+
+        Args:
+            batch_size: Number of images in the batch.
+        """
+        ht, wd, dim = 4, 4, 8
+        memory = torch.randn(batch_size, ht * wd, dim)
+        spatial_shapes = [(ht, wd)]
+
+        output_memory, output_proposals = gen_encoder_output_proposals(
+            memory, memory_padding_mask=None, spatial_shapes=spatial_shapes
+        )
+
+        assert output_memory.shape == (batch_size, ht * wd, dim)
+        assert output_proposals.shape == (batch_size, ht * wd, 4)
+
+    def test_proposals_semantically_equivalent_across_batch_sizes(self) -> None:
+        """Proposals for batch=1 and batch=4 must be identical per image.
+
+        Regression: if batch_size were baked as a constant, repeating the same image
+        N times would produce different proposals for each copy.
+        """
+        ht, wd, dim = 4, 4, 8
+        memory_single = torch.randn(1, ht * wd, dim)
+        memory_multi = memory_single.expand(4, -1, -1).contiguous()
+        spatial_shapes = [(ht, wd)]
+
+        _, proposals_single = gen_encoder_output_proposals(
+            memory_single, memory_padding_mask=None, spatial_shapes=spatial_shapes
+        )
+        _, proposals_multi = gen_encoder_output_proposals(
+            memory_multi, memory_padding_mask=None, spatial_shapes=spatial_shapes
+        )
+
+        torch.testing.assert_close(proposals_single.expand(4, -1, -1), proposals_multi)
+
+    @pytest.mark.parametrize("batch_size", [1, 4])
+    def test_output_shape_invariant_with_padding_mask(self, batch_size: int) -> None:
+        """Output shapes must be correct when memory_padding_mask is provided with varying batch sizes.
+
+        Regression for PR #950 / issue #949: the masked branch used .reshape(-1, h, w, 1) to infer the batch dimension
+        dynamically; this test verifies the branch handles varying batch sizes without error.
+
+        Args:
+            batch_size: Number of images in the batch.
+        """
+        ht, wd, dim = 4, 4, 8
+        total_hw = ht * wd
+        memory = torch.randn(batch_size, total_hw, dim)
+        # Mask shape: (batch, sum_hw) — True means padding (invalid position)
+        memory_padding_mask = torch.zeros(batch_size, total_hw, dtype=torch.bool)
+        spatial_shapes = [(ht, wd)]
+
+        output_memory, output_proposals = gen_encoder_output_proposals(
+            memory, memory_padding_mask=memory_padding_mask, spatial_shapes=spatial_shapes
+        )
+
+        assert output_memory.shape == (batch_size, total_hw, dim)
+        assert output_proposals.shape == (batch_size, total_hw, 4)
+
+    @pytest.mark.parametrize("batch_size", [1, 4, 8])
+    def test_onnx_export_with_dynamic_batch_axis(self, batch_size: int) -> None:
+        """ONNX export with dynamic batch axis must run inference for batch sizes other than the trace batch.
+
+        Regression for issue #949: exporting with a fixed trace batch baked `Reshape([8,...])` as a constant ONNX node,
+        causing TRT engines to fail at inference for any batch != 8. Skipped when onnx or onnxruntime is not installed.
+        """
+        pytest.importorskip("onnx")
+        onnxruntime = pytest.importorskip("onnxruntime")
+
+        ht, wd, dim = 4, 4, 8
+        spatial_shapes_list = [(ht, wd)]
+
+        class _ProposalModule(torch.nn.Module):
+            """Thin wrapper to export gen_encoder_output_proposals via torch.onnx."""
+
+            def forward(self, memory: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+                """Forward pass delegating to gen_encoder_output_proposals."""
+                return gen_encoder_output_proposals(
+                    memory, memory_padding_mask=None, spatial_shapes=spatial_shapes_list
+                )
+
+        module = _ProposalModule()
+        trace_memory = torch.randn(2, ht * wd, dim)
+
+        buf = io.BytesIO()
+        torch.onnx.export(
+            module,
+            (trace_memory,),
+            buf,
+            input_names=["memory"],
+            output_names=["output_memory", "output_proposals"],
+            dynamic_axes={"memory": {0: "batch"}},
+            opset_version=17,
+        )
+        buf.seek(0)
+        onnx_bytes = buf.read()
+
+        session = onnxruntime.InferenceSession(onnx_bytes, providers=["CPUExecutionProvider"])
+        memory_np = np.random.randn(batch_size, ht * wd, dim).astype(np.float32)
+        out_memory, out_proposals = session.run(None, {"memory": memory_np})
+        assert out_memory.shape == (batch_size, ht * wd, dim), f"wrong memory shape for batch={batch_size}"
+        assert out_proposals.shape == (batch_size, ht * wd, 4), f"wrong proposals shape for batch={batch_size}"
+
+
+def test_ms_deform_attn_core_pytorch_export_compatible() -> None:
+    """torch.export.export must succeed on a module using ms_deform_attn_core_pytorch with hw param.
+
+    Regression test for the FakeTensor tracing failure: iterating over spatial_shapes and using the scalar elements as
+    split/view sizes fails during torch.export.export because FakeTensor data is not allocated. Passing
+    value_spatial_shapes_hw (concrete Python ints from a module attribute) bypasses the tensor iteration entirely.
+    """
+    levels: list[tuple[int, int]] = [(4, 4), (2, 2)]
+    bsz, n_heads, head_dim = 1, 2, 4
+    total_hw = sum(ht * wd for ht, wd in levels)
+    len_q, nlvl, npts = 3, len(levels), 1
+
+    class _MinimalDeformAttn(torch.nn.Module):
+        """Minimal wrapper to test torch.export.export on the hw-param code path."""
+
+        def __init__(self, hw: list[tuple[int, int]]) -> None:
+            super().__init__()
+            self.hw = hw
+
+        def forward(
+            self,
+            value: torch.Tensor,
+            spatial_shapes: torch.Tensor,
+            sampling_locations: torch.Tensor,
+            attention_weights: torch.Tensor,
+        ) -> torch.Tensor:
+            """Forward using concrete Python int pairs for export compatibility."""
+            return ms_deform_attn_core_pytorch(
+                value,
+                spatial_shapes,
+                sampling_locations,
+                attention_weights,
+                value_spatial_shapes_hw=self.hw,
+            )
+
+    value = torch.randn(bsz, n_heads, head_dim, total_hw)
+    spatial_shapes = torch.tensor(levels, dtype=torch.long)
+    sampling_locations = torch.rand(bsz, len_q, n_heads, nlvl, npts, 2)
+    attention_weights = torch.softmax(torch.randn(bsz, len_q, n_heads, nlvl * npts), dim=-1)
+
+    module = _MinimalDeformAttn(hw=levels)
+
+    exported = torch.export.export(module, args=(value, spatial_shapes, sampling_locations, attention_weights))
+    assert exported is not None

tests/models/test_transformer_keypoints.py

@@ -0,0 +1,250 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests for GroupPose-oriented transformer streams."""
+
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+
+from rfdetr.models.transformer import Transformer, TransformerDecoder, TransformerDecoderLayer, build_transformer
+
+
+def _build_transformer_inputs(
+    batch_size: int = 2,
+    hidden_dim: int = 16,
+    num_levels: int = 2,
+) -> tuple[list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], torch.Tensor, torch.Tensor]:
+    """Build minimal synthetic multi-scale inputs for `Transformer.forward`.
+
+    Args:
+        batch_size: Mini-batch size.
+        hidden_dim: Transformer and input channel size.
+        num_levels: Number of feature pyramid levels.
+
+    Returns:
+        srcs, masks, pos_embeds, refpoint_embed, query_feat.
+    """
+    spatial_shapes = [(4, 4), (2, 2)]
+    srcs = [
+        torch.randn(batch_size, hidden_dim, spatial_shapes[idx][0], spatial_shapes[idx][1]) for idx in range(num_levels)
+    ]
+    masks = [torch.zeros(batch_size, h, w, dtype=torch.bool) for h, w in spatial_shapes[:num_levels]]
+    pos_embeds = [
+        torch.randn(batch_size, hidden_dim, spatial_shapes[idx][0], spatial_shapes[idx][1]) for idx in range(num_levels)
+    ]
+    refpoint_embed = torch.rand(6, 4)
+    query_feat = torch.randn(6, hidden_dim)
+    return srcs, masks, pos_embeds, refpoint_embed, query_feat
+
+
+def test_transformer_keypoint_disabled_matches_default_contract() -> None:
+    """Transformer without GroupPose should keep the 4-item return contract."""
+    srcs, masks, pos_embeds, refpoint_embed, query_feat = _build_transformer_inputs()
+    transformer = Transformer(
+        d_model=16,
+        num_queries=6,
+        num_decoder_layers=1,
+        sa_nhead=4,
+        ca_nhead=4,
+        num_feature_levels=2,
+        dec_n_points=1,
+        return_intermediate_dec=True,
+        lite_refpoint_refine=True,
+        use_grouppose_keypoints=False,
+    )
+
+    outputs = transformer(srcs, masks, pos_embeds, refpoint_embed, query_feat, cross_attn_srcs=None)
+    assert len(outputs) == 4, f"Expected 4 outputs, got {len(outputs)}"
+    hs, references, memory_ts, boxes_ts = outputs
+    assert hs is not None and references is not None
+    assert memory_ts is None and boxes_ts is None
+
+
+def test_transformer_keypoint_enabled_shapes() -> None:
+    """GroupPose path should emit keypoint hidden states and encoder keypoint slots."""
+    srcs, masks, pos_embeds, refpoint_embed, query_feat = _build_transformer_inputs()
+    transformer = Transformer(
+        d_model=16,
+        num_queries=6,
+        num_decoder_layers=1,
+        sa_nhead=4,
+        ca_nhead=4,
+        num_feature_levels=2,
+        dec_n_points=1,
+        return_intermediate_dec=True,
+        lite_refpoint_refine=True,
+        two_stage=True,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=[17],
+    )
+    transformer.enc_out_class_embed = nn.ModuleList([nn.Linear(16, 2)])
+    transformer.enc_out_bbox_embed = nn.ModuleList([nn.Linear(16, 4)])
+
+    outputs = transformer(srcs, masks, pos_embeds, refpoint_embed, query_feat, cross_attn_srcs=None)
+    assert len(outputs) == 7, f"Expected 7 outputs, got {len(outputs)}"
+    hs, references, memory_ts, boxes_ts, keypoint_hs, enc_kp_predictions, keypoint_memory_ts = outputs
+
+    assert isinstance(hs, torch.Tensor)
+    assert hs.shape[-1] == 16
+    assert references.shape[-1] == 4
+    assert memory_ts is not None and boxes_ts is not None
+    assert keypoint_hs is not None
+    assert keypoint_hs.shape[3] == 17
+    assert enc_kp_predictions is not None
+    assert enc_kp_predictions.shape[-1] == 16
+    assert keypoint_memory_ts is not None
+
+
+def test_build_transformer_defaults_inter_instance_keypoint_attention_to_config_default() -> None:
+    """Older args objects without `inter_instance_kp_attn` should keep the preview topology default."""
+    args = SimpleNamespace(
+        hidden_dim=16,
+        sa_nheads=4,
+        ca_nheads=4,
+        num_queries=6,
+        dropout=0.0,
+        dim_feedforward=32,
+        dec_layers=1,
+        group_detr=1,
+        num_feature_levels=2,
+        dec_n_points=1,
+        lite_refpoint_refine=True,
+        decoder_norm="LN",
+        bbox_reparam=False,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=[17],
+    )
+
+    transformer = build_transformer(args)
+
+    decoder_layer = transformer.decoder.layers[0]
+    assert isinstance(decoder_layer, TransformerDecoderLayer)
+    assert decoder_layer.enable_keypoint_processing
+    assert not decoder_layer.inter_instance_kp_attn
+
+
+def test_keypoint_class_mask_person_only() -> None:
+    """Person-only schema `[17]` should build a keypoint class mask with only self-class tokens."""
+    layer = TransformerDecoderLayer(
+        d_model=16,
+        sa_nhead=4,
+        ca_nhead=4,
+        dim_feedforward=32,
+        num_feature_levels=2,
+        enable_keypoint_processing=True,
+        grouppose_keypoint_dim_downscale=1,
+        keypoint_cross_attn=False,
+        inter_instance_kp_attn=False,
+    )
+    decoder = TransformerDecoder(
+        decoder_layer=layer,
+        num_layers=1,
+        return_intermediate=True,
+        d_model=16,
+        lite_refpoint_refine=True,
+        enable_keypoint_processing=True,
+        num_keypoints_per_class=[17],
+        grouppose_keypoint_dim_downscale=1,
+    )
+
+    assert decoder.keypoint_pos_embed is not None
+    assert decoder.keypoint_class_mask.shape == (18, 18)
+    assert decoder.keypoint_class_mask.dtype == torch.bool
+    assert not decoder.keypoint_class_mask.any()
+
+
+def test_enc_keypoint_embed_eval_uses_only_head_zero() -> None:
+    """Encoder keypoint path must use a single head (head 0) in eval mode.
+
+    Regression: ``group_detr = len(self.enc_out_keypoint_embed)`` without a
+    ``self.training`` guard caused eval mode to split ``num_queries`` across all
+    group heads instead of routing every query through head 0. Fix:
+    ``group_detr = len(...) if self.training else 1``.
+
+    Strategy: zero all head weights/biases; set head-0 last-layer bias to
+    ``sentinel``. In eval mode every query routes through head 0, so
+    ``kp_pred[..., 2:]`` (the pure-delta dims unaffected by ref_xy/wh) must all
+    equal ``sentinel``. In training mode only the first 1/group_detr queries go
+    through head 0 (the rest equal 0.0).
+    """
+    group_detr = 3
+    num_queries = 6  # divisible by group_detr
+    hidden_dim = 16
+    batch_size = 1
+    sentinel = 50.0
+
+    srcs, masks, pos_embeds, _, _ = _build_transformer_inputs(batch_size=batch_size, hidden_dim=hidden_dim)
+    refpoint_embed = torch.rand(num_queries, 4)
+    query_feat = torch.randn(num_queries, hidden_dim)
+
+    transformer = Transformer(
+        d_model=hidden_dim,
+        num_queries=num_queries,
+        num_decoder_layers=1,
+        sa_nhead=4,
+        ca_nhead=4,
+        num_feature_levels=2,
+        dec_n_points=1,
+        return_intermediate_dec=True,
+        lite_refpoint_refine=True,
+        two_stage=True,
+        group_detr=group_detr,
+        use_grouppose_keypoints=True,
+        num_keypoints_per_class=[2],
+    )
+    transformer.enc_out_class_embed = nn.ModuleList([nn.Linear(hidden_dim, 2) for _ in range(group_detr)])
+    transformer.enc_out_bbox_embed = nn.ModuleList([nn.Linear(hidden_dim, 4) for _ in range(group_detr)])
+
+    # Zero all keypoint head weights and biases; give head 0 a distinctive output.
+    with torch.no_grad():
+        for _, head in enumerate(transformer.enc_out_keypoint_embed):
+            for layer in head.layers:
+                layer.weight.zero_()
+                layer.bias.zero_()
+        transformer.enc_out_keypoint_embed[0].layers[-1].bias.fill_(sentinel)
+
+    transformer.eval()
+    with torch.no_grad():
+        outputs = transformer(srcs, masks, pos_embeds, refpoint_embed, query_feat, cross_attn_srcs=None)
+
+    _, _, _, _, _, enc_kp_predictions, _ = outputs
+    assert enc_kp_predictions is not None, "enc_kp_predictions should not be None in keypoint mode"
+
+    # kp_pred = [kp_xy(2 dims), kp_delta[2:]]; dims 2: are pure MLP output unaffected by ref_xy/wh.
+    kp_beyond_xy = enc_kp_predictions[..., 2:]
+    assert (kp_beyond_xy == sentinel).all(), (
+        f"Eval mode must route all {num_queries} queries through head 0 (bias={sentinel}). "
+        f"Got min={kp_beyond_xy.min().item():.2f}, max={kp_beyond_xy.max().item():.2f}. "
+        "Bug: group_detr not guarded by self.training in enc_out_keypoint_embed loop."
+    )
+
+
+def test_cross_attn_srcs_none_backward_compat() -> None:
+    """`cross_attn_srcs=None` must remain equivalent to passing the primary feature stream."""
+    srcs, masks, pos_embeds, refpoint_embed, query_feat = _build_transformer_inputs()
+
+    transformer = Transformer(
+        d_model=16,
+        num_queries=6,
+        num_decoder_layers=1,
+        sa_nhead=4,
+        ca_nhead=4,
+        num_feature_levels=2,
+        dec_n_points=1,
+        return_intermediate_dec=True,
+        lite_refpoint_refine=True,
+        use_grouppose_keypoints=False,
+    )
+    outputs_default = transformer(srcs, masks, pos_embeds, refpoint_embed, query_feat, cross_attn_srcs=None)
+    outputs_explicit = transformer(srcs, masks, pos_embeds, refpoint_embed, query_feat, cross_attn_srcs=srcs)
+
+    assert len(outputs_default) == len(outputs_explicit) == 4
+    for default_part, explicit_part in zip(outputs_default, outputs_explicit):
+        if default_part is None:
+            assert explicit_part is None
+        else:
+            torch.testing.assert_close(default_part, explicit_part)

tests/models/test_transformer_onnx_spatial_shapes.py

@@ -0,0 +1,409 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""ONNX-export regression tests for the shared Transformer's spatial_shapes path.
+
+These guard the fix that builds ``spatial_shapes`` from symbolic feature-map shapes (``Shape`` -> ``Concat``) instead of
+``torch.empty`` + in-place index assignment. The latter (added in #871 to keep the trace symbolic for dynamic-batch
+export) emitted a ``ScatterND`` that fed a shape tensor, which TensorRT rejects ("IScatterLayer cannot be used to
+compute a shape tensor"). The constant-baking ``torch.as_tensor`` alternative avoids the ScatterND but regresses the
+symbolic trace back to a baked constant.
+
+The Transformer is shared by detection, segmentation and keypoint models, so a single low-level export here covers the
+spatial_shapes path for all of them.
+"""
+
+import inspect
+import io
+
+import numpy as np
+import pytest
+import torch
+from torch import nn
+
+onnx = pytest.importorskip("onnx", reason="onnx not installed; skip ONNX export tests")
+
+from rfdetr.models.transformer import Transformer  # noqa: E402
+
+# CI guard: torch._shape_as_tensor is a private ATen API used on the live forward path in
+# Transformer.forward(). If a future PyTorch upgrade removes it, this assertion fails
+# loudly in CI before users hit an AttributeError at inference time.
+assert hasattr(torch, "_shape_as_tensor"), (
+    "torch._shape_as_tensor not found — update spatial_shapes construction in "
+    "Transformer.forward() for the new PyTorch API."
+)
+
+# dynamo kwarg was added in PyTorch 2.1; our minimum is >=2.2.0, so it is always present.
+# Check via inspect so that a future signature removal surfaces here rather than as a
+# confusing TypeError inside torch.onnx.export.
+_DYNAMO_KWARG: dict[str, bool] = (
+    {"dynamo": False} if "dynamo" in inspect.signature(torch.onnx.export).parameters else {}
+)
+
+
+class _TransformerExportWrapper(nn.Module):
+    """Wrap Transformer.forward with flat tensor args for 2-level ONNX export.
+
+    ``torch.onnx.export`` (TorchScript/non-dynamo path) cannot trace Python list arguments; this wrapper flattens the
+    list args of ``Transformer.forward`` into positional tensor arguments.
+    """
+
+    def __init__(self, transformer: Transformer) -> None:
+        super().__init__()
+        self.transformer = transformer
+
+    def forward(
+        self,
+        s0: torch.Tensor,
+        s1: torch.Tensor,
+        p0: torch.Tensor,
+        p1: torch.Tensor,
+        m0: torch.Tensor,
+        m1: torch.Tensor,
+        refpoint_embed: torch.Tensor,
+        query_feat: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run Transformer with two feature levels and return the first decoder output.
+
+        Args:
+            s0: First-level feature map of shape ``(B, C, H0, W0)``.
+            s1: Second-level feature map of shape ``(B, C, H1, W1)``.
+            p0: Positional embeddings for level 0, same shape as ``s0``.
+            p1: Positional embeddings for level 1, same shape as ``s1``.
+            m0: Boolean padding mask for level 0 of shape ``(B, H0, W0)``.
+            m1: Boolean padding mask for level 1 of shape ``(B, H1, W1)``.
+            refpoint_embed: Reference point embeddings of shape ``(num_queries, 4)``.
+            query_feat: Query feature embeddings of shape ``(num_queries, C)``.
+
+        Returns:
+            First intermediate decoder output tensor.
+        """
+        outputs = self.transformer([s0, s1], [m0, m1], [p0, p1], refpoint_embed, query_feat, cross_attn_srcs=None)
+        return outputs[0]
+
+
+class _TransformerExportWrapper1Level(nn.Module):
+    """Wrap Transformer.forward with flat tensor args for 1-level ONNX export.
+
+    Production RF-DETR models (Base, Small, Nano, Medium, Large) use a single feature level
+    (``projector_scale=["P4"]``). This wrapper exercises that path.
+    """
+
+    def __init__(self, transformer: Transformer) -> None:
+        super().__init__()
+        self.transformer = transformer
+
+    def forward(
+        self,
+        s0: torch.Tensor,
+        p0: torch.Tensor,
+        m0: torch.Tensor,
+        refpoint_embed: torch.Tensor,
+        query_feat: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run Transformer with one feature level and return the first decoder output.
+
+        Args:
+            s0: Feature map of shape ``(B, C, H0, W0)``.
+            p0: Positional embeddings for level 0, same shape as ``s0``.
+            m0: Boolean padding mask for level 0 of shape ``(B, H0, W0)``.
+            refpoint_embed: Reference point embeddings of shape ``(num_queries, 4)``.
+            query_feat: Query feature embeddings of shape ``(num_queries, C)``.
+
+        Returns:
+            First intermediate decoder output tensor.
+        """
+        outputs = self.transformer([s0], [m0], [p0], refpoint_embed, query_feat, cross_attn_srcs=None)
+        return outputs[0]
+
+
+# ---------------------------------------------------------------------------
+# Session-scoped fixtures — build and export once per test session
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="session")
+def transformer_wrapper_2lvl() -> _TransformerExportWrapper:
+    """Build a 2-level Transformer and wrap it for ONNX export.
+
+    Returns:
+        Eval-mode ``_TransformerExportWrapper`` with ``d_model=16``.
+    """
+    transformer = Transformer(
+        d_model=16,
+        num_queries=6,
+        num_decoder_layers=1,
+        sa_nhead=4,
+        ca_nhead=4,
+        num_feature_levels=2,
+        dec_n_points=1,
+        return_intermediate_dec=True,
+        lite_refpoint_refine=True,
+        use_grouppose_keypoints=False,
+    )
+    return _TransformerExportWrapper(transformer.eval()).eval()
+
+
+@pytest.fixture(scope="session")
+def transformer_wrapper_1lvl() -> _TransformerExportWrapper1Level:
+    """Build a 1-level Transformer and wrap it for ONNX export.
+
+    Returns:
+        Eval-mode ``_TransformerExportWrapper1Level`` with ``d_model=16``.
+    """
+    transformer = Transformer(
+        d_model=16,
+        num_queries=6,
+        num_decoder_layers=1,
+        sa_nhead=4,
+        ca_nhead=4,
+        num_feature_levels=1,
+        dec_n_points=1,
+        return_intermediate_dec=True,
+        lite_refpoint_refine=True,
+        use_grouppose_keypoints=False,
+    )
+    return _TransformerExportWrapper1Level(transformer.eval()).eval()
+
+
+@pytest.fixture(scope="session")
+def example_inputs_2lvl() -> tuple[torch.Tensor, ...]:
+    """Return fixed 2-level example input tensors for ``_TransformerExportWrapper``.
+
+    Level 0: spatial size 4×4. Level 1: spatial size 2×2. Batch size 1, ``d_model=16``.
+
+    Returns:
+        8-tuple ``(s0, s1, p0, p1, m0, m1, refpoint_embed, query_feat)``.
+    """
+    return (
+        torch.randn(1, 16, 4, 4),
+        torch.randn(1, 16, 2, 2),
+        torch.randn(1, 16, 4, 4),
+        torch.randn(1, 16, 2, 2),
+        torch.zeros(1, 4, 4, dtype=torch.bool),
+        torch.zeros(1, 2, 2, dtype=torch.bool),
+        torch.rand(6, 4),
+        torch.randn(6, 16),
+    )
+
+
+@pytest.fixture(scope="session")
+def example_inputs_1lvl() -> tuple[torch.Tensor, ...]:
+    """Return fixed 1-level example input tensors for ``_TransformerExportWrapper1Level``.
+
+    Level 0: spatial size 4×4. Batch size 1, ``d_model=16``.
+
+    Returns:
+        5-tuple ``(s0, p0, m0, refpoint_embed, query_feat)``.
+    """
+    return (
+        torch.randn(1, 16, 4, 4),
+        torch.randn(1, 16, 4, 4),
+        torch.zeros(1, 4, 4, dtype=torch.bool),
+        torch.rand(6, 4),
+        torch.randn(6, 16),
+    )
+
+
+@pytest.fixture(scope="session")
+def exported_static_onnx(
+    tmp_path_factory: pytest.TempPathFactory,
+    transformer_wrapper_2lvl: _TransformerExportWrapper,
+    example_inputs_2lvl: tuple[torch.Tensor, ...],
+) -> "onnx.ModelProto":
+    """Export the 2-level Transformer to ONNX (static batch) and return the loaded proto.
+
+    Returns:
+        Loaded ``onnx.ModelProto`` for structural graph assertions.
+    """
+    out = tmp_path_factory.mktemp("onnx_static") / "transformer.onnx"
+    torch.onnx.export(
+        transformer_wrapper_2lvl,
+        example_inputs_2lvl,
+        str(out),
+        input_names=["s0", "s1", "p0", "p1", "m0", "m1", "refpoint_embed", "query_feat"],
+        output_names=["hs"],
+        opset_version=17,
+        **_DYNAMO_KWARG,
+    )
+    return onnx.load(str(out))
+
+
+@pytest.fixture(scope="session")
+def exported_dynamic_onnx_bytes(
+    transformer_wrapper_2lvl: _TransformerExportWrapper,
+    example_inputs_2lvl: tuple[torch.Tensor, ...],
+) -> bytes:
+    """Export the 2-level Transformer with dynamic batch axis and return ONNX bytes.
+
+    Returns:
+        Serialized ONNX model bytes for onnxruntime inference with variable batch sizes.
+    """
+    buf = io.BytesIO()
+    torch.onnx.export(
+        transformer_wrapper_2lvl,
+        example_inputs_2lvl,
+        buf,
+        input_names=["s0", "s1", "p0", "p1", "m0", "m1", "refpoint_embed", "query_feat"],
+        output_names=["hs"],
+        dynamic_axes={
+            "s0": {0: "batch"},
+            "s1": {0: "batch"},
+            "p0": {0: "batch"},
+            "p1": {0: "batch"},
+            "m0": {0: "batch"},
+            "m1": {0: "batch"},
+        },
+        opset_version=17,
+        **_DYNAMO_KWARG,
+    )
+    return buf.getvalue()
+
+
+@pytest.fixture(scope="session")
+def exported_1lvl_onnx(
+    tmp_path_factory: pytest.TempPathFactory,
+    transformer_wrapper_1lvl: _TransformerExportWrapper1Level,
+    example_inputs_1lvl: tuple[torch.Tensor, ...],
+) -> "onnx.ModelProto":
+    """Export the 1-level Transformer to ONNX and return the loaded proto.
+
+    Returns:
+        Loaded ``onnx.ModelProto`` for structural graph assertions.
+    """
+    out = tmp_path_factory.mktemp("onnx_1lvl") / "transformer_1lvl.onnx"
+    torch.onnx.export(
+        transformer_wrapper_1lvl,
+        example_inputs_1lvl,
+        str(out),
+        input_names=["s0", "p0", "m0", "refpoint_embed", "query_feat"],
+        output_names=["hs"],
+        opset_version=17,
+        **_DYNAMO_KWARG,
+    )
+    return onnx.load(str(out))
+
+
+# ---------------------------------------------------------------------------
+# Tests: structural ONNX graph assertions (2-level static export)
+# ---------------------------------------------------------------------------
+
+
+def test_spatial_shapes_export_has_no_scatternd(exported_static_onnx: "onnx.ModelProto") -> None:
+    """The exported Transformer must not contain a ScatterND (TRT shape-tensor killer)."""
+    op_types = [n.op_type for n in exported_static_onnx.graph.node]
+    assert "ScatterND" not in op_types, (
+        "ScatterND reintroduced in Transformer export — spatial_shapes is no longer "
+        "built from symbolic Shape ops; this breaks TensorRT engine building."
+    )
+
+
+def test_spatial_shapes_export_is_shape_derived(exported_static_onnx: "onnx.ModelProto") -> None:
+    """Sanity-check that the exported 2-level Transformer graph contains Shape ops.
+
+    Note: ``spatial_shapes`` itself traces as a ``Constant`` node in TorchScript ONNX export —
+    the tracer records concrete H,W values at trace time, not ``Shape`` ops. The ``Shape`` ops
+    present here originate from other model computations (e.g. batch-size extraction). The true
+    regression guards are ``test_spatial_shapes_export_has_no_scatternd`` (no ScatterND) and
+    ``test_spatial_shapes_dynamic_batch_inference`` (runtime correctness at variable batch size).
+    """
+    op_types = [n.op_type for n in exported_static_onnx.graph.node]
+    assert "Shape" in op_types, "No Shape op in 2-level Transformer ONNX graph — unexpected graph structure change."
+
+
+# ---------------------------------------------------------------------------
+# Tests: dynamic-batch export (2-level)
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("batch_size", [pytest.param(1, id="batch1"), pytest.param(2, id="batch2")])
+def test_spatial_shapes_dynamic_batch_inference(
+    exported_dynamic_onnx_bytes: bytes,
+    batch_size: int,
+) -> None:
+    """Dynamic-batch Transformer ONNX must run onnxruntime inference at batch != trace batch.
+
+    Regression guard: a baked batch constant in ``spatial_shapes`` or any upstream tensor
+    would cause shape mismatches at runtime for any batch size other than the trace batch (1).
+    """
+    onnxruntime = pytest.importorskip("onnxruntime", reason="onnxruntime not installed")
+    session = onnxruntime.InferenceSession(exported_dynamic_onnx_bytes, providers=["CPUExecutionProvider"])
+    # The TorchScript tracer may constant-fold positional embeddings (p0, p1) into the
+    # graph; query the actual session inputs rather than assuming all 8 are present.
+    actual_names = {inp.name for inp in session.get_inputs()}
+    candidate_feeds = {
+        "s0": np.random.randn(batch_size, 16, 4, 4).astype(np.float32),
+        "s1": np.random.randn(batch_size, 16, 2, 2).astype(np.float32),
+        "p0": np.random.randn(batch_size, 16, 4, 4).astype(np.float32),
+        "p1": np.random.randn(batch_size, 16, 2, 2).astype(np.float32),
+        "m0": np.zeros((batch_size, 4, 4), dtype=bool),
+        "m1": np.zeros((batch_size, 2, 2), dtype=bool),
+        "refpoint_embed": np.random.rand(6, 4).astype(np.float32),
+        "query_feat": np.random.randn(6, 16).astype(np.float32),
+    }
+    feeds = {k: v for k, v in candidate_feeds.items() if k in actual_names}
+    (hs,) = session.run(None, feeds)
+    assert hs.shape[1] == batch_size, f"Expected batch dim=={batch_size} at index 1, got shape {hs.shape}"
+
+
+# ---------------------------------------------------------------------------
+# Tests: single-level export (production models use 1 feature level)
+# ---------------------------------------------------------------------------
+
+
+def test_spatial_shapes_single_level_export_has_no_scatternd(
+    exported_1lvl_onnx: "onnx.ModelProto",
+) -> None:
+    """Single-level Transformer ONNX export must not contain ScatterND."""
+    op_types = [n.op_type for n in exported_1lvl_onnx.graph.node]
+    assert "ScatterND" not in op_types, (
+        "ScatterND present in single-level Transformer export — torch.stack on a "
+        "one-element list must still produce a Shape-derived result."
+    )
+
+
+def test_spatial_shapes_single_level_export_is_shape_derived(
+    exported_1lvl_onnx: "onnx.ModelProto",
+) -> None:
+    """Sanity-check that the exported 1-level Transformer graph contains Shape ops.
+
+    Note: ``spatial_shapes`` itself traces as a ``Constant`` node in TorchScript ONNX export —
+    the tracer records the concrete H,W value at trace time. The ``Shape`` ops present here
+    originate from other model computations. The true regression guards are
+    ``test_spatial_shapes_single_level_export_has_no_scatternd`` and the dynamic-batch
+    inference test.
+    """
+    op_types = [n.op_type for n in exported_1lvl_onnx.graph.node]
+    assert "Shape" in op_types, "No Shape op in 1-level Transformer ONNX graph — unexpected graph structure change."
+
+
+# ---------------------------------------------------------------------------
+# Tests: level_start_index numerical correctness
+# ---------------------------------------------------------------------------
+
+
+def test_level_start_index_correctness_two_levels() -> None:
+    """spatial_shapes construction must extract correct (H, W) for non-square feature maps.
+
+    Uses H0=8, W0=6 and H1=4, W1=3 — non-square and H≠W so a transposed [2:4] slice would produce wrong values. Verifies
+    that the ``torch.stack(_shape_as_tensor)`` formula in ``Transformer.forward()`` yields
+    ``spatial_shapes=[[8,6],[4,3]]`` and ``level_start_index=[0, 48]`` (cumulative H*W: 8*6=48, then 48+4*3=60 but index
+    stops at level boundaries → [0, 48]).
+    """
+    s0 = torch.randn(1, 16, 8, 6)
+    s1 = torch.randn(1, 16, 4, 3)
+    srcs = [s0, s1]
+
+    spatial_shapes = torch.stack([torch._shape_as_tensor(src)[2:4] for src in srcs]).to(
+        device=s0.device, dtype=torch.long
+    )
+    level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]))
+
+    assert torch.equal(spatial_shapes, torch.tensor([[8, 6], [4, 3]], dtype=torch.long)), (
+        f"spatial_shapes mismatch: got {spatial_shapes.tolist()}"
+    )
+    assert torch.equal(level_start_index, torch.tensor([0, 48], dtype=torch.long)), (
+        f"level_start_index expected [0, 48], got {level_start_index.tolist()}"
+    )

tests/models/test_validate_shape_dims.py

@@ -0,0 +1,152 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for :func:`rfdetr.detr._validate_shape_dims` and :func:`rfdetr.detr._resolve_patch_size`.
+
+Tests call each helper directly so each validation path has a single focused test without the export/predict scaffolding
+overhead.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+
+from rfdetr.detr import _resolve_patch_size, _validate_shape_dims
+
+
+class TestValidateShapeDimsHappyPath:
+    """_validate_shape_dims returns normalised (height, width) for valid inputs."""
+
+    def test_exact_plain_ints(self) -> None:
+        """Plain int dims divisible by block_size are returned unchanged."""
+        assert _validate_shape_dims((56, 112), 14, 14, 1) == (56, 112)
+
+    def test_returns_plain_int_tuple(self) -> None:
+        """Return type is always a tuple of plain Python int."""
+        h, w = _validate_shape_dims((56, 56), 14, 14, 1)
+        assert type(h) is int and type(w) is int
+
+    def test_numpy_int_accepted(self) -> None:
+        """numpy.int64 dims are accepted via operator.index and normalised."""
+        import numpy as np
+
+        h, w = _validate_shape_dims((np.int64(56), np.int64(112)), 14, 14, 1)
+        assert (h, w) == (56, 112)
+        assert type(h) is int and type(w) is int
+
+    def test_non_square_shape(self) -> None:
+        """Non-square (H != W) shapes are returned correctly."""
+        assert _validate_shape_dims((112, 224), 14, 14, 1) == (112, 224)
+
+    def test_block_size_from_num_windows(self) -> None:
+        """block_size = patch_size * num_windows; both dims divisible by it."""
+        # patch_size=16, num_windows=2 → block_size=32
+        assert _validate_shape_dims((64, 128), 32, 16, 2) == (64, 128)
+
+
+class TestValidateShapeDimsArityErrors:
+    """_validate_shape_dims raises ValueError for non-two-element shapes."""
+
+    def test_one_element_raises(self) -> None:
+        """Single-element tuple must raise ValueError."""
+        with pytest.raises(ValueError, match="shape must be a sequence"):
+            _validate_shape_dims((56,), 14, 14, 1)
+
+    def test_three_element_raises(self) -> None:
+        """Three-element tuple must raise ValueError."""
+        with pytest.raises(ValueError, match="shape must be a sequence"):
+            _validate_shape_dims((56, 56, 3), 14, 14, 1)
+
+    def test_scalar_raises(self) -> None:
+        """Bare scalar (not a sequence) must raise ValueError."""
+        with pytest.raises(ValueError, match="shape must be a sequence"):
+            _validate_shape_dims(56, 14, 14, 1)  # type: ignore[arg-type]
+
+
+class TestValidateShapeDimsInvalidDim:
+    """_validate_shape_dims rejects bool, float, and non-positive dimension values."""
+
+    @pytest.mark.parametrize("shape,match", [((True, 56), "height"), ((56, False), "width")])
+    def test_bool_dim_raises(self, shape: tuple, match: str) -> None:
+        """Bool dims must raise ValueError even though bool is an int subtype."""
+        with pytest.raises(ValueError, match=f"{match} must be an integer"):
+            _validate_shape_dims(shape, 14, 14, 1)  # type: ignore[arg-type]
+
+    @pytest.mark.parametrize("shape", [(56.0, 56.0), (56.0, 56), (56, 56.0)])
+    def test_float_dim_raises(self, shape: tuple) -> None:
+        """Float dims must raise ValueError (operator.index rejects them)."""
+        with pytest.raises(ValueError, match="must be an integer"):
+            _validate_shape_dims(shape, 14, 14, 1)
+
+    @pytest.mark.parametrize("shape", [(0, 56), (56, 0), (-14, 56), (56, -14)])
+    def test_non_positive_dim_raises(self, shape: tuple[int, int]) -> None:
+        """Zero or negative dims must raise ValueError."""
+        with pytest.raises(ValueError, match="positive integers"):
+            _validate_shape_dims(shape, 14, 14, 1)
+
+
+class TestValidateShapeDimsDivisibilityCheck:
+    """_validate_shape_dims enforces divisibility by block_size."""
+
+    @pytest.mark.parametrize(
+        "shape, block_size",
+        [
+            pytest.param((55, 56), 14, id="height_not_divisible"),
+            pytest.param((56, 55), 14, id="width_not_divisible"),
+            pytest.param((48, 64), 32, id="height_not_divisible_large_block"),
+        ],
+    )
+    def test_indivisible_shape_raises(self, shape: tuple[int, int], block_size: int) -> None:
+        """Shapes not divisible by block_size must raise ValueError."""
+        with pytest.raises(ValueError, match=f"divisible by {block_size}"):
+            _validate_shape_dims(shape, block_size, 14, 1)
+
+    def test_error_message_includes_patch_size_and_num_windows(self) -> None:
+        """Error message must name patch_size and num_windows for debuggability."""
+        with pytest.raises(ValueError, match="patch_size=16") as exc_info:
+            _validate_shape_dims((48, 64), 32, 16, 2)
+        assert "num_windows=2" in str(exc_info.value)
+
+
+class TestResolvePatchSize:
+    """_resolve_patch_size resolves and validates patch_size for export()/predict()."""
+
+    def _cfg(self, patch_size: int) -> SimpleNamespace:
+        """Return a minimal model_config stub with the given patch_size."""
+        return SimpleNamespace(patch_size=patch_size)
+
+    def test_none_reads_from_model_config(self) -> None:
+        """patch_size=None resolves to model_config.patch_size."""
+        assert _resolve_patch_size(None, self._cfg(16), "export") == 16
+
+    def test_none_falls_back_to_14_when_config_missing(self) -> None:
+        """patch_size=None falls back to 14 when model_config has no patch_size."""
+        assert _resolve_patch_size(None, SimpleNamespace(), "export") == 14
+
+    def test_explicit_matching_config_accepted(self) -> None:
+        """Providing patch_size equal to model_config.patch_size succeeds."""
+        assert _resolve_patch_size(14, self._cfg(14), "export") == 14
+
+    def test_explicit_mismatch_raises(self) -> None:
+        """Providing patch_size != model_config.patch_size must raise ValueError."""
+        with pytest.raises(ValueError, match="does not match"):
+            _resolve_patch_size(16, self._cfg(14), "export")
+
+    def test_mismatch_error_includes_caller_name(self) -> None:
+        """Mismatch error message includes the caller name for context."""
+        with pytest.raises(ValueError, match="predict"):
+            _resolve_patch_size(16, self._cfg(14), "predict")
+
+    @pytest.mark.parametrize("bad", [0, -1, True, False])
+    def test_invalid_explicit_patch_size_raises(self, bad: int) -> None:
+        """Non-positive-int patch_size must raise ValueError before the mismatch check."""
+        cfg = SimpleNamespace(patch_size=bad)
+        with pytest.raises(ValueError, match="patch_size must be a positive integer"):
+            _resolve_patch_size(bad, cfg, "export")
+
+    def test_invalid_config_patch_size_raises(self) -> None:
+        """Bad patch_size in model_config (when caller passes None) must raise ValueError."""
+        with pytest.raises(ValueError, match="patch_size must be a positive integer"):
+            _resolve_patch_size(None, SimpleNamespace(patch_size=0), "export")

tests/run_smoke_all_models.py

@@ -0,0 +1,290 @@
+#!/usr/bin/env python3
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Smoke-test model instantiation and basic inference with all available weights.
+
+Tests detection, segmentation, and keypoint-preview model classes from rf-detr by importing and instantiating them.
+Validates: imports, download, MD5 hash, model instantiation, from_checkpoint round-trip, and basic inference smoke
+checks.
+
+Usage:
+    python tests/run_smoke_all_models.py
+"""
+
+import argparse
+import os
+import sys
+import tempfile
+from functools import partial
+
+import torch
+from tqdm.auto import tqdm
+
+import rfdetr as _rfdetr
+from rfdetr import (
+    RFDETRKeypointPreview,
+    RFDETRLarge,
+    RFDETRMedium,
+    RFDETRNano,
+    RFDETRSeg2XLarge,
+    RFDETRSegLarge,
+    RFDETRSegMedium,
+    RFDETRSegNano,
+    RFDETRSegSmall,
+    RFDETRSegXLarge,
+    RFDETRSmall,
+)
+
+try:
+    from rfdetr import RFDETR2XLarge, RFDETRXLarge
+except ImportError:
+    RFDETR2XLarge = None
+    RFDETRXLarge = None
+
+# Explicitly list all models to validate
+MODELS_TO_TEST = [
+    # Detection Models
+    RFDETRNano,
+    RFDETRSmall,
+    RFDETRMedium,
+    RFDETRLarge,
+    # Keypoint Models
+    RFDETRKeypointPreview,
+    # Segmentation Models
+    RFDETRSegNano,
+    RFDETRSegSmall,
+    RFDETRSegMedium,
+    RFDETRSegLarge,
+    RFDETRSegXLarge,
+    RFDETRSeg2XLarge,
+]
+
+if RFDETRXLarge is not None:
+    MODELS_TO_TEST.append(partial(RFDETRXLarge, accept_platform_model_license=True))
+if RFDETR2XLarge is not None:
+    MODELS_TO_TEST.append(partial(RFDETR2XLarge, accept_platform_model_license=True))
+
+# 1008 = LCM(12, 16) × 21: valid for all patch sizes (PE=63 for det ÷16,
+# PE=84 for seg ÷12). Each model is tested at its default resolution and at
+# 1008 (regression #1038).
+#
+# Note on Base: ``RFDETRBaseConfig.positional_encoding_size = 37`` is *not*
+# formula-derived (see test_load_pretrain_weights.py:TestLoadPretrainWeightsPEInterpolation
+# ::test_base_config_non_formula_pe_is_interpolated_from_smaller_checkpoint),
+# so this `÷16` description applies only to Nano/Small/Medium/Large.
+_CUSTOM_RESOLUTION = 1008
+
+# Plus models (XLarge / 2XLarge) are heavy enough that running them at
+# resolution=1008 risks the 15-min CI timeout on windows-latest / macos-latest
+# runners.  Smaller models still exercise the 1008 path for #1038 coverage.
+_HEAVY_MODEL_NAMES = {
+    "xlarge",
+    "2xlarge",
+    "xxlarge",
+    "seg-xlarge",
+    "seg-2xlarge",
+    "seg-xxlarge",
+    "rfdetr-xlarge",
+    "rfdetr-2xlarge",
+    "rfdetr-xxlarge",
+    "rfdetr-keypoint-preview",
+    "rfdetr-seg-xlarge",
+    "rfdetr-seg-2xlarge",
+    "rfdetr-seg-xxlarge",
+}
+
+
+def _test_from_checkpoint(
+    model_instance: object, actual_cls: type, extra_kwargs: dict, *, test_starter: bool = True
+) -> None:
+    """Round-trip a model through from_checkpoint using a temp training checkpoint.
+
+    Saves the instantiated model's weights into a minimal training-style checkpoint (``{"args": ..., "model":
+    state_dict}``), calls ``rfdetr.from_checkpoint`` on it, and asserts the returned object is an instance of
+    *actual_cls*.
+
+    Args:
+        model_instance: An already-loaded RFDETR model instance.
+        actual_cls: The expected model class (e.g. ``RFDETRSmall``).
+        extra_kwargs: Extra kwargs to pass to ``from_checkpoint`` (e.g.
+            ``{"accept_platform_model_license": True}`` for plus models).
+        test_starter: When ``True`` (default) also run the starter-like
+            checkpoint round-trip.  Pass ``False`` for non-default resolutions
+            to avoid running the same resolution-independent test multiple times.
+
+    Raises:
+        AssertionError: If the recovered model is not an instance of *actual_cls*.
+        Exception: Propagates any error from ``from_checkpoint`` to the caller.
+    """
+    # Build a minimal training-style checkpoint. The pretrain_weights value only
+    # needs to contain the model-size substring that from_checkpoint matches on
+    # (e.g. "small", "seg-large").  Using cls.size directly satisfies this.
+    fake_pretrain_name = f"{actual_cls.size}.pth"
+    num_classes = model_instance.model.args.num_classes
+    ckpt = {
+        "args": argparse.Namespace(
+            pretrain_weights=fake_pretrain_name,
+            num_classes=num_classes,
+        ),
+        "model": model_instance.model.model.state_dict(),
+    }
+    starter_like_ckpt = {
+        "args": argparse.Namespace(
+            pretrain_weights="none",
+            num_classes=num_classes,
+        ),
+        "model": model_instance.model.model.state_dict(),
+    }
+
+    tmp_fd, tmp_path = tempfile.mkstemp(suffix=".pth")
+    os.close(tmp_fd)
+    try:
+        torch.save(ckpt, tmp_path)
+        recovered = _rfdetr.from_checkpoint(tmp_path, **extra_kwargs)
+        assert recovered is not None, "from_checkpoint returned None"
+        assert hasattr(recovered, "model"), "from_checkpoint result missing 'model' attribute"
+        assert isinstance(recovered, actual_cls), (
+            f"from_checkpoint returned {type(recovered).__name__}, expected {actual_cls.__name__}"
+        )
+    finally:
+        os.unlink(tmp_path)
+
+    if test_starter:
+        starter_tmp_fd, starter_path = tempfile.mkstemp(prefix=f"{actual_cls.size}-starter-", suffix=".pth")
+        os.close(starter_tmp_fd)
+        try:
+            torch.save(starter_like_ckpt, starter_path)
+            starter_recovered = _rfdetr.from_checkpoint(starter_path, **extra_kwargs)
+            assert starter_recovered is not None, "from_checkpoint returned None for starter-like checkpoint"
+            assert hasattr(starter_recovered, "model"), "starter-like from_checkpoint result missing 'model' attribute"
+            got = type(starter_recovered).__name__
+            assert isinstance(starter_recovered, actual_cls), (
+                f"starter-like from_checkpoint returned {got}, expected {actual_cls.__name__}"
+            )
+        finally:
+            os.unlink(starter_path)
+
+
+def _test_coco_class_name_mapping(model_instance: object) -> None:
+    """Verify predict() uses sparse COCO category-ID → class-name mapping.
+
+    Issue #988: RFDETRSegSmall returned "sheep" for class_id=18 instead of "dog" because 0-indexed
+    ``COCO_CLASS_NAMES[18]`` was used instead of the sparse-dict lookup ``COCO_CLASSES[18]``.  threshold=0 forces all
+    top-k queries through so every class ID in the output is covered.  Covers both detection and segmentation nano
+    variants (RFDETRNano, RFDETRSegNano).
+
+    Args:
+        model_instance: An already-loaded pretrained COCO model instance (det or seg).
+
+    Raises:
+        AssertionError: On any class-name mapping failure.
+    """
+    import PIL.Image
+
+    from rfdetr.assets.coco_classes import COCO_CLASS_NAMES, COCO_CLASSES
+
+    # Sanity-check model properties required for the pretrained COCO branch.
+    class_names = model_instance.class_names
+    assert class_names is not None, "Pretrained COCO model must have class_names set"
+    assert len(class_names) == len(COCO_CLASS_NAMES), (
+        f"Expected {len(COCO_CLASS_NAMES)} COCO class names, got {len(class_names)}"
+    )
+    assert class_names == list(COCO_CLASS_NAMES), "model.class_names must equal COCO_CLASS_NAMES"
+    num_classes = model_instance.model.args.num_classes
+    assert num_classes == 90, f"Pretrained COCO model must have num_classes=90, got {num_classes}"
+
+    # Run at threshold=0 to exercise all top-k output slots.
+    img = PIL.Image.new("RGB", (640, 640), color=(128, 128, 128))
+    detections = model_instance.predict(img, threshold=0.0)
+
+    assert "class_name" in detections.data, "data['class_name'] must be present after predict()"
+
+    # For every detection whose class_id is a valid COCO category, class_name must
+    # use sparse-ID lookup (COCO_CLASSES[class_id]), not 0-indexed lookup.
+    # Canonical regression case: class_id=18 → "dog", NOT "sheep" (COCO_CLASS_NAMES[18]).
+    for class_id, class_name in zip(detections.class_id, detections.data["class_name"]):
+        cid = int(class_id)
+        if cid in COCO_CLASSES:
+            expected = COCO_CLASSES[cid]
+            assert class_name == expected, (
+                f"Sparse COCO ID mapping broken (issue #988): "
+                f"class_id={cid} must map to '{expected}', got '{class_name}'"
+            )
+
+    # Regression for PR #1051 HIGH-1: no COCO-pretrained detection may carry
+    # '__background__' — background is implicit (below threshold), never a sentinel label.
+    background_labeled = [
+        (int(cid), name)
+        for cid, name in zip(detections.class_id, detections.data["class_name"])
+        if name == "__background__"
+    ]
+    assert not background_labeled, (
+        "COCO-pretrained predict() must never produce '__background__' class names "
+        f"(PR #1051 HIGH-1 regression); found: {background_labeled[:3]}"
+    )
+
+
+def main() -> None:
+    """Download, validate, instantiate all models, and test from_checkpoint round-trip."""
+    print("Model Instantiation & Download Validation\n")
+
+    succeeded = 0
+    pbar = tqdm(MODELS_TO_TEST, desc="Testing models", unit="model")
+    for model_class in pbar:
+        actual_cls = model_class.func if isinstance(model_class, partial) else model_class
+        extra_kwargs = model_class.keywords if isinstance(model_class, partial) else {}
+        base_name = actual_cls.size
+
+        for res in (None, _CUSTOM_RESOLUTION):
+            # Skip the 1008-resolution variant for heavyweight Plus models — they
+            # risk the 15-min CI timeout on windows-latest / macos-latest runners.
+            if res == _CUSTOM_RESOLUTION and base_name in _HEAVY_MODEL_NAMES:
+                continue
+
+            model_name = base_name if res is None else f"{base_name}@{res}"
+            # Build the kwargs once so `_test_from_checkpoint` and the
+            # instantiation call share the same parameter set (avoids the
+            # `functools.partial.size` AttributeError seen in the previous form).
+            instantiate_kwargs = dict(extra_kwargs)
+            if res is not None:
+                instantiate_kwargs["resolution"] = res
+
+            pbar.set_description(f"Testing {model_name}")
+            try:
+                # Instantiate model class - triggers download, MD5 validation, and loading
+                model_instance = actual_cls(**instantiate_kwargs)
+
+                # Verify model was created
+                assert model_instance is not None, "Model instance is None"
+                assert hasattr(model_instance, "model"), "Model missing 'model' attribute"
+
+                # from_checkpoint round-trip: save a training-style checkpoint and reload it.
+                # Pass the real class (not a partial) so `_test_from_checkpoint` can read
+                # `.size` and `.__name__` and run `isinstance(recovered, actual_cls)`.
+                _test_from_checkpoint(model_instance, actual_cls, instantiate_kwargs, test_starter=(res is None))
+
+                # Inference class-name regression for issue #988 — run on all
+                # nano-sized pretrained COCO models at default resolution only.
+                if "nano" in base_name.lower() and res is None:
+                    _test_coco_class_name_mapping(model_instance)
+
+                succeeded += 1
+            except Exception as ex:
+                # Fail-fast: surface the first failing model directly so CI logs the
+                # root cause cleanly instead of burying it under later cascade failures.
+                pbar.close()
+                print(f"\n[FAIL] {model_name}: {ex}")
+                raise
+
+    pbar.close()
+    print("\nResults:")
+    print(f"\tSucceeded:\t{succeeded}")
+    print("\n[OK] All models validated successfully")
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()

tests/training/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/training/callbacks/__init__.py

@@ -0,0 +1,5 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------

tests/training/callbacks/test_drop_path_callback.py

@@ -0,0 +1,293 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for :class:`rfdetr.training.callbacks.drop_schedule.DropPathCallback`."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+
+from rfdetr.training.callbacks.drop_schedule import DropPathCallback
+from rfdetr.training.drop_schedule import drop_scheduler
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_mock_trainer(global_step: int = 0, estimated_stepping_batches: int = 50) -> MagicMock:
+    """Create a minimal mock Trainer with controllable step metadata."""
+    trainer = MagicMock()
+    trainer.global_step = global_step
+    trainer.estimated_stepping_batches = estimated_stepping_batches
+    return trainer
+
+
+def _make_mock_pl_module(epochs: int = 5) -> MagicMock:
+    """Create a minimal mock RFDETRModule with ``train_config.epochs``."""
+    pl_module = MagicMock()
+    pl_module.train_config.epochs = epochs
+    return pl_module
+
+
+# ---------------------------------------------------------------------------
+# TestDropPathCallbackInit
+# ---------------------------------------------------------------------------
+
+
+class TestDropPathCallbackInit:
+    """Verify constructor defaults."""
+
+    def test_default_args(self) -> None:
+        """Default rates are zero and vit_encoder_num_layers is 12."""
+        cb = DropPathCallback()
+        assert cb._drop_path == 0.0
+        assert cb._dropout == 0.0
+        assert cb._vit_encoder_num_layers == 12
+        assert cb._dp_schedule is None
+        assert cb._do_schedule is None
+
+
+# ---------------------------------------------------------------------------
+# TestOnTrainStart
+# ---------------------------------------------------------------------------
+
+
+class TestOnTrainStart:
+    """Verify schedule arrays built in ``on_train_start``."""
+
+    def test_dp_schedule_matches_drop_scheduler_standard(self) -> None:
+        """drop_path schedule matches ``drop_scheduler`` for standard mode."""
+        cb = DropPathCallback(drop_path=0.3)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        expected = drop_scheduler(0.3, 5, 10)
+        assert cb._dp_schedule is not None
+        np.testing.assert_array_equal(cb._dp_schedule, expected)
+
+    def test_do_schedule_matches_drop_scheduler_standard(self) -> None:
+        """Dropout schedule matches ``drop_scheduler`` for standard mode."""
+        cb = DropPathCallback(dropout=0.1)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        expected = drop_scheduler(0.1, 5, 10)
+        assert cb._do_schedule is not None
+        np.testing.assert_array_equal(cb._do_schedule, expected)
+
+    def test_no_dp_schedule_when_rate_zero(self) -> None:
+        """drop_path=0.0 leaves ``_dp_schedule`` as None."""
+        cb = DropPathCallback(drop_path=0.0)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        assert cb._dp_schedule is None
+
+    def test_dp_schedule_early_mode(self) -> None:
+        """Early mode: rates at step 0 and step 30 match ``drop_scheduler``."""
+        cb = DropPathCallback(drop_path=0.3, cutoff_epoch=2, mode="early")
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        expected = drop_scheduler(0.3, 5, 10, 2, "early")
+        assert cb._dp_schedule is not None
+        assert cb._dp_schedule[0] == expected[0]
+        assert cb._dp_schedule[30] == expected[30]
+
+    def test_dp_schedule_late_mode(self) -> None:
+        """Late mode: rates at step 0 and step 30 match ``drop_scheduler``."""
+        cb = DropPathCallback(drop_path=0.3, cutoff_epoch=2, mode="late")
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        expected = drop_scheduler(0.3, 5, 10, 2, "late")
+        assert cb._dp_schedule is not None
+        assert cb._dp_schedule[0] == expected[0]
+        assert cb._dp_schedule[30] == expected[30]
+
+
+# ---------------------------------------------------------------------------
+# TestOnTrainBatchStart
+# ---------------------------------------------------------------------------
+
+
+class TestOnTrainBatchStart:
+    """Verify model update calls in ``on_train_batch_start``."""
+
+    def test_update_drop_path_called_with_correct_rate(self) -> None:
+        """``update_drop_path`` is called with the schedule value at step 0."""
+        cb = DropPathCallback(drop_path=0.3, vit_encoder_num_layers=6)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        trainer.global_step = 0
+        cb.on_train_batch_start(trainer, pl_module, batch=None, batch_idx=0)
+
+        assert cb._dp_schedule is not None
+        pl_module.model.update_drop_path.assert_called_once_with(cb._dp_schedule[0], 6)
+
+    def test_update_dropout_called_with_correct_rate(self) -> None:
+        """``update_dropout`` is called with the schedule value at step 0."""
+        cb = DropPathCallback(dropout=0.1)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        trainer.global_step = 0
+        cb.on_train_batch_start(trainer, pl_module, batch=None, batch_idx=0)
+
+        assert cb._do_schedule is not None
+        pl_module.model.update_dropout.assert_called_once_with(cb._do_schedule[0])
+
+    def test_no_update_when_step_out_of_bounds(self) -> None:
+        """No model updates when ``global_step`` exceeds schedule length."""
+        cb = DropPathCallback(drop_path=0.3, dropout=0.1)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        trainer.global_step = 9999
+        cb.on_train_batch_start(trainer, pl_module, batch=None, batch_idx=0)
+
+        pl_module.model.update_drop_path.assert_not_called()
+        pl_module.model.update_dropout.assert_not_called()
+
+    @pytest.mark.parametrize(
+        "step",
+        [
+            pytest.param(0, id="first_step"),
+            pytest.param(5, id="mid_step"),
+            pytest.param(9, id="last_of_first_epoch"),
+        ],
+    )
+    def test_drop_rates_at_multiple_steps_match_schedule(self, step: int) -> None:
+        """Each step uses the correct value from the pre-computed schedule."""
+        cb = DropPathCallback(drop_path=0.3, vit_encoder_num_layers=6)
+        trainer = _make_mock_trainer(estimated_stepping_batches=50)
+        pl_module = _make_mock_pl_module(epochs=5)
+
+        cb.on_train_start(trainer, pl_module)
+
+        trainer.global_step = step
+        cb.on_train_batch_start(trainer, pl_module, batch=None, batch_idx=0)
+
+        assert cb._dp_schedule is not None
+        pl_module.model.update_drop_path.assert_called_once_with(cb._dp_schedule[step], 6)
+
+
+# ---------------------------------------------------------------------------
+# TestDropSchedulerValidation
+# ---------------------------------------------------------------------------
+
+
+class TestDropSchedulerValidation:
+    """Verify drop_scheduler raises for invalid inputs."""
+
+    @pytest.mark.parametrize(
+        "cutoff_epoch",
+        [
+            pytest.param(6, id="above_epochs"),
+            pytest.param(-1, id="negative"),
+        ],
+    )
+    def test_raises_for_invalid_cutoff_epoch(self, cutoff_epoch: int) -> None:
+        """drop_scheduler raises ValueError when cutoff_epoch is outside [0, epochs]."""
+        with pytest.raises(ValueError, match="cutoff_epoch must be in"):
+            drop_scheduler(0.3, 5, 10, cutoff_epoch=cutoff_epoch, mode="early")
+
+    @pytest.mark.parametrize(
+        ("epochs", "niter_per_ep", "match"),
+        [
+            pytest.param(0, 10, "epochs must be >= 1", id="epochs_zero"),
+            pytest.param(5, 0, "niter_per_ep must be >= 1", id="niter_per_ep_zero"),
+        ],
+    )
+    def test_raises_for_invalid_epoch_counts(self, epochs: int, niter_per_ep: int, match: str) -> None:
+        """drop_scheduler raises ValueError when epochs or niter_per_ep is less than 1."""
+        with pytest.raises(ValueError, match=match):
+            drop_scheduler(0.3, epochs, niter_per_ep)
+
+
+# ---------------------------------------------------------------------------
+# TestDropSchedulerBoundary
+# ---------------------------------------------------------------------------
+
+
+class TestDropSchedulerBoundary:
+    """Verify drop_scheduler with cutoff_epoch at the inclusive boundaries 0 and epochs."""
+
+    @pytest.mark.parametrize(
+        ("cutoff_epoch", "mode", "expected_first", "expected_last"),
+        [
+            pytest.param(0, "early", 0.0, 0.0, id="early_cutoff_zero_all_zeros"),
+            pytest.param(5, "early", 0.3, 0.3, id="early_cutoff_full_all_rate"),
+            pytest.param(0, "late", 0.3, 0.3, id="late_cutoff_zero_all_rate"),
+            pytest.param(5, "late", 0.0, 0.0, id="late_cutoff_full_all_zeros"),
+        ],
+    )
+    def test_boundary_cutoff_epoch(
+        self,
+        cutoff_epoch: int,
+        mode: str,
+        expected_first: float,
+        expected_last: float,
+    ) -> None:
+        """Boundary cutoff_epoch values (0 and epochs) produce correct first and last rates."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=cutoff_epoch, mode=mode)
+        assert schedule[0] == expected_first
+        assert schedule[-1] == expected_last
+
+
+# ---------------------------------------------------------------------------
+# TestDropSchedulerLinear
+# ---------------------------------------------------------------------------
+
+
+class TestDropSchedulerLinear:
+    """Verify drop_scheduler with schedule='linear' in early mode."""
+
+    def test_linear_early_starts_at_drop_rate(self) -> None:
+        """Linear early schedule first value equals drop_rate."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        assert schedule[0] == pytest.approx(0.3, abs=1e-9)
+
+    def test_linear_early_ends_early_phase_at_zero(self) -> None:
+        """Linear early schedule last value of the early phase equals 0."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        assert schedule[19] == pytest.approx(0.0, abs=1e-9)
+
+    def test_linear_early_late_phase_is_zero(self) -> None:
+        """Linear early schedule: all values after cutoff_epoch are zero."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        np.testing.assert_array_equal(schedule[20:], 0.0)
+
+    def test_linear_early_decreases_monotonically(self) -> None:
+        """Linear early schedule values decrease monotonically during the early phase."""
+        schedule = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        assert np.all(np.diff(schedule[:20]) <= 0)
+
+    def test_linear_same_shape_as_constant(self) -> None:
+        """Schedule='linear' output has the same length as schedule='constant'."""
+        linear = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="linear")
+        constant = drop_scheduler(0.3, 5, 10, cutoff_epoch=2, mode="early", schedule="constant")
+        assert linear.shape == constant.shape

tests/training/callbacks/test_ema_callback.py

@@ -0,0 +1,260 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit and parity tests for RFDETREMACallback."""
+
+from __future__ import annotations
+
+import math
+import warnings
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from torch import nn
+from torch.optim.swa_utils import AveragedModel
+
+from rfdetr.training.callbacks.ema import RFDETREMACallback
+from rfdetr.training.model_ema import ModelEma
+
+
+class _EMAContainerModule(nn.Module):
+    """Minimal module with `.model` to mirror RFDETRModelModule shape."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.model = nn.Linear(4, 2)
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+
+class TestAvgFnDecayFormula:
+    """Verify the tau / no-tau decay formula matches ModelEma."""
+
+    @pytest.mark.parametrize(
+        "num_averaged",
+        [
+            pytest.param(0, id="step-0"),
+            pytest.param(5, id="step-5"),
+            pytest.param(99, id="step-99"),
+        ],
+    )
+    def test_tau_zero_uses_fixed_decay(self, num_averaged: int) -> None:
+        """With tau=0 the effective decay equals the base decay at every step."""
+        decay = 0.99
+        cb = RFDETREMACallback(decay=decay, tau=0)
+        ema_val = torch.tensor(1.0)
+        model_val = torch.tensor(2.0)
+
+        result = cb._avg_fn(ema_val, model_val, num_averaged)
+
+        expected = ema_val * decay + model_val * (1.0 - decay)
+        assert torch.allclose(result, expected, atol=1e-7)
+
+    def test_tau_warmup_at_step_1(self) -> None:
+        """At the first call (num_averaged=0) with tau>0 the effective decay uses updates=1 matching ModelEma's
+        1-indexed counter."""
+        decay = 0.993
+        tau = 100
+        cb = RFDETREMACallback(decay=decay, tau=tau)
+        ema_val = torch.tensor(1.0)
+        model_val = torch.tensor(2.0)
+
+        result = cb._avg_fn(ema_val, model_val, num_averaged=0)
+
+        updates = 1  # num_averaged + 1
+        effective_decay = decay * (1 - math.exp(-updates / tau))
+        expected = ema_val * effective_decay + model_val * (1.0 - effective_decay)
+        assert torch.allclose(result, expected, atol=1e-7)
+
+
+class TestModelEmaParity:
+    """Ensure N-step EMA weights match ModelEma exactly."""
+
+    def test_avg_fn_matches_modelema_weight_parity(self) -> None:
+        """Simulate 500 update steps and compare final EMA weights with ModelEma.module to confirm numerical parity."""
+        torch.manual_seed(42)
+        n_steps = 500
+        decay = 0.993
+        tau = 100
+
+        model = nn.Linear(4, 4)
+        model_ema = ModelEma(model, decay=decay, tau=tau)
+        cb = RFDETREMACallback(decay=decay, tau=tau)
+
+        # Initialise manual EMA state from model (same as ModelEma deepcopy)
+        ema_weights: dict[str, torch.Tensor] = {name: p.clone() for name, p in model.named_parameters()}
+
+        for step in range(n_steps):
+            # Perturb model parameters
+            with torch.no_grad():
+                for p in model.parameters():
+                    p.add_(torch.randn_like(p) * 0.01)
+
+            # Update legacy ModelEma
+            model_ema.update(model)
+
+            # Replicate update via callback avg_fn
+            model_weights = {name: p.clone() for name, p in model.named_parameters()}
+            for name in ema_weights:
+                ema_weights[name] = cb._avg_fn(ema_weights[name], model_weights[name], step)
+
+        # Compare
+        legacy_state = dict(model_ema.module.named_parameters())
+        for name, cb_val in ema_weights.items():
+            assert torch.allclose(cb_val, legacy_state[name], atol=1e-5), (
+                f"Parity failed for {name}: max diff = {(cb_val - legacy_state[name]).abs().max().item()}"
+            )
+
+
+class TestShouldUpdate:
+    """Verify should_update triggers on steps and epochs."""
+
+    def test_should_update_on_step(self) -> None:
+        cb = RFDETREMACallback()
+        assert cb.should_update(step_idx=42) is True
+
+    def test_should_update_on_epoch(self) -> None:
+        cb = RFDETREMACallback()
+        assert cb.should_update(epoch_idx=3) is True
+
+    def test_should_update_neither(self) -> None:
+        cb = RFDETREMACallback()
+        assert cb.should_update() is False
+
+
+class TestInit:
+    """Construction and EMA-state access behavior."""
+
+    def test_init_emits_no_user_warning(self) -> None:
+        """Instantiation should not emit runtime UserWarnings."""
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            RFDETREMACallback()
+        user_warns = [w for w in caught if issubclass(w.category, UserWarning)]
+        assert not user_warns
+
+    def test_get_ema_model_state_dict_none_before_setup(self) -> None:
+        """EMA state accessor returns None before averaged model is created."""
+        cb = RFDETREMACallback()
+        assert cb.get_ema_model_state_dict() is None
+
+    def test_get_ema_model_state_dict_returns_model_weights(self) -> None:
+        """EMA state accessor returns the wrapped `.model` state dict."""
+
+        class _Container(nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+                self.model = nn.Linear(4, 2)
+
+        cb = RFDETREMACallback()
+        container = _Container()
+        cb._average_model = AveragedModel(container, avg_fn=cb._avg_fn)
+
+        state = cb.get_ema_model_state_dict()
+
+        assert state is not None
+        assert "weight" in state
+        assert "bias" in state
+
+
+class TestUpdateInterval:
+    """Verify update_interval_steps throttles EMA updates on step hooks."""
+
+    def test_updates_only_on_interval_steps(self) -> None:
+        """update_interval_steps=2 updates on steps 2, 4, ...
+
+        only.
+        """
+        cb = RFDETREMACallback(update_interval_steps=2)
+        cb._average_model = MagicMock()
+
+        trainer = MagicMock()
+        pl_module = MagicMock()
+
+        for step in (1, 2, 3, 4):
+            trainer.global_step = step
+            cb.on_train_batch_end(trainer, pl_module, outputs=None, batch=None, batch_idx=step - 1)
+
+        assert cb._average_model.update_parameters.call_count == 2
+
+
+class TestLegacyEMAResume:
+    """Legacy checkpoint EMA payload is consumed by the callback setup path."""
+
+    def test_setup_loads_pending_legacy_ema_state_into_average_model(self) -> None:
+        """`_pending_legacy_ema_state` must initialize EMA weights at fit setup."""
+        cb = RFDETREMACallback()
+        pl_module = _EMAContainerModule()
+        trainer = MagicMock()
+
+        legacy_ema_state = {k: torch.full_like(v, 2.0) for k, v in pl_module.model.state_dict().items()}
+        pl_module._pending_legacy_ema_state = legacy_ema_state
+
+        cb.setup(trainer, pl_module, stage="fit")
+
+        assert cb._average_model is not None
+        restored = cb._average_model.module.model.state_dict()
+        for key, expected in legacy_ema_state.items():
+            assert torch.allclose(restored[key], expected)
+        assert not hasattr(pl_module, "_pending_legacy_ema_state")
+
+
+class TestSuppressTestSwap:
+    """suppress_test_swap must disable the test-time EMA weight swap while leaving defaults unchanged."""
+
+    @staticmethod
+    def _make_swap_scenario() -> tuple[RFDETREMACallback, _EMAContainerModule]:
+        """Build a module at weight 7.0 with an EMA average model captured at weight 5.0."""
+        cb = RFDETREMACallback()
+        pl_module = _EMAContainerModule()
+        with torch.no_grad():
+            for p in pl_module.parameters():
+                p.fill_(5.0)
+        cb._average_model = AveragedModel(model=pl_module, use_buffers=True, avg_fn=cb._avg_fn)
+        with torch.no_grad():
+            for p in pl_module.parameters():
+                p.fill_(7.0)
+        return cb, pl_module
+
+    def test_default_flag_is_false(self) -> None:
+        """The suppression flag defaults to False so standalone trainer.test() keeps EMA evaluation."""
+        cb = RFDETREMACallback()
+        assert cb.suppress_test_swap is False
+
+    def test_on_test_epoch_start_swaps_by_default(self) -> None:
+        """Without suppression, the test hooks swap live weights (7.0) for EMA weights (5.0)."""
+        cb, pl_module = self._make_swap_scenario()
+        trainer = MagicMock()
+
+        cb.on_test_epoch_start(trainer, pl_module)
+
+        weight = pl_module.model.weight.detach()
+        assert torch.allclose(weight, torch.full_like(weight, 5.0))
+
+    def test_on_test_epoch_start_suppressed_keeps_live_weights(self) -> None:
+        """With suppress_test_swap=True the live weights (7.0) must stay in place during test."""
+        cb, pl_module = self._make_swap_scenario()
+        cb.suppress_test_swap = True
+        trainer = MagicMock()
+
+        cb.on_test_epoch_start(trainer, pl_module)
+
+        weight = pl_module.model.weight.detach()
+        assert torch.allclose(weight, torch.full_like(weight, 7.0))
+
+    def test_on_test_epoch_end_suppressed_does_not_swap(self) -> None:
+        """With suppression active, on_test_epoch_end must not swap EMA weights in unpaired."""
+        cb, pl_module = self._make_swap_scenario()
+        cb.suppress_test_swap = True
+        trainer = MagicMock()
+
+        cb.on_test_epoch_start(trainer, pl_module)
+        cb.on_test_epoch_end(trainer, pl_module)
+
+        weight = pl_module.model.weight.detach()
+        assert torch.allclose(weight, torch.full_like(weight, 7.0))

tests/training/conftest.py

@@ -0,0 +1,132 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Package-level pytest fixtures for tests/training/.
+
+Provides cross-test cleanup that prevents class-level state from leaking between individual tests in the training/ test
+package, plus shared config factory fixtures used across multiple test modules.
+"""
+
+import pytest
+
+from rfdetr.config import RFDETRBaseConfig, SegmentationTrainConfig, TrainConfig
+
+# ---------------------------------------------------------------------------
+# Shared config factory fixtures (used by test_module, test_datamodule,
+# test_args — avoids duplicate fixture definitions across files)
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture
+def base_model_config():
+    """Factory fixture — call with **overrides to get a minimal RFDETRBaseConfig."""
+
+    def _make(**overrides):
+        defaults = dict(pretrain_weights=None, device="cpu", num_classes=5)
+        defaults.update(overrides)
+        return RFDETRBaseConfig(**defaults)
+
+    return _make
+
+
+@pytest.fixture
+def base_train_config(tmp_path):
+    """Factory fixture — call with **overrides to get a minimal TrainConfig.
+
+    tmp_path is injected automatically so test methods do not need to declare it.
+    """
+
+    def _make(**overrides):
+        defaults = dict(
+            dataset_dir=str(tmp_path / "dataset"),
+            output_dir=str(tmp_path / "output"),
+            epochs=10,
+            lr=1e-4,
+            lr_encoder=1.5e-4,
+            batch_size=2,
+            weight_decay=1e-4,
+            lr_drop=8,
+            warmup_epochs=1.0,
+            drop_path=0.0,
+            multi_scale=False,
+            expanded_scales=False,
+            do_random_resize_via_padding=False,
+            grad_accum_steps=1,
+            num_workers=0,
+            tensorboard=False,
+        )
+        defaults.update(overrides)
+        return TrainConfig(**defaults)
+
+    return _make
+
+
+@pytest.fixture
+def seg_train_config(tmp_path):
+    """Factory fixture — call with **overrides to get a minimal SegmentationTrainConfig.
+
+    tmp_path is injected automatically so test methods do not need to declare it.
+    """
+
+    def _make(**overrides):
+        defaults = dict(
+            dataset_dir=str(tmp_path / "dataset"),
+            output_dir=str(tmp_path / "output"),
+            epochs=10,
+            batch_size=2,
+            multi_scale=False,
+            expanded_scales=False,
+            do_random_resize_via_padding=False,
+            grad_accum_steps=1,
+            drop_path=0.0,
+            num_workers=0,
+            tensorboard=False,
+        )
+        defaults.update(overrides)
+        return SegmentationTrainConfig(**defaults)
+
+    return _make
+
+
+# ---------------------------------------------------------------------------
+# Class-level isolation
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(autouse=True)
+def _restore_rfdetr_module_trainer_property():
+    """Restore RFDETRModelModule.trainer to the LightningModule parent property after each test.
+
+    Several unit tests in test_module_model.py patch the ``trainer`` property directly on the ``RFDETRModelModule``
+    class (``type(module).trainer = property(...)``). Without cleanup this mutates the class for the remainder of the
+    session and breaks ``Trainer.fit()`` calls in smoke tests (PTL cannot set ``.trainer`` on the module because the
+    patched property has no setter).
+
+    This fixture deletes any class-level override from ``RFDETRModelModule.__dict__`` after every test, so the next test
+    starts with a clean class that inherits PTL's read/write ``trainer`` descriptor from ``LightningModule``.
+    """
+    yield
+    # Lazy import so the fixture does not force module import at collection time.
+    from rfdetr.training.module_model import RFDETRModelModule
+
+    if "trainer" in RFDETRModelModule.__dict__:
+        delattr(RFDETRModelModule, "trainer")
+
+
+@pytest.fixture(autouse=True)
+def _restore_rfdetr_datamodule_trainer_property():
+    """Restore RFDETRDataModule.trainer to the LightningDataModule parent property after each test.
+
+    Tests that mock the ``trainer`` property on ``RFDETRDataModule`` (e.g. for ``on_after_batch_transfer`` tests) patch
+    it at the class level.  Without cleanup this mutates the class for the remainder of the session.
+
+    This fixture deletes any class-level override from ``RFDETRDataModule.__dict__`` after every test, mirroring the
+    ``_restore_rfdetr_module_trainer_property`` pattern above.
+    """
+    yield
+    from rfdetr.training.module_data import RFDETRDataModule
+
+    if "trainer" in RFDETRDataModule.__dict__:
+        delattr(RFDETRDataModule, "trainer")

tests/training/helpers.py

@@ -0,0 +1,126 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Shared test helpers for the rfdetr.training test suite.
+
+Plain classes and functions (not pytest fixtures) shared across multiple test modules to avoid verbatim duplication.
+Import with a relative import::
+
+    from .helpers import _FakeCriterion, _FakeDataset, _TinyModel
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.utils.data
+
+
+class _TinyModel(nn.Module):
+    """Minimal real nn.Module satisfying the RFDETRModule model contract.
+
+    Has a single trainable parameter so the optimizer has something to update and the loss has a gradient path back
+    through the model.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.dummy = nn.Parameter(torch.zeros(1))
+
+    def forward(self, samples, targets=None):
+        return {"dummy": self.dummy}
+
+    def update_drop_path(self, *args, **kwargs) -> None:
+        pass
+
+    def update_dropout(self, *args, **kwargs) -> None:
+        pass
+
+    def reinitialize_detection_head(self, *args, **kwargs) -> None:
+        pass
+
+
+class _FakeCriterion:
+    """Callable criterion that returns a loss connected to the model output.
+
+    Keeps a gradient path from the loss back to _TinyModel.dummy so that ``loss.backward()`` does not error when the
+    Trainer calls it.
+    """
+
+    weight_dict = {"loss_ce": 1.0}
+
+    def num_boxes_for_targets(self, outputs, targets):
+        dummy = outputs.get("dummy", torch.zeros(1))
+        return torch.ones((), dtype=dummy.dtype, device=dummy.device)
+
+    def __call__(self, outputs, targets, num_boxes=None):
+        dummy = outputs.get("dummy", torch.zeros(1))
+        denominator = self.num_boxes_for_targets(outputs, targets) if num_boxes is None else num_boxes
+        return {"loss_ce": dummy.mean() / denominator}
+
+
+class _FakeDataset(torch.utils.data.Dataset):
+    """Dataset with ``(image, target)`` pairs for detection.
+
+    The image is a ``(3, 32, 32)`` float tensor; the target dict includes the fields expected by RFDETRModule:
+    ``boxes``, ``labels``, ``image_id``, ``orig_size``, ``size``.
+    """
+
+    def __init__(self, length: int = 20) -> None:
+        self._length = length
+
+    def __len__(self) -> int:
+        return self._length
+
+    def __getitem__(self, idx):
+        image = torch.randn(3, 32, 32)
+        target = {
+            "boxes": torch.tensor([[0.5, 0.5, 0.1, 0.1]]),
+            "labels": torch.tensor([1]),
+            "image_id": torch.tensor(idx),
+            "orig_size": torch.tensor([32, 32]),
+            "size": torch.tensor([32, 32]),
+        }
+        return image, target
+
+
+class _FakeDatasetWithMasks(_FakeDataset):
+    """Like _FakeDataset but includes binary instance masks (for segmentation)."""
+
+    def __getitem__(self, idx):
+        image, target = super().__getitem__(idx)
+        target["masks"] = torch.zeros(1, 32, 32, dtype=torch.bool)
+        return image, target
+
+
+class _FakePostProcess:
+    """Picklable postprocessor for ddp_spawn tests.
+
+    ``MagicMock`` is not picklable and cannot survive the subprocess boundary that ``ddp_spawn`` creates.  This plain
+    class is a drop-in replacement.
+
+    Delegates to ``_fake_postprocess``; keep both in sync if the fake output format changes.
+    """
+
+    def __call__(self, outputs, orig_sizes):
+        return _fake_postprocess(outputs, orig_sizes)
+
+
+def _fake_postprocess(outputs, orig_sizes):
+    """Return one non-empty prediction per image so COCOEvalCallback has something to score."""
+    n = orig_sizes.shape[0]
+    return [
+        {
+            "boxes": torch.tensor([[5.0, 5.0, 20.0, 20.0]]),
+            "scores": torch.tensor([0.9]),
+            "labels": torch.tensor([1]),
+        }
+        for _ in range(n)
+    ]
+
+
+def _make_param_dicts(model: nn.Module) -> list[dict]:
+    """Build a minimal param-dict list for AdamW from all trainable parameters."""
+    return [{"params": p, "lr": 1e-4} for p in model.parameters() if p.requires_grad]

tests/training/test_args.py

@@ -0,0 +1,156 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Unit tests for _namespace_from_configs — the canonical config-to-Namespace mapping."""
+
+import pytest
+
+from rfdetr._namespace import _namespace_from_configs
+
+
+class TestNamespaceFromConfigs:
+    """_namespace_from_configs maps ModelConfig + TrainConfig to a Namespace."""
+
+    def test_forwards_model_config_fields(self, base_model_config, base_train_config):
+        """All key ModelConfig fields are faithfully mapped."""
+        mc = base_model_config(num_classes=7)
+        args = _namespace_from_configs(mc, base_train_config())
+
+        assert args.encoder == mc.encoder
+        assert args.num_classes == 7
+        assert args.hidden_dim == mc.hidden_dim
+        assert args.resolution == mc.resolution
+        assert args.patch_size == mc.patch_size
+        assert args.num_windows == mc.num_windows
+        assert args.segmentation_head == mc.segmentation_head
+        assert args.positional_encoding_size == mc.positional_encoding_size
+
+    def test_forwards_train_config_fields(self, base_model_config, base_train_config):
+        """All key TrainConfig fields are faithfully mapped."""
+        tc = base_train_config(
+            lr=3e-4,
+            epochs=20,
+            weight_decay=5e-5,
+            batch_size=4,
+            num_workers=0,
+            eval_interval=3,
+            log_per_class_metrics=False,
+            train_log_sync_dist=True,
+            train_log_on_step=True,
+            compute_val_loss=False,
+            compute_test_loss=False,
+            ema_update_interval=2,
+            prefetch_factor=4,
+        )
+        args = _namespace_from_configs(base_model_config(), tc)
+
+        assert args.lr == pytest.approx(3e-4)
+        assert args.epochs == 20
+        assert args.weight_decay == pytest.approx(5e-5)
+        assert args.batch_size == 4
+        assert args.num_workers == 0
+        assert args.eval_interval == 3
+        assert args.log_per_class_metrics is False
+        assert args.train_log_sync_dist is True
+        assert args.train_log_on_step is True
+        assert args.compute_val_loss is False
+        assert args.compute_test_loss is False
+        assert args.ema_update_interval == 2
+        assert args.prefetch_factor == 4
+
+    def test_forwards_promoted_train_fields(self, base_model_config, base_train_config):
+        """Promoted TrainConfig fields are forwarded to the namespace."""
+        tc = base_train_config(clip_max_norm=0.35, seed=123, sync_bn=True, fp16_eval=True)
+        args = _namespace_from_configs(base_model_config(), tc)
+
+        assert args.clip_max_norm == pytest.approx(0.35)
+        assert args.seed == 123
+        assert args.sync_bn is True
+        assert args.fp16_eval is True
+
+    def test_seed_falls_back_to_legacy_default_when_unset(self, base_model_config, base_train_config):
+        """Seed defaults to 42 in the namespace when TrainConfig.seed is None."""
+        tc = base_train_config(seed=None)
+        args = _namespace_from_configs(base_model_config(), tc)
+        assert args.seed == 42
+
+    def test_forwards_dataset_fields(self, base_model_config, base_train_config):
+        """Dataset-routing fields are forwarded to the Namespace."""
+        tc = base_train_config(multi_scale=True, expanded_scales=True, dataset_file="coco")
+        args = _namespace_from_configs(base_model_config(), tc)
+
+        assert args.multi_scale is True
+        assert args.expanded_scales is True
+        assert args.dataset_file == "coco"
+
+    def test_num_queries_from_subclass_config(self, base_model_config, base_train_config):
+        """num_queries is read from subclass config attributes."""
+        mc = base_model_config()  # RFDETRBaseConfig has num_queries=300
+        args = _namespace_from_configs(mc, base_train_config())
+        assert args.num_queries == 300
+
+    def test_resume_none_becomes_empty_string(self, base_model_config, base_train_config):
+        """Resume=None (the default) is converted to '' for the Namespace."""
+        tc = base_train_config()
+        assert tc.resume is None
+        args = _namespace_from_configs(base_model_config(), tc)
+        assert args.resume == ""
+
+    def test_segmentation_extras_forwarded_from_seg_config(self, base_model_config, seg_train_config):
+        """SegmentationTrainConfig mask loss coefficients are forwarded."""
+        mc = base_model_config(segmentation_head=True)
+        tc = seg_train_config()
+        args = _namespace_from_configs(mc, tc)
+
+        assert args.mask_ce_loss_coef == pytest.approx(5.0)
+        assert args.mask_dice_loss_coef == pytest.approx(5.0)
+
+    def test_segmentation_cls_loss_default_matches_pre_1_7_effective_weight(self, base_model_config, seg_train_config):
+        """Default segmentation classification loss weight must stay at the pre-1.7 effective value."""
+        mc = base_model_config(segmentation_head=True)
+        tc = seg_train_config()
+        args = _namespace_from_configs(mc, tc)
+
+        assert args.cls_loss_coef == pytest.approx(1.0)
+
+    def test_segmentation_cls_loss_explicit_override_is_forwarded(self, base_model_config, seg_train_config):
+        """Explicit segmentation classification loss weight overrides are preserved."""
+        mc = base_model_config(segmentation_head=True)
+        tc = seg_train_config(cls_loss_coef=5.0)
+        args = _namespace_from_configs(mc, tc)
+
+        assert args.cls_loss_coef == pytest.approx(5.0)
+
+    def test_segmentation_num_select_none_falls_back_to_model_config(self, base_model_config, seg_train_config) -> None:
+        """SegmentationTrainConfig(num_select=None) must not overwrite ModelConfig.num_select."""
+        mc = base_model_config(segmentation_head=True, num_select=200)
+        # Explicitly passing num_select=None triggers the deprecation warning (Item #3).
+        with pytest.warns(DeprecationWarning, match="TrainConfig.num_select is deprecated"):
+            tc = seg_train_config(num_select=None)
+
+        args = _namespace_from_configs(mc, tc)
+
+        assert args.num_select == 200
+
+    def test_segmentation_extras_default_for_plain_config(self, base_model_config, base_train_config):
+        """mask_* attributes default to 5.0 for a plain TrainConfig (not segmentation)."""
+        args = _namespace_from_configs(base_model_config(), base_train_config())
+        assert args.mask_ce_loss_coef == pytest.approx(5.0)
+        assert args.mask_dice_loss_coef == pytest.approx(5.0)
+
+    def test_segmentation_head_flag_forwarded(self, base_model_config, base_train_config):
+        """segmentation_head=True from ModelConfig reaches the Namespace."""
+        mc = base_model_config(segmentation_head=True)
+        args = _namespace_from_configs(mc, base_train_config())
+        assert args.segmentation_head is True
+
+    def test_build_namespace_emits_deprecation_warning(
+        self, base_model_config, base_train_config, reset_build_namespace_warning_state
+    ):
+        """build_namespace() must emit a DeprecationWarning on every call."""
+        from rfdetr._namespace import build_namespace
+
+        with pytest.warns(FutureWarning, match="build_namespace"):
+            build_namespace(base_model_config(), base_train_config())

tests/training/test_auto_batch.py

@@ -0,0 +1,228 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from rfdetr.detr import RFDETR
+from rfdetr.training import auto_batch
+from rfdetr.training.auto_batch import AutoBatchResult
+
+
+class _TinyModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.w = torch.nn.Parameter(torch.ones(1))
+
+
+def test_recommend_grad_accum_steps_rounds_up():
+    assert auto_batch.recommend_grad_accum_steps(3, 16) == 6
+
+
+def test_probe_max_micro_batch_uses_exponential_then_binary_search():
+    model = _TinyModule()
+    criterion = _TinyModule()
+    threshold = 7
+
+    def _fake_probe(*args, **kwargs):
+        micro_batch_size = args[2]
+        return micro_batch_size <= threshold
+
+    with (
+        patch("rfdetr.training.auto_batch._probe_step", side_effect=_fake_probe),
+        patch("rfdetr.training.auto_batch.torch.cuda.empty_cache"),
+    ):
+        safe = auto_batch.probe_max_micro_batch(
+            model=model,
+            criterion=criterion,
+            resolution=64,
+            device=torch.device("cuda"),
+            num_classes=5,
+            amp=False,
+            safety_margin=1.0,
+            max_micro_batch=32,
+        )
+    assert safe == threshold
+
+
+def test_probe_max_micro_batch_raises_if_one_is_not_safe():
+    model = _TinyModule()
+    criterion = _TinyModule()
+
+    with (
+        patch("rfdetr.training.auto_batch._probe_step", return_value=False),
+        patch("rfdetr.training.auto_batch.torch.cuda.empty_cache"),
+        pytest.raises(RuntimeError, match="micro_batch_size=1"),
+    ):
+        auto_batch.probe_max_micro_batch(
+            model=model,
+            criterion=criterion,
+            resolution=64,
+            device=torch.device("cuda"),
+            num_classes=5,
+            amp=False,
+        )
+
+
+def test_probe_step_raises_when_loss_keys_do_not_overlap_weight_keys():
+    """_probe_step must fail fast when weighted loss would be empty."""
+
+    class _DummyCriterion(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.weight_dict = {"loss_bbox": 1.0}
+
+        def forward(self, outputs, targets):
+            return {"loss_ce": torch.tensor(1.0)}
+
+    class _DummyModel(torch.nn.Module):
+        def forward(self, samples, targets):
+            return {}
+
+    model = _DummyModel()
+    criterion = _DummyCriterion()
+
+    with (
+        patch(
+            "rfdetr.training.auto_batch._make_synthetic_batch",
+            return_value=(MagicMock(), []),
+        ),
+        pytest.raises(RuntimeError, match="no overlap between criterion loss_dict and weight_dict keys"),
+    ):
+        auto_batch._probe_step(
+            model=model,
+            criterion=criterion,
+            micro_batch_size=1,
+            resolution=64,
+            device=torch.device("cpu"),
+            num_classes=5,
+            amp=False,
+        )
+
+
+def test_resolve_auto_batch_config_requires_cuda():
+    model_context = SimpleNamespace(device=torch.device("cpu"), model=MagicMock())
+    model_config = SimpleNamespace(resolution=64, num_classes=5, amp=False, segmentation_head=False)
+    train_config = SimpleNamespace(batch_size="auto", auto_batch_target_effective=16)
+
+    with (
+        patch("rfdetr.training.auto_batch.torch.cuda.is_available", return_value=False),
+        pytest.raises(RuntimeError, match="requires a CUDA device"),
+    ):
+        auto_batch.resolve_auto_batch_config(model_context, model_config, train_config)
+
+
+def test_resolve_auto_batch_config_returns_expected_values():
+    model_context = SimpleNamespace(device=torch.device("cuda"), model=MagicMock())
+    model_config = SimpleNamespace(resolution=64, num_classes=5, amp=False, segmentation_head=True)
+    train_config = SimpleNamespace(batch_size="auto", auto_batch_target_effective=16)
+    criterion = MagicMock()
+    criterion.to.return_value = criterion
+
+    with (
+        patch("rfdetr.training.auto_batch.torch.cuda.is_available", return_value=True),
+        patch("rfdetr.training.auto_batch.build_criterion_from_config", return_value=(criterion, None)),
+        patch("rfdetr.training.auto_batch.probe_max_micro_batch", return_value=5),
+        patch("rfdetr.training.auto_batch.torch.cuda.get_device_name", return_value="Fake GPU"),
+    ):
+        result = auto_batch.resolve_auto_batch_config(model_context, model_config, train_config)
+
+    assert isinstance(result, AutoBatchResult)
+    assert result.safe_micro_batch == 5
+    assert result.recommended_grad_accum_steps == 4
+    assert result.effective_batch_size == 20
+    assert result.device_name == "Fake GPU"
+
+
+@patch("rfdetr.detr.is_main_process", return_value=False)
+@patch("rfdetr.training.auto_batch.resolve_auto_batch_config")
+@patch("rfdetr.training.build_trainer")
+@patch("rfdetr.training.RFDETRDataModule")
+@patch("rfdetr.training.RFDETRModelModule")
+@patch("rfdetr.detr._move_model_context_to_device")
+def test_train_auto_batch_ensures_model_on_device_before_resolve(
+    mock_move: MagicMock,
+    _mock_module: MagicMock,
+    _mock_data_module: MagicMock,
+    _mock_build_trainer: MagicMock,
+    mock_resolve: MagicMock,
+    _mock_is_main: MagicMock,
+) -> None:
+    """Model weights must be moved before resolve_auto_batch_config when batch_size='auto'."""
+    auto_result = SimpleNamespace(safe_micro_batch=4, recommended_grad_accum_steps=1, effective_batch_size=4)
+    call_order: list[str] = []
+
+    def _move_side_effect(model: object) -> None:
+        call_order.append("ensure")
+
+    def _resolve_side_effect(**_kwargs: object) -> object:
+        call_order.append("resolve")
+        return auto_result
+
+    mock_move.side_effect = _move_side_effect
+    mock_resolve.side_effect = _resolve_side_effect
+
+    train_config = SimpleNamespace(
+        batch_size="auto",
+        grad_accum_steps=99,
+        dataset_dir=None,
+        resume=None,
+        class_names=None,
+        save_dataset_grids=False,
+    )
+    mock_self = MagicMock()
+    mock_self.model_config = SimpleNamespace(model_name=None)
+    mock_self.get_train_config.return_value = train_config
+
+    RFDETR.train(mock_self)
+
+    assert train_config.batch_size == 4
+    assert train_config.grad_accum_steps == 1
+    mock_move.assert_called_once_with(mock_self.model)
+    mock_resolve.assert_called_once_with(
+        model_context=mock_self.model,
+        model_config=mock_self.model_config,
+        train_config=train_config,
+    )
+    assert call_order == ["ensure", "resolve"]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for segmentation probe")
+def test_probe_step_with_real_segmentation_criterion(tmp_path):
+    """Run one probe step with real segmentation model and criterion so loss_masks and t['masks'] are exercised."""
+    from rfdetr._namespace import _namespace_from_configs
+    from rfdetr.config import RFDETRSegNanoConfig, SegmentationTrainConfig
+    from rfdetr.models.lwdetr import build_criterion_and_postprocessors, build_model
+
+    mc = RFDETRSegNanoConfig(pretrain_weights=None, device="cuda", num_classes=2)
+    tc = SegmentationTrainConfig(
+        dataset_dir=str(tmp_path / "ds"),
+        output_dir=str(tmp_path / "out"),
+        batch_size=2,
+        grad_accum_steps=1,
+        tensorboard=False,
+    )
+    args = _namespace_from_configs(mc, tc)
+    model = build_model(args)
+    criterion, _ = build_criterion_and_postprocessors(args)
+    device = torch.device("cuda")
+    model = model.to(device)
+    criterion = criterion.to(device)
+
+    ok = auto_batch._probe_step(
+        model=model,
+        criterion=criterion,
+        micro_batch_size=1,
+        resolution=mc.resolution,
+        device=device,
+        num_classes=mc.num_classes,
+        amp=False,
+        segmentation_head=True,
+    )
+    assert ok is True

tests/training/test_cli.py

@@ -0,0 +1,117 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Tests for RFDETRCli — PTL Ch4/T4.
+
+Verifies that the CLI module is correctly structured: importable, subclasses LightningCLI, overrides
+add_arguments_to_parser, and exposes a callable main() entry point.  CLI integration / smoke tests (--help subprocess,
+YAML roundtrip) live in T4-7.
+"""
+
+import pytest
+
+# ---------------------------------------------------------------------------
+# Structure and importability
+# ---------------------------------------------------------------------------
+
+
+class TestRFDETRCliStructure:
+    """RFDETRCli is correctly structured and importable."""
+
+    def test_cli_module_importable(self):
+        """rfdetr.training.cli imports without error."""
+        import rfdetr.training.cli  # noqa: F401
+
+    def test_rfdetr_cli_importable(self):
+        """RFDETRCli can be imported from rfdetr.training.cli."""
+        from rfdetr.training.cli import RFDETRCli  # noqa: F401
+
+    def test_main_importable(self):
+        """Main() can be imported from rfdetr.training.cli."""
+        from rfdetr.training.cli import main  # noqa: F401
+
+    def test_rfdetr_cli_is_lightning_cli_subclass(self):
+        """RFDETRCli must subclass pytorch_lightning LightningCLI."""
+        from pytorch_lightning.cli import LightningCLI
+
+        from rfdetr.training.cli import RFDETRCli
+
+        assert issubclass(RFDETRCli, LightningCLI)
+
+    def test_main_is_callable(self):
+        """Main must be a callable (function, not e.g. a string)."""
+        from rfdetr.training.cli import main
+
+        assert callable(main)
+
+    def test_add_arguments_to_parser_is_overridden(self):
+        """RFDETRCli overrides add_arguments_to_parser from LightningCLI."""
+        from pytorch_lightning.cli import LightningCLI
+
+        from rfdetr.training.cli import RFDETRCli
+
+        assert RFDETRCli.add_arguments_to_parser is not LightningCLI.add_arguments_to_parser
+
+    def test_exported_from_lit_package(self):
+        """RFDETRCli is exported from rfdetr.training (appears in __all__)."""
+        import rfdetr.training as lit
+
+        assert hasattr(lit, "RFDETRCli")
+        assert "RFDETRCli" in lit.__all__
+
+
+# ---------------------------------------------------------------------------
+# Argument linking
+# ---------------------------------------------------------------------------
+
+
+class TestRFDETRCliArgumentLinking:
+    """add_arguments_to_parser registers the expected argument links."""
+
+    def _collect_links(self):
+        """Instantiate a minimal parser and collect registered link sources."""
+        import unittest.mock as mock
+
+        from rfdetr.training.cli import RFDETRCli
+
+        captured = []
+
+        class _FakeParser:
+            def link_arguments(self, source, target, **kwargs):
+                captured.append({"source": source, "target": target, **kwargs})
+
+            # LightningArgumentParser methods that may be called during setup
+            def __getattr__(self, name):
+                return mock.MagicMock()
+
+        cli = RFDETRCli.__new__(RFDETRCli)
+        cli.add_arguments_to_parser(_FakeParser())
+        return captured
+
+    def test_model_config_link_registered(self):
+        """model.model_config is linked to data.model_config."""
+        links = self._collect_links()
+        sources = [lnk["source"] for lnk in links]
+        assert "model.model_config" in sources
+
+    def test_train_config_link_registered(self):
+        """model.train_config is linked to data.train_config."""
+        links = self._collect_links()
+        sources = [lnk["source"] for lnk in links]
+        assert "model.train_config" in sources
+
+    @pytest.mark.parametrize(
+        "source, expected_target",
+        [
+            pytest.param("model.model_config", "data.model_config", id="model_config"),
+            pytest.param("model.train_config", "data.train_config", id="train_config"),
+        ],
+    )
+    def test_link_target(self, source, expected_target):
+        """Each link points to the correct data.* target."""
+        links = self._collect_links()
+        match = next((lnk for lnk in links if lnk["source"] == source), None)
+        assert match is not None, f"No link registered for source {source!r}"
+        assert match["target"] == expected_target

tests/training/test_drop_path.py

@@ -0,0 +1,203 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+import logging
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+import torch.nn as nn
+
+import rfdetr.models.backbone.dinov2 as dinov2_module
+from rfdetr._namespace import _namespace_from_configs
+from rfdetr.config import RFDETRNanoConfig, TrainConfig
+from rfdetr.models import build_model
+from rfdetr.models.backbone.dinov2 import DinoV2
+from rfdetr.models.backbone.dinov2_with_windowed_attn import (
+    Dinov2WithRegistersDropPath,
+    WindowedDinov2WithRegistersBackbone,
+)
+from rfdetr.models.lwdetr import LWDETR
+
+
+@pytest.fixture
+def model_with_drop_path(monkeypatch: pytest.MonkeyPatch) -> LWDETR:
+    """Create RF-DETR Nano LWDETR with drop_path enabled."""
+    monkeypatch.setattr(
+        WindowedDinov2WithRegistersBackbone,
+        "from_pretrained",
+        classmethod(lambda cls, name, config: cls(config)),
+    )
+    mc = RFDETRNanoConfig(num_classes=3, pretrain_weights=None)
+    tc = TrainConfig(
+        dataset_dir=".",
+        output_dir=".",
+        drop_path=0.1,
+    )
+    args = _namespace_from_configs(mc, tc)
+    return build_model(args)
+
+
+@pytest.fixture
+def model_without_drop_path(monkeypatch: pytest.MonkeyPatch) -> LWDETR:
+    """Create RF-DETR Nano LWDETR without drop_path."""
+    monkeypatch.setattr(
+        WindowedDinov2WithRegistersBackbone,
+        "from_pretrained",
+        classmethod(lambda cls, name, config: cls(config)),
+    )
+    mc = RFDETRNanoConfig(num_classes=3, pretrain_weights=None)
+    tc = TrainConfig(
+        dataset_dir=".",
+        output_dir=".",
+        drop_path=0.0,
+    )
+    args = _namespace_from_configs(mc, tc)
+    return build_model(args)
+
+
+def test_get_backbone_encoder_layers_dinov2(model_with_drop_path: LWDETR) -> None:
+    """Verify _get_backbone_encoder_layers() returns encoder.encoder.layer for DinoV2."""
+    model = model_with_drop_path
+
+    layers = model._get_backbone_encoder_layers()
+    assert layers is not None
+
+    enc = model.backbone[0].encoder
+    assert hasattr(enc, "encoder"), "DinoV2 encoder should have encoder attribute"
+    assert hasattr(enc.encoder, "encoder"), "DinoV2 encoder.encoder should have encoder attribute"
+    assert hasattr(enc.encoder.encoder, "layer"), "DinoV2 encoder.encoder.encoder should have layer attribute"
+    assert layers is enc.encoder.encoder.layer, "Should return encoder.encoder.encoder.layer"
+
+    assert len(layers) > 0, "Should have at least one layer"
+    for layer in layers:
+        assert hasattr(layer, "drop_path"), "Each layer should have drop_path attribute"
+
+
+def test_update_drop_path_dinov2(model_with_drop_path: LWDETR) -> None:
+    """Verify update_drop_path() sets drop_prob values correctly with linear schedule."""
+    model = model_with_drop_path
+
+    layers = model._get_backbone_encoder_layers()
+    assert layers is not None
+
+    num_layers = len(layers)
+    drop_path_rate = 0.1
+
+    model.update_drop_path(drop_path_rate, num_layers)
+
+    # All layers must be Dinov2WithRegistersDropPath (drop_path_rate=0.1 > 0 at model build time).
+    expected_rates = [x.item() for x in torch.linspace(0, drop_path_rate, num_layers)]
+    for i, layer in enumerate(layers):
+        assert isinstance(layer.drop_path, Dinov2WithRegistersDropPath), (
+            f"Layer {i} drop_path should be Dinov2WithRegistersDropPath, got {type(layer.drop_path)}"
+        )
+        actual_prob = layer.drop_path.drop_prob
+        assert abs(actual_prob - expected_rates[i]) < 1e-6, (
+            f"Layer {i} drop_prob should be {expected_rates[i]}, got {actual_prob}"
+        )
+
+    assert abs(layers[0].drop_path.drop_prob - 0.0) < 1e-6, "First layer should have drop_prob = 0"
+    assert abs(layers[-1].drop_path.drop_prob - drop_path_rate) < 1e-6, (
+        f"Last layer should have drop_prob = {drop_path_rate}"
+    )
+
+
+def test_drop_path_initialization(model_with_drop_path: LWDETR, model_without_drop_path: LWDETR) -> None:
+    """Verify drop_path initialization: Dinov2WithRegistersDropPath vs Identity based on rate."""
+    layers_with_dp = model_with_drop_path._get_backbone_encoder_layers()
+    layers_without_dp = model_without_drop_path._get_backbone_encoder_layers()
+
+    assert layers_with_dp is not None
+    assert layers_without_dp is not None
+
+    # drop_path_rate=0.1 -> every layer initialised as Dinov2WithRegistersDropPath
+    for i, layer in enumerate(layers_with_dp):
+        assert hasattr(layer, "drop_path"), "Layer should have drop_path attribute"
+        assert isinstance(layer.drop_path, Dinov2WithRegistersDropPath), (
+            f"Layer {i}: expected Dinov2WithRegistersDropPath, got {type(layer.drop_path)}"
+        )
+
+    # drop_path_rate=0.0 -> every layer initialised as nn.Identity
+    for i, layer in enumerate(layers_without_dp):
+        assert hasattr(layer, "drop_path"), "Layer should have drop_path attribute"
+        assert isinstance(layer.drop_path, torch.nn.Identity), (
+            f"Layer {i}: expected nn.Identity for zero drop_path, got {type(layer.drop_path)}"
+        )
+
+
+def test_update_drop_path_handles_missing_layers(model_with_drop_path: LWDETR, monkeypatch: pytest.MonkeyPatch) -> None:
+    """Verify update_drop_path() handles models without recognizable layer structure gracefully."""
+    model = model_with_drop_path
+
+    monkeypatch.setattr(model, "_get_backbone_encoder_layers", lambda: None)
+
+    # Should not raise an error, just return early
+    model.update_drop_path(0.1, 12)
+
+
+def test_update_drop_path_partial_layers(model_with_drop_path: LWDETR) -> None:
+    """Verify min() guard prevents IndexError when vit_encoder_num_layers > len(layers)."""
+    model = model_with_drop_path
+
+    layers = model._get_backbone_encoder_layers()
+    assert layers is not None
+    actual_num_layers = len(layers)
+
+    # Request more layers than exist in the backbone
+    requested_num_layers = actual_num_layers + 4
+    drop_path_rate = 0.2
+
+    # Should not raise IndexError
+    model.update_drop_path(drop_path_rate, requested_num_layers)
+
+    # Each updated layer gets a rate from 0 to drop_path_rate (shorter, capped linspace)
+    expected_rates = [x.item() for x in torch.linspace(0, drop_path_rate, actual_num_layers)]
+    for i in range(actual_num_layers):
+        actual_prob = layers[i].drop_path.drop_prob
+        assert abs(actual_prob - expected_rates[i]) < 1e-6, (
+            f"Layer {i} drop_prob should be {expected_rates[i]}, got {actual_prob}"
+        )
+
+
+def test_non_windowed_drop_path_warns(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Verify a warning is emitted when drop_path_rate > 0 with non-windowed backbone."""
+    mock_backbone = MagicMock()
+    monkeypatch.setattr(dinov2_module, "AutoBackbone", MagicMock(from_pretrained=MagicMock(return_value=mock_backbone)))
+
+    # The rf-detr logger sets propagate=False, so intercept warning() directly.
+    warning_messages: list[str] = []
+    rf_detr_logger = logging.getLogger("rf-detr")
+    monkeypatch.setattr(rf_detr_logger, "warning", lambda msg, *args, **kwargs: warning_messages.append(msg))
+
+    DinoV2(size="base", use_windowed_attn=False, drop_path_rate=0.1)
+
+    assert any("drop_path_rate" in msg and "ignored" in msg for msg in warning_messages), (
+        "Expected warning about drop_path_rate being ignored for non-windowed backbone"
+    )
+
+
+def test_get_backbone_encoder_layers_blocks_path() -> None:
+    """Verify _get_backbone_encoder_layers() returns enc.blocks for standard ViT backbones."""
+    mock_blocks = nn.ModuleList([nn.Linear(1, 1) for _ in range(3)])
+    # SimpleNamespace gives only the attributes we define, so hasattr checks work correctly.
+    mock_encoder = SimpleNamespace(blocks=mock_blocks)
+    mock_self = SimpleNamespace(backbone=[SimpleNamespace(encoder=mock_encoder)])
+
+    result = LWDETR._get_backbone_encoder_layers(mock_self)  # type: ignore[arg-type]
+    assert result is mock_blocks, "Should return enc.blocks for standard ViT backbone"
+
+
+def test_get_backbone_encoder_layers_trunk_blocks_path() -> None:
+    """Verify _get_backbone_encoder_layers() returns enc.trunk.blocks for aimv2 backbones."""
+    mock_blocks = nn.ModuleList([nn.Linear(1, 1) for _ in range(3)])
+    mock_trunk = SimpleNamespace(blocks=mock_blocks)
+    mock_encoder = SimpleNamespace(trunk=mock_trunk)  # no 'blocks' at top level
+    mock_self = SimpleNamespace(backbone=[SimpleNamespace(encoder=mock_encoder)])
+
+    result = LWDETR._get_backbone_encoder_layers(mock_self)  # type: ignore[arg-type]
+    assert result is mock_blocks, "Should return enc.trunk.blocks for aimv2 backbone"

tests/training/test_load_pretrain_weights.py

@@ -0,0 +1,870 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Regression tests for fine-tuned checkpoint weight destruction.
+
+When a user loads a fine-tuned N-class checkpoint but has ``num_classes`` configured to a LARGER value (e.g. default
+90), the second reinit in ``load_pretrain_weights`` (models/weights.py) must NOT erroneously resize the detection head
+to ``num_classes + 1``, destroying the loaded weights.
+
+The fix changes the second reinit condition from:
+    ``checkpoint_num_classes != args.num_classes + 1``
+to the user-override-aware logic that auto-aligns to the checkpoint when the user did not explicitly set
+``num_classes``.
+
+These tests exercise ``rfdetr.models.weights.load_pretrain_weights`` directly, which is the unified function that
+replaced the two prior separate implementations (``detr.py:_load_pretrain_weights_into`` and
+``module_model.py:RFDETRModelModule._load_pretrain_weights``).
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, call
+
+import pytest
+import torch
+
+from rfdetr.config import (
+    RFDETRBaseConfig,
+    RFDETRKeypointPreviewConfig,
+    RFDETRLargeConfig,
+    RFDETRMediumConfig,
+    RFDETRNanoConfig,
+    RFDETRSeg2XLargeConfig,
+    RFDETRSegLargeConfig,
+    RFDETRSegMediumConfig,
+    RFDETRSegNanoConfig,
+    RFDETRSegPreviewConfig,
+    RFDETRSegSmallConfig,
+    RFDETRSegXLargeConfig,
+    RFDETRSmallConfig,
+    TrainConfig,
+)
+from rfdetr.models.weights import load_pretrain_weights
+
+# ---------------------------------------------------------------------------
+# Shared helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_checkpoint(num_classes=91, num_queries=300, group_detr=13):
+    """Build a minimal checkpoint dict with the given class count.
+
+    Args:
+        num_classes: Total classes including background (bias shape).
+        num_queries: Number of object queries per group.
+        group_detr: Number of groups.
+    """
+    total_queries = num_queries * group_detr
+    state = {
+        "class_embed.weight": torch.randn(num_classes, 256),
+        "class_embed.bias": torch.randn(num_classes),
+        "refpoint_embed.weight": torch.randn(total_queries, 4),
+        "query_feat.weight": torch.randn(total_queries, 256),
+        "other_layer.weight": torch.randn(10, 10),
+    }
+    ckpt_args = SimpleNamespace(
+        segmentation_head=False,
+        patch_size=14,
+        class_names=[],
+    )
+    return {"model": state, "args": ckpt_args}
+
+
+def _make_train_config():
+    """Return a minimal TrainConfig for use in load_pretrain_weights.
+
+    Returns:
+        Minimal TrainConfig with placeholder dataset and output dirs.
+    """
+    return TrainConfig(
+        dataset_dir="/nonexistent/dataset",
+        output_dir="/nonexistent/output",
+        epochs=10,
+        lr=1e-4,
+        lr_encoder=1.5e-4,
+        batch_size=2,
+        weight_decay=1e-4,
+        lr_drop=8,
+        warmup_epochs=1.0,
+        drop_path=0.0,
+        multi_scale=False,
+        expanded_scales=False,
+        do_random_resize_via_padding=False,
+        grad_accum_steps=1,
+        tensorboard=False,
+    )
+
+
+def _suppress_pretrain_io(monkeypatch) -> None:
+    """Suppress download/validate/file-existence side effects on the canonical load path."""
+    monkeypatch.setattr("rfdetr.models.weights.download_pretrain_weights", lambda *a, **kw: None)
+    monkeypatch.setattr("rfdetr.models.weights.validate_pretrain_weights", lambda *a, **kw: None)
+    monkeypatch.setattr("rfdetr.models.weights.validate_checkpoint_compatibility", lambda *a, **kw: None)
+    monkeypatch.setattr("rfdetr.models.weights.os.path.isfile", lambda _: True)
+
+
+# ---------------------------------------------------------------------------
+# Regression tests: load_pretrain_weights (models/weights.py)
+# ---------------------------------------------------------------------------
+
+
+class TestLoadPretrainWeightsSecondReinit:
+    """Regression tests for ``load_pretrain_weights`` in ``rfdetr.models.weights``.
+
+    Validates that the second reinitialize_detection_head call only fires when the checkpoint has MORE classes than
+    configured (backbone pretrain scenario), not when it has fewer (fine-tuned checkpoint scenario).
+    """
+
+    @pytest.fixture(autouse=True)
+    def _patch_download(self, monkeypatch):
+        """Suppress all download and file-existence side effects."""
+        _suppress_pretrain_io(monkeypatch)
+
+    def test_finetune_checkpoint_preserves_weights(self, monkeypatch):
+        """Fine-tuned checkpoint (fewer classes) must NOT trigger second reinit.
+
+        Scenario: 2-class fine-tuned checkpoint (bias shape [3]) loaded with
+        default num_classes=90. The first reinit correctly resizes the head to 3 so load_state_dict works. The second
+        reinit must NOT resize to 91 — that would destroy the loaded fine-tuned weights.
+        """
+        from rfdetr.models.weights import load_pretrain_weights
+
+        mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu")
+        checkpoint = _make_checkpoint(num_classes=3)
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        calls = fake_model.reinitialize_detection_head.call_args_list
+        assert calls[0] == call(3), f"First reinit should resize to checkpoint size 3, got {calls[0]}"
+        assert len(calls) == 1, (
+            f"Expected exactly 1 reinit call (to checkpoint size), but got {len(calls)}: "
+            f"{calls}. The second reinit to 91 destroys loaded weights."
+        )
+        assert mc.num_classes == 2, (
+            f"mc.num_classes must be auto-aligned to 2 (checkpoint_logits - 1), got {mc.num_classes}"
+        )
+
+    def test_no_mismatch_no_reinit(self, monkeypatch):
+        """Checkpoint class count matches config — no reinit at all.
+
+        Scenario: COCO checkpoint (91 classes) with num_classes=90. 91 == 90 + 1, so no reinit should fire.
+        """
+        from rfdetr.models.weights import load_pretrain_weights
+
+        mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu", num_classes=90)
+        checkpoint = _make_checkpoint(num_classes=91)
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        fake_model.reinitialize_detection_head.assert_not_called()
+
+    def test_backbone_pretrain_still_reinits(self, monkeypatch):
+        """Backbone pretrain (more classes in checkpoint) must still reinit.
+
+        Scenario: COCO 91-class checkpoint loaded for 2-class fine-tuning
+        (num_classes=2). Both reinits are correct here: first to 91 for load_state_dict, second to 3 for the configured
+        class count.
+        """
+        from rfdetr.models.weights import load_pretrain_weights
+
+        mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu", num_classes=2)
+        checkpoint = _make_checkpoint(num_classes=91)
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        calls = fake_model.reinitialize_detection_head.call_args_list
+        assert calls == [call(91), call(3)], f"Expected reinit to [91, 3] (expand then trim), got {calls}"
+
+    def test_user_override_larger_than_checkpoint_reexpands_head(self, monkeypatch):
+        """Explicit larger num_classes must be restored after checkpoint load.
+
+        Scenario: 91-class checkpoint loaded with explicit num_classes=93.
+        Loader must temporarily match checkpoint size for load_state_dict, then expand to 94 logits and keep
+        args.num_classes unchanged.
+        """
+        from rfdetr.models.weights import load_pretrain_weights
+
+        mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu", num_classes=93)
+        checkpoint = _make_checkpoint(num_classes=91)
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        calls = fake_model.reinitialize_detection_head.call_args_list
+        assert calls == [call(91), call(94)], f"Expected reinit to [91, 94] (load then expand), got {calls}"
+        assert mc.num_classes == 93, "Explicitly configured num_classes must not be overwritten."
+
+    # All non-deprecated model configs (RFDETRLargeDeprecatedConfig and
+    # RFDETRBaseConfig are excluded; the former is deprecated, the latter
+    # serves as the base class for the concrete variants below).
+    @pytest.mark.parametrize(
+        "config_cls",
+        [
+            pytest.param(RFDETRNanoConfig, id="nano"),
+            pytest.param(RFDETRSmallConfig, id="small"),
+            pytest.param(RFDETRMediumConfig, id="medium"),
+            pytest.param(RFDETRLargeConfig, id="large"),
+            pytest.param(RFDETRSegNanoConfig, id="seg_nano"),
+            pytest.param(RFDETRSegSmallConfig, id="seg_small"),
+            pytest.param(RFDETRSegMediumConfig, id="seg_medium"),
+            pytest.param(RFDETRSegLargeConfig, id="seg_large"),
+            pytest.param(RFDETRSegXLargeConfig, id="seg_xlarge"),
+            pytest.param(RFDETRSeg2XLargeConfig, id="seg_2xlarge"),
+        ],
+    )
+    def test_eight_class_finetune_checkpoint_auto_aligns_num_classes_and_reinits_once(self, monkeypatch, config_cls):
+        """Auto-align ``mc.num_classes`` and avoid a second reinit for 8-class checkpoints.
+
+        Scenario (from user bug report): user trains on 8 categories (IDs 0–7).
+        The checkpoint stores ``class_embed.bias`` with shape [9] (8 user classes + 1 background). Loading without
+        specifying ``num_classes`` must NOT trigger a second reinit to 91 after temporarily matching the checkpoint size
+        for ``load_state_dict``.
+
+        This test asserts the loader auto-aligns ``mc.num_classes`` to 8 (9 - 1) and fires exactly one reinit call — to
+        9 (the checkpoint size).
+        """
+        # 8 dataset categories → training builds a model with 8+1=9 logits.
+        checkpoint = _make_checkpoint(num_classes=9)
+        mc = config_cls(pretrain_weights="/fake/weights.pth", device="cpu")
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        calls = fake_model.reinitialize_detection_head.call_args_list
+        assert len(calls) == 1, (
+            f"Expected exactly 1 reinit call (to checkpoint size 9), but got {len(calls)}: "
+            f"{calls}. A second reinit to 91 would produce OOB class IDs like 73."
+        )
+        assert calls[0] == call(9), f"Reinit must resize to checkpoint's 9 logits, got {calls[0]}"
+        assert mc.num_classes == 8, (
+            f"mc.num_classes must be auto-aligned to 8 (checkpoint_logits - 1), got {mc.num_classes}"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Regression #960: PE interpolation for custom resolution
+# ---------------------------------------------------------------------------
+
+PE_KEY = "backbone.0.encoder.encoder.embeddings.position_embeddings"
+
+
+class TestLoadPretrainWeightsPEInterpolation:
+    """Regression tests for #960 — PE must be interpolated when resolution changes.
+
+    ``load_pretrain_weights`` must bicubic-interpolate the checkpoint's DINOv2 positional embeddings to match the
+    model's ``positional_encoding_size`` before calling ``load_state_dict``.  Without this, any custom ``resolution``
+    that changes the PE grid size causes a ``RuntimeError: size mismatch``.
+    """
+
+    @pytest.fixture(autouse=True)
+    def _patch_download(self, monkeypatch):
+        """Suppress all download and file-existence side effects."""
+        _suppress_pretrain_io(monkeypatch)
+
+    @pytest.mark.parametrize(
+        "src_pe_size, tgt_resolution, patch_size, expected_tgt_pe_size",
+        [
+            pytest.param(24, 640, 16, 40, id="nano_24x24_upscale_to_40x40"),
+            pytest.param(40, 384, 16, 24, id="nano_40x40_downscale_to_24x24"),
+            pytest.param(32, 640, 16, 40, id="small_32x32_upscale_to_40x40"),
+        ],
+    )
+    def test_pe_in_checkpoint_is_interpolated_to_model_resolution(
+        self, monkeypatch, src_pe_size, tgt_resolution, patch_size, expected_tgt_pe_size
+    ):
+        """Checkpoint PE is bicubic-interpolated to match model_config.positional_encoding_size.
+
+        Regression for #960: ``load_pretrain_weights`` must not raise ``RuntimeError`` when model resolution differs
+        from checkpoint resolution.  The PE tensor in the checkpoint must be resized in-place before ``load_state_dict``
+        is called.
+        """
+        mc = RFDETRNanoConfig(
+            pretrain_weights="/fake/weights.pth",
+            device="cpu",
+            resolution=tgt_resolution,
+            patch_size=patch_size,
+        )
+        assert mc.positional_encoding_size == tgt_resolution // patch_size
+
+        dim = 384
+        src_n = src_pe_size * src_pe_size + 1  # patches + class token
+        checkpoint = _make_checkpoint(num_classes=91)
+        checkpoint["model"][PE_KEY] = torch.randn(1, src_n, dim).half()  # float16 to verify dtype round-trip
+
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        pe = checkpoint["model"][PE_KEY]
+        expected_n = expected_tgt_pe_size * expected_tgt_pe_size + 1
+        assert pe.shape == torch.Size([1, expected_n, dim]), (
+            f"Expected PE shape [1, {expected_n}, {dim}], got {tuple(pe.shape)}. "
+            f"PE was not interpolated from {src_pe_size}x{src_pe_size} "
+            f"to {expected_tgt_pe_size}x{expected_tgt_pe_size}."
+        )
+        assert pe.dtype == torch.float16, f"Dtype must be preserved after interpolation, got {pe.dtype}"
+
+    def test_matching_pe_shape_is_not_modified(self, monkeypatch):
+        """When checkpoint PE matches model expectations, the tensor is not changed.
+
+        Ensures PE interpolation is a no-op for same-resolution checkpoints so that normal weight loading is unaffected.
+        """
+        mc = RFDETRNanoConfig(pretrain_weights="/fake/weights.pth", device="cpu")
+        # Default: positional_encoding_size=24 → PE = [1, 24*24+1, 384] = [1, 577, 384]
+
+        dim = 384
+        original_pe = torch.randn(1, 577, dim)
+        checkpoint = _make_checkpoint(num_classes=91)
+        checkpoint["model"][PE_KEY] = original_pe.clone()
+
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        pe = checkpoint["model"][PE_KEY]
+        assert pe.shape == torch.Size([1, 577, dim]), "Matching PE shape must not be modified."
+        assert torch.equal(pe, original_pe), "Matching PE tensor values must not be modified."
+
+    def test_base_config_non_formula_pe_is_interpolated_from_smaller_checkpoint(self, monkeypatch):
+        """RFDETRBaseConfig PE=37 (not formula-derived) is interpolated when checkpoint differs.
+
+        RFDETRBaseConfig.positional_encoding_size=37 is not updated by ``_sync_pe_with_resolution`` because 37 ≠
+        560//16=35 (not formula-derived). Loading a checkpoint with a smaller PE grid (e.g., 24×24) must still trigger
+        interpolation to the model's fixed PE=37×37 target.
+        """
+        mc = RFDETRBaseConfig(pretrain_weights="/fake/weights.pth", device="cpu")
+        assert mc.positional_encoding_size == 37, "RFDETRBaseConfig PE must remain 37 (not formula-derived)"
+
+        dim = 384
+        src_pe_size = 24
+        src_n = src_pe_size * src_pe_size + 1
+        checkpoint = _make_checkpoint(num_classes=91)
+        checkpoint["model"][PE_KEY] = torch.randn(1, src_n, dim)
+
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        pe = checkpoint["model"][PE_KEY]
+        expected_n = 37 * 37 + 1
+        assert pe.shape == torch.Size([1, expected_n, dim]), (
+            f"Expected PE shape [1, {expected_n}, {dim}] (37×37 grid), got {tuple(pe.shape)}. "
+            "BaseConfig's non-formula-derived PE must be the interpolation target."
+        )
+
+    def test_non_square_source_pe_logs_warning_and_is_not_modified(self, monkeypatch):
+        """Non-square source PE grids are skipped with a warning and left unchanged.
+
+        When ``n_source`` is not a perfect square the interpolation is skipped to avoid producing malformed embeddings.
+        The tensor must remain untouched and a warning must be emitted via the weights module logger.
+        """
+        mc = RFDETRNanoConfig(pretrain_weights="/fake/weights.pth", device="cpu")
+        # positional_encoding_size=24 → n_target=576 (perfect square, so the
+        # target-side guard does not trigger; only the source-side guard fires)
+
+        dim = 384
+        # 17 is not a perfect square: isqrt(17)=4, 4*4=16 ≠ 17
+        non_square_n_source = 17
+        original_pe = torch.randn(1, non_square_n_source + 1, dim)
+        checkpoint = _make_checkpoint(num_classes=91)
+        checkpoint["model"][PE_KEY] = original_pe.clone()
+
+        warning_calls: list[tuple] = []
+        monkeypatch.setattr("rfdetr.models.weights.logger.warning", lambda *a, **kw: warning_calls.append(a))
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = MagicMock()
+        load_pretrain_weights(fake_model, mc)
+
+        pe = checkpoint["model"][PE_KEY]
+        assert torch.equal(pe, original_pe), "Non-square source PE must not be modified."
+        assert any("not a perfect square" in str(args) for args in warning_calls), (
+            f"Expected a 'not a perfect square' warning; got calls: {warning_calls}"
+        )
+
+
+class TestL1FacadePEInterpolationEndToEnd:
+    """Regression for instantiating an RF-DETR L1 facade variant with a custom ``resolution`` and a checkpoint trained
+    at the variant's default resolution must not raise ``RuntimeError`` from a PE shape mismatch.
+
+    In v1.6.5 the L1 facade (``RFDETRLarge``, ``RFDETRNano``, ...) used a private ``_load_pretrain_weights_into`` helper
+    in ``detr.py`` that bypassed the PE bicubic-interpolation added to ``models.weights.load_pretrain_weights`` Code
+    that wired the L1 facade through the unified loader landed later (``inference._build_model_context`` calling
+    ``load_pretrain_weights`` from ``models.weights``).  This test pins that wiring so a future refactor cannot
+    reintroduce a divergent loader path that silently skips PE interpolation.
+
+    Current coverage: ``RFDETRNano`` (detection) and ``RFDETRSegNano`` (segmentation), upward-interpolation only.  When
+    a third L1 facade variant is added, collapse both methods to a single ``@pytest.mark.parametrize`` over
+    ``(variant_class, default_pe_grid, patch_size, new_resolution)``. Downward-interpolation (high-res checkpoint →
+    lower-res model) is not currently exercised; add a reverse-direction parametrize row when refactoring.
+    """
+
+    def test_rfdetr_nano_loads_default_pe_checkpoint_at_custom_resolution(self, tmp_path):
+        """Saving an RFDETRNano state_dict at default resolution and loading at a higher resolution must succeed via PE
+        interpolation in the L1 facade.
+
+        Mirrors the user-reported scenario in https://github.com/roboflow/rf-detr/issues/990 (PE size mismatch ``[1,
+        1937, 384]`` vs ``[1, 6401, 384]`` raised from ``LWDETR.load_state_dict``), reduced to RFDETRNano for test
+        speed.
+        """
+        from rfdetr import RFDETRNano
+
+        # 1. Build a default-resolution RFDETRNano (no pretrain, on CPU) so that
+        #    it produces a state_dict with the variant's native PE grid.
+        default_model = RFDETRNano(pretrain_weights=None, num_classes=2, device="cpu")
+        default_pe_grid = default_model.model_config.positional_encoding_size
+        assert default_pe_grid == 24, "RFDETRNano default PE grid must be 24×24"
+        patch_size = default_model.model_config.patch_size
+        default_state = default_model.model.model.state_dict()
+        default_pe = default_state[PE_KEY]
+        pe_dim = default_pe.shape[-1]
+        assert default_pe.shape == torch.Size([1, default_pe_grid * default_pe_grid + 1, pe_dim])
+
+        # 2. Persist as a checkpoint that mimics what `model.train()` saves —
+        #    a top-level "model" key plus a SimpleNamespace "args" block.
+        ckpt_path = tmp_path / "user_finetuned.pth"
+        torch.save(
+            {
+                "model": dict(default_state),
+                "args": SimpleNamespace(class_names=["a", "b"], patch_size=patch_size),
+            },
+            ckpt_path,
+        )
+
+        # 3. Re-instantiate at a NEW resolution.  Without PE interpolation in
+        #    the L1 facade path this raises ``RuntimeError: size mismatch for
+        #    backbone.0.encoder.encoder.embeddings.position_embeddings`` from
+        #    LWDETR.load_state_dict — exactly the user-reported failure.
+        new_resolution = 640
+        loaded = RFDETRNano(
+            pretrain_weights=str(ckpt_path),
+            resolution=new_resolution,
+            num_classes=2,
+            device="cpu",
+        )
+
+        # 4. The model_config validator must update PE proportionally to the
+        #    new resolution, and the loaded backbone PE parameter must have the
+        #    interpolated target shape (40 × 40 + 1 = 1601 tokens).
+        expected_pe_grid = new_resolution // patch_size
+        assert expected_pe_grid == 40
+        assert loaded.model_config.positional_encoding_size == expected_pe_grid
+        loaded_pe = loaded.model.model.state_dict()[PE_KEY]
+        assert loaded_pe.shape == torch.Size([1, expected_pe_grid * expected_pe_grid + 1, pe_dim]), (
+            f"Backbone PE was not interpolated to the requested resolution; "
+            f"got shape {tuple(loaded_pe.shape)}, expected [1, {expected_pe_grid**2 + 1}, {pe_dim}]."
+        )
+
+    def test_rfdetr_seg_nano_loads_default_pe_checkpoint_at_custom_resolution(self, tmp_path):
+        """Saving an RFDETRSegNano state_dict at default resolution and loading at a higher resolution must succeed via
+        PE interpolation in the L1 facade.
+
+        Regression for https://github.com/roboflow/rf-detr/issues/1023 — the segmentation model variant
+        (``RFDETRSegNano``) raised ``RuntimeError: size mismatch for
+        backbone.0.encoder.encoder.embeddings.position_embeddings`` when instantiated with a non-default ``resolution``
+        because the L1 facade's checkpoint-loading path did not interpolate positional embeddings for segmentation
+        models.
+        """
+        from rfdetr import RFDETRSegNano
+
+        # 1. Build a default-resolution RFDETRSegNano (no pretrain, on CPU).
+        #    Uses 90 classes to mimic an official COCO-pretrained checkpoint so
+        #    the load path at step 3 exercises both head-reinit (90 → 2 classes)
+        #    and PE interpolation simultaneously.
+        default_model = RFDETRSegNano(pretrain_weights=None, num_classes=90, device="cpu")
+        default_pe_grid = default_model.model_config.positional_encoding_size
+        assert default_pe_grid == 26, "RFDETRSegNano default PE grid must be 26×26 (312 // 12)"
+        patch_size = default_model.model_config.patch_size
+        assert patch_size == 12, "RFDETRSegNano patch_size must be 12"
+        default_state = default_model.model.model.state_dict()
+        default_pe = default_state[PE_KEY]
+        pe_dim = default_pe.shape[-1]
+        assert default_pe.shape == torch.Size([1, default_pe_grid * default_pe_grid + 1, pe_dim])
+
+        # 2. Persist as a checkpoint that mimics the official pretrain weights
+        #    format.  Saved as .pth (not .pt) so the tmp_path fixture path does
+        #    not trigger ModelWeights registry / MD5 lookup.  Top-level keys
+        #    match the real checkpoint: "model" (state_dict) and "args" with
+        #    segmentation_head=True and patch_size=12.
+        ckpt_path = tmp_path / "rf-detr-seg-nano.pth"
+        torch.save(
+            {
+                "model": dict(default_state),
+                "args": SimpleNamespace(
+                    class_names=[],
+                    patch_size=patch_size,
+                    segmentation_head=True,
+                ),
+            },
+            ckpt_path,
+        )
+
+        # 3. Re-instantiate at a custom resolution with fewer classes.  Without
+        #    PE interpolation this raises
+        #    ``RuntimeError: size mismatch for
+        #    backbone.0.encoder.encoder.embeddings.position_embeddings``
+        #    from LWDETR.load_state_dict — exactly the user-reported failure.
+        # resolution=1008 (user-reported in #1023, 84×84=7057 tokens) is deferred to follow-up parametrization.
+        new_resolution = 624  # 2× the default 312; divisible by patch_size=12
+        loaded = RFDETRSegNano(
+            pretrain_weights=str(ckpt_path),
+            resolution=new_resolution,
+            num_classes=2,
+            device="cpu",
+        )
+
+        # 4. The model_config validator must update PE proportionally to the
+        #    new resolution, and the loaded backbone PE parameter must have the
+        #    interpolated target shape (52 × 52 + 1 = 2705 tokens).
+        expected_pe_grid = new_resolution // patch_size
+        assert expected_pe_grid == 52
+        assert loaded.model_config.positional_encoding_size == expected_pe_grid
+        loaded_pe = loaded.model.model.state_dict()[PE_KEY]
+        assert loaded_pe.shape == torch.Size([1, expected_pe_grid * expected_pe_grid + 1, pe_dim]), (
+            f"Backbone PE was not interpolated to the requested resolution; "
+            f"got shape {tuple(loaded_pe.shape)}, expected [1, {expected_pe_grid**2 + 1}, {pe_dim}]."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Deprecation: train_config argument
+# ---------------------------------------------------------------------------
+
+
+class TestLoadPretrainWeightsDeprecation:
+    """Passing train_config must emit a DeprecationWarning."""
+
+    def test_emits_deprecation_warning_when_train_config_passed(self, monkeypatch):
+        """Any non-None train_config triggers a DeprecationWarning."""
+        from rfdetr.models.weights import load_pretrain_weights
+
+        mc = RFDETRBaseConfig(pretrain_weights=None, device="cpu")
+        tc = _make_train_config()
+
+        with pytest.warns(FutureWarning, match="train_config.*deprecated"):
+            load_pretrain_weights(MagicMock(), mc, tc)
+
+
+# ---------------------------------------------------------------------------
+# Regression #1038: PE interpolation for custom resolution — training path
+# ---------------------------------------------------------------------------
+
+
+class TestModuleLoadPretrainWeightsPEInterpolationCustomResolution:
+    """Regression for #1038 — PE interpolation must fire through ``RFDETRModelModule.__init__``.
+
+    The L2 training entry path (``RFDETRSegLarge(resolution=1008).train(...)``) constructs an
+    :class:`~rfdetr.training.module_model.RFDETRModelModule` whose ``__init__`` delegates to
+    :func:`~rfdetr.models.weights.load_pretrain_weights`. That helper must bicubic-interpolate the checkpoint's DINOv2
+    positional embeddings to match ``model_config.positional_encoding_size`` before calling ``load_state_dict``.
+    Without this, any ``model.train()`` call with a custom ``resolution`` that changes the PE grid raises::
+
+        RuntimeError: Error(s) in loading state_dict for LWDETR:
+            size mismatch for backbone.0.encoder.encoder.embeddings.position_embeddings
+
+    These tests exercise the construction path end-to-end (mocking only the heavy ``build_model_from_config`` /
+    ``build_criterion_from_config`` calls and disk I/O), so the regression cannot reappear if the in-init delegation to
+    ``load_pretrain_weights`` is removed.
+    """
+
+    @pytest.fixture(autouse=True)
+    def _patch_download(self, monkeypatch):
+        """Suppress download/validate side effects on the canonical load path."""
+        _suppress_pretrain_io(monkeypatch)
+
+    def _construct_module(self, mc, checkpoint, monkeypatch, tmp_path):
+        """Construct an RFDETRModelModule with all heavy work mocked.
+
+        Returns the constructed module and the fake nn_model whose ``load_state_dict`` receives the (now-interpolated)
+        state dict.
+        """
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = MagicMock()
+        # Pretend the head was already aligned so the canonical loader does not
+        # try to introspect the MagicMock's internals.
+        fake_model.num_classes = mc.num_classes
+        monkeypatch.setattr(
+            "rfdetr.training.module_model.build_model_from_config",
+            lambda *a, **kw: fake_model,
+        )
+        monkeypatch.setattr(
+            "rfdetr.training.module_model.build_criterion_from_config",
+            lambda *a, **kw: (MagicMock(), MagicMock()),
+        )
+
+        tc = TrainConfig(
+            dataset_dir=str(tmp_path / "dataset"),
+            output_dir=str(tmp_path / "output"),
+            epochs=1,
+            lr=1e-4,
+            lr_encoder=1.5e-4,
+            batch_size=2,
+            weight_decay=1e-4,
+            lr_drop=1,
+            warmup_epochs=0.0,
+            drop_path=0.0,
+            multi_scale=False,
+            expanded_scales=False,
+            do_random_resize_via_padding=False,
+            grad_accum_steps=1,
+            tensorboard=False,
+        )
+        from rfdetr.training.module_model import RFDETRModelModule
+
+        module = RFDETRModelModule(mc, tc)
+        return module, fake_model
+
+    @pytest.mark.parametrize(
+        "config_cls, src_pe_size, tgt_pe_size",
+        [
+            # All 7 segmentation variants listed in #1038, plus detection up/downscale.
+            pytest.param(RFDETRSegNanoConfig, 26, 84, id="seg_nano_upscale_26_to_84"),
+            pytest.param(RFDETRSegSmallConfig, 32, 84, id="seg_small_upscale_32_to_84"),
+            pytest.param(RFDETRSegMediumConfig, 36, 84, id="seg_medium_upscale_36_to_84"),
+            pytest.param(RFDETRSegLargeConfig, 42, 84, id="seg_large_upscale_42_to_84"),
+            pytest.param(RFDETRSegPreviewConfig, 24, 84, id="seg_preview_upscale_24_to_84"),
+            pytest.param(RFDETRSegXLargeConfig, 52, 84, id="seg_xlarge_upscale_52_to_84"),
+            pytest.param(RFDETRSeg2XLargeConfig, 64, 84, id="seg_2xlarge_upscale_64_to_84"),
+            pytest.param(RFDETRNanoConfig, 24, 40, id="nano_upscale_24_to_40"),
+            pytest.param(RFDETRNanoConfig, 40, 24, id="nano_downscale_40_to_24"),
+        ],
+    )
+    def test_pe_interpolated_in_training_path(self, monkeypatch, config_cls, src_pe_size, tgt_pe_size, tmp_path):
+        """Module construction interpolates PE to match ``positional_encoding_size``.
+
+        Regression for #1038 — ``RFDETRModelModule.__init__`` must trigger PE interpolation through the canonical loader
+        so ``load_state_dict`` does not raise ``RuntimeError: size mismatch`` at custom training resolutions.
+        """
+        mc = config_cls(
+            pretrain_weights="/fake/weights.pth",
+            device="cpu",
+            positional_encoding_size=tgt_pe_size,
+        )
+
+        dim = 384
+        src_n = src_pe_size * src_pe_size + 1
+        checkpoint = _make_checkpoint(num_classes=mc.num_classes + 1)
+        checkpoint["model"][PE_KEY] = torch.randn(1, src_n, dim)
+
+        _, fake_model = self._construct_module(mc, checkpoint, monkeypatch, tmp_path)
+
+        pe = checkpoint["model"][PE_KEY]
+        expected_n = tgt_pe_size * tgt_pe_size + 1
+        assert pe.shape == torch.Size([1, expected_n, dim]), (
+            f"Expected PE shape [1, {expected_n}, {dim}] after interpolation from "
+            f"{src_pe_size}x{src_pe_size} to {tgt_pe_size}x{tgt_pe_size}, got {tuple(pe.shape)}. "
+            "PE interpolation must fire during RFDETRModelModule.__init__ via canonical load_pretrain_weights."
+        )
+        # load_state_dict was called on the model with the interpolated state dict.
+        fake_model.load_state_dict.assert_called_once()
+
+    def test_matching_pe_not_modified_in_training_path(self, monkeypatch, tmp_path):
+        """Same-resolution checkpoint PE is untouched in the training path."""
+        pe_size = 24
+        mc = RFDETRNanoConfig(
+            pretrain_weights="/fake/weights.pth",
+            device="cpu",
+            positional_encoding_size=pe_size,
+        )
+
+        dim = 384
+        original_pe = torch.randn(1, pe_size * pe_size + 1, dim)
+        checkpoint = _make_checkpoint(num_classes=mc.num_classes + 1)
+        checkpoint["model"][PE_KEY] = original_pe.clone()
+
+        self._construct_module(mc, checkpoint, monkeypatch, tmp_path)
+
+        pe = checkpoint["model"][PE_KEY]
+        assert pe.shape == torch.Size([1, pe_size * pe_size + 1, dim]), "Matching PE must not be modified."
+        assert torch.equal(pe, original_pe), "Matching PE values must not be modified."
+
+
+# ---------------------------------------------------------------------------
+# Regression: keypoint schema auto-align from checkpoint _kp_active_mask
+# ---------------------------------------------------------------------------
+
+
+def _make_kp_checkpoint(kp_schema: list[int], num_queries: int = 300, group_detr: int = 13) -> dict:
+    """Build a minimal keypoint checkpoint encoding *kp_schema* via ``_kp_active_mask``.
+
+    Args:
+        kp_schema: Keypoints-per-class list, e.g. ``[0, 17]`` for bg-first or ``[17]`` for active-first.
+        num_queries: Number of queries per group.
+        group_detr: Number of decoder groups.
+
+    Returns:
+        Checkpoint dict with ``model``, ``args``, and schema-derived ``_kp_active_mask``.
+    """
+    num_classes = len(kp_schema)
+    total_queries = num_queries * group_detr
+    max_kp = max(kp_schema) if kp_schema else 0
+    mask = torch.zeros(num_classes, max_kp, dtype=torch.bool)
+    for idx, n_kp in enumerate(kp_schema):
+        mask[idx, :n_kp] = True
+    state = {
+        "class_embed.weight": torch.randn(num_classes + 1, 256),
+        "class_embed.bias": torch.randn(num_classes + 1),
+        "refpoint_embed.weight": torch.randn(total_queries, 4),
+        "query_feat.weight": torch.randn(total_queries, 256),
+        "_kp_active_mask": mask,
+    }
+    ckpt_args = SimpleNamespace(segmentation_head=False, patch_size=14, class_names=[])
+    return {"model": state, "args": ckpt_args}
+
+
+def _make_kp_fake_model(initial_schema: list[int]) -> MagicMock:
+    """Build a fake nn_model mock that mimics the GroupPose keypoint model interface.
+
+    Args:
+        initial_schema: Keypoints-per-class schema the model was built with.
+
+    Returns:
+        Configured MagicMock with keypoint schema methods and ``_kp_active_mask`` state.
+    """
+    max_kp = max(initial_schema) if initial_schema else 0
+    current_schema = list(initial_schema)
+    mask = torch.zeros(len(initial_schema), max_kp, dtype=torch.bool)
+    for idx, n in enumerate(initial_schema):
+        mask[idx, :n] = True
+
+    fake_model = MagicMock()
+    fake_model.state_dict.return_value = {"_kp_active_mask": mask.clone()}
+
+    def _reinit_kp(schema):
+        current_schema[:] = schema
+        new_max = max(schema) if schema else 0
+        new_mask = torch.zeros(len(schema), new_max, dtype=torch.bool)
+        for i, n in enumerate(schema):
+            new_mask[i, :n] = True
+        fake_model.state_dict.return_value = {"_kp_active_mask": new_mask}
+        fake_model.get_num_keypoints_per_class.return_value = list(schema)
+
+    fake_model.reinitialize_keypoint_head.side_effect = _reinit_kp
+    fake_model.get_num_keypoints_per_class.return_value = list(initial_schema)
+    fake_model.get_num_keypoints_per_class_from_checkpoint = lambda sd: (
+        [int(n) for n in sd["_kp_active_mask"].sum(dim=1).tolist()] if "_kp_active_mask" in sd else None
+    )
+    return fake_model
+
+
+class TestLoadPretrainWeightsKeypointSchemaAutoAlign:
+    """Regression for AP=0 when loading bg-first checkpoint into active-first default config.
+
+    Before the fix, ``load_pretrain_weights`` would restore ``mc.num_keypoints_per_class`` to the config default
+    ``[17]`` after loading a ``[0, 17]`` pretrained checkpoint.  The detection-head trimming kept rows 0 (background)
+    and 1 (person) from the checkpoint but the active-first schema treats row 0 as person, producing AP≈0.
+
+    The fix auto-aligns ``mc.num_keypoints_per_class`` from the checkpoint ``_kp_active_mask`` when the user did not
+    explicitly set the field, mirroring the existing ``num_classes`` auto-align pattern.
+    """
+
+    @pytest.fixture(autouse=True)
+    def _patch_io(self, monkeypatch):
+        """Suppress all download and file-existence side effects."""
+        _suppress_pretrain_io(monkeypatch)
+
+    def test_bg_first_checkpoint_auto_aligns_active_first_config(self, monkeypatch):
+        """Loading bg-first [0,17] checkpoint into active-first [17] config auto-aligns schema."""
+        mc = RFDETRKeypointPreviewConfig(pretrain_weights="/fake/kp.pth", device="cpu")
+        assert mc.num_keypoints_per_class == [17], "Precondition: config default is active-first [17]."
+
+        checkpoint = _make_kp_checkpoint([0, 17])
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = _make_kp_fake_model([17])
+
+        load_pretrain_weights(fake_model, mc)
+
+        assert mc.num_keypoints_per_class == [0, 17], (
+            f"expected auto-align to [0, 17], got {mc.num_keypoints_per_class}"
+        )
+        assert mc.num_classes == 2, (
+            f"mc.num_classes must be auto-aligned to 2 (checkpoint has 3 logit slots), got {mc.num_classes}"
+        )
+
+    def test_matching_schema_no_change(self, monkeypatch):
+        """Config and checkpoint with same schema leaves mc.num_keypoints_per_class unchanged."""
+        mc = RFDETRKeypointPreviewConfig(pretrain_weights="/fake/kp.pth", device="cpu")
+        mc.num_keypoints_per_class = [17]
+
+        checkpoint = _make_kp_checkpoint([17])
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = _make_kp_fake_model([17])
+
+        load_pretrain_weights(fake_model, mc)
+
+        assert mc.num_keypoints_per_class == [17]
+
+    def test_user_explicit_schema_not_overridden(self, monkeypatch):
+        """Explicit num_keypoints_per_class from user survives even when checkpoint differs."""
+        mc = RFDETRKeypointPreviewConfig(pretrain_weights="/fake/kp.pth", device="cpu")
+        # Simulate user explicitly providing num_keypoints_per_class
+        mc.num_keypoints_per_class = [0, 33]
+        mc.model_fields_set.add("num_keypoints_per_class")
+
+        checkpoint = _make_kp_checkpoint([0, 17])
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = _make_kp_fake_model([0, 33])
+
+        load_pretrain_weights(fake_model, mc)
+
+        assert mc.num_keypoints_per_class == [0, 33], (
+            "Explicit user schema must not be overridden by checkpoint auto-align."
+        )
+
+    def test_1d_kp_active_mask_skips_auto_align_with_warning(self, monkeypatch):
+        """Malformed 1-D _kp_active_mask skips auto-align and emits a logger.warning."""
+        mc = RFDETRKeypointPreviewConfig(pretrain_weights="/fake/kp.pth", device="cpu")
+        assert mc.num_keypoints_per_class == [17], "Precondition: config default is active-first [17]."
+
+        checkpoint = _make_kp_checkpoint([0, 17])
+        # Replace the 2-D mask with a malformed 1-D tensor.
+        checkpoint["model"]["_kp_active_mask"] = torch.ones(17, dtype=torch.bool)
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = _make_kp_fake_model([17])
+        # The fake model's get_num_keypoints_per_class_from_checkpoint calls sum(dim=1),
+        # which raises IndexError on a 1-D tensor; return None to skip that code path.
+        fake_model.get_num_keypoints_per_class_from_checkpoint = lambda sd: None
+
+        warned: list[str] = []
+        monkeypatch.setattr("rfdetr.models.weights.logger.warning", lambda msg, *args: warned.append(msg % args))
+
+        load_pretrain_weights(fake_model, mc)
+
+        assert mc.num_keypoints_per_class == [17], (
+            "Auto-align must not fire for a 1-D _kp_active_mask; config default should be unchanged."
+        )
+        assert any("unexpected shape" in msg for msg in warned), (
+            f"Expected a warning mentioning 'unexpected shape'; got: {warned}"
+        )
+
+    def test_all_zero_kp_active_mask_skips_auto_align_with_warning(self, monkeypatch):
+        """All-zero 2-D _kp_active_mask (degenerate) skips auto-align and emits a logger.warning."""
+        mc = RFDETRKeypointPreviewConfig(pretrain_weights="/fake/kp.pth", device="cpu")
+        assert mc.num_keypoints_per_class == [17], "Precondition: config default is active-first [17]."
+
+        checkpoint = _make_kp_checkpoint([0, 17])
+        # Replace mask with a valid 2-D shape but all zeros (no active slots).
+        checkpoint["model"]["_kp_active_mask"] = torch.zeros(2, 17, dtype=torch.bool)
+        monkeypatch.setattr("rfdetr.models.weights.torch.load", lambda *a, **kw: checkpoint)
+        fake_model = _make_kp_fake_model([17])
+
+        warned: list[str] = []
+        monkeypatch.setattr("rfdetr.models.weights.logger.warning", lambda msg, *args: warned.append(msg % args))
+
+        load_pretrain_weights(fake_model, mc)
+
+        assert mc.num_keypoints_per_class == [17], (
+            "Auto-align must not fire for an all-zero _kp_active_mask; config default should be unchanged."
+        )
+        assert any("no active slots" in msg for msg in warned), (
+            f"Expected a warning mentioning 'no active slots'; got: {warned}"
+        )

tests/training/test_metrics_csv.py

@@ -0,0 +1,159 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+"""Integration tests: metrics.csv contains all columns used by plot_metrics().
+
+Runs a minimal PTL training loop (1 epoch, 2 batches each) using mocked model internals so no real dataset or GPU is
+required.  After training, reads the CSVLogger output and asserts that every metric column that ``plot_metrics()`` needs
+is present and has at least one non-NaN value.
+
+Also verifies that ``train/loss`` is logged at the same scale as ``val/loss`` (i.e. NOT divided by ``grad_accum_steps``
+before logging).
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pandas as pd
+import torch
+
+from rfdetr.config import RFDETRBaseConfig, TrainConfig
+from rfdetr.training import build_trainer
+from rfdetr.training.module_data import RFDETRDataModule
+from rfdetr.training.module_model import RFDETRModelModule
+
+from .helpers import _fake_postprocess, _FakeCriterion, _FakeDataset, _make_param_dicts, _TinyModel
+
+# ---------------------------------------------------------------------------
+# Helpers local to this module
+# ---------------------------------------------------------------------------
+
+
+def _fit_and_read_csv(mc: RFDETRBaseConfig, tc: TrainConfig, criterion=None) -> pd.DataFrame:
+    """Run 1 epoch (2 train + 2 val batches) and return the resulting metrics.csv."""
+    fake_criterion = criterion or _FakeCriterion()
+    with (
+        patch("rfdetr.training.module_model.build_model_from_config", return_value=_TinyModel()),
+        patch(
+            "rfdetr.training.module_model.build_criterion_from_config",
+            return_value=(fake_criterion, MagicMock(side_effect=_fake_postprocess)),
+        ),
+        patch("rfdetr.training.module_data.build_dataset", return_value=_FakeDataset(length=20)),
+        patch(
+            "rfdetr.training.module_model.get_param_dict",
+            side_effect=lambda args, model: _make_param_dicts(model),
+        ),
+    ):
+        module = RFDETRModelModule(mc, tc)
+        datamodule = RFDETRDataModule(mc, tc)
+        trainer = build_trainer(
+            tc,
+            mc,
+            accelerator="cpu",
+            max_epochs=1,
+            limit_train_batches=2,
+            limit_val_batches=2,
+            log_every_n_steps=1,
+        )
+        trainer.fit(module, datamodule=datamodule)
+
+    csv_path = Path(tc.output_dir) / "metrics.csv"
+    assert csv_path.exists(), "CSVLogger must write metrics.csv to output_dir"
+    return pd.read_csv(csv_path)
+
+
+# ---------------------------------------------------------------------------
+# Expected columns (must exist and have ≥1 non-NaN row after one epoch)
+# ---------------------------------------------------------------------------
+
+_REQUIRED_DETECTION = frozenset(
+    {
+        "train/loss",
+        "train/lr",
+        "val/loss",
+        "val/mAP_50",
+        "val/mAP_50_95",
+        "val/mAR",
+    }
+)
+
+_REQUIRED_DETECTION_EMA = _REQUIRED_DETECTION | frozenset(
+    {
+        "val/ema_mAP_50",
+        "val/ema_mAP_50_95",
+        "val/ema_mAR",
+    }
+)
+
+
+# ---------------------------------------------------------------------------
+# Tests
+# ---------------------------------------------------------------------------
+
+
+class TestDetectionMetricsCSV:
+    """metrics.csv contains all columns that plot_metrics() needs for detection."""
+
+    def test_base_metrics_present_without_ema(self, base_model_config, base_train_config):
+        """Without EMA all core val/* columns must appear in metrics.csv with non-NaN data."""
+        mc = base_model_config()
+        tc = base_train_config(use_ema=False, run_test=False)
+        df = _fit_and_read_csv(mc, tc)
+
+        missing = _REQUIRED_DETECTION - set(df.columns)
+        assert not missing, f"Missing columns in metrics.csv: {sorted(missing)}"
+
+        all_nan = {c for c in _REQUIRED_DETECTION if df[c].isna().all()}
+        assert not all_nan, f"Columns with all-NaN values: {sorted(all_nan)}"
+
+    def test_ema_metrics_present_with_ema_enabled(self, base_model_config, base_train_config):
+        """With use_ema=True the ema_* aliases must also appear in metrics.csv."""
+        mc = base_model_config()
+        tc = base_train_config(use_ema=True, run_test=False)
+        df = _fit_and_read_csv(mc, tc)
+
+        missing = _REQUIRED_DETECTION_EMA - set(df.columns)
+        assert not missing, f"Missing EMA columns in metrics.csv: {sorted(missing)}"
+
+        all_nan = {c for c in _REQUIRED_DETECTION_EMA if df[c].isna().all()}
+        assert not all_nan, f"EMA columns with all-NaN values: {sorted(all_nan)}"
+
+    def test_train_loss_is_unscaled(self, base_model_config, base_train_config):
+        """Train/loss must be logged at the raw criterion scale, not divided by grad_accum_steps.
+
+        With grad_accum_steps=4 the old code divided the logged value by 4, making train/loss ~4× smaller than val/loss.
+        After the fix the logged value equals the raw weighted criterion output so both losses are on the same scale.
+        """
+        fixed_loss_value = 5.0
+        grad_accum_steps = 4
+
+        class _FixedCriterion:
+            weight_dict = {"loss_ce": 1.0}
+
+            def num_boxes_for_targets(self, outputs, targets):
+                dummy = outputs.get("dummy", torch.zeros(1))
+                return torch.ones((), dtype=dummy.dtype, device=dummy.device)
+
+            def __call__(self, outputs, targets, num_boxes=None):
+                # Loss is always fixed_loss_value, connected to model params for gradient.
+                dummy = outputs.get("dummy", torch.zeros(1))
+                denominator = self.num_boxes_for_targets(outputs, targets) if num_boxes is None else num_boxes
+                return {"loss_ce": (dummy.mean() * 0 + fixed_loss_value) / denominator}
+
+        mc = base_model_config()
+        tc = base_train_config(use_ema=False, run_test=False, grad_accum_steps=grad_accum_steps)
+        df = _fit_and_read_csv(mc, tc, criterion=_FixedCriterion())
+
+        logged = df["train/loss"].dropna().mean()
+        expected_unscaled = fixed_loss_value
+        expected_if_divided = fixed_loss_value / grad_accum_steps
+
+        assert abs(logged - expected_unscaled) < abs(logged - expected_if_divided), (
+            f"train/loss={logged:.4f} is closer to the grad-accum-divided value "
+            f"({expected_if_divided:.4f}) than the raw criterion output "
+            f"({expected_unscaled:.4f}). The division must have been removed."
+        )