RF-DETR

The base RF-DETR class implements the core methods for training RF-DETR models, running inference on the models, optimising models, and uploading trained models for deployment.

Source code in src/rfdetr/detr.py

class RFDETR:
    """The base RF-DETR class implements the core methods for training RF-DETR models, running inference on the models,
    optimising models, and uploading trained models for deployment."""

    means = [0.485, 0.456, 0.406]
    stds = [0.229, 0.224, 0.225]
    size = None
    _model_config_class: type[ModelConfig] = ModelConfig
    _train_config_class: type[TrainConfig] = TrainConfig

    def __init__(self, **kwargs):
        self.model_config = self.get_model_config(**kwargs)
        self.maybe_download_pretrain_weights()
        self.model = self.get_model(self.model_config)
        self.callbacks = defaultdict(list)

        # repeat means and stds for non-rgb images
        if self.model_config.num_channels != 3:
            from itertools import cycle

            self.means = [val for _, val in zip(range(self.model_config.num_channels), cycle(self.means))]
            self.stds = [val for _, val in zip(range(self.model_config.num_channels), cycle(self.stds))]

        self.model.inference_model = None
        self._is_optimized_for_inference = False
        self._has_warned_about_not_being_optimized_for_inference = False
        self._optimized_has_been_compiled = False
        self._optimized_batch_size = None
        self._optimized_resolution = None
        self._optimized_dtype = None

    def maybe_download_pretrain_weights(self):
        """Download pre-trained weights if they are not already downloaded.

        Bare filenames (no directory component, e.g. ``rf-detr-base.pth``) are resolved to the model cache directory —
        set the ``RF_HOME`` environment variable to override the location (default: ``~/.roboflow/models``). Resolution
        happens in ``ModelConfig.expand_path`` for explicitly-provided values, and here as a fallback for field defaults
        (which Pydantic does not validate by default).

        Paths that already contain a directory component are used as-is; the parent directory is created if it does not
        yet exist.
        """
        pretrain_weights = self.model_config.pretrain_weights
        if pretrain_weights is None:
            return
        if not os.path.dirname(pretrain_weights):
            # Field default was not processed by expand_path — resolve to cache dir.
            cache_dir = get_model_cache_dir()
            os.makedirs(cache_dir, exist_ok=True)
            pretrain_weights = os.path.join(cache_dir, pretrain_weights)
        else:
            os.makedirs(os.path.dirname(pretrain_weights), exist_ok=True)
        self.model_config.pretrain_weights = pretrain_weights
        download_pretrain_weights(self.model_config.pretrain_weights)

    def get_model_config(self, **kwargs) -> ModelConfig:
        """Retrieve the configuration parameters used by the model."""
        return self._model_config_class(**kwargs)

    @classmethod
    def from_checkpoint(cls, path: str | os.PathLike[str], **kwargs: Any) -> RFDETR:
        """Load an RF-DETR model from a training checkpoint, automatically inferring the model class.

        The correct subclass is resolved in order of preference:

        1. ``model_name`` key in the checkpoint (written by the PTL training
           stack since v1.7.0).
        2. ``pretrain_weights`` field in the checkpoint's ``args`` entry
           (legacy fallback for older checkpoints).
        3. The **filename** of *path* itself, used as a last resort when
           ``pretrain_weights`` is absent or an unset-like sentinel value
           (empty string, ``"none"``, or ``"null"``).  Starter weights
           published by Roboflow store ``pretrain_weights="none"`` in their
           ``args``; passing the canonical filename (e.g.
           ``rf-detr-small.pth``) lets ``from_checkpoint`` infer the class
           automatically.

        Both legacy ``argparse.Namespace`` checkpoints (produced by ``engine.py``) and dict-style checkpoints (produced
        by the PTL training stack) are supported.

        Args:
            path: Path to a checkpoint file (e.g. ``checkpoint_best_total.pth``).
            **kwargs: Additional keyword arguments forwarded to the model
                constructor (e.g. ``accept_platform_model_license=True`` for XLarge / 2XLarge models).
                If ``num_classes`` is not supplied here, the value stored in the checkpoint is
                used when present; otherwise the constructor default applies.  In either case the
                field is not recorded as a user-set override, so :meth:`train` can still adapt the
                detection head to the training dataset's class count.  Pass an explicit
                ``num_classes=N`` to pin the head and prevent adaptation, even when ``N`` equals
                the class default (e.g. ``num_classes=90`` on a COCO-pretrained checkpoint).

        Returns:
            An instance of the appropriate :class:`RFDETR` subclass loaded from the checkpoint.

        Warning:
            This method calls ``torch.load`` with ``weights_only=False``, which
            unpickles arbitrary Python objects. Only load checkpoints from trusted sources.

        Raises:
            FileNotFoundError: If *path* does not exist.
            OSError: If *path* exists but cannot be read.
            KeyError: If the checkpoint does not contain an ``"args"`` key.
            ValueError: If the model class cannot be inferred from ``model_name``,
                ``pretrain_weights``, or the checkpoint filename.

        Examples:
            >>> model = RFDETR.from_checkpoint("checkpoint_best_total.pth")  # doctest: +SKIP
            >>> model = RFDETRSmall.from_checkpoint("checkpoint_best_total.pth")  # doctest: +SKIP
        """
        # Local import breaks the variants → detr import cycle.
        import rfdetr.variants as rfdetr_variants

        _plus_available = False
        _plus_symbols: dict[str, type[RFDETR]] = {}
        _plus_entries: list[tuple[str, type[RFDETR]]] = []
        from rfdetr.platform import _IS_RFDETR_PLUS_AVAILABLE

        if _IS_RFDETR_PLUS_AVAILABLE:
            try:
                import rfdetr.platform.models as platform_models

                for class_symbol in _CHECKPOINT_PLUS_MODEL_NAME_CLASS_SYMBOLS:
                    plus_obj = getattr(platform_models, class_symbol)
                    _plus_symbols[class_symbol] = plus_obj
                _plus_entries = [
                    (name, _plus_symbols[class_symbol]) for name, class_symbol in _CHECKPOINT_PLUS_MODEL_MAP_ENTRIES
                ]
                _plus_available = True
            except ModuleNotFoundError as ex:
                if ex.name not in {"rfdetr_plus", "rfdetr_plus.models"}:
                    raise

        # weights_only=False is required because legacy checkpoints embed
        # argparse.Namespace objects that cannot be deserialised with
        # weights_only=True.
        ckpt: dict[str, Any] = torch.load(path, map_location="cpu", weights_only=False)
        args = ckpt["args"]

        _variant_name_to_class: dict[str, type[RFDETR]] = {
            getattr(variant_obj, "__name__", symbol): variant_obj
            for symbol in dir(rfdetr_variants)
            if symbol.startswith("RFDETR")
            for variant_obj in [getattr(rfdetr_variants, symbol)]
        }
        _variant_symbols: dict[str, type[RFDETR]] = {
            class_symbol: _variant_name_to_class[class_symbol] for class_symbol in _CHECKPOINT_MODEL_NAME_CLASS_SYMBOLS
        }
        # Build in three explicit segments: seg-* entries, then plus-model entries
        # (xlarge/2xlarge), then base entries — order determines lookup priority.
        _seg_map: list[tuple[str, type[RFDETR]]] = [
            (name, _variant_symbols[class_symbol])
            for name, class_symbol in _CHECKPOINT_MODEL_MAP_ENTRIES
            if name.startswith("seg-")
        ]
        _keypoint_map: list[tuple[str, type[RFDETR]]] = [
            (name, _variant_symbols[class_symbol])
            for name, class_symbol in _CHECKPOINT_MODEL_MAP_ENTRIES
            if "keypoint" in name
        ]
        _base_map: list[tuple[str, type[RFDETR]]] = [
            (name, _variant_symbols[class_symbol])
            for name, class_symbol in _CHECKPOINT_MODEL_MAP_ENTRIES
            if not name.startswith("seg-") and "keypoint" not in name
        ]
        _model_map: list[tuple[str, type[RFDETR]]] = _seg_map + _keypoint_map + _plus_entries + _base_map

        # New checkpoints store model_name directly — use it when available.
        _name_map: dict[str, type[RFDETR]] = dict(_variant_symbols)
        # Plus-model classes are resolved only when rfdetr_plus is installed.
        if _plus_available:
            _name_map.update(_plus_symbols)
        saved_model_name = ckpt.get("model_name")
        model_cls: type[RFDETR] | None = None
        if isinstance(saved_model_name, str):
            normalized_name = saved_model_name.strip()
            if normalized_name:
                model_cls = _name_map.get(normalized_name)
        else:
            normalized_name = ""

        # Fall back to pretrain_weights (legacy) or, when unset-like, the checkpoint filename.
        if isinstance(args, dict):
            weights_name = str(args.get("pretrain_weights", "")).strip().lower()
        else:
            weights_name = str(getattr(args, "pretrain_weights", "")).strip().lower()
        # The sentinel set {"", "none", "null"} covers unset-like checkpoint values:
        #   ""     — pretrain_weights key absent entirely
        #   "none" — checkpoint value was None or the literal string "none";
        #            after str(...).strip().lower() both normalize to the same sentinel.
        #            This is NOT an intentional "no pretraining" flag (see
        #            test_pretrain_weights_none_warns, which operates at the config
        #            level, not the checkpoint level)
        #   "null" — checkpoint stored the literal string "null" (for example from a
        #            YAML-originated value), which is also treated as unset-like here
        _filename_fallback = False
        if weights_name in {"", "none", "null"}:
            weights_name = os.path.basename(os.fspath(path)).lower()
            _filename_fallback = True

        if model_cls is None:
            # Guard: plus-only checkpoints should raise an actionable install error
            # when rfdetr_plus is missing, regardless of whether class inference
            # relies on model_name (new format) or pretrain_weights (legacy format).
            plus_by_model_name = normalized_name in _CHECKPOINT_PLUS_MODEL_NAME_CLASS_SYMBOLS
            plus_by_weights_name = (
                "xlarge" in weights_name and "seg-" not in weights_name and "keypoint-preview" not in weights_name
            )
            if not _plus_available and (plus_by_model_name or plus_by_weights_name):
                from rfdetr.platform import _INSTALL_MSG

                raise ImportError(
                    f"Checkpoint model_name={saved_model_name!r}, pretrain_weights={weights_name!r} requires the "
                    f"rfdetr_plus package. " + _INSTALL_MSG.format(name="platform model downloads")
                )

            for name, klass in _model_map:
                if name in weights_name:
                    model_cls = klass
                    break

            if _filename_fallback and model_cls is not None:
                logger.info(
                    "pretrain_weights unset in checkpoint %r; inferred model class %s from filename %r",
                    path,
                    getattr(model_cls, "__name__", repr(model_cls)),
                    weights_name,
                )

        if model_cls is None:
            raise ValueError(
                f"Could not infer model class from checkpoint at {path!r} "
                f"(model_name={saved_model_name!r}, pretrain_weights={weights_name!r}). "
                f"Please instantiate the model class directly."
            )

        if isinstance(args, dict):
            num_classes: int | None = args.get("num_classes")
        else:
            num_classes = getattr(args, "num_classes", None)

        constructor_kwargs: dict[str, Any] = {}
        checkpoint_config_keys: set[str] = set()  # keys injected from checkpoint, not from caller
        saved_model_config = ckpt.get("model_config")
        if isinstance(saved_model_config, dict):
            model_config_class = getattr(model_cls, "_model_config_class", None)
            model_fields = getattr(model_config_class, "model_fields", None)
            if not isinstance(model_fields, dict):
                model_fields = getattr(model_config_class, "__fields__", None)
            if not isinstance(model_fields, dict):
                model_fields = {}
            for key, value in saved_model_config.items():
                if key == "pretrain_weights":
                    continue
                if not model_fields or key in model_fields:
                    constructor_kwargs[key] = value
                    checkpoint_config_keys.add(key)

        if num_classes is not None and "num_classes" not in kwargs:
            constructor_kwargs["num_classes"] = num_classes
            checkpoint_config_keys.add("num_classes")
        constructor_kwargs.update(kwargs)
        # pretrain_weights is placed after **kwargs so it always wins even if
        # a caller accidentally passes pretrain_weights inside kwargs.
        constructor_kwargs["pretrain_weights"] = str(path)

        # Fields injected from the checkpoint but not supplied by the caller must not be
        # treated as explicit user overrides in Pydantic's model_fields_set.  Downstream
        # alignment guards (e.g. _align_num_classes_from_dataset,
        # _align_keypoint_schema_from_dataset, load_pretrain_weights) all read
        # model_fields_set to decide whether to adapt model internals to the training
        # dataset — leaving checkpoint-derived fields marked as user-set breaks them.
        checkpoint_derived_keys = checkpoint_config_keys - set(kwargs)

        model = model_cls(**constructor_kwargs)

        if checkpoint_derived_keys:
            loaded_config = getattr(model, "model_config", None)
            # model_fields_set is the public API and returns the live backing set
            # in Pydantic v2; fall back to the private attribute only if that changes.
            fields_set = getattr(loaded_config, "model_fields_set", None)
            if fields_set is None:
                fields_set = getattr(loaded_config, "__pydantic_fields_set__", None)
            if fields_set is not None:
                fields_set.difference_update(checkpoint_derived_keys)
            # Verify num_classes specifically — if Pydantic ever returns a snapshot instead
            # of the live backing set, this assertion will catch the silent regression before
            # it causes a training-time head-adaptation failure.
            if "num_classes" in checkpoint_derived_keys:
                assert "num_classes" not in getattr(loaded_config, "model_fields_set", set()), (
                    "num_classes still in model_fields_set after checkpoint load; "
                    "Pydantic may return a snapshot rather than the live backing set — "
                    "switch to model_construct(_fields_set=...) for Pydantic v3 compatibility."
                )

        return model

    @staticmethod
    def _resolve_trainer_device_kwargs(device: Any) -> tuple[str | None, list[int] | None]:
        """Map a torch-style device specifier to PTL ``accelerator``/``devices`` kwargs.

        Args:
            device: A device specifier accepted by ``torch.device``.

        Returns:
            ``(accelerator, devices)`` where ``devices`` is ``None`` unless an explicit device index is provided (for
            example ``cuda:1``).

        Raises:
            ValueError: If ``device`` is not a valid torch device specifier.
        """
        if device is None:
            return None, None
        try:
            resolved_device = torch.device(device)
        except (TypeError, ValueError, RuntimeError) as exc:
            raise ValueError(
                f"Invalid device specifier for train(): {device!r}. "
                "Expected values like 'cpu', 'cuda', 'cuda:0', or torch.device(...).",
            ) from exc

        if resolved_device.type == "cpu":
            return "cpu", None
        if resolved_device.type == "cuda":
            return "gpu", [resolved_device.index] if resolved_device.index is not None else None
        if resolved_device.type == "mps":
            return "mps", [resolved_device.index] if resolved_device.index is not None else None

        warnings.warn(
            f"Device type {resolved_device.type!r} is not explicitly mapped to a PyTorch Lightning "
            "accelerator; falling back to PTL auto-detection. Training may use an unexpected device.",
            UserWarning,
            stacklevel=2,
        )
        return None, None

    def train(self, **kwargs):
        """Train an RF-DETR model via the PyTorch Lightning stack.

        All keyword arguments are forwarded to :meth:`get_train_config` to build a :class:`~rfdetr.config.TrainConfig`.
        Several kwargs are absorbed and handled specially so that existing call-sites do not break:

        * ``resolution`` — updates the model's input resolution by mutating
          :attr:`model_config.resolution` in place before the train config is built. This change persists on
          :attr:`model_config` after :meth:`train` returns. The value must be a positive integer divisible by
          ``patch_size * num_windows`` for the model variant; a :class:`ValueError` is raised otherwise.
          :attr:`model_config.positional_encoding_size` is also updated when the config derives it formulaically (``PE
          == resolution // patch_size``); configs with a pretrained-specific PE value (e.g. ``RFDETRBase`` uses DINOv2's
          PE=37 at 560 px) are left unchanged to preserve checkpoint compatibility.
        * ``device`` — normalized via :class:`torch.device` and mapped to PyTorch
          Lightning trainer arguments. ``"cpu"`` becomes ``accelerator="cpu"``; ``"cuda"`` and ``"cuda:N"`` become
          ``accelerator="gpu"`` and optionally ``devices=[N]``; ``"mps"`` becomes ``accelerator="mps"``. Other valid
          torch device types fall back to PTL auto-detection and emit a :class:`UserWarning`.
        * ``callbacks`` — if the dict contains any non-empty lists a
          :class:`DeprecationWarning` is emitted; the dict is then discarded. Use PTL
          :class:`~pytorch_lightning.Callback` objects passed via :func:`~rfdetr.training.build_trainer` instead.
        * ``start_epoch`` — emits :class:`DeprecationWarning` and is dropped.
        * ``do_benchmark`` — emits :class:`DeprecationWarning` and is dropped.
        * ``notes`` — optional user-defined metadata (string, dict, list, or
          any JSON-serialisable value) stored under the ``"notes"`` key in every ``.pth`` checkpoint produced during
          training.  The value is also available inside ``args["notes"]`` for full provenance.  Pass the same value to
          :meth:`export` to embed it in the ONNX file as well.

        After training completes the underlying ``nn.Module`` is synced back onto ``self.model.model`` so that
        :meth:`predict` and :meth:`export` continue to work without reloading the checkpoint.

        Raises:
            ImportError: If training dependencies are not installed. Install with
                ``pip install "rfdetr[train,loggers]"``.
            ValueError: If ``resolution`` is not a positive integer or is not
                divisible by ``patch_size * num_windows`` for the model variant.
        """
        # Both imports are grouped in a single try block because they both live in
        # the `rfdetr[train]` extras group — a missing `pytorch_lightning` (or any
        # other training-extras package) causes either import to fail, and the
        # remediation is identical: `pip install "rfdetr[train,loggers]"`.
        try:
            from rfdetr.training import RFDETRDataModule, RFDETRModelModule, build_trainer
            from rfdetr.training.auto_batch import resolve_auto_batch_config
        except ModuleNotFoundError as exc:
            # Preserve internal import errors so packaging/regression issues in
            # rfdetr.* are not misreported as missing optional extras.
            if exc.name and exc.name.startswith("rfdetr."):
                raise
            raise ImportError(
                "RF-DETR training dependencies are missing. "
                'Install them with `pip install "rfdetr[train,loggers]"` and try again.',
            ) from exc

        # Absorb legacy `callbacks` dict — warn if non-empty, then discard.
        callbacks_dict = kwargs.pop("callbacks", None)
        if callbacks_dict and any(callbacks_dict.values()):
            warnings.warn(
                "Custom callbacks dict is not forwarded to PTL. "
                "Deprecated since v1.7.0, will be removed in v1.9.0. "
                "Use PTL Callback objects instead.",
                DeprecationWarning,
                stacklevel=2,
            )

        # Parse `device` kwarg and map it to PTL accelerator/devices.
        # Supports torch-style strings and torch.device (e.g. "cuda:1").
        _device = kwargs.pop("device", None)
        _accelerator, _devices = RFDETR._resolve_trainer_device_kwargs(_device)

        # Absorb legacy `start_epoch` — PTL resumes automatically via ckpt_path.
        if "start_epoch" in kwargs:
            warnings.warn(
                "`start_epoch` is deprecated since v1.7.0 and will be removed in v1.9.0; "
                "PTL resumes automatically via `resume`.",
                DeprecationWarning,
                stacklevel=2,
            )
            kwargs.pop("start_epoch")

        # Pop `do_benchmark`; benchmarking via `.train()` is deprecated.
        run_benchmark = bool(kwargs.pop("do_benchmark", False))
        if run_benchmark:
            warnings.warn(
                "`do_benchmark` in `.train()` is deprecated since v1.7.0 and will be removed in v1.9.0; "
                "use `rfdetr benchmark`.",
                DeprecationWarning,
                stacklevel=2,
            )

        # Apply resolution override to model_config before building the train config.
        # resolution is a ModelConfig field, not a TrainConfig field, so we pop it
        # here to avoid it being silently ignored by TrainConfig.
        _resolution = kwargs.pop("resolution", None)
        if _resolution is not None:
            if isinstance(_resolution, bool):
                raise ValueError("resolution must be a positive integer")
            try:
                _resolution = operator.index(_resolution)
            except TypeError as error:
                raise ValueError("resolution must be a positive integer") from error
            if _resolution <= 0:
                raise ValueError("resolution must be a positive integer")
            block_size = self.model_config.patch_size * self.model_config.num_windows
            if _resolution % block_size != 0:
                raise ValueError(
                    f"resolution={_resolution} is not divisible by "
                    f"patch_size ({self.model_config.patch_size}) * num_windows "
                    f"({self.model_config.num_windows}) = {block_size}. "
                    f"Choose a resolution that is a multiple of {block_size}."
                )
            # Smart PE update: only recompute positional_encoding_size when the
            # current config derives it formulaically (PE == resolution // patch_size).
            # Configs with a pretrained-specific PE (e.g. RFDETRBase uses DINOv2's
            # PE=37 at 518 px, training at 560 px) must not have PE silently changed
            # — doing so causes shape mismatches when loading pretrained checkpoints.
            _current_pe = self.model_config.positional_encoding_size
            _derived_pe = self.model_config.resolution // self.model_config.patch_size
            if _current_pe == _derived_pe:
                # Formula-derived: update PE proportionally to the new resolution.
                new_pe = _resolution // self.model_config.patch_size
                self.model_config.positional_encoding_size = new_pe
            else:
                # Pretrained-specific PE; leave it unchanged.
                new_pe = _current_pe
            self.model_config.resolution = _resolution

            # Keep the cached inference/export context in sync with model_config so
            # predict()/export()/deployment all see the same resolution metadata.
            if hasattr(self, "model") and self.model is not None:
                if hasattr(self.model, "resolution"):
                    self.model.resolution = _resolution
                model_args = getattr(self.model, "args", None)
                if model_args is not None:
                    if hasattr(model_args, "resolution"):
                        model_args.resolution = _resolution
                    if hasattr(model_args, "positional_encoding_size"):
                        model_args.positional_encoding_size = new_pe
        config = self.get_train_config(**kwargs)
        if config.batch_size == "auto":
            # Auto-batch probing runs forward/backward on the actual model, which
            # must be on the target device (typically CUDA).  Lazy placement keeps
            # the model on CPU until first use — move it now.
            _move_model_context_to_device(self.model)
            auto_batch = resolve_auto_batch_config(
                model_context=self.model,
                model_config=self.model_config,
                train_config=config,
            )
            config.batch_size = auto_batch.safe_micro_batch
            config.grad_accum_steps = auto_batch.recommended_grad_accum_steps
            logger.info(
                "[auto-batch] resolved train config: batch_size=%s grad_accum_steps=%s effective_batch_size=%s",
                config.batch_size,
                config.grad_accum_steps,
                auto_batch.effective_batch_size,
            )
        self.model_config.model_name = type(self).__name__

        # Auto-detect num_classes from the training dataset and align model_config.
        # This must run before RFDETRModelModule is constructed so that weight loading
        # inside the module uses the correct (dataset-derived) class count.
        dataset_dir = getattr(config, "dataset_dir", None)
        if dataset_dir:
            self._align_keypoint_schema_from_dataset(config)
            self._align_num_classes_from_dataset(dataset_dir)

        module = RFDETRModelModule(self.model_config, config)
        datamodule = RFDETRDataModule(self.model_config, config)

        # Guard with LOCAL_RANK env var rather than is_main_process() because torch.distributed
        # is not yet initialized here (it is set up inside trainer.fit()).  In Lightning DDP
        # subprocesses, LOCAL_RANK is set by the launcher before the subprocess calls train(),
        # so this correctly identifies rank 0 even before dist.init_process_group() runs.
        if config.save_dataset_grids and os.environ.get("LOCAL_RANK", "0") == "0":
            try:
                from rfdetr.datasets.save_grids import DatasetGridSaver

                datamodule.setup("fit")
                grids_output_dir = Path(config.output_dir) / "dataset_grids"
                DatasetGridSaver(datamodule.train_dataloader(), grids_output_dir, dataset_type="train").save_grid()
                DatasetGridSaver(datamodule.val_dataloader(), grids_output_dir, dataset_type="val").save_grid()
            except Exception:
                logger.warning(
                    "Failed to save dataset grids; training will continue without them.",
                    exc_info=True,
                )

        trainer_kwargs = {"accelerator": _accelerator}
        if _devices is not None:
            trainer_kwargs["devices"] = _devices
        trainer = build_trainer(config, self.model_config, **trainer_kwargs)
        trainer.fit(module, datamodule, ckpt_path=config.resume or None)

        # Sync the trained weights back so predict() / export() see the updated model.
        self.model.model = module.model
        # Sync class names: prefer explicit config.class_names, otherwise fall back to dataset (#509).
        config_class_names = getattr(config, "class_names", None)
        if config_class_names is not None:
            self.model.class_names = config_class_names
        else:
            dataset_class_names = getattr(datamodule, "class_names", None)
            if dataset_class_names is not None:
                self.model.class_names = dataset_class_names

        # Save complete training configuration to disk for reproducibility.
        # Guard to main process only to avoid races in distributed/multi-GPU training.
        if is_main_process():
            complete_config = {
                "train_config": config.model_dump(),
                "model_config": self.model_config.model_dump(),
                "model_config_type": self.model_config.__class__.__name__,
                "class_names": self.model.class_names,
                "num_classes": len(self.model.class_names) if self.model.class_names else 0,
            }
            try:
                os.makedirs(config.output_dir, exist_ok=True)
                with open(os.path.join(config.output_dir, "training_config.json"), "w") as f:
                    json.dump(complete_config, f, indent=2, default=str)
            except OSError as exc:
                logger.warning("Could not save training_config.json to %s: %s", config.output_dir, exc)

    @_ensure_model_on_device
    def optimize_for_inference(
        self, compile: bool = True, batch_size: int = 1, dtype: torch.dtype | str = torch.float32
    ) -> None:
        """Optimize the model for inference with optional JIT compilation and dtype casting.

        Operations are wrapped in the correct CUDA device context to prevent context leaks on multi-GPU setups. When
        ``compile=True`` the model is traced with ``torch.jit.trace`` using a dummy input of ``batch_size`` images at
        the model's current resolution.

        Args:
            compile: If ``True``, trace the model with ``torch.jit.trace`` to obtain
                a JIT-compiled ``ScriptModule``. Set to ``False`` for broader compatibility (e.g. models with dynamic
                control flow).
            batch_size: Number of images the traced model will be optimized for. Ignored when ``compile=False``.
            dtype: Target floating-point dtype for the inference model. Accepts a
                ``torch.dtype`` directly (e.g. ``torch.float16``) or its string name (e.g. ``"float16"``). Defaults to
                ``torch.float32``.

        Raises:
            TypeError: If ``dtype`` is not a ``torch.dtype``, or if ``dtype`` is a
                string that does not correspond to a valid ``torch.dtype`` attribute.

        Examples:
            >>> from types import SimpleNamespace
            >>> import torch
            >>> class _TinyModel(torch.nn.Module):
            ...     def __init__(self):
            ...         super().__init__()
            ...         self.linear = torch.nn.Linear(1, 1)
            ...     def forward(self, x):
            ...         return {"pred_boxes": self.linear(x[:, :1, :1, :1].squeeze(-1).squeeze(-1))}
            ...     def export(self):
            ...         return None
            >>> class _TinyContext:
            ...     def __init__(self):
            ...         self.device = torch.device("cpu")
            ...         self.resolution = 28
            ...         self.model = _TinyModel()
            ...         self.inference_model = None
            >>> model = object.__new__(RFDETR)
            >>> model.model_config = SimpleNamespace(num_channels=3)
            >>> model.model = _TinyContext()
            >>> model._is_optimized_for_inference = False
            >>> model._has_warned_about_not_being_optimized_for_inference = False
            >>> model._optimized_has_been_compiled = False
            >>> model._optimized_batch_size = None
            >>> model._optimized_resolution = None
            >>> model._optimized_dtype = None
            >>> model.optimize_for_inference(compile=False, dtype="float16")
            >>> model._is_optimized_for_inference
            True
            >>> model._optimized_dtype
            torch.float16
        """
        if isinstance(dtype, str):
            try:
                dtype = getattr(torch, dtype)
            except AttributeError:
                raise TypeError(f"dtype must be a torch.dtype or a string name of a dtype, got {dtype!r}") from None
        if not isinstance(dtype, torch.dtype):
            raise TypeError(f"dtype must be a torch.dtype or a string name of a dtype, got {type(dtype)!r}")

        # Clear any previously optimized state before starting a new optimization run.
        self.remove_optimized_model()

        device = self.model.device
        cuda_ctx = torch.cuda.device(device) if device.type == "cuda" else contextlib.nullcontext()

        try:
            with cuda_ctx:
                self.model.inference_model = deepcopy(self.model.model)
                self.model.inference_model.eval()
                self.model.inference_model.export()

                self.model.inference_model = self.model.inference_model.to(dtype=dtype)

                if compile:
                    self.model.inference_model = torch.jit.trace(
                        self.model.inference_model,
                        torch.randn(
                            batch_size,
                            self.model_config.num_channels,
                            self.model.resolution,
                            self.model.resolution,
                            device=self.model.device,
                            dtype=dtype,
                        ),
                    )
                    self._optimized_has_been_compiled = True
                    self._optimized_batch_size = batch_size

                # Set success flags only after all operations complete.
                self._optimized_resolution = self.model.resolution
                self._is_optimized_for_inference = True
                self._optimized_dtype = dtype
        except Exception:
            # Ensure the object is left in a consistent, unoptimized state if optimization fails.
            with contextlib.suppress(Exception):
                self.remove_optimized_model()
            raise

    def remove_optimized_model(self) -> None:
        """Remove the optimized inference model and reset all optimization flags.

        Clears ``model.inference_model`` and resets all internal state set by :meth:`optimize_for_inference`. Safe to
        call even if the model has not been optimized.

        Examples:
            >>> from types import SimpleNamespace
            >>> import torch
            >>> class _TinyModel(torch.nn.Module):
            ...     def __init__(self):
            ...         super().__init__()
            ...         self.linear = torch.nn.Linear(1, 1)
            ...     def forward(self, x):
            ...         return {"pred_boxes": self.linear(x[:, :1, :1, :1].squeeze(-1).squeeze(-1))}
            ...     def export(self):
            ...         return None
            >>> class _TinyContext:
            ...     def __init__(self):
            ...         self.device = torch.device("cpu")
            ...         self.resolution = 28
            ...         self.model = _TinyModel()
            ...         self.inference_model = None
            >>> model = object.__new__(RFDETR)
            >>> model.model_config = SimpleNamespace(num_channels=3)
            >>> model.model = _TinyContext()
            >>> model._is_optimized_for_inference = False
            >>> model._has_warned_about_not_being_optimized_for_inference = False
            >>> model._optimized_has_been_compiled = False
            >>> model._optimized_batch_size = None
            >>> model._optimized_resolution = None
            >>> model._optimized_dtype = None
            >>> model.optimize_for_inference(compile=False)
            >>> model.remove_optimized_model()
            >>> model._is_optimized_for_inference
            False
        """
        self.model.inference_model = None
        self._is_optimized_for_inference = False
        self._optimized_has_been_compiled = False
        self._optimized_batch_size = None
        self._optimized_resolution = None
        self._optimized_dtype = None

    def export(
        self,
        output_dir: str = "output",
        infer_dir: str | None = None,
        backbone_only: bool = False,
        opset_version: int = 17,
        verbose: bool = True,
        shape: tuple[int, int] | None = None,
        batch_size: int = 1,
        dynamic_batch: bool = False,
        patch_size: int | None = None,
        format: str = "onnx",
        quantization: str | None = None,
        calibration_data: str | np.ndarray | None = None,
        max_images: int = 100,
        *,
        notes: object = None,
    ) -> Path:
        """Export the trained model to ONNX or TFLite format.

        See the `export documentation <https://rfdetr.roboflow.com/learn/export/>`_ for more information.

        Args:
            output_dir: Directory to write the exported model to.
            infer_dir: Optional directory of sample images for dynamic-axes inference.
            backbone_only: Export only the backbone (feature extractor).
            opset_version: ONNX opset version to target.
            verbose: Print export progress information.
            shape: ``(height, width)`` tuple; defaults to square at model resolution.
                Both dimensions must be divisible by ``patch_size * num_windows``.
            batch_size: Static batch size to bake into the ONNX graph.
            dynamic_batch: If True, export with a dynamic batch dimension
                so the ONNX model accepts variable batch sizes at runtime.
            patch_size: Backbone patch size. Defaults to the value stored in
                ``model_config.patch_size`` (typically 14 or 16). When provided explicitly it must match the
                instantiated model's patch size. Shape divisibility is validated against ``patch_size * num_windows``.
            format: Export format — ``"onnx"`` (default) or ``"tflite"``.
                When ``"tflite"`` is selected the model is first exported to ONNX then converted to TFLite via
                ``onnx2tf``.  Requires ``pip install rfdetr[onnx,tflite]``.

                .. warning::
                    TFLite export is experimental and subject to change; upstream dependency instabilities (``onnx2tf``,
                    ``ai_edge_litert``) may affect results.
            quantization: TFLite quantization mode (ignored when
                ``format="onnx"``).  One of ``None``, ``"fp32"``, ``"fp16"``, ``"int8"``.  ``None`` / ``"fp32"`` /
                ``"fp16"`` produce FP32 + FP16 ``.tflite`` files; ``"int8"`` additionally produces an INT8-quantized
                model.
            calibration_data: Representative images for INT8 calibration and ``onnx2tf`` output validation.  Accepts:

                * ``None`` — auto-generate random data (sufficient for fp32/fp16; warns for int8).
                * A **directory path** (``str``) containing JPEG/PNG
                  images — the converter automatically loads, resizes, and prepares them.  This is the simplest
                  approach.
                * A path (``str``) to a ``.npy`` file of shape ``(N, H, W, 3)``, dtype float32, values in ``[0, 1]``.
                * A :class:`numpy.ndarray` with the same format.

                For INT8 quantization, provide 20–100 representative images from your training/validation set for best
                accuracy.
            max_images: Maximum number of images to load from a calibration directory.  Defaults to ``100``.  Only used
                when *calibration_data* is a directory path.
            notes: Optional user-defined metadata (string, dict, list, or
                any JSON-serialisable value) to embed in the exported ONNX model under the ``"rfdetr_notes"`` metadata
                property.  When ``None`` no metadata entry is written.  String values are stored verbatim; all other
                types are JSON-encoded so consumers must call ``json.loads()`` to recover a dict or list.  The same
                value can be passed to :meth:`train` so the checkpoint and the ONNX file share the same provenance
                information.

        Returns:
            Path to the exported model file (``.onnx`` or ``.tflite``).
        """
        logger.info("Exporting model to ONNX format")
        _valid_formats = ("onnx", "tflite")
        if format not in _valid_formats:
            raise ValueError(f"Unsupported export format {format!r}. Choose from: {_valid_formats}")
        try:
            from rfdetr.export.main import export_onnx, make_infer_image
        except ImportError:
            logger.error(
                "It seems some dependencies for ONNX export are missing."
                " Please run `pip install rfdetr[onnx]` and try again.",
            )
            raise

        device = self.model.device
        # Move the live model to CPU before deepcopying and keep it there during export. ``nn.Module.to(...)`` mutates
        # in place, so this frees GPU memory for the local export copy, ONNX tracing, TFLite conversion, and any
        # calibration tensors. The ``finally`` block restores the live model even if export or conversion raises.
        self.model.model = self.model.model.to("cpu")
        model = deepcopy(self.model.model)
        model.to(device)
        try:
            os.makedirs(output_dir, exist_ok=True)
            output_dir_path = Path(output_dir)
            patch_size = _resolve_patch_size(patch_size, self.model_config, "export")
            num_windows = getattr(self.model_config, "num_windows", 1)
            if isinstance(num_windows, bool) or not isinstance(num_windows, int) or num_windows <= 0:
                raise ValueError(f"num_windows must be a positive integer, got {num_windows!r}")
            block_size = patch_size * num_windows
            if shape is None:
                shape = (self.model.resolution, self.model.resolution)
                if shape[0] % block_size != 0:
                    raise ValueError(
                        f"Model's default resolution ({self.model.resolution}) is not divisible by "
                        f"block_size={block_size} (patch_size={patch_size} * num_windows={num_windows}). "
                        f"Provide an explicit shape divisible by {block_size}.",
                    )
            else:
                shape = _validate_shape_dims(shape, block_size, patch_size, num_windows)

            input_tensors = make_infer_image(
                infer_dir, shape, batch_size, device, num_channels=self.model_config.num_channels
            ).to(device)
            input_names = ["input"]
            if backbone_only:
                output_names = ["features"]
            elif self.model_config.segmentation_head:
                output_names = ["dets", "labels", "masks"]
            elif self.model_config.use_grouppose_keypoints:
                output_names = ["dets", "labels", "keypoints"]
            else:
                output_names = ["dets", "labels"]

            if dynamic_batch:
                dynamic_axes = {name: {0: "batch"} for name in input_names + output_names}
            else:
                dynamic_axes = None
            model.eval()
            with torch.no_grad():
                if backbone_only:
                    features = model(input_tensors)
                    logger.debug(f"PyTorch inference output shape: {features.shape}")
                elif self.model_config.segmentation_head:
                    outputs = model(input_tensors)
                    dets = outputs["pred_boxes"]
                    labels = outputs["pred_logits"]
                    masks = outputs["pred_masks"]
                    if isinstance(masks, torch.Tensor):
                        logger.debug(
                            f"PyTorch inference output shapes - Boxes: {dets.shape}, Labels: {labels.shape}, "
                            f"Masks: {masks.shape}",
                        )
                    else:
                        logger.debug(f"PyTorch inference output shapes - Boxes: {dets.shape}, Labels: {labels.shape}")
                elif self.model_config.use_grouppose_keypoints:
                    outputs = model(input_tensors)
                    dets = outputs["pred_boxes"]
                    labels = outputs["pred_logits"]
                    keypoints = outputs["pred_keypoints"]
                    logger.debug(
                        f"PyTorch inference output shapes - Boxes: {dets.shape}, Labels: {labels.shape}, "
                        f"Keypoints: {keypoints.shape}",
                    )
                else:
                    outputs = model(input_tensors)
                    dets = outputs["pred_boxes"]
                    labels = outputs["pred_logits"]
                    logger.debug(f"PyTorch inference output shapes - Boxes: {dets.shape}, Labels: {labels.shape}")

            model.cpu()
            input_tensors = input_tensors.cpu()

            output_file = export_onnx(
                output_dir=str(output_dir_path),
                model=model,
                input_names=input_names,
                input_tensors=input_tensors,
                output_names=output_names,
                dynamic_axes=dynamic_axes,
                backbone_only=backbone_only,
                verbose=verbose,
                opset_version=opset_version,
                variant_name=getattr(self, "size", None),
                notes=notes,
            )

            logger.info(f"Successfully exported ONNX model to: {output_file}")

            if format == "tflite":
                warnings.warn(
                    "TFLite export is experimental and work-in-progress. "
                    "Upstream dependency instabilities (onnx2tf, ai_edge_litert) may affect results.",
                    UserWarning,
                    stacklevel=2,
                )
                try:
                    from rfdetr.export._tflite.converter import export_tflite
                except ImportError:
                    logger.error(
                        "It seems some dependencies for TFLite export are missing."
                        " Please run `pip install rfdetr[onnx,tflite]` and try again.",
                    )
                    raise

                tflite_path = export_tflite(
                    onnx_path=output_file,
                    output_dir=str(output_dir_path),
                    quantization=quantization,
                    calibration_data=calibration_data,
                    verbosity="info" if verbose else "error",
                    max_images=max_images,
                    verbose=verbose,
                )
                logger.info(f"Successfully exported TFLite model to: {tflite_path}")
                return tflite_path

            logger.info("Export completed successfully")
            return Path(output_file)
        finally:
            self.model.model = self.model.model.to(device)

    @staticmethod
    def _load_classes(dataset_dir: str) -> list[str]:
        """Load class names from a COCO or YOLO dataset directory."""
        if is_valid_coco_dataset(dataset_dir):
            coco_path = os.path.join(dataset_dir, "train", "_annotations.coco.json")
            with open(coco_path, encoding="utf-8") as f:
                anns = json.load(f)
            categories = sorted(anns["categories"], key=lambda category: category.get("id", float("inf")))

            # Catch possible placeholders for no supercategory
            placeholders = {"", "none", "null", None}

            # If no meaningful supercategory exists anywhere, treat as flat dataset
            has_any_sc = any(c.get("supercategory", "none") not in placeholders for c in categories)
            if not has_any_sc:
                return [c["name"] for c in categories]

            # Mixed/Hierarchical: keep only categories that are not parents of other categories.
            # Both leaves (with a real supercategory) and standalone top-level nodes (supercategory is a
            # placeholder) satisfy this condition — neither appears as another category's supercategory.
            parents = {c.get("supercategory") for c in categories if c.get("supercategory", "none") not in placeholders}
            has_children = {c["name"] for c in categories if c["name"] in parents}

            class_names = [c["name"] for c in categories if c["name"] not in has_children]
            # Safety fallback for pathological inputs
            return class_names or [c["name"] for c in categories]

        # list all YAML files in the folder
        if is_valid_yolo_dataset(dataset_dir):
            yaml_paths = glob.glob(os.path.join(dataset_dir, "*.yaml")) + glob.glob(os.path.join(dataset_dir, "*.yml"))
            # any YAML file starting with data e.g. data.yaml, dataset.yaml
            yaml_data_files = [yp for yp in yaml_paths if os.path.basename(yp).startswith("data")]
            yaml_path = yaml_data_files[0]
            with open(yaml_path) as f:
                data = yaml.safe_load(f)
            if "names" in data:
                if isinstance(data["names"], dict):
                    return [data["names"][i] for i in sorted(data["names"].keys())]
                return data["names"]
            raise ValueError(f"Found {yaml_path} but it does not contain 'names' field.")
        raise FileNotFoundError(
            f"Could not find class names in {dataset_dir}."
            " Checked for COCO (train/_annotations.coco.json) and YOLO (data.yaml, data.yml) styles.",
        )

    @staticmethod
    def _detect_num_classes_for_training(dataset_dir: str, *, use_grouppose_keypoints: bool = False) -> int:
        """Detect the class count using the same category basis as training labels.

        For COCO-style datasets this counts all categories by ``id`` from ``train/_annotations.coco.json`` (matching the
        remapping based on ``coco.cats`` used by the training datamodule). In keypoint mode it instead counts the
        inferred RF-DETR keypoint label slots, where slot ``0`` may be reserved for classes without keypoints. For YOLO-
        style datasets it falls back to ``_load_classes``.
        """
        if is_valid_coco_dataset(dataset_dir):
            coco_path = os.path.join(dataset_dir, "train", "_annotations.coco.json")
            if use_grouppose_keypoints:
                return len(infer_coco_keypoint_schema(coco_path).class_names)
            with open(coco_path, encoding="utf-8") as f:
                anns = json.load(f)
            categories = anns["categories"]
            cat_by_id = {category["id"]: category for category in categories}
            return len(cat_by_id)

        return len(RFDETR._load_classes(dataset_dir))

    def _align_num_classes_from_dataset(self, dataset_dir: str) -> None:
        """Auto-detect the dataset class count and align ``model_config.num_classes`` in-place.

        Must be called before ``RFDETRModelModule`` is constructed so that weight loading inside the module uses the
        correct (dataset-derived) class count.

        When the user did **not** explicitly set ``num_classes`` (it is left unset, e.g. inferred from a
        checkpoint), ``model_config.num_classes`` and ``self.model.args.num_classes`` are updated to match the dataset.
        When the user *did* set ``num_classes`` explicitly — to any value, including the class default — and it differs
        from the dataset, the configured value is preserved and a warning is emitted.

        Failures from ``_detect_num_classes_for_training`` are caught and logged at DEBUG level so that training is
        never blocked by detection errors.

        When ``model_config.use_grouppose_keypoints`` is True and
        ``model_config.num_keypoints_per_class`` is shorter than the adjusted
        ``num_classes``, the schema is zero-padded in-place so that
        ``len(num_keypoints_per_class) == num_classes``.  Both ``model_config``
        and ``model.args`` (if present) are updated.  Appended classes receive
        zero keypoints and contribute no class-logit boost.

        Args:
            dataset_dir: Path to the training dataset root directory.
        """
        try:
            dataset_num_classes = RFDETR._detect_num_classes_for_training(
                dataset_dir,
                use_grouppose_keypoints=self.model_config.use_grouppose_keypoints,
            )
        except (FileNotFoundError, ValueError, KeyError, OSError) as exc:
            # Best-effort only; do not block training if detection fails.
            logger.debug("Could not auto-detect num_classes from dataset '%s': %s", dataset_dir, exc)
            return

        # Hoist so both branches below can reference the schema without re-fetching.
        keypoint_schema: list[int] = []
        if self.model_config.use_grouppose_keypoints:
            # Older configs may omit the schema; absence means no schema-based class-count expansion.
            keypoint_schema = list(getattr(self.model_config, "num_keypoints_per_class", []) or [])
            if keypoint_schema:
                dataset_num_classes = max(dataset_num_classes, len(keypoint_schema))

        model_num_classes = self.model_config.num_classes

        if dataset_num_classes == model_num_classes:
            return

        # Determine whether the user explicitly set num_classes.  "num_classes" in
        # model_fields_set is True only when the field was explicitly provided at construction
        # (or assigned afterwards); an explicit value is honored regardless of whether it equals
        # the class default, so an intentional num_classes is never silently overridden by the
        # dataset count.  A checkpoint-derived num_classes is cleared from model_fields_set by
        # ``from_checkpoint`` (see PR #1106 / issue #1092), so it correctly counts as "not set" here.
        user_overrode = "num_classes" in getattr(self.model_config, "model_fields_set", set())

        if not user_overrode:
            logger.debug(
                "Detected %d classes in dataset '%s'; auto-adjusting model num_classes from %d to %d.",
                dataset_num_classes,
                dataset_dir,
                model_num_classes,
                dataset_num_classes,
            )
            self.model_config.num_classes = dataset_num_classes
            # Keep serialized checkpoint metadata in sync with the updated class count.
            model_args = getattr(self.model, "args", None)
            if model_args is not None:
                model_args.num_classes = dataset_num_classes
            # Pad keypoint schema with zeros so len(num_keypoints_per_class) == num_classes.
            # Without this, _aggregate_keypoint_class_logits emits a one-time mismatch
            # warning per model instance and the config state is inconsistent with the
            # detection head width.
            if keypoint_schema and len(keypoint_schema) < dataset_num_classes:
                padded_schema = keypoint_schema + [0] * (dataset_num_classes - len(keypoint_schema))
                self.model_config.num_keypoints_per_class = padded_schema
                if model_args is not None:
                    model_args.num_keypoints_per_class = padded_schema
        else:
            logger.warning(
                "Dataset '%s' has %d classes but model was initialized with num_classes=%d. "
                "Using the model's configured value (%d). If this is unintentional, "
                "reinitialize the model with num_classes=%d.",
                dataset_dir,
                dataset_num_classes,
                model_num_classes,
                model_num_classes,
                dataset_num_classes,
            )
            # Also pad schema when the user-configured num_classes exceeds the schema length,
            # to prevent the _aggregate_keypoint_class_logits mismatch warning in this path too.
            if keypoint_schema and len(keypoint_schema) < model_num_classes:
                padded_schema = keypoint_schema + [0] * (model_num_classes - len(keypoint_schema))
                self.model_config.num_keypoints_per_class = padded_schema
                model_args = getattr(self.model, "args", None)
                if model_args is not None:
                    model_args.num_keypoints_per_class = padded_schema

    @staticmethod
    def _roboflow_keypoint_annotation_path(dataset_dir: str) -> Path | None:
        """Return the Roboflow COCO train annotation path when it exists.

        Args:
            dataset_dir: Path to the Roboflow dataset root.

        Returns:
            Train split annotation path, or ``None`` when the dataset is not Roboflow COCO style.

        Raises:
            This helper does not raise.

        Example:
            >>> RFDETR._roboflow_keypoint_annotation_path("/missing") is None
            True
        """
        if not is_valid_coco_dataset(dataset_dir):
            return None
        annotation_path = Path(dataset_dir) / "train" / "_annotations.coco.json"
        return annotation_path if annotation_path.exists() else None

    def _align_keypoint_schema_from_dataset(self, config: TrainConfig) -> None:
        """Infer or validate keypoint schema from Roboflow COCO metadata.

        Args:
            config: Training configuration containing dataset location and format.

        Returns:
            ``None``. The model config is updated in-place when dataset metadata is available.

        Raises:
            This method does not raise for missing or malformed metadata; later dataset construction still validates
            keypoint-mode requirements.

        Example:
            >>> from rfdetr.config import RFDETRKeypointPreviewConfig, TrainConfig
            >>> model = object.__new__(RFDETR)
            >>> model.model_config = RFDETRKeypointPreviewConfig(pretrain_weights=None)
            >>> model.model = type("Context", (), {"args": None})()
            >>> model._align_keypoint_schema_from_dataset(TrainConfig(dataset_dir="/missing", tensorboard=False))
        """

        if not self.model_config.use_grouppose_keypoints:
            return
        if getattr(config, "dataset_file", None) != "roboflow":
            return
        dataset_dir = getattr(config, "dataset_dir", None)
        if not dataset_dir:
            return
        annotation_path = RFDETR._roboflow_keypoint_annotation_path(dataset_dir)
        if annotation_path is None:
            return

        try:
            inferred = infer_coco_keypoint_schema(annotation_path)
        except (FileNotFoundError, ValueError, KeyError, OSError) as exc:
            logger.info("Could not infer keypoint schema from dataset '%s': %s", dataset_dir, exc)
            return

        inferred_schema = inferred.num_keypoints_per_class
        # Older configs may omit the schema; absence lets dataset inference populate it.
        current_schema = list(getattr(self.model_config, "num_keypoints_per_class", []) or [])
        user_set_schema = "num_keypoints_per_class" in getattr(self.model_config, "model_fields_set", set())

        if user_set_schema and active_keypoint_counts(current_schema) == active_keypoint_counts(inferred_schema):
            return

        if current_schema != inferred_schema:
            if user_set_schema:
                logger.warning(
                    "Configured num_keypoints_per_class=%s does not match dataset keypoint metadata %s from '%s'. "
                    "Using dataset metadata as the source of truth.",
                    current_schema,
                    inferred_schema,
                    annotation_path,
                )
            else:
                logger.info(
                    "Inferred num_keypoints_per_class=%s from Roboflow COCO keypoint metadata at '%s'.",
                    inferred_schema,
                    annotation_path,
                )
            self.model_config.num_keypoints_per_class = inferred_schema
            model_args = getattr(self.model, "args", None)
            if model_args is not None:
                model_args.num_keypoints_per_class = inferred_schema

    def get_train_config(self, **kwargs) -> TrainConfig:
        """Retrieve the configuration parameters that will be used for training."""
        return self._train_config_class(**kwargs)

    def get_model(self, config: ModelConfig) -> ModelContext:
        """Retrieve a model context from the provided architecture configuration.

        Args:
            config: Architecture configuration.

        Returns:
            ModelContext with model, postprocess, device, resolution, args, and class_names attributes.
        """
        return _build_model_context(config)

    @property
    def class_names(self) -> list[str]:
        """Retrieve the class names supported by the loaded model.

        Returns:
            A list of class name strings, 0-indexed.  When no custom class names are embedded in the checkpoint, returns
            the standard 80 COCO class names.
        """
        if hasattr(self.model, "class_names") and self.model.class_names is not None:
            return list(self.model.class_names)

        return list(COCO_CLASS_NAMES)

    @torch.no_grad()
    @_ensure_model_on_device
    def predict(
        self,
        images: str | Image.Image | np.ndarray | torch.Tensor | list[str | np.ndarray | Image.Image | torch.Tensor],
        threshold: float = 0.5,
        shape: tuple[int, int] | None = None,
        patch_size: int | None = None,
        include_source_image: bool = True,
        **kwargs: Any,
    ) -> Detections | KeyPoints | list[Detections | KeyPoints]:
        """Performs model inference on the input images.

        This method accepts a single image or a list of images in various formats (file path, image url, PIL Image,
        NumPy array, or torch.Tensor). The images should be in RGB channel order. If a torch.Tensor is provided, it must
        already be normalized to values in the [0, 1] range and have the shape (C, H, W).

        Args:
            images:
                A single image or a list of images to process. Images can be provided
                as file paths, PIL Images, NumPy arrays, or torch.Tensors.
            threshold:
                The minimum confidence score needed to consider a detected bounding box valid.
            shape:
                Optional ``(height, width)`` tuple to resize images to before inference. When provided, overrides the
                model's default inference resolution. The tuple should match the resolution used when exporting the
                model (typically a square shape). Both dimensions must be positive integers divisible by ``patch_size *
                num_windows``. Defaults to ``(model.resolution, model.resolution)`` when not set.
            patch_size:
                Backbone patch size used for shape divisibility validation. Defaults to ``model_config.patch_size``
                (typically 14 for large models, 16 for smaller ones). Divisibility is checked against ``patch_size *
                num_windows``.
            include_source_image:
                Whether to attach the original image to the returned prediction. Detection and segmentation outputs use
                ``detections.metadata["source_image"]``. Keypoint outputs use per-object
                ``key_points.data["source_image"]`` because Supervision ``KeyPoints`` currently has no collection-level
                metadata field. Defaults to ``True``. Set to ``False`` to reduce memory use when source images are not
                needed.
            **kwargs:
                Additional keyword arguments.

        Returns:
            A single or multiple Supervision prediction objects. Detection and segmentation models return
            :class:`~supervision.Detections`. Keypoint models return :class:`~supervision.KeyPoints`, with keypoint
            coordinates in ``xy``. Keypoint predictions preserve the detection-level fields produced by RF-DETR:
            ``key_points.detection_confidence`` is the per-object score used by ``threshold``. For keypoint models this
            is the postprocessed detection score and, by default, includes keypoint uncertainty fusion controlled by
            ``model_config.postprocess_trace_alpha``. ``key_points.keypoint_confidence`` is separate: it is a
            ``(num_detections, num_keypoints)`` array of per-keypoint findability scores decoded from the keypoint head,
            not a repeated copy of the detection score. When RF-DETR emits keypoint precision parameters,
            ``key_points.data["covariance"]`` stores per-keypoint pixel-space covariance matrices with shape
            ``(num_detections, num_keypoints, 2, 2)``. ``key_points.data["xyxy"]`` stores the corresponding detection
            boxes as a ``(num_detections, 4)`` array in the same row order as ``key_points.xy`` because Supervision
            ``KeyPoints`` does not have a native bounding-box field. The ``data`` dict also contains ``class_name`` and
            ``source_shape`` as per-object arrays. When ``include_source_image=True`` for keypoint models,
            ``source_image`` is stored as per-object data until Supervision exposes collection-level metadata for
            ``KeyPoints``.

        Note:
            For ``Detections`` outputs, ``source_image`` moved from ``detections.data`` to ``detections.metadata``.
            Update detection callers reading ``detections.data["source_image"]`` to use
            ``detections.metadata["source_image"]``.

        Note:
            ``class_name`` mapping uses one of two modes depending on the checkpoint. For pretrained COCO checkpoints
            (detected when ``model.args.num_classes > len(class_names)`` and ``class_names`` matches
            ``COCO_CLASS_NAMES``), raw COCO category IDs (1–90, sparse) are looked up by category ID rather than by
            position — so ``class_id=18`` yields ``"dog"``, not ``class_names[18]``. For fine-tuned models, ``class_id``
            is a 0-based index into ``class_names``.

        Raises:
            ValueError: If ``shape`` cannot be unpacked as a two-element sequence,
                if either dimension does not support the ``__index__`` protocol (e.g. ``float``) or is a ``bool``, if
                either dimension is zero or negative, if either dimension is not divisible by ``patch_size *
                num_windows``, or if ``patch_size`` is not a positive integer.
        """
        from supervision import Detections, KeyPoints

        patch_size = _resolve_patch_size(patch_size, self.model_config, "predict")
        num_windows = getattr(self.model_config, "num_windows", 1)
        if isinstance(num_windows, bool) or not isinstance(num_windows, int) or num_windows <= 0:
            raise ValueError(f"model_config.num_windows must be a positive integer, got {num_windows!r}")
        block_size = patch_size * num_windows

        if shape is None:
            default_res = self.model.resolution
            if default_res % block_size != 0:
                raise ValueError(
                    f"Model's default resolution ({default_res}) is not divisible by "
                    f"block_size={block_size} (patch_size={patch_size} * num_windows={num_windows}). "
                    f"Provide an explicit shape divisible by {block_size}.",
                )
        else:
            shape = _validate_shape_dims(shape, block_size, patch_size, num_windows)

        if not self._is_optimized_for_inference and not self._has_warned_about_not_being_optimized_for_inference:
            logger.warning(
                "Model is not optimized for inference. Latency may be higher than expected."
                " You can optimize the model for inference by calling model.optimize_for_inference().",
            )
            self._has_warned_about_not_being_optimized_for_inference = True

            self.model.model.eval()

        if not isinstance(images, list):
            images = [images]

        orig_sizes = []
        processed_images = []
        source_images = [] if include_source_image else None

        for img in images:
            if isinstance(img, str):
                if img.startswith("http"):
                    img = requests.get(img, stream=True).raw
                img = Image.open(img)

            if not isinstance(img, torch.Tensor):
                if include_source_image:
                    src = np.array(img)
                    if src.dtype != np.uint8:
                        src = (src * 255).clip(0, 255).astype(np.uint8)
                    source_images.append(src)
                img = F.to_tensor(img)
            elif include_source_image:
                source_images.append((img.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))

            if (img > 1).any():
                raise ValueError(
                    "Image has pixel values above 1. Please ensure the image is normalized (scaled to [0, 1]).",
                )
            if (img < 0).any():
                raise ValueError(
                    "Image has pixel values below 0. Please ensure the image is normalized (scaled to [0, 1]).",
                )
            if img.shape[0] != self.model_config.num_channels:
                raise ValueError(
                    "Invalid tensor image shape. Tensor inputs to `predict()` must be in (C, H, W) format "
                    f"with C matching the model configuration ({self.model_config.num_channels} channels). "
                    f"Received tensor with shape {tuple(img.shape)}."
                )
            img_tensor = img

            h, w = img_tensor.shape[1:]
            orig_sizes.append((h, w))

            img_tensor = img_tensor.to(self.model.device)
            resize_to = list(shape) if shape is not None else [self.model.resolution, self.model.resolution]
            img_tensor = F.resize(img_tensor, resize_to)
            img_tensor = F.normalize(img_tensor, self.means, self.stds)

            processed_images.append(img_tensor)

        batch_tensor = torch.stack(processed_images)

        if self._is_optimized_for_inference:
            if (
                self._optimized_resolution != batch_tensor.shape[2]
                or self._optimized_resolution != batch_tensor.shape[3]
            ):
                # this could happen if someone manually changes self.model.resolution after optimizing the model,
                # or if predict(shape=...) is used with a shape that doesn't match the compiled square resolution.
                raise ValueError(
                    f"Resolution mismatch. "
                    f"Model was optimized for resolution {self._optimized_resolution}x{self._optimized_resolution}, "
                    f"but got {batch_tensor.shape[2]}x{batch_tensor.shape[3]}."
                    " You can explicitly remove the optimized model by calling model.remove_optimized_model().",
                )
            if self._optimized_has_been_compiled:
                if self._optimized_batch_size != batch_tensor.shape[0]:
                    raise ValueError(
                        f"Batch size mismatch. "
                        f"Optimized model was compiled for batch size {self._optimized_batch_size}, "
                        f"but got {batch_tensor.shape[0]}."
                        " You can explicitly remove the optimized model by calling model.remove_optimized_model()."
                        " Alternatively, you can recompile the optimized model for a different batch size"
                        " by calling model.optimize_for_inference(batch_size=<new_batch_size>).",
                    )

        if self._is_optimized_for_inference:
            predictions = self.model.inference_model(batch_tensor.to(dtype=self._optimized_dtype))
        else:
            predictions = self.model.model(batch_tensor)
        if isinstance(predictions, tuple):
            return_predictions = {
                "pred_logits": predictions[1],
                "pred_boxes": predictions[0],
            }
            if len(predictions) == 3:
                # Distinguish optional keypoint vs mask tuple output for legacy compiled/export shims.
                if getattr(getattr(self.model, "model_config", None), "use_grouppose_keypoints", False):
                    return_predictions["pred_keypoints"] = predictions[2]
                else:
                    return_predictions["pred_masks"] = predictions[2]
            predictions = return_predictions
        target_sizes = torch.tensor(orig_sizes, device=self.model.device)
        results = self.model.postprocess(predictions, target_sizes=target_sizes)

        model_class_names = self.class_names
        n = len(model_class_names)
        # Pretrained COCO models use COCO category IDs (1–90, with gaps) as class_ids,
        # while class_names is a flat 0-indexed list of 80 entries. Detected when
        # args.num_classes > len(class_names) AND class_names == COCO_CLASS_NAMES.
        # Fine-tuned models remap category IDs to 0-based contiguous indices, so
        # class_id i maps directly to class_names[i].
        _model_args = getattr(self.model, "args", None)
        if _model_args is None and model_class_names == list(COCO_CLASS_NAMES):
            logger.warning_once(
                "predict(): model has no 'args' attribute — COCO sparse-ID mapping cannot activate; "
                "class_ids are treated as 0-indexed (may be wrong for pretrained COCO checkpoints)"
            )
        num_logit_slots: int = getattr(_model_args, "num_classes", n)
        _is_coco_pretrained = num_logit_slots > n and model_class_names == list(COCO_CLASS_NAMES)
        if _is_coco_pretrained:
            _class_id_to_name: dict[int, str] = {
                coco_id: model_class_names[i] for i, coco_id in enumerate(COCO_CLASSES) if i < n
            }
        else:
            _class_id_to_name = dict(enumerate(model_class_names))
        predictions_list: list[Detections | KeyPoints] = []
        for i, result in enumerate(results):
            scores = result["scores"]
            labels = result["labels"]
            boxes = result["boxes"]

            keep = scores > threshold
            scores = scores[keep]
            labels = labels[keep]
            boxes = boxes[keep]
            keypoints_array = None
            if "keypoints" in result:
                keypoints = result["keypoints"][keep]
                keypoints_array = keypoints.float().cpu().numpy()
            has_keypoints = keypoints_array is not None

            if "masks" in result:
                masks = result["masks"]
                masks = masks[keep]

                detections = Detections(
                    xyxy=boxes.float().cpu().numpy(),
                    confidence=scores.float().cpu().numpy(),
                    class_id=labels.cpu().numpy(),
                    mask=masks.squeeze(1).cpu().numpy(),
                )
            else:
                detections = Detections(
                    xyxy=boxes.float().cpu().numpy(),
                    confidence=scores.float().cpu().numpy(),
                    class_id=labels.cpu().numpy(),
                )
            if "keypoint_precision_cholesky" in result:
                keypoint_precision = result["keypoint_precision_cholesky"][keep]
                detections.data["keypoint_precision_cholesky"] = keypoint_precision.float().cpu().numpy()

            if include_source_image:
                detections.metadata["source_image"] = source_images[i]
            detections.data["source_shape"] = np.tile(np.array(orig_sizes[i], dtype=np.int64), (len(detections), 1))

            # Attach class names so callers can map class_id → name without a
            # separate lookup. Always set data["class_name"] for a consistent interface.
            #
            # For fine-tuned models, logit index num_logit_slots is the no-object slot —
            # map it to "__background__" without warning. For COCO-pretrained models,
            # background is implicit (filtered by threshold); class ID 90 is "toothbrush".
            # IDs not in _class_id_to_name are genuinely unexpected and produce an empty
            # string with a one-time warning.
            class_ids = detections.class_id if detections.class_id is not None else np.array([], dtype=int)
            truly_oob = [cid for cid in class_ids if cid not in _class_id_to_name and cid != num_logit_slots]
            if truly_oob:
                logger.warning_once(
                    "predict() encountered unmapped class_id(s): %s — mapping to empty string",
                    truly_oob[:5],
                )
            if _is_coco_pretrained:
                class_names = [_class_id_to_name.get(cid, "") for cid in class_ids]
            else:
                class_names = [
                    "__background__" if cid == num_logit_slots else _class_id_to_name.get(cid, "") for cid in class_ids
                ]
            detections.data["class_name"] = np.array(class_names, dtype=object)

            if has_keypoints and keypoints_array is not None:
                keypoint_data = dict(detections.data)
                keypoint_data["xyxy"] = detections.xyxy.astype(np.float32)
                if include_source_image:
                    keypoint_data["source_image"] = [source_images[i] for _ in range(len(detections))]
                raw_precision = keypoint_data.get("keypoint_precision_cholesky")
                raw_source_shape = keypoint_data.get("source_shape")
                if raw_precision is not None and raw_source_shape is not None and len(detections) > 0:
                    precision = np.asarray(raw_precision, dtype=np.float32)
                    source_shape = np.asarray(raw_source_shape, dtype=np.float32)
                    if precision.shape[:2] == keypoints_array.shape[:2] and source_shape.shape == (len(detections), 2):
                        keypoint_data["covariance"] = precision_cholesky_to_pixel_covariance(
                            precision_cholesky=precision, source_shape=source_shape
                        )
                keypoints_array = keypoints_array.astype(np.float32, copy=False)
                keypoint_confidence = keypoints_array[:, :, 2]
                key_points = KeyPoints(
                    xy=keypoints_array[:, :, :2],
                    keypoint_confidence=keypoint_confidence,
                    detection_confidence=detections.confidence.astype(np.float32)
                    if detections.confidence is not None
                    else None,
                    class_id=detections.class_id.astype(int) if detections.class_id is not None else None,
                    visible=keypoint_confidence > 0,
                    data=keypoint_data,
                )
                predictions_list.append(key_points)
            else:
                predictions_list.append(detections)

        return predictions_list if len(predictions_list) > 1 else predictions_list[0]

    def deploy_to_roboflow(
        self,
        workspace: str,
        project_id: str,
        version: int | str,
        api_key: str | None = None,
        size: str | None = None,
    ) -> None:
        """Deploy the trained RF-DETR model to Roboflow.

        Deploying with Roboflow will create a Serverless API to which you can make requests.

        You can also download weights into a Roboflow Inference deployment for use in Roboflow Workflows and on-device
        deployment.

        Args:
            workspace: The name of the Roboflow workspace to deploy to.
            project_id: The project ID to which the model will be deployed.
            version: The project version to which the model will be deployed.
            api_key: Your Roboflow API key. If not provided,
                it will be read from the environment variable `ROBOFLOW_API_KEY`.
            size: The size of the model to deploy. If not provided,
                it will default to the size of the model being trained (e.g., "rfdetr-base", "rfdetr-large", etc.).

        Raises:
            ValueError: If the `api_key` is not provided and not found in the
                environment variable `ROBOFLOW_API_KEY`, or if the `size` is not set for custom architectures.

        Note:
            Bundle creation is delegated to :meth:`export_for_roboflow`, which can be called independently
            to write ``weights.pt`` and ``class_names.txt`` without a network round-trip.
        """
        from roboflow import Roboflow

        if api_key is None:
            api_key = os.getenv("ROBOFLOW_API_KEY")
            if api_key is None:
                raise ValueError("Set api_key=<KEY> in deploy_to_roboflow or export ROBOFLOW_API_KEY=<KEY>")

        rf = Roboflow(api_key=api_key)
        workspace = rf.workspace(workspace)

        if self.size is None and size is None:
            raise ValueError("Must set size for custom architectures")

        size = self.size or size
        with tempfile.TemporaryDirectory(prefix="roboflow_upload_") as tmp_out_dir:
            self.export_for_roboflow(tmp_out_dir)
            project = workspace.project(project_id)
            project_version = project.version(version)
            project_version.deploy(model_type=size, model_path=tmp_out_dir, filename="weights.pt")

    def export_for_roboflow(self, output_dir: str | os.PathLike[str]) -> None:
        """Write a Roboflow upload bundle (``weights.pt`` + ``class_names.txt``) into *output_dir*.

        This is the network-free core of :meth:`deploy_to_roboflow`: it serialises the model state and
        training args into ``weights.pt``, always embedding ``class_names`` into a copy of the args so
        the bundle is self-contained, and writes the class labels to ``class_names.txt``.  The Roboflow
        SDK uses this format to adapt raw PyTorch-Lightning checkpoints into a deploy-ready bundle.

        Args:
            output_dir: Directory into which ``weights.pt`` and ``class_names.txt`` are written.  Created
                if it does not exist.  Existing files are silently overwritten.

        Raises:
            PermissionError: If the process lacks write access to *output_dir* or its parent directory.
            OSError: On disk-full, invalid path, or other filesystem failure during directory creation,
                file write, or ``torch.save``.
        """
        os.makedirs(output_dir, exist_ok=True)
        # Write class_names.txt so the Roboflow upload pipeline can discover
        # the class labels without relying on args.class_names in the checkpoint.
        class_names_path = os.path.join(output_dir, "class_names.txt")
        with open(class_names_path, "w", encoding="utf-8", newline="\n") as f:
            f.write("\n".join(self.class_names))

        # Embed class_names in a shallow copy of args so the saved bundle is
        # self-contained (roboflow-python's second fallback reads args.class_names
        # directly from the checkpoint).  Using a copy leaves self.model.args
        # unmodified — each export call is independent regardless of call order.
        args = copy(self.model.args)
        if not hasattr(args, "class_names") or args.class_names is None:
            args.class_names = self.class_names

        outpath = os.path.join(output_dir, "weights.pt")
        torch.save({"model": self.model.model.state_dict(), "args": args}, outpath)

