chore: import upstream snapshot with attribution

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