chore: import upstream snapshot with attribution

tests/benchmarks/conftest.py

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

tests/benchmarks/test_develop_downloads.py

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

tests/benchmarks/test_develop_urls.py

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

tests/benchmarks/test_inference_coco.py

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

tests/benchmarks/test_training_coco.py

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

tests/benchmarks/test_training_synthetic.py

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