Attributes

class_names property

Retrieve the class names supported by the loaded model.

Returns:

Type	Description
list[str]	A list of class name strings, 0-indexed. When no custom class names are embedded in the checkpoint, returns
list[str]	the standard 80 COCO class names.

Functions

deploy_to_roboflow(workspace, project_id, version, api_key=None, size=None)

Deploy the trained RF-DETR model to Roboflow.

Deploying with Roboflow will create a Serverless API to which you can make requests.

You can also download weights into a Roboflow Inference deployment for use in Roboflow Workflows and on-device deployment.

Parameters:

Name	Type	Description	Default
workspace	str	The name of the Roboflow workspace to deploy to.	required
project_id	str	The project ID to which the model will be deployed.	required
version	int | str	The project version to which the model will be deployed.	required
api_key	str | None	Your Roboflow API key. If not provided, it will be read from the environment variable ROBOFLOW_API_KEY.	None
size	str | None	The size of the model to deploy. If not provided, it will default to the size of the model being trained (e.g., "rfdetr-base", "rfdetr-large", etc.).	None

Raises:

Type	Description
ValueError	If the api_key is not provided and not found in the environment variable ROBOFLOW_API_KEY, or if the size is not set for custom architectures.

Note

Bundle creation is delegated to :meth:export_for_roboflow, which can be called independently to write weights.pt and class_names.txt without a network round-trip.

