/tests/run_smoke_all_models.py

#!/usr/bin/env python3
# ------------------------------------------------------------------------
# RF-DETR
# Copyright (c) 2025 Roboflow. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------
"""Smoke-test model instantiation and basic inference with all available weights.

Tests detection, segmentation, and keypoint-preview model classes from rf-detr by importing and instantiating them.
Validates: imports, download, MD5 hash, model instantiation, from_checkpoint round-trip, and basic inference smoke
checks.

Usage:
    python tests/run_smoke_all_models.py
"""

import argparse
import os
import sys
import tempfile
from functools import partial

import torch
from tqdm.auto import tqdm

import rfdetr as _rfdetr
from rfdetr import (
    RFDETRKeypointPreview,
    RFDETRLarge,
    RFDETRMedium,
    RFDETRNano,
    RFDETRSeg2XLarge,
    RFDETRSegLarge,
    RFDETRSegMedium,
    RFDETRSegNano,
    RFDETRSegSmall,
    RFDETRSegXLarge,
    RFDETRSmall,
)

try:
    from rfdetr import RFDETR2XLarge, RFDETRXLarge
except ImportError:
    RFDETR2XLarge = None
    RFDETRXLarge = None

# Explicitly list all models to validate
MODELS_TO_TEST = [
    # Detection Models
    RFDETRNano,
    RFDETRSmall,
    RFDETRMedium,
    RFDETRLarge,
    # Keypoint Models
    RFDETRKeypointPreview,
    # Segmentation Models
    RFDETRSegNano,
    RFDETRSegSmall,
    RFDETRSegMedium,
    RFDETRSegLarge,
    RFDETRSegXLarge,
    RFDETRSeg2XLarge,
]

if RFDETRXLarge is not None:
    MODELS_TO_TEST.append(partial(RFDETRXLarge, accept_platform_model_license=True))
if RFDETR2XLarge is not None:
    MODELS_TO_TEST.append(partial(RFDETR2XLarge, accept_platform_model_license=True))

# 1008 = LCM(12, 16) × 21: valid for all patch sizes (PE=63 for det ÷16,
# PE=84 for seg ÷12). Each model is tested at its default resolution and at
# 1008 (regression #1038).
#
# Note on Base: ``RFDETRBaseConfig.positional_encoding_size = 37`` is *not*
# formula-derived (see test_load_pretrain_weights.py:TestLoadPretrainWeightsPEInterpolation
# ::test_base_config_non_formula_pe_is_interpolated_from_smaller_checkpoint),
# so this `÷16` description applies only to Nano/Small/Medium/Large.
_CUSTOM_RESOLUTION = 1008

# Plus models (XLarge / 2XLarge) are heavy enough that running them at
# resolution=1008 risks the 15-min CI timeout on windows-latest / macos-latest
# runners.  Smaller models still exercise the 1008 path for #1038 coverage.
_HEAVY_MODEL_NAMES = {
    "xlarge",
    "2xlarge",
    "xxlarge",
    "seg-xlarge",
    "seg-2xlarge",
    "seg-xxlarge",
    "rfdetr-xlarge",
    "rfdetr-2xlarge",
    "rfdetr-xxlarge",
    "rfdetr-keypoint-preview",
    "rfdetr-seg-xlarge",
    "rfdetr-seg-2xlarge",
    "rfdetr-seg-xxlarge",
}