Source code in src/rfdetr/detr.py

def deploy_to_roboflow(
    self,
    workspace: str,
    project_id: str,
    version: int | str,
    api_key: str | None = None,
    size: str | None = None,
) -> None:
    """Deploy the trained RF-DETR model to Roboflow.

    Deploying with Roboflow will create a Serverless API to which you can make requests.

    You can also download weights into a Roboflow Inference deployment for use in Roboflow Workflows and on-device
    deployment.

    Args:
        workspace: The name of the Roboflow workspace to deploy to.
        project_id: The project ID to which the model will be deployed.
        version: The project version to which the model will be deployed.
        api_key: Your Roboflow API key. If not provided,
            it will be read from the environment variable `ROBOFLOW_API_KEY`.
        size: The size of the model to deploy. If not provided,
            it will default to the size of the model being trained (e.g., "rfdetr-base", "rfdetr-large", etc.).

    Raises:
        ValueError: If the `api_key` is not provided and not found in the
            environment variable `ROBOFLOW_API_KEY`, or if the `size` is not set for custom architectures.

    Note:
        Bundle creation is delegated to :meth:`export_for_roboflow`, which can be called independently
        to write ``weights.pt`` and ``class_names.txt`` without a network round-trip.
    """
    from roboflow import Roboflow

    if api_key is None:
        api_key = os.getenv("ROBOFLOW_API_KEY")
        if api_key is None:
            raise ValueError("Set api_key=<KEY> in deploy_to_roboflow or export ROBOFLOW_API_KEY=<KEY>")

    rf = Roboflow(api_key=api_key)
    workspace = rf.workspace(workspace)

    if self.size is None and size is None:
        raise ValueError("Must set size for custom architectures")

    size = self.size or size
    with tempfile.TemporaryDirectory(prefix="roboflow_upload_") as tmp_out_dir:
        self.export_for_roboflow(tmp_out_dir)
        project = workspace.project(project_id)
        project_version = project.version(version)
        project_version.deploy(model_type=size, model_path=tmp_out_dir, filename="weights.pt")

export(output_dir='output', infer_dir=None, backbone_only=False, opset_version=17, verbose=True, shape=None, batch_size=1, dynamic_batch=False, patch_size=None, format='onnx', quantization=None, calibration_data=None, max_images=100, *, notes=None)

Export the trained model to ONNX or TFLite format.

See the export documentation <https://rfdetr.roboflow.com/learn/export/>_ for more information.

Parameters:

Name	Type	Description	Default
output_dir	str	Directory to write the exported model to.	'output'
infer_dir	str | None	Optional directory of sample images for dynamic-axes inference.	None
backbone_only	bool	Export only the backbone (feature extractor).	False
opset_version	int	ONNX opset version to target.	17
verbose	bool	Print export progress information.	True
shape	tuple[int, int] | None	(height, width) tuple; defaults to square at model resolution. Both dimensions must be divisible by patch_size * num_windows.	None
batch_size	int	Static batch size to bake into the ONNX graph.	1
dynamic_batch	bool	If True, export with a dynamic batch dimension so the ONNX model accepts variable batch sizes at runtime.	False
patch_size	int | None	Backbone patch size. Defaults to the value stored in model_config.patch_size (typically 14 or 16). When provided explicitly it must match the instantiated model's patch size. Shape divisibility is validated against patch_size * num_windows.	None
format	str	Export format — "onnx" (default) or "tflite". When "tflite" is selected the model is first exported to ONNX then converted to TFLite via onnx2tf. Requires pip install rfdetr[onnx,tflite]. .. warning:: TFLite export is experimental and subject to change; upstream dependency instabilities (onnx2tf, ai_edge_litert) may affect results.	'onnx'
quantization	str | None	TFLite quantization mode (ignored when format="onnx"). One of None, "fp32", "fp16", "int8". None / "fp32" / "fp16" produce FP32 + FP16 .tflite files; "int8" additionally produces an INT8-quantized model.	None
calibration_data	str | ndarray | None	Representative images for INT8 calibration and onnx2tf output validation. Accepts: None — auto-generate random data (sufficient for fp32/fp16; warns for int8). A directory path (str) containing JPEG/PNG images — the converter automatically loads, resizes, and prepares them. This is the simplest approach. A path (str) to a .npy file of shape (N, H, W, 3), dtype float32, values in [0, 1]. A :class:numpy.ndarray with the same format. For INT8 quantization, provide 20–100 representative images from your training/validation set for best accuracy.	None
max_images	int	Maximum number of images to load from a calibration directory. Defaults to 100. Only used when calibration_data is a directory path.	100
notes	object	Optional user-defined metadata (string, dict, list, or any JSON-serialisable value) to embed in the exported ONNX model under the "rfdetr_notes" metadata property. When None no metadata entry is written. String values are stored verbatim; all other types are JSON-encoded so consumers must call json.loads() to recover a dict or list. The same value can be passed to :meth:train so the checkpoint and the ONNX file share the same provenance information.	None