def _test_from_checkpoint(
    model_instance: object, actual_cls: type, extra_kwargs: dict, *, test_starter: bool = True
) -> None:
    """Round-trip a model through from_checkpoint using a temp training checkpoint.

    Saves the instantiated model's weights into a minimal training-style checkpoint (``{"args": ..., "model":
    state_dict}``), calls ``rfdetr.from_checkpoint`` on it, and asserts the returned object is an instance of
    *actual_cls*.

    Args:
        model_instance: An already-loaded RFDETR model instance.
        actual_cls: The expected model class (e.g. ``RFDETRSmall``).
        extra_kwargs: Extra kwargs to pass to ``from_checkpoint`` (e.g.
            ``{"accept_platform_model_license": True}`` for plus models).
        test_starter: When ``True`` (default) also run the starter-like
            checkpoint round-trip.  Pass ``False`` for non-default resolutions
            to avoid running the same resolution-independent test multiple times.

    Raises:
        AssertionError: If the recovered model is not an instance of *actual_cls*.
        Exception: Propagates any error from ``from_checkpoint`` to the caller.
    """
    # Build a minimal training-style checkpoint. The pretrain_weights value only
    # needs to contain the model-size substring that from_checkpoint matches on
    # (e.g. "small", "seg-large").  Using cls.size directly satisfies this.
    fake_pretrain_name = f"{actual_cls.size}.pth"
    num_classes = model_instance.model.args.num_classes
    ckpt = {
        "args": argparse.Namespace(
            pretrain_weights=fake_pretrain_name,
            num_classes=num_classes,
        ),
        "model": model_instance.model.model.state_dict(),
    }
    starter_like_ckpt = {
        "args": argparse.Namespace(
            pretrain_weights="none",
            num_classes=num_classes,
        ),
        "model": model_instance.model.model.state_dict(),
    }

    tmp_fd, tmp_path = tempfile.mkstemp(suffix=".pth")
    os.close(tmp_fd)
    try:
        torch.save(ckpt, tmp_path)
        recovered = _rfdetr.from_checkpoint(tmp_path, **extra_kwargs)
        assert recovered is not None, "from_checkpoint returned None"
        assert hasattr(recovered, "model"), "from_checkpoint result missing 'model' attribute"
        assert isinstance(recovered, actual_cls), (
            f"from_checkpoint returned {type(recovered).__name__}, expected {actual_cls.__name__}"
        )
    finally:
        os.unlink(tmp_path)

    if test_starter:
        starter_tmp_fd, starter_path = tempfile.mkstemp(prefix=f"{actual_cls.size}-starter-", suffix=".pth")
        os.close(starter_tmp_fd)
        try:
            torch.save(starter_like_ckpt, starter_path)
            starter_recovered = _rfdetr.from_checkpoint(starter_path, **extra_kwargs)
            assert starter_recovered is not None, "from_checkpoint returned None for starter-like checkpoint"
            assert hasattr(starter_recovered, "model"), "starter-like from_checkpoint result missing 'model' attribute"
            got = type(starter_recovered).__name__
            assert isinstance(starter_recovered, actual_cls), (
                f"starter-like from_checkpoint returned {got}, expected {actual_cls.__name__}"
            )
        finally:
            os.unlink(starter_path)