Returns:

Type	Description
Path	Path to the exported model file (.onnx or .tflite).

Source code in src/rfdetr/detr.py

def export(
    self,
    output_dir: str = "output",
    infer_dir: str | None = None,
    backbone_only: bool = False,
    opset_version: int = 17,
    verbose: bool = True,
    shape: tuple[int, int] | None = None,
    batch_size: int = 1,
    dynamic_batch: bool = False,
    patch_size: int | None = None,
    format: str = "onnx",
    quantization: str | None = None,
    calibration_data: str | np.ndarray | None = None,
    max_images: int = 100,
    *,
    notes: object = None,
) -> Path:
    """Export the trained model to ONNX or TFLite format.

    See the `export documentation <https://rfdetr.roboflow.com/learn/export/>`_ for more information.

    Args:
        output_dir: Directory to write the exported model to.
        infer_dir: Optional directory of sample images for dynamic-axes inference.
        backbone_only: Export only the backbone (feature extractor).
        opset_version: ONNX opset version to target.
        verbose: Print export progress information.
        shape: ``(height, width)`` tuple; defaults to square at model resolution.
            Both dimensions must be divisible by ``patch_size * num_windows``.
        batch_size: Static batch size to bake into the ONNX graph.
        dynamic_batch: If True, export with a dynamic batch dimension
            so the ONNX model accepts variable batch sizes at runtime.
        patch_size: Backbone patch size. Defaults to the value stored in
            ``model_config.patch_size`` (typically 14 or 16). When provided explicitly it must match the
            instantiated model's patch size. Shape divisibility is validated against ``patch_size * num_windows``.
        format: Export format — ``"onnx"`` (default) or ``"tflite"``.
            When ``"tflite"`` is selected the model is first exported to ONNX then converted to TFLite via
            ``onnx2tf``.  Requires ``pip install rfdetr[onnx,tflite]``.

            .. warning::
                TFLite export is experimental and subject to change; upstream dependency instabilities (``onnx2tf``,
                ``ai_edge_litert``) may affect results.
        quantization: TFLite quantization mode (ignored when
            ``format="onnx"``).  One of ``None``, ``"fp32"``, ``"fp16"``, ``"int8"``.  ``None`` / ``"fp32"`` /
            ``"fp16"`` produce FP32 + FP16 ``.tflite`` files; ``"int8"`` additionally produces an INT8-quantized
            model.
        calibration_data: Representative images for INT8 calibration and ``onnx2tf`` output validation.  Accepts:

            * ``None`` — auto-generate random data (sufficient for fp32/fp16; warns for int8).
            * A **directory path** (``str``) containing JPEG/PNG
              images — the converter automatically loads, resizes, and prepares them.  This is the simplest
              approach.
            * A path (``str``) to a ``.npy`` file of shape ``(N, H, W, 3)``, dtype float32, values in ``[0, 1]``.
            * A :class:`numpy.ndarray` with the same format.

            For INT8 quantization, provide 20–100 representative images from your training/validation set for best
            accuracy.
        max_images: Maximum number of images to load from a calibration directory.  Defaults to ``100``.  Only used
            when *calibration_data* is a directory path.
        notes: Optional user-defined metadata (string, dict, list, or
            any JSON-serialisable value) to embed in the exported ONNX model under the ``"rfdetr_notes"`` metadata
            property.  When ``None`` no metadata entry is written.  String values are stored verbatim; all other
            types are JSON-encoded so consumers must call ``json.loads()`` to recover a dict or list.  The same
            value can be passed to :meth:`train` so the checkpoint and the ONNX file share the same provenance
            information.

    Returns:
        Path to the exported model file (``.onnx`` or ``.tflite``).
    """
    logger.info("Exporting model to ONNX format")
    _valid_formats = ("onnx", "tflite")
    if format not in _valid_formats:
        raise ValueError(f"Unsupported export format {format!r}. Choose from: {_valid_formats}")
    try:
        from rfdetr.export.main import export_onnx, make_infer_image
    except ImportError:
        logger.error(
            "It seems some dependencies for ONNX export are missing."
            " Please run `pip install rfdetr[onnx]` and try again.",
        )
        raise

    device = self.model.device
    # Move the live model to CPU before deepcopying and keep it there during export. ``nn.Module.to(...)`` mutates
    # in place, so this frees GPU memory for the local export copy, ONNX tracing, TFLite conversion, and any
    # calibration tensors. The ``finally`` block restores the live model even if export or conversion raises.
    self.model.model = self.model.model.to("cpu")
    model = deepcopy(self.model.model)
    model.to(device)
    try:
        os.makedirs(output_dir, exist_ok=True)
        output_dir_path = Path(output_dir)
        patch_size = _resolve_patch_size(patch_size, self.model_config, "export")
        num_windows = getattr(self.model_config, "num_windows", 1)
        if isinstance(num_windows, bool) or not isinstance(num_windows, int) or num_windows <= 0:
            raise ValueError(f"num_windows must be a positive integer, got {num_windows!r}")
        block_size = patch_size * num_windows
        if shape is None:
            shape = (self.model.resolution, self.model.resolution)
            if shape[0] % block_size != 0:
                raise ValueError(
                    f"Model's default resolution ({self.model.resolution}) is not divisible by "
                    f"block_size={block_size} (patch_size={patch_size} * num_windows={num_windows}). "
                    f"Provide an explicit shape divisible by {block_size}.",
                )
        else:
            shape = _validate_shape_dims(shape, block_size, patch_size, num_windows)

        input_tensors = make_infer_image(
            infer_dir, shape, batch_size, device, num_channels=self.model_config.num_channels
        ).to(device)
        input_names = ["input"]
        if backbone_only:
            output_names = ["features"]
        elif self.model_config.segmentation_head:
            output_names = ["dets", "labels", "masks"]
        elif self.model_config.use_grouppose_keypoints:
            output_names = ["dets", "labels", "keypoints"]
        else:
            output_names = ["dets", "labels"]

        if dynamic_batch:
            dynamic_axes = {name: {0: "batch"} for name in input_names + output_names}
        else:
            dynamic_axes = None
        model.eval()
        with torch.no_grad():
            if backbone_only:
                features = model(input_tensors)
                logger.debug(f"PyTorch inference output shape: {features.shape}")
            elif self.model_config.segmentation_head:
                outputs = model(input_tensors)
                dets = outputs["pred_boxes"]
                labels = outputs["pred_logits"]
                masks = outputs["pred_masks"]
                if isinstance(masks, torch.Tensor):
                    logger.debug(
                        f"PyTorch inference output shapes - Boxes: {dets.shape}, Labels: {labels.shape}, "
                        f"Masks: {masks.shape}",
                    )
                else:
                    logger.debug(f"PyTorch inference output shapes - Boxes: {dets.shape}, Labels: {labels.shape}")
            elif self.model_config.use_grouppose_keypoints:
                outputs = model(input_tensors)
                dets = outputs["pred_boxes"]
                labels = outputs["pred_logits"]
                keypoints = outputs["pred_keypoints"]
                logger.debug(
                    f"PyTorch inference output shapes - Boxes: {dets.shape}, Labels: {labels.shape}, "
                    f"Keypoints: {keypoints.shape}",
                )
            else:
                outputs = model(input_tensors)
                dets = outputs["pred_boxes"]
                labels = outputs["pred_logits"]
                logger.debug(f"PyTorch inference output shapes - Boxes: {dets.shape}, Labels: {labels.shape}")

        model.cpu()
        input_tensors = input_tensors.cpu()

        output_file = export_onnx(
            output_dir=str(output_dir_path),
            model=model,
            input_names=input_names,
            input_tensors=input_tensors,
            output_names=output_names,
            dynamic_axes=dynamic_axes,
            backbone_only=backbone_only,
            verbose=verbose,
            opset_version=opset_version,
            variant_name=getattr(self, "size", None),
            notes=notes,
        )

        logger.info(f"Successfully exported ONNX model to: {output_file}")

        if format == "tflite":
            warnings.warn(
                "TFLite export is experimental and work-in-progress. "
                "Upstream dependency instabilities (onnx2tf, ai_edge_litert) may affect results.",
                UserWarning,
                stacklevel=2,
            )
            try:
                from rfdetr.export._tflite.converter import export_tflite
            except ImportError:
                logger.error(
                    "It seems some dependencies for TFLite export are missing."
                    " Please run `pip install rfdetr[onnx,tflite]` and try again.",
                )
                raise

            tflite_path = export_tflite(
                onnx_path=output_file,
                output_dir=str(output_dir_path),
                quantization=quantization,
                calibration_data=calibration_data,
                verbosity="info" if verbose else "error",
                max_images=max_images,
                verbose=verbose,
            )
            logger.info(f"Successfully exported TFLite model to: {tflite_path}")
            return tflite_path

        logger.info("Export completed successfully")
        return Path(output_file)
    finally:
        self.model.model = self.model.model.to(device)

export_for_roboflow(output_dir)

Write a Roboflow upload bundle (weights.pt + class_names.txt) into output_dir.

This is the network-free core of :meth:deploy_to_roboflow: it serialises the model state and training args into weights.pt, always embedding class_names into a copy of the args so the bundle is self-contained, and writes the class labels to class_names.txt. The Roboflow SDK uses this format to adapt raw PyTorch-Lightning checkpoints into a deploy-ready bundle.

Parameters:

Name	Type	Description	Default
output_dir	str | PathLike[str]	Directory into which weights.pt and class_names.txt are written. Created if it does not exist. Existing files are silently overwritten.	required

Raises:

Type	Description
PermissionError	If the process lacks write access to output_dir or its parent directory.
OSError	On disk-full, invalid path, or other filesystem failure during directory creation, file write, or torch.save.

Source code in src/rfdetr/detr.py

def export_for_roboflow(self, output_dir: str | os.PathLike[str]) -> None:
    """Write a Roboflow upload bundle (``weights.pt`` + ``class_names.txt``) into *output_dir*.

    This is the network-free core of :meth:`deploy_to_roboflow`: it serialises the model state and
    training args into ``weights.pt``, always embedding ``class_names`` into a copy of the args so
    the bundle is self-contained, and writes the class labels to ``class_names.txt``.  The Roboflow
    SDK uses this format to adapt raw PyTorch-Lightning checkpoints into a deploy-ready bundle.

    Args:
        output_dir: Directory into which ``weights.pt`` and ``class_names.txt`` are written.  Created
            if it does not exist.  Existing files are silently overwritten.

    Raises:
        PermissionError: If the process lacks write access to *output_dir* or its parent directory.
        OSError: On disk-full, invalid path, or other filesystem failure during directory creation,
            file write, or ``torch.save``.
    """
    os.makedirs(output_dir, exist_ok=True)
    # Write class_names.txt so the Roboflow upload pipeline can discover
    # the class labels without relying on args.class_names in the checkpoint.
    class_names_path = os.path.join(output_dir, "class_names.txt")
    with open(class_names_path, "w", encoding="utf-8", newline="\n") as f:
        f.write("\n".join(self.class_names))

    # Embed class_names in a shallow copy of args so the saved bundle is
    # self-contained (roboflow-python's second fallback reads args.class_names
    # directly from the checkpoint).  Using a copy leaves self.model.args
    # unmodified — each export call is independent regardless of call order.
    args = copy(self.model.args)
    if not hasattr(args, "class_names") or args.class_names is None:
        args.class_names = self.class_names

    outpath = os.path.join(output_dir, "weights.pt")
    torch.save({"model": self.model.model.state_dict(), "args": args}, outpath)

from_checkpoint(path, **kwargs) classmethod

Load an RF-DETR model from a training checkpoint, automatically inferring the model class.

The correct subclass is resolved in order of preference:

1. model_name key in the checkpoint (written by the PTL training stack since v1.7.0).
2. pretrain_weights field in the checkpoint's args entry (legacy fallback for older checkpoints).
3. The filename of path itself, used as a last resort when pretrain_weights is absent or an unset-like sentinel value (empty string, "none", or "null"). Starter weights published by Roboflow store pretrain_weights="none" in their args; passing the canonical filename (e.g. rf-detr-small.pth) lets from_checkpoint infer the class automatically.

Both legacy argparse.Namespace checkpoints (produced by engine.py) and dict-style checkpoints (produced by the PTL training stack) are supported.

Parameters:

Name	Type	Description	Default
path	str | PathLike[str]	Path to a checkpoint file (e.g. checkpoint_best_total.pth).	required
**kwargs	Any	Additional keyword arguments forwarded to the model constructor (e.g. accept_platform_model_license=True for XLarge / 2XLarge models). If num_classes is not supplied here, the value stored in the checkpoint is used when present; otherwise the constructor default applies. In either case the field is not recorded as a user-set override, so :meth:train can still adapt the detection head to the training dataset's class count. Pass an explicit num_classes=N to pin the head and prevent adaptation, even when N equals the class default (e.g. num_classes=90 on a COCO-pretrained checkpoint).	{}

Returns:

Type	Description
RFDETR	An instance of the appropriate :class:RFDETR subclass loaded from the checkpoint.

Warning

This method calls torch.load with weights_only=False, which unpickles arbitrary Python objects. Only load checkpoints from trusted sources.

Raises:

Type	Description
FileNotFoundError	If path does not exist.
OSError	If path exists but cannot be read.
KeyError	If the checkpoint does not contain an "args" key.
ValueError	If the model class cannot be inferred from model_name, pretrain_weights, or the checkpoint filename.

Examples:

>>> model = RFDETR.from_checkpoint("checkpoint_best_total.pth")
>>> model = RFDETRSmall.from_checkpoint("checkpoint_best_total.pth")

Source code in src/rfdetr/detr.py