def _test_coco_class_name_mapping(model_instance: object) -> None:
    """Verify predict() uses sparse COCO category-ID → class-name mapping.

    Issue #988: RFDETRSegSmall returned "sheep" for class_id=18 instead of "dog" because 0-indexed
    ``COCO_CLASS_NAMES[18]`` was used instead of the sparse-dict lookup ``COCO_CLASSES[18]``.  threshold=0 forces all
    top-k queries through so every class ID in the output is covered.  Covers both detection and segmentation nano
    variants (RFDETRNano, RFDETRSegNano).

    Args:
        model_instance: An already-loaded pretrained COCO model instance (det or seg).

    Raises:
        AssertionError: On any class-name mapping failure.
    """
    import PIL.Image

    from rfdetr.assets.coco_classes import COCO_CLASS_NAMES, COCO_CLASSES

    # Sanity-check model properties required for the pretrained COCO branch.
    class_names = model_instance.class_names
    assert class_names is not None, "Pretrained COCO model must have class_names set"
    assert len(class_names) == len(COCO_CLASS_NAMES), (
        f"Expected {len(COCO_CLASS_NAMES)} COCO class names, got {len(class_names)}"
    )
    assert class_names == list(COCO_CLASS_NAMES), "model.class_names must equal COCO_CLASS_NAMES"
    num_classes = model_instance.model.args.num_classes
    assert num_classes == 90, f"Pretrained COCO model must have num_classes=90, got {num_classes}"

    # Run at threshold=0 to exercise all top-k output slots.
    img = PIL.Image.new("RGB", (640, 640), color=(128, 128, 128))
    detections = model_instance.predict(img, threshold=0.0)

    assert "class_name" in detections.data, "data['class_name'] must be present after predict()"

    # For every detection whose class_id is a valid COCO category, class_name must
    # use sparse-ID lookup (COCO_CLASSES[class_id]), not 0-indexed lookup.
    # Canonical regression case: class_id=18 → "dog", NOT "sheep" (COCO_CLASS_NAMES[18]).
    for class_id, class_name in zip(detections.class_id, detections.data["class_name"]):
        cid = int(class_id)
        if cid in COCO_CLASSES:
            expected = COCO_CLASSES[cid]
            assert class_name == expected, (
                f"Sparse COCO ID mapping broken (issue #988): "
                f"class_id={cid} must map to '{expected}', got '{class_name}'"
            )

    # Regression for PR #1051 HIGH-1: no COCO-pretrained detection may carry
    # '__background__' — background is implicit (below threshold), never a sentinel label.
    background_labeled = [
        (int(cid), name)
        for cid, name in zip(detections.class_id, detections.data["class_name"])
        if name == "__background__"
    ]
    assert not background_labeled, (
        "COCO-pretrained predict() must never produce '__background__' class names "
        f"(PR #1051 HIGH-1 regression); found: {background_labeled[:3]}"
    )


def main() -> None:
    """Download, validate, instantiate all models, and test from_checkpoint round-trip."""
    print("Model Instantiation & Download Validation\n")

    succeeded = 0
    pbar = tqdm(MODELS_TO_TEST, desc="Testing models", unit="model")
    for model_class in pbar:
        actual_cls = model_class.func if isinstance(model_class, partial) else model_class
        extra_kwargs = model_class.keywords if isinstance(model_class, partial) else {}
        base_name = actual_cls.size

        for res in (None, _CUSTOM_RESOLUTION):
            # Skip the 1008-resolution variant for heavyweight Plus models — they
            # risk the 15-min CI timeout on windows-latest / macos-latest runners.
            if res == _CUSTOM_RESOLUTION and base_name in _HEAVY_MODEL_NAMES:
                continue

            model_name = base_name if res is None else f"{base_name}@{res}"
            # Build the kwargs once so `_test_from_checkpoint` and the
            # instantiation call share the same parameter set (avoids the
            # `functools.partial.size` AttributeError seen in the previous form).
            instantiate_kwargs = dict(extra_kwargs)
            if res is not None:
                instantiate_kwargs["resolution"] = res

            pbar.set_description(f"Testing {model_name}")
            try:
                # Instantiate model class - triggers download, MD5 validation, and loading
                model_instance = actual_cls(**instantiate_kwargs)

                # Verify model was created
                assert model_instance is not None, "Model instance is None"
                assert hasattr(model_instance, "model"), "Model missing 'model' attribute"

                # from_checkpoint round-trip: save a training-style checkpoint and reload it.
                # Pass the real class (not a partial) so `_test_from_checkpoint` can read
                # `.size` and `.__name__` and run `isinstance(recovered, actual_cls)`.
                _test_from_checkpoint(model_instance, actual_cls, instantiate_kwargs, test_starter=(res is None))

                # Inference class-name regression for issue #988 — run on all
                # nano-sized pretrained COCO models at default resolution only.
                if "nano" in base_name.lower() and res is None:
                    _test_coco_class_name_mapping(model_instance)

                succeeded += 1
            except Exception as ex:
                # Fail-fast: surface the first failing model directly so CI logs the
                # root cause cleanly instead of burying it under later cascade failures.
                pbar.close()
                print(f"\n[FAIL] {model_name}: {ex}")
                raise

    pbar.close()
    print("\nResults:")
    print(f"\tSucceeded:\t{succeeded}")
    print("\n[OK] All models validated successfully")
    sys.exit(0)


if __name__ == "__main__":
    main()