@classmethod
def from_checkpoint(cls, path: str | os.PathLike[str], **kwargs: Any) -> RFDETR:
    """Load an RF-DETR model from a training checkpoint, automatically inferring the model class.

    The correct subclass is resolved in order of preference:

    1. ``model_name`` key in the checkpoint (written by the PTL training
       stack since v1.7.0).
    2. ``pretrain_weights`` field in the checkpoint's ``args`` entry
       (legacy fallback for older checkpoints).
    3. The **filename** of *path* itself, used as a last resort when
       ``pretrain_weights`` is absent or an unset-like sentinel value
       (empty string, ``"none"``, or ``"null"``).  Starter weights
       published by Roboflow store ``pretrain_weights="none"`` in their
       ``args``; passing the canonical filename (e.g.
       ``rf-detr-small.pth``) lets ``from_checkpoint`` infer the class
       automatically.

    Both legacy ``argparse.Namespace`` checkpoints (produced by ``engine.py``) and dict-style checkpoints (produced
    by the PTL training stack) are supported.

    Args:
        path: Path to a checkpoint file (e.g. ``checkpoint_best_total.pth``).
        **kwargs: Additional keyword arguments forwarded to the model
            constructor (e.g. ``accept_platform_model_license=True`` for XLarge / 2XLarge models).
            If ``num_classes`` is not supplied here, the value stored in the checkpoint is
            used when present; otherwise the constructor default applies.  In either case the
            field is not recorded as a user-set override, so :meth:`train` can still adapt the
            detection head to the training dataset's class count.  Pass an explicit
            ``num_classes=N`` to pin the head and prevent adaptation, even when ``N`` equals
            the class default (e.g. ``num_classes=90`` on a COCO-pretrained checkpoint).

    Returns:
        An instance of the appropriate :class:`RFDETR` subclass loaded from the checkpoint.

    Warning:
        This method calls ``torch.load`` with ``weights_only=False``, which
        unpickles arbitrary Python objects. Only load checkpoints from trusted sources.

    Raises:
        FileNotFoundError: If *path* does not exist.
        OSError: If *path* exists but cannot be read.
        KeyError: If the checkpoint does not contain an ``"args"`` key.
        ValueError: If the model class cannot be inferred from ``model_name``,
            ``pretrain_weights``, or the checkpoint filename.

    Examples:
        >>> model = RFDETR.from_checkpoint("checkpoint_best_total.pth")  # doctest: +SKIP
        >>> model = RFDETRSmall.from_checkpoint("checkpoint_best_total.pth")  # doctest: +SKIP
    """
    # Local import breaks the variants → detr import cycle.
    import rfdetr.variants as rfdetr_variants

    _plus_available = False
    _plus_symbols: dict[str, type[RFDETR]] = {}
    _plus_entries: list[tuple[str, type[RFDETR]]] = []
    from rfdetr.platform import _IS_RFDETR_PLUS_AVAILABLE

    if _IS_RFDETR_PLUS_AVAILABLE:
        try:
            import rfdetr.platform.models as platform_models

            for class_symbol in _CHECKPOINT_PLUS_MODEL_NAME_CLASS_SYMBOLS:
                plus_obj = getattr(platform_models, class_symbol)
                _plus_symbols[class_symbol] = plus_obj
            _plus_entries = [
                (name, _plus_symbols[class_symbol]) for name, class_symbol in _CHECKPOINT_PLUS_MODEL_MAP_ENTRIES
            ]
            _plus_available = True
        except ModuleNotFoundError as ex:
            if ex.name not in {"rfdetr_plus", "rfdetr_plus.models"}:
                raise

    # weights_only=False is required because legacy checkpoints embed
    # argparse.Namespace objects that cannot be deserialised with
    # weights_only=True.
    ckpt: dict[str, Any] = torch.load(path, map_location="cpu", weights_only=False)
    args = ckpt["args"]

    _variant_name_to_class: dict[str, type[RFDETR]] = {
        getattr(variant_obj, "__name__", symbol): variant_obj
        for symbol in dir(rfdetr_variants)
        if symbol.startswith("RFDETR")
        for variant_obj in [getattr(rfdetr_variants, symbol)]
    }
    _variant_symbols: dict[str, type[RFDETR]] = {
        class_symbol: _variant_name_to_class[class_symbol] for class_symbol in _CHECKPOINT_MODEL_NAME_CLASS_SYMBOLS
    }
    # Build in three explicit segments: seg-* entries, then plus-model entries
    # (xlarge/2xlarge), then base entries — order determines lookup priority.
    _seg_map: list[tuple[str, type[RFDETR]]] = [
        (name, _variant_symbols[class_symbol])
        for name, class_symbol in _CHECKPOINT_MODEL_MAP_ENTRIES
        if name.startswith("seg-")
    ]
    _keypoint_map: list[tuple[str, type[RFDETR]]] = [
        (name, _variant_symbols[class_symbol])
        for name, class_symbol in _CHECKPOINT_MODEL_MAP_ENTRIES
        if "keypoint" in name
    ]
    _base_map: list[tuple[str, type[RFDETR]]] = [
        (name, _variant_symbols[class_symbol])
        for name, class_symbol in _CHECKPOINT_MODEL_MAP_ENTRIES
        if not name.startswith("seg-") and "keypoint" not in name
    ]
    _model_map: list[tuple[str, type[RFDETR]]] = _seg_map + _keypoint_map + _plus_entries + _base_map

    # New checkpoints store model_name directly — use it when available.
    _name_map: dict[str, type[RFDETR]] = dict(_variant_symbols)
    # Plus-model classes are resolved only when rfdetr_plus is installed.
    if _plus_available:
        _name_map.update(_plus_symbols)
    saved_model_name = ckpt.get("model_name")
    model_cls: type[RFDETR] | None = None
    if isinstance(saved_model_name, str):
        normalized_name = saved_model_name.strip()
        if normalized_name:
            model_cls = _name_map.get(normalized_name)
    else:
        normalized_name = ""

    # Fall back to pretrain_weights (legacy) or, when unset-like, the checkpoint filename.
    if isinstance(args, dict):
        weights_name = str(args.get("pretrain_weights", "")).strip().lower()
    else:
        weights_name = str(getattr(args, "pretrain_weights", "")).strip().lower()
    # The sentinel set {"", "none", "null"} covers unset-like checkpoint values:
    #   ""     — pretrain_weights key absent entirely
    #   "none" — checkpoint value was None or the literal string "none";
    #            after str(...).strip().lower() both normalize to the same sentinel.
    #            This is NOT an intentional "no pretraining" flag (see
    #            test_pretrain_weights_none_warns, which operates at the config
    #            level, not the checkpoint level)
    #   "null" — checkpoint stored the literal string "null" (for example from a
    #            YAML-originated value), which is also treated as unset-like here
    _filename_fallback = False
    if weights_name in {"", "none", "null"}:
        weights_name = os.path.basename(os.fspath(path)).lower()
        _filename_fallback = True

    if model_cls is None:
        # Guard: plus-only checkpoints should raise an actionable install error
        # when rfdetr_plus is missing, regardless of whether class inference
        # relies on model_name (new format) or pretrain_weights (legacy format).
        plus_by_model_name = normalized_name in _CHECKPOINT_PLUS_MODEL_NAME_CLASS_SYMBOLS
        plus_by_weights_name = (
            "xlarge" in weights_name and "seg-" not in weights_name and "keypoint-preview" not in weights_name
        )
        if not _plus_available and (plus_by_model_name or plus_by_weights_name):
            from rfdetr.platform import _INSTALL_MSG

            raise ImportError(
                f"Checkpoint model_name={saved_model_name!r}, pretrain_weights={weights_name!r} requires the "
                f"rfdetr_plus package. " + _INSTALL_MSG.format(name="platform model downloads")
            )

        for name, klass in _model_map:
            if name in weights_name:
                model_cls = klass
                break

        if _filename_fallback and model_cls is not None:
            logger.info(
                "pretrain_weights unset in checkpoint %r; inferred model class %s from filename %r",
                path,
                getattr(model_cls, "__name__", repr(model_cls)),
                weights_name,
            )

    if model_cls is None:
        raise ValueError(
            f"Could not infer model class from checkpoint at {path!r} "
            f"(model_name={saved_model_name!r}, pretrain_weights={weights_name!r}). "
            f"Please instantiate the model class directly."
        )

    if isinstance(args, dict):
        num_classes: int | None = args.get("num_classes")
    else:
        num_classes = getattr(args, "num_classes", None)

    constructor_kwargs: dict[str, Any] = {}
    checkpoint_config_keys: set[str] = set()  # keys injected from checkpoint, not from caller
    saved_model_config = ckpt.get("model_config")
    if isinstance(saved_model_config, dict):
        model_config_class = getattr(model_cls, "_model_config_class", None)
        model_fields = getattr(model_config_class, "model_fields", None)
        if not isinstance(model_fields, dict):
            model_fields = getattr(model_config_class, "__fields__", None)
        if not isinstance(model_fields, dict):
            model_fields = {}
        for key, value in saved_model_config.items():
            if key == "pretrain_weights":
                continue
            if not model_fields or key in model_fields:
                constructor_kwargs[key] = value
                checkpoint_config_keys.add(key)

    if num_classes is not None and "num_classes" not in kwargs:
        constructor_kwargs["num_classes"] = num_classes
        checkpoint_config_keys.add("num_classes")
    constructor_kwargs.update(kwargs)
    # pretrain_weights is placed after **kwargs so it always wins even if
    # a caller accidentally passes pretrain_weights inside kwargs.
    constructor_kwargs["pretrain_weights"] = str(path)

    # Fields injected from the checkpoint but not supplied by the caller must not be
    # treated as explicit user overrides in Pydantic's model_fields_set.  Downstream
    # alignment guards (e.g. _align_num_classes_from_dataset,
    # _align_keypoint_schema_from_dataset, load_pretrain_weights) all read
    # model_fields_set to decide whether to adapt model internals to the training
    # dataset — leaving checkpoint-derived fields marked as user-set breaks them.
    checkpoint_derived_keys = checkpoint_config_keys - set(kwargs)

    model = model_cls(**constructor_kwargs)

    if checkpoint_derived_keys:
        loaded_config = getattr(model, "model_config", None)
        # model_fields_set is the public API and returns the live backing set
        # in Pydantic v2; fall back to the private attribute only if that changes.
        fields_set = getattr(loaded_config, "model_fields_set", None)
        if fields_set is None:
            fields_set = getattr(loaded_config, "__pydantic_fields_set__", None)
        if fields_set is not None:
            fields_set.difference_update(checkpoint_derived_keys)
        # Verify num_classes specifically — if Pydantic ever returns a snapshot instead
        # of the live backing set, this assertion will catch the silent regression before
        # it causes a training-time head-adaptation failure.
        if "num_classes" in checkpoint_derived_keys:
            assert "num_classes" not in getattr(loaded_config, "model_fields_set", set()), (
                "num_classes still in model_fields_set after checkpoint load; "
                "Pydantic may return a snapshot rather than the live backing set — "
                "switch to model_construct(_fields_set=...) for Pydantic v3 compatibility."
            )

    return model

get_model(config)

Retrieve a model context from the provided architecture configuration.

Parameters:

Name	Type	Description	Default
config	ModelConfig	Architecture configuration.	required

Returns:

Type	Description
ModelContext	ModelContext with model, postprocess, device, resolution, args, and class_names attributes.

Source code in src/rfdetr/detr.py

def get_model(self, config: ModelConfig) -> ModelContext:
    """Retrieve a model context from the provided architecture configuration.

    Args:
        config: Architecture configuration.

    Returns:
        ModelContext with model, postprocess, device, resolution, args, and class_names attributes.
    """
    return _build_model_context(config)

get_model_config(**kwargs)

Retrieve the configuration parameters used by the model.

Source code in src/rfdetr/detr.py

def get_model_config(self, **kwargs) -> ModelConfig:
    """Retrieve the configuration parameters used by the model."""
    return self._model_config_class(**kwargs)

get_train_config(**kwargs)

Retrieve the configuration parameters that will be used for training.

Source code in src/rfdetr/detr.py

def get_train_config(self, **kwargs) -> TrainConfig:
    """Retrieve the configuration parameters that will be used for training."""
    return self._train_config_class(**kwargs)

maybe_download_pretrain_weights()

Download pre-trained weights if they are not already downloaded.

Bare filenames (no directory component, e.g. rf-detr-base.pth) are resolved to the model cache directory — set the RF_HOME environment variable to override the location (default: ~/.roboflow/models). Resolution happens in ModelConfig.expand_path for explicitly-provided values, and here as a fallback for field defaults (which Pydantic does not validate by default).

Paths that already contain a directory component are used as-is; the parent directory is created if it does not yet exist.

Source code in src/rfdetr/detr.py

def maybe_download_pretrain_weights(self):
    """Download pre-trained weights if they are not already downloaded.

    Bare filenames (no directory component, e.g. ``rf-detr-base.pth``) are resolved to the model cache directory —
    set the ``RF_HOME`` environment variable to override the location (default: ``~/.roboflow/models``). Resolution
    happens in ``ModelConfig.expand_path`` for explicitly-provided values, and here as a fallback for field defaults
    (which Pydantic does not validate by default).

    Paths that already contain a directory component are used as-is; the parent directory is created if it does not
    yet exist.
    """
    pretrain_weights = self.model_config.pretrain_weights
    if pretrain_weights is None:
        return
    if not os.path.dirname(pretrain_weights):
        # Field default was not processed by expand_path — resolve to cache dir.
        cache_dir = get_model_cache_dir()
        os.makedirs(cache_dir, exist_ok=True)
        pretrain_weights = os.path.join(cache_dir, pretrain_weights)
    else:
        os.makedirs(os.path.dirname(pretrain_weights), exist_ok=True)
    self.model_config.pretrain_weights = pretrain_weights
    download_pretrain_weights(self.model_config.pretrain_weights)

optimize_for_inference(compile=True, batch_size=1, dtype=torch.float32)

Optimize the model for inference with optional JIT compilation and dtype casting.

Operations are wrapped in the correct CUDA device context to prevent context leaks on multi-GPU setups. When compile=True the model is traced with torch.jit.trace using a dummy input of batch_size images at the model's current resolution.

Parameters:

Name	Type	Description	Default
compile	bool	If True, trace the model with torch.jit.trace to obtain a JIT-compiled ScriptModule. Set to False for broader compatibility (e.g. models with dynamic control flow).	True
batch_size	int	Number of images the traced model will be optimized for. Ignored when compile=False.	1
dtype	dtype | str	Target floating-point dtype for the inference model. Accepts a torch.dtype directly (e.g. torch.float16) or its string name (e.g. "float16"). Defaults to torch.float32.	float32

Raises:

Type	Description
TypeError	If dtype is not a torch.dtype, or if dtype is a string that does not correspond to a valid torch.dtype attribute.

Examples:

>>> from types import SimpleNamespace
>>> import torch
>>> class _TinyModel(torch.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.linear = torch.nn.Linear(1, 1)
...     def forward(self, x):
...         return {"pred_boxes": self.linear(x[:, :1, :1, :1].squeeze(-1).squeeze(-1))}
...     def export(self):
...         return None
>>> class _TinyContext:
...     def __init__(self):
...         self.device = torch.device("cpu")
...         self.resolution = 28
...         self.model = _TinyModel()
...         self.inference_model = None
>>> model = object.__new__(RFDETR)
>>> model.model_config = SimpleNamespace(num_channels=3)
>>> model.model = _TinyContext()
>>> model._is_optimized_for_inference = False
>>> model._has_warned_about_not_being_optimized_for_inference = False
>>> model._optimized_has_been_compiled = False
>>> model._optimized_batch_size = None
>>> model._optimized_resolution = None
>>> model._optimized_dtype = None
>>> model.optimize_for_inference(compile=False, dtype="float16")
>>> model._is_optimized_for_inference
True
>>> model._optimized_dtype
torch.float16

Source code in src/rfdetr/detr.py

@_ensure_model_on_device
def optimize_for_inference(
    self, compile: bool = True, batch_size: int = 1, dtype: torch.dtype | str = torch.float32
) -> None:
    """Optimize the model for inference with optional JIT compilation and dtype casting.

    Operations are wrapped in the correct CUDA device context to prevent context leaks on multi-GPU setups. When
    ``compile=True`` the model is traced with ``torch.jit.trace`` using a dummy input of ``batch_size`` images at
    the model's current resolution.

    Args:
        compile: If ``True``, trace the model with ``torch.jit.trace`` to obtain
            a JIT-compiled ``ScriptModule``. Set to ``False`` for broader compatibility (e.g. models with dynamic
            control flow).
        batch_size: Number of images the traced model will be optimized for. Ignored when ``compile=False``.
        dtype: Target floating-point dtype for the inference model. Accepts a
            ``torch.dtype`` directly (e.g. ``torch.float16``) or its string name (e.g. ``"float16"``). Defaults to
            ``torch.float32``.

    Raises:
        TypeError: If ``dtype`` is not a ``torch.dtype``, or if ``dtype`` is a
            string that does not correspond to a valid ``torch.dtype`` attribute.

    Examples:
        >>> from types import SimpleNamespace
        >>> import torch
        >>> class _TinyModel(torch.nn.Module):
        ...     def __init__(self):
        ...         super().__init__()
        ...         self.linear = torch.nn.Linear(1, 1)
        ...     def forward(self, x):
        ...         return {"pred_boxes": self.linear(x[:, :1, :1, :1].squeeze(-1).squeeze(-1))}
        ...     def export(self):
        ...         return None
        >>> class _TinyContext:
        ...     def __init__(self):
        ...         self.device = torch.device("cpu")
        ...         self.resolution = 28
        ...         self.model = _TinyModel()
        ...         self.inference_model = None
        >>> model = object.__new__(RFDETR)
        >>> model.model_config = SimpleNamespace(num_channels=3)
        >>> model.model = _TinyContext()
        >>> model._is_optimized_for_inference = False
        >>> model._has_warned_about_not_being_optimized_for_inference = False
        >>> model._optimized_has_been_compiled = False
        >>> model._optimized_batch_size = None
        >>> model._optimized_resolution = None
        >>> model._optimized_dtype = None
        >>> model.optimize_for_inference(compile=False, dtype="float16")
        >>> model._is_optimized_for_inference
        True
        >>> model._optimized_dtype
        torch.float16
    """
    if isinstance(dtype, str):
        try:
            dtype = getattr(torch, dtype)
        except AttributeError:
            raise TypeError(f"dtype must be a torch.dtype or a string name of a dtype, got {dtype!r}") from None
    if not isinstance(dtype, torch.dtype):
        raise TypeError(f"dtype must be a torch.dtype or a string name of a dtype, got {type(dtype)!r}")

    # Clear any previously optimized state before starting a new optimization run.
    self.remove_optimized_model()

    device = self.model.device
    cuda_ctx = torch.cuda.device(device) if device.type == "cuda" else contextlib.nullcontext()

    try:
        with cuda_ctx:
            self.model.inference_model = deepcopy(self.model.model)
            self.model.inference_model.eval()
            self.model.inference_model.export()

            self.model.inference_model = self.model.inference_model.to(dtype=dtype)

            if compile:
                self.model.inference_model = torch.jit.trace(
                    self.model.inference_model,
                    torch.randn(
                        batch_size,
                        self.model_config.num_channels,
                        self.model.resolution,
                        self.model.resolution,
                        device=self.model.device,
                        dtype=dtype,
                    ),
                )
                self._optimized_has_been_compiled = True
                self._optimized_batch_size = batch_size

            # Set success flags only after all operations complete.
            self._optimized_resolution = self.model.resolution
            self._is_optimized_for_inference = True
            self._optimized_dtype = dtype
    except Exception:
        # Ensure the object is left in a consistent, unoptimized state if optimization fails.
        with contextlib.suppress(Exception):
            self.remove_optimized_model()
        raise

predict(images, threshold=0.5, shape=None, patch_size=None, include_source_image=True, **kwargs)

Performs model inference on the input images.

This method accepts a single image or a list of images in various formats (file path, image url, PIL Image, NumPy array, or torch.Tensor). The images should be in RGB channel order. If a torch.Tensor is provided, it must already be normalized to values in the [0, 1] range and have the shape (C, H, W).

Parameters:

Name	Type	Description	Default
images	str | Image | ndarray | Tensor | list[str | ndarray | Image | Tensor]	A single image or a list of images to process. Images can be provided as file paths, PIL Images, NumPy arrays, or torch.Tensors.	required
threshold	float	The minimum confidence score needed to consider a detected bounding box valid.	0.5
shape	tuple[int, int] | None	Optional (height, width) tuple to resize images to before inference. When provided, overrides the model's default inference resolution. The tuple should match the resolution used when exporting the model (typically a square shape). Both dimensions must be positive integers divisible by patch_size * num_windows. Defaults to (model.resolution, model.resolution) when not set.	None
patch_size	int | None	Backbone patch size used for shape divisibility validation. Defaults to model_config.patch_size (typically 14 for large models, 16 for smaller ones). Divisibility is checked against patch_size * num_windows.	None
include_source_image	bool	Whether to attach the original image to the returned prediction. Detection and segmentation outputs use detections.metadata["source_image"]. Keypoint outputs use per-object key_points.data["source_image"] because Supervision KeyPoints currently has no collection-level metadata field. Defaults to True. Set to False to reduce memory use when source images are not needed.	True
**kwargs	Any	Additional keyword arguments.	{}

Returns:

Type	Description
Detections | KeyPoints | list[Detections | KeyPoints]	A single or multiple Supervision prediction objects. Detection and segmentation models return
Detections | KeyPoints | list[Detections | KeyPoints]	class:~supervision.Detections. Keypoint models return :class:~supervision.KeyPoints, with keypoint
Detections | KeyPoints | list[Detections | KeyPoints]	coordinates in xy. Keypoint predictions preserve the detection-level fields produced by RF-DETR:
Detections | KeyPoints | list[Detections | KeyPoints]	key_points.detection_confidence is the per-object score used by threshold. For keypoint models this
Detections | KeyPoints | list[Detections | KeyPoints]	is the postprocessed detection score and, by default, includes keypoint uncertainty fusion controlled by
Detections | KeyPoints | list[Detections | KeyPoints]	model_config.postprocess_trace_alpha. key_points.keypoint_confidence is separate: it is a
Detections | KeyPoints | list[Detections | KeyPoints]	(num_detections, num_keypoints) array of per-keypoint findability scores decoded from the keypoint head,
Detections | KeyPoints | list[Detections | KeyPoints]	not a repeated copy of the detection score. When RF-DETR emits keypoint precision parameters,
Detections | KeyPoints | list[Detections | KeyPoints]	key_points.data["covariance"] stores per-keypoint pixel-space covariance matrices with shape
Detections | KeyPoints | list[Detections | KeyPoints]	(num_detections, num_keypoints, 2, 2). key_points.data["xyxy"] stores the corresponding detection
Detections | KeyPoints | list[Detections | KeyPoints]	boxes as a (num_detections, 4) array in the same row order as key_points.xy because Supervision
Detections | KeyPoints | list[Detections | KeyPoints]	KeyPoints does not have a native bounding-box field. The data dict also contains class_name and
Detections | KeyPoints | list[Detections | KeyPoints]	source_shape as per-object arrays. When include_source_image=True for keypoint models,
Detections | KeyPoints | list[Detections | KeyPoints]	source_image is stored as per-object data until Supervision exposes collection-level metadata for
Detections | KeyPoints | list[Detections | KeyPoints]	KeyPoints.

Note

For Detections outputs, source_image moved from detections.data to detections.metadata. Update detection callers reading detections.data["source_image"] to use detections.metadata["source_image"].

Note

class_name mapping uses one of two modes depending on the checkpoint. For pretrained COCO checkpoints (detected when model.args.num_classes > len(class_names) and class_names matches COCO_CLASS_NAMES), raw COCO category IDs (1–90, sparse) are looked up by category ID rather than by position — so class_id=18 yields "dog", not class_names[18]. For fine-tuned models, class_id is a 0-based index into class_names.

Raises:

Type	Description
ValueError	If shape cannot be unpacked as a two-element sequence, if either dimension does not support the __index__ protocol (e.g. float) or is a bool, if either dimension is zero or negative, if either dimension is not divisible by patch_size * num_windows, or if patch_size is not a positive integer.

Source code in src/rfdetr/detr.py

@torch.no_grad()
@_ensure_model_on_device
def predict(
    self,
    images: str | Image.Image | np.ndarray | torch.Tensor | list[str | np.ndarray | Image.Image | torch.Tensor],
    threshold: float = 0.5,
    shape: tuple[int, int] | None = None,
    patch_size: int | None = None,
    include_source_image: bool = True,
    **kwargs: Any,
) -> Detections | KeyPoints | list[Detections | KeyPoints]:
    """Performs model inference on the input images.

    This method accepts a single image or a list of images in various formats (file path, image url, PIL Image,
    NumPy array, or torch.Tensor). The images should be in RGB channel order. If a torch.Tensor is provided, it must
    already be normalized to values in the [0, 1] range and have the shape (C, H, W).

    Args:
        images:
            A single image or a list of images to process. Images can be provided
            as file paths, PIL Images, NumPy arrays, or torch.Tensors.
        threshold:
            The minimum confidence score needed to consider a detected bounding box valid.
        shape:
            Optional ``(height, width)`` tuple to resize images to before inference. When provided, overrides the
            model's default inference resolution. The tuple should match the resolution used when exporting the
            model (typically a square shape). Both dimensions must be positive integers divisible by ``patch_size *
            num_windows``. Defaults to ``(model.resolution, model.resolution)`` when not set.
        patch_size:
            Backbone patch size used for shape divisibility validation. Defaults to ``model_config.patch_size``
            (typically 14 for large models, 16 for smaller ones). Divisibility is checked against ``patch_size *
            num_windows``.
        include_source_image:
            Whether to attach the original image to the returned prediction. Detection and segmentation outputs use
            ``detections.metadata["source_image"]``. Keypoint outputs use per-object
            ``key_points.data["source_image"]`` because Supervision ``KeyPoints`` currently has no collection-level
            metadata field. Defaults to ``True``. Set to ``False`` to reduce memory use when source images are not
            needed.
        **kwargs:
            Additional keyword arguments.

    Returns:
        A single or multiple Supervision prediction objects. Detection and segmentation models return
        :class:`~supervision.Detections`. Keypoint models return :class:`~supervision.KeyPoints`, with keypoint
        coordinates in ``xy``. Keypoint predictions preserve the detection-level fields produced by RF-DETR:
        ``key_points.detection_confidence`` is the per-object score used by ``threshold``. For keypoint models this
        is the postprocessed detection score and, by default, includes keypoint uncertainty fusion controlled by
        ``model_config.postprocess_trace_alpha``. ``key_points.keypoint_confidence`` is separate: it is a
        ``(num_detections, num_keypoints)`` array of per-keypoint findability scores decoded from the keypoint head,
        not a repeated copy of the detection score. When RF-DETR emits keypoint precision parameters,
        ``key_points.data["covariance"]`` stores per-keypoint pixel-space covariance matrices with shape
        ``(num_detections, num_keypoints, 2, 2)``. ``key_points.data["xyxy"]`` stores the corresponding detection
        boxes as a ``(num_detections, 4)`` array in the same row order as ``key_points.xy`` because Supervision
        ``KeyPoints`` does not have a native bounding-box field. The ``data`` dict also contains ``class_name`` and
        ``source_shape`` as per-object arrays. When ``include_source_image=True`` for keypoint models,
        ``source_image`` is stored as per-object data until Supervision exposes collection-level metadata for
        ``KeyPoints``.

    Note:
        For ``Detections`` outputs, ``source_image`` moved from ``detections.data`` to ``detections.metadata``.
        Update detection callers reading ``detections.data["source_image"]`` to use
        ``detections.metadata["source_image"]``.

    Note:
        ``class_name`` mapping uses one of two modes depending on the checkpoint. For pretrained COCO checkpoints
        (detected when ``model.args.num_classes > len(class_names)`` and ``class_names`` matches
        ``COCO_CLASS_NAMES``), raw COCO category IDs (1–90, sparse) are looked up by category ID rather than by
        position — so ``class_id=18`` yields ``"dog"``, not ``class_names[18]``. For fine-tuned models, ``class_id``
        is a 0-based index into ``class_names``.

    Raises:
        ValueError: If ``shape`` cannot be unpacked as a two-element sequence,
            if either dimension does not support the ``__index__`` protocol (e.g. ``float``) or is a ``bool``, if
            either dimension is zero or negative, if either dimension is not divisible by ``patch_size *
            num_windows``, or if ``patch_size`` is not a positive integer.
    """
    from supervision import Detections, KeyPoints

    patch_size = _resolve_patch_size(patch_size, self.model_config, "predict")
    num_windows = getattr(self.model_config, "num_windows", 1)
    if isinstance(num_windows, bool) or not isinstance(num_windows, int) or num_windows <= 0:
        raise ValueError(f"model_config.num_windows must be a positive integer, got {num_windows!r}")
    block_size = patch_size * num_windows

    if shape is None:
        default_res = self.model.resolution
        if default_res % block_size != 0:
            raise ValueError(
                f"Model's default resolution ({default_res}) is not divisible by "
                f"block_size={block_size} (patch_size={patch_size} * num_windows={num_windows}). "
                f"Provide an explicit shape divisible by {block_size}.",
            )
    else:
        shape = _validate_shape_dims(shape, block_size, patch_size, num_windows)

    if not self._is_optimized_for_inference and not self._has_warned_about_not_being_optimized_for_inference:
        logger.warning(
            "Model is not optimized for inference. Latency may be higher than expected."
            " You can optimize the model for inference by calling model.optimize_for_inference().",
        )
        self._has_warned_about_not_being_optimized_for_inference = True

        self.model.model.eval()

    if not isinstance(images, list):
        images = [images]

    orig_sizes = []
    processed_images = []
    source_images = [] if include_source_image else None

    for img in images:
        if isinstance(img, str):
            if img.startswith("http"):
                img = requests.get(img, stream=True).raw
            img = Image.open(img)

        if not isinstance(img, torch.Tensor):
            if include_source_image:
                src = np.array(img)
                if src.dtype != np.uint8:
                    src = (src * 255).clip(0, 255).astype(np.uint8)
                source_images.append(src)
            img = F.to_tensor(img)
        elif include_source_image:
            source_images.append((img.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8))

        if (img > 1).any():
            raise ValueError(
                "Image has pixel values above 1. Please ensure the image is normalized (scaled to [0, 1]).",
            )
        if (img < 0).any():
            raise ValueError(
                "Image has pixel values below 0. Please ensure the image is normalized (scaled to [0, 1]).",
            )
        if img.shape[0] != self.model_config.num_channels:
            raise ValueError(
                "Invalid tensor image shape. Tensor inputs to `predict()` must be in (C, H, W) format "
                f"with C matching the model configuration ({self.model_config.num_channels} channels). "
                f"Received tensor with shape {tuple(img.shape)}."
            )
        img_tensor = img

        h, w = img_tensor.shape[1:]
        orig_sizes.append((h, w))

        img_tensor = img_tensor.to(self.model.device)
        resize_to = list(shape) if shape is not None else [self.model.resolution, self.model.resolution]
        img_tensor = F.resize(img_tensor, resize_to)
        img_tensor = F.normalize(img_tensor, self.means, self.stds)

        processed_images.append(img_tensor)

    batch_tensor = torch.stack(processed_images)

    if self._is_optimized_for_inference:
        if (
            self._optimized_resolution != batch_tensor.shape[2]
            or self._optimized_resolution != batch_tensor.shape[3]
        ):
            # this could happen if someone manually changes self.model.resolution after optimizing the model,
            # or if predict(shape=...) is used with a shape that doesn't match the compiled square resolution.
            raise ValueError(
                f"Resolution mismatch. "
                f"Model was optimized for resolution {self._optimized_resolution}x{self._optimized_resolution}, "
                f"but got {batch_tensor.shape[2]}x{batch_tensor.shape[3]}."
                " You can explicitly remove the optimized model by calling model.remove_optimized_model().",
            )
        if self._optimized_has_been_compiled:
            if self._optimized_batch_size != batch_tensor.shape[0]:
                raise ValueError(
                    f"Batch size mismatch. "
                    f"Optimized model was compiled for batch size {self._optimized_batch_size}, "
                    f"but got {batch_tensor.shape[0]}."
                    " You can explicitly remove the optimized model by calling model.remove_optimized_model()."
                    " Alternatively, you can recompile the optimized model for a different batch size"
                    " by calling model.optimize_for_inference(batch_size=<new_batch_size>).",
                )

    if self._is_optimized_for_inference:
        predictions = self.model.inference_model(batch_tensor.to(dtype=self._optimized_dtype))
    else:
        predictions = self.model.model(batch_tensor)
    if isinstance(predictions, tuple):
        return_predictions = {
            "pred_logits": predictions[1],
            "pred_boxes": predictions[0],
        }
        if len(predictions) == 3:
            # Distinguish optional keypoint vs mask tuple output for legacy compiled/export shims.
            if getattr(getattr(self.model, "model_config", None), "use_grouppose_keypoints", False):
                return_predictions["pred_keypoints"] = predictions[2]
            else:
                return_predictions["pred_masks"] = predictions[2]
        predictions = return_predictions
    target_sizes = torch.tensor(orig_sizes, device=self.model.device)
    results = self.model.postprocess(predictions, target_sizes=target_sizes)

    model_class_names = self.class_names
    n = len(model_class_names)
    # Pretrained COCO models use COCO category IDs (1–90, with gaps) as class_ids,
    # while class_names is a flat 0-indexed list of 80 entries. Detected when
    # args.num_classes > len(class_names) AND class_names == COCO_CLASS_NAMES.
    # Fine-tuned models remap category IDs to 0-based contiguous indices, so
    # class_id i maps directly to class_names[i].
    _model_args = getattr(self.model, "args", None)
    if _model_args is None and model_class_names == list(COCO_CLASS_NAMES):
        logger.warning_once(
            "predict(): model has no 'args' attribute — COCO sparse-ID mapping cannot activate; "
            "class_ids are treated as 0-indexed (may be wrong for pretrained COCO checkpoints)"
        )
    num_logit_slots: int = getattr(_model_args, "num_classes", n)
    _is_coco_pretrained = num_logit_slots > n and model_class_names == list(COCO_CLASS_NAMES)
    if _is_coco_pretrained:
        _class_id_to_name: dict[int, str] = {
            coco_id: model_class_names[i] for i, coco_id in enumerate(COCO_CLASSES) if i < n
        }
    else:
        _class_id_to_name = dict(enumerate(model_class_names))
    predictions_list: list[Detections | KeyPoints] = []
    for i, result in enumerate(results):
        scores = result["scores"]
        labels = result["labels"]
        boxes = result["boxes"]

        keep = scores > threshold
        scores = scores[keep]
        labels = labels[keep]
        boxes = boxes[keep]
        keypoints_array = None
        if "keypoints" in result:
            keypoints = result["keypoints"][keep]
            keypoints_array = keypoints.float().cpu().numpy()
        has_keypoints = keypoints_array is not None

        if "masks" in result:
            masks = result["masks"]
            masks = masks[keep]

            detections = Detections(
                xyxy=boxes.float().cpu().numpy(),
                confidence=scores.float().cpu().numpy(),
                class_id=labels.cpu().numpy(),
                mask=masks.squeeze(1).cpu().numpy(),
            )
        else:
            detections = Detections(
                xyxy=boxes.float().cpu().numpy(),
                confidence=scores.float().cpu().numpy(),
                class_id=labels.cpu().numpy(),
            )
        if "keypoint_precision_cholesky" in result:
            keypoint_precision = result["keypoint_precision_cholesky"][keep]
            detections.data["keypoint_precision_cholesky"] = keypoint_precision.float().cpu().numpy()

        if include_source_image:
            detections.metadata["source_image"] = source_images[i]
        detections.data["source_shape"] = np.tile(np.array(orig_sizes[i], dtype=np.int64), (len(detections), 1))

        # Attach class names so callers can map class_id → name without a
        # separate lookup. Always set data["class_name"] for a consistent interface.
        #
        # For fine-tuned models, logit index num_logit_slots is the no-object slot —
        # map it to "__background__" without warning. For COCO-pretrained models,
        # background is implicit (filtered by threshold); class ID 90 is "toothbrush".
        # IDs not in _class_id_to_name are genuinely unexpected and produce an empty
        # string with a one-time warning.
        class_ids = detections.class_id if detections.class_id is not None else np.array([], dtype=int)
        truly_oob = [cid for cid in class_ids if cid not in _class_id_to_name and cid != num_logit_slots]
        if truly_oob:
            logger.warning_once(
                "predict() encountered unmapped class_id(s): %s — mapping to empty string",
                truly_oob[:5],
            )
        if _is_coco_pretrained:
            class_names = [_class_id_to_name.get(cid, "") for cid in class_ids]
        else:
            class_names = [
                "__background__" if cid == num_logit_slots else _class_id_to_name.get(cid, "") for cid in class_ids
            ]
        detections.data["class_name"] = np.array(class_names, dtype=object)

        if has_keypoints and keypoints_array is not None:
            keypoint_data = dict(detections.data)
            keypoint_data["xyxy"] = detections.xyxy.astype(np.float32)
            if include_source_image:
                keypoint_data["source_image"] = [source_images[i] for _ in range(len(detections))]
            raw_precision = keypoint_data.get("keypoint_precision_cholesky")
            raw_source_shape = keypoint_data.get("source_shape")
            if raw_precision is not None and raw_source_shape is not None and len(detections) > 0:
                precision = np.asarray(raw_precision, dtype=np.float32)
                source_shape = np.asarray(raw_source_shape, dtype=np.float32)
                if precision.shape[:2] == keypoints_array.shape[:2] and source_shape.shape == (len(detections), 2):
                    keypoint_data["covariance"] = precision_cholesky_to_pixel_covariance(
                        precision_cholesky=precision, source_shape=source_shape
                    )
            keypoints_array = keypoints_array.astype(np.float32, copy=False)
            keypoint_confidence = keypoints_array[:, :, 2]
            key_points = KeyPoints(
                xy=keypoints_array[:, :, :2],
                keypoint_confidence=keypoint_confidence,
                detection_confidence=detections.confidence.astype(np.float32)
                if detections.confidence is not None
                else None,
                class_id=detections.class_id.astype(int) if detections.class_id is not None else None,
                visible=keypoint_confidence > 0,
                data=keypoint_data,
            )
            predictions_list.append(key_points)
        else:
            predictions_list.append(detections)

    return predictions_list if len(predictions_list) > 1 else predictions_list[0]

remove_optimized_model()

Remove the optimized inference model and reset all optimization flags.

Clears model.inference_model and resets all internal state set by :meth:optimize_for_inference. Safe to call even if the model has not been optimized.

Examples:

>>> from types import SimpleNamespace
>>> import torch
>>> class _TinyModel(torch.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.linear = torch.nn.Linear(1, 1)
...     def forward(self, x):
...         return {"pred_boxes": self.linear(x[:, :1, :1, :1].squeeze(-1).squeeze(-1))}
...     def export(self):
...         return None
>>> class _TinyContext:
...     def __init__(self):
...         self.device = torch.device("cpu")
...         self.resolution = 28
...         self.model = _TinyModel()
...         self.inference_model = None
>>> model = object.__new__(RFDETR)
>>> model.model_config = SimpleNamespace(num_channels=3)
>>> model.model = _TinyContext()
>>> model._is_optimized_for_inference = False
>>> model._has_warned_about_not_being_optimized_for_inference = False
>>> model._optimized_has_been_compiled = False
>>> model._optimized_batch_size = None
>>> model._optimized_resolution = None
>>> model._optimized_dtype = None
>>> model.optimize_for_inference(compile=False)
>>> model.remove_optimized_model()
>>> model._is_optimized_for_inference
False

Source code in src/rfdetr/detr.py

def remove_optimized_model(self) -> None:
    """Remove the optimized inference model and reset all optimization flags.

    Clears ``model.inference_model`` and resets all internal state set by :meth:`optimize_for_inference`. Safe to
    call even if the model has not been optimized.

    Examples:
        >>> from types import SimpleNamespace
        >>> import torch
        >>> class _TinyModel(torch.nn.Module):
        ...     def __init__(self):
        ...         super().__init__()
        ...         self.linear = torch.nn.Linear(1, 1)
        ...     def forward(self, x):
        ...         return {"pred_boxes": self.linear(x[:, :1, :1, :1].squeeze(-1).squeeze(-1))}
        ...     def export(self):
        ...         return None
        >>> class _TinyContext:
        ...     def __init__(self):
        ...         self.device = torch.device("cpu")
        ...         self.resolution = 28
        ...         self.model = _TinyModel()
        ...         self.inference_model = None
        >>> model = object.__new__(RFDETR)
        >>> model.model_config = SimpleNamespace(num_channels=3)
        >>> model.model = _TinyContext()
        >>> model._is_optimized_for_inference = False
        >>> model._has_warned_about_not_being_optimized_for_inference = False
        >>> model._optimized_has_been_compiled = False
        >>> model._optimized_batch_size = None
        >>> model._optimized_resolution = None
        >>> model._optimized_dtype = None
        >>> model.optimize_for_inference(compile=False)
        >>> model.remove_optimized_model()
        >>> model._is_optimized_for_inference
        False
    """
    self.model.inference_model = None
    self._is_optimized_for_inference = False
    self._optimized_has_been_compiled = False
    self._optimized_batch_size = None
    self._optimized_resolution = None
    self._optimized_dtype = None

train(**kwargs)

Train an RF-DETR model via the PyTorch Lightning stack.

All keyword arguments are forwarded to :meth:get_train_config to build a :class:~rfdetr.config.TrainConfig. Several kwargs are absorbed and handled specially so that existing call-sites do not break:

• resolution — updates the model's input resolution by mutating :attr:model_config.resolution in place before the train config is built. This change persists on :attr:model_config after :meth:train returns. The value must be a positive integer divisible by patch_size * num_windows for the model variant; a :class:ValueError is raised otherwise. :attr:model_config.positional_encoding_size is also updated when the config derives it formulaically (PE == resolution // patch_size); configs with a pretrained-specific PE value (e.g. RFDETRBase uses DINOv2's PE=37 at 560 px) are left unchanged to preserve checkpoint compatibility.
• device — normalized via :class:torch.device and mapped to PyTorch Lightning trainer arguments. "cpu" becomes accelerator="cpu"; "cuda" and "cuda:N" become accelerator="gpu" and optionally devices=[N]; "mps" becomes accelerator="mps". Other valid torch device types fall back to PTL auto-detection and emit a :class:UserWarning.
• callbacks — if the dict contains any non-empty lists a :class:DeprecationWarning is emitted; the dict is then discarded. Use PTL :class:~pytorch_lightning.Callback objects passed via :func:~rfdetr.training.build_trainer instead.
• start_epoch — emits :class:DeprecationWarning and is dropped.
• do_benchmark — emits :class:DeprecationWarning and is dropped.
• notes — optional user-defined metadata (string, dict, list, or any JSON-serialisable value) stored under the "notes" key in every .pth checkpoint produced during training. The value is also available inside args["notes"] for full provenance. Pass the same value to :meth:export to embed it in the ONNX file as well.

After training completes the underlying nn.Module is synced back onto self.model.model so that :meth:predict and :meth:export continue to work without reloading the checkpoint.

Raises:

Type	Description
ImportError	If training dependencies are not installed. Install with pip install "rfdetr[train,loggers]".
ValueError	If resolution is not a positive integer or is not divisible by patch_size * num_windows for the model variant.

Source code in src/rfdetr/detr.py

def train(self, **kwargs):
    """Train an RF-DETR model via the PyTorch Lightning stack.

    All keyword arguments are forwarded to :meth:`get_train_config` to build a :class:`~rfdetr.config.TrainConfig`.
    Several kwargs are absorbed and handled specially so that existing call-sites do not break:

    * ``resolution`` — updates the model's input resolution by mutating
      :attr:`model_config.resolution` in place before the train config is built. This change persists on
      :attr:`model_config` after :meth:`train` returns. The value must be a positive integer divisible by
      ``patch_size * num_windows`` for the model variant; a :class:`ValueError` is raised otherwise.
      :attr:`model_config.positional_encoding_size` is also updated when the config derives it formulaically (``PE
      == resolution // patch_size``); configs with a pretrained-specific PE value (e.g. ``RFDETRBase`` uses DINOv2's
      PE=37 at 560 px) are left unchanged to preserve checkpoint compatibility.
    * ``device`` — normalized via :class:`torch.device` and mapped to PyTorch
      Lightning trainer arguments. ``"cpu"`` becomes ``accelerator="cpu"``; ``"cuda"`` and ``"cuda:N"`` become
      ``accelerator="gpu"`` and optionally ``devices=[N]``; ``"mps"`` becomes ``accelerator="mps"``. Other valid
      torch device types fall back to PTL auto-detection and emit a :class:`UserWarning`.
    * ``callbacks`` — if the dict contains any non-empty lists a
      :class:`DeprecationWarning` is emitted; the dict is then discarded. Use PTL
      :class:`~pytorch_lightning.Callback` objects passed via :func:`~rfdetr.training.build_trainer` instead.
    * ``start_epoch`` — emits :class:`DeprecationWarning` and is dropped.
    * ``do_benchmark`` — emits :class:`DeprecationWarning` and is dropped.
    * ``notes`` — optional user-defined metadata (string, dict, list, or
      any JSON-serialisable value) stored under the ``"notes"`` key in every ``.pth`` checkpoint produced during
      training.  The value is also available inside ``args["notes"]`` for full provenance.  Pass the same value to
      :meth:`export` to embed it in the ONNX file as well.

    After training completes the underlying ``nn.Module`` is synced back onto ``self.model.model`` so that
    :meth:`predict` and :meth:`export` continue to work without reloading the checkpoint.

    Raises:
        ImportError: If training dependencies are not installed. Install with
            ``pip install "rfdetr[train,loggers]"``.
        ValueError: If ``resolution`` is not a positive integer or is not
            divisible by ``patch_size * num_windows`` for the model variant.
    """
    # Both imports are grouped in a single try block because they both live in
    # the `rfdetr[train]` extras group — a missing `pytorch_lightning` (or any
    # other training-extras package) causes either import to fail, and the
    # remediation is identical: `pip install "rfdetr[train,loggers]"`.
    try:
        from rfdetr.training import RFDETRDataModule, RFDETRModelModule, build_trainer
        from rfdetr.training.auto_batch import resolve_auto_batch_config
    except ModuleNotFoundError as exc:
        # Preserve internal import errors so packaging/regression issues in
        # rfdetr.* are not misreported as missing optional extras.
        if exc.name and exc.name.startswith("rfdetr."):
            raise
        raise ImportError(
            "RF-DETR training dependencies are missing. "
            'Install them with `pip install "rfdetr[train,loggers]"` and try again.',
        ) from exc

    # Absorb legacy `callbacks` dict — warn if non-empty, then discard.
    callbacks_dict = kwargs.pop("callbacks", None)
    if callbacks_dict and any(callbacks_dict.values()):
        warnings.warn(
            "Custom callbacks dict is not forwarded to PTL. "
            "Deprecated since v1.7.0, will be removed in v1.9.0. "
            "Use PTL Callback objects instead.",
            DeprecationWarning,
            stacklevel=2,
        )

    # Parse `device` kwarg and map it to PTL accelerator/devices.
    # Supports torch-style strings and torch.device (e.g. "cuda:1").
    _device = kwargs.pop("device", None)
    _accelerator, _devices = RFDETR._resolve_trainer_device_kwargs(_device)

    # Absorb legacy `start_epoch` — PTL resumes automatically via ckpt_path.
    if "start_epoch" in kwargs:
        warnings.warn(
            "`start_epoch` is deprecated since v1.7.0 and will be removed in v1.9.0; "
            "PTL resumes automatically via `resume`.",
            DeprecationWarning,
            stacklevel=2,
        )
        kwargs.pop("start_epoch")

    # Pop `do_benchmark`; benchmarking via `.train()` is deprecated.
    run_benchmark = bool(kwargs.pop("do_benchmark", False))
    if run_benchmark:
        warnings.warn(
            "`do_benchmark` in `.train()` is deprecated since v1.7.0 and will be removed in v1.9.0; "
            "use `rfdetr benchmark`.",
            DeprecationWarning,
            stacklevel=2,
        )

    # Apply resolution override to model_config before building the train config.
    # resolution is a ModelConfig field, not a TrainConfig field, so we pop it
    # here to avoid it being silently ignored by TrainConfig.
    _resolution = kwargs.pop("resolution", None)
    if _resolution is not None:
        if isinstance(_resolution, bool):
            raise ValueError("resolution must be a positive integer")
        try:
            _resolution = operator.index(_resolution)
        except TypeError as error:
            raise ValueError("resolution must be a positive integer") from error
        if _resolution <= 0:
            raise ValueError("resolution must be a positive integer")
        block_size = self.model_config.patch_size * self.model_config.num_windows
        if _resolution % block_size != 0:
            raise ValueError(
                f"resolution={_resolution} is not divisible by "
                f"patch_size ({self.model_config.patch_size}) * num_windows "
                f"({self.model_config.num_windows}) = {block_size}. "
                f"Choose a resolution that is a multiple of {block_size}."
            )
        # Smart PE update: only recompute positional_encoding_size when the
        # current config derives it formulaically (PE == resolution // patch_size).
        # Configs with a pretrained-specific PE (e.g. RFDETRBase uses DINOv2's
        # PE=37 at 518 px, training at 560 px) must not have PE silently changed
        # — doing so causes shape mismatches when loading pretrained checkpoints.
        _current_pe = self.model_config.positional_encoding_size
        _derived_pe = self.model_config.resolution // self.model_config.patch_size
        if _current_pe == _derived_pe:
            # Formula-derived: update PE proportionally to the new resolution.
            new_pe = _resolution // self.model_config.patch_size
            self.model_config.positional_encoding_size = new_pe
        else:
            # Pretrained-specific PE; leave it unchanged.
            new_pe = _current_pe
        self.model_config.resolution = _resolution

        # Keep the cached inference/export context in sync with model_config so
        # predict()/export()/deployment all see the same resolution metadata.
        if hasattr(self, "model") and self.model is not None:
            if hasattr(self.model, "resolution"):
                self.model.resolution = _resolution
            model_args = getattr(self.model, "args", None)
            if model_args is not None:
                if hasattr(model_args, "resolution"):
                    model_args.resolution = _resolution
                if hasattr(model_args, "positional_encoding_size"):
                    model_args.positional_encoding_size = new_pe
    config = self.get_train_config(**kwargs)
    if config.batch_size == "auto":
        # Auto-batch probing runs forward/backward on the actual model, which
        # must be on the target device (typically CUDA).  Lazy placement keeps
        # the model on CPU until first use — move it now.
        _move_model_context_to_device(self.model)
        auto_batch = resolve_auto_batch_config(
            model_context=self.model,
            model_config=self.model_config,
            train_config=config,
        )
        config.batch_size = auto_batch.safe_micro_batch
        config.grad_accum_steps = auto_batch.recommended_grad_accum_steps
        logger.info(
            "[auto-batch] resolved train config: batch_size=%s grad_accum_steps=%s effective_batch_size=%s",
            config.batch_size,
            config.grad_accum_steps,
            auto_batch.effective_batch_size,
        )
    self.model_config.model_name = type(self).__name__

    # Auto-detect num_classes from the training dataset and align model_config.
    # This must run before RFDETRModelModule is constructed so that weight loading
    # inside the module uses the correct (dataset-derived) class count.
    dataset_dir = getattr(config, "dataset_dir", None)
    if dataset_dir:
        self._align_keypoint_schema_from_dataset(config)
        self._align_num_classes_from_dataset(dataset_dir)

    module = RFDETRModelModule(self.model_config, config)
    datamodule = RFDETRDataModule(self.model_config, config)

    # Guard with LOCAL_RANK env var rather than is_main_process() because torch.distributed
    # is not yet initialized here (it is set up inside trainer.fit()).  In Lightning DDP
    # subprocesses, LOCAL_RANK is set by the launcher before the subprocess calls train(),
    # so this correctly identifies rank 0 even before dist.init_process_group() runs.
    if config.save_dataset_grids and os.environ.get("LOCAL_RANK", "0") == "0":
        try:
            from rfdetr.datasets.save_grids import DatasetGridSaver

            datamodule.setup("fit")
            grids_output_dir = Path(config.output_dir) / "dataset_grids"
            DatasetGridSaver(datamodule.train_dataloader(), grids_output_dir, dataset_type="train").save_grid()
            DatasetGridSaver(datamodule.val_dataloader(), grids_output_dir, dataset_type="val").save_grid()
        except Exception:
            logger.warning(
                "Failed to save dataset grids; training will continue without them.",
                exc_info=True,
            )

    trainer_kwargs = {"accelerator": _accelerator}
    if _devices is not None:
        trainer_kwargs["devices"] = _devices
    trainer = build_trainer(config, self.model_config, **trainer_kwargs)
    trainer.fit(module, datamodule, ckpt_path=config.resume or None)

    # Sync the trained weights back so predict() / export() see the updated model.
    self.model.model = module.model
    # Sync class names: prefer explicit config.class_names, otherwise fall back to dataset (#509).
    config_class_names = getattr(config, "class_names", None)
    if config_class_names is not None:
        self.model.class_names = config_class_names
    else:
        dataset_class_names = getattr(datamodule, "class_names", None)
        if dataset_class_names is not None:
            self.model.class_names = dataset_class_names

    # Save complete training configuration to disk for reproducibility.
    # Guard to main process only to avoid races in distributed/multi-GPU training.
    if is_main_process():
        complete_config = {
            "train_config": config.model_dump(),
            "model_config": self.model_config.model_dump(),
            "model_config_type": self.model_config.__class__.__name__,
            "class_names": self.model.class_names,
            "num_classes": len(self.model.class_names) if self.model.class_names else 0,
        }
        try:
            os.makedirs(config.output_dir, exist_ok=True)
            with open(os.path.join(config.output_dir, "training_config.json"), "w") as f:
                json.dump(complete_config, f, indent=2, default=str)
        except OSError as exc:
            logger.warning("Could not save training_config.json to %s: %s", config.output_dir, exc)