DFINE

The public, backend-agnostic detector class: build from typed params, load released weights, and predict / train / val / export.

dfine.model.DFINE

DFINE(size: str | None = None, *, config: DFINEConfig | None = None, weights: str | PathLike | None = None, device: str | device | None = None, **params)

Config-first D-FINE detector with an ultralytics-style predict.

Source code in dfine/model.py

def __init__(
    self,
    size: str | None = None,
    *,
    config: DFINEConfig | None = None,
    weights: str | os.PathLike | None = None,
    device: str | torch.device | None = None,
    **params,
):
    if config is not None:
        if size is not None or params:
            raise ValueError("Pass either `config=` or `size=`/kwargs, not both.")
        self.config = config
    else:
        self.config = DFINEConfig.preset(size, **params) if size else DFINEConfig(**params)
    self.device = _resolve_device(device)
    self.names = _build_names(self.config)

    # Backend is imported here (not at module top) so it's only pulled in when a
    # model is actually built.
    from .backends.native import DFINE as _NativeDFINE
    from .backends.native import DFINEPostProcessor

    self.model = _NativeDFINE.from_config(self.config).to(self.device).eval()
    self.postprocessor = DFINEPostProcessor.from_config(self.config).to(self.device).eval()

    if weights is not None:
        self.load(weights)

from_pretrained classmethod

from_pretrained(name: str, device: str | device | None = None, **overrides) -> DFINE

Build a model matching a released checkpoint and load its weights.

name is a catalogue entry ("dfine-s", "dfine-l-obj365" ...); the size and num_classes are taken from it. See dfine models.

Source code in dfine/model.py

@classmethod
def from_pretrained(
    cls, name: str, device: str | torch.device | None = None, **overrides
) -> DFINE:
    """Build a model matching a released checkpoint and load its weights.

    ``name`` is a catalogue entry (``"dfine-s"``, ``"dfine-l-obj365"`` ...); the
    size and ``num_classes`` are taken from it. See ``dfine models``.
    """
    from .registry import resolve

    spec = resolve(name)
    # Skip the ImageNet backbone fetch — the checkpoint overwrites it anyway.
    model = cls(
        size=spec.size,
        device=device,
        num_classes=spec.num_classes,
        backbone_pretrained=False,
        **overrides,
    )
    model.load(name)
    return model

load

load(weights: str | PathLike, use_ema: bool = True) -> DFINE

Load weights into the model, in place.

weights is either a catalogue name (downloaded + cached) or a local .pth path. Returns self for chaining: DFINE(size="s").load("dfine-s").

Source code in dfine/model.py

def load(self, weights: str | os.PathLike, use_ema: bool = True) -> DFINE:
    """Load weights into the model, in place.

    ``weights`` is either a catalogue name (downloaded + cached) or a local
    ``.pth`` path. Returns ``self`` for chaining: ``DFINE(size="s").load("dfine-s")``.
    """
    from .backends.native.loader import load_checkpoint
    from .downloads import download_weights
    from .registry import CHECKPOINTS

    if isinstance(weights, str) and weights.lower() in CHECKPOINTS:
        path = download_weights(weights.lower())
    else:
        path = Path(weights)
        if not path.exists():
            raise FileNotFoundError(
                f"{weights!r} is neither a known checkpoint name nor an existing file."
            )
    load_checkpoint(self.model, path, use_ema=use_ema, strict=True)
    self.model.to(self.device)
    return self

predict

predict(source, conf: float = 0.25, imgsz: int | None = None, mask_thresh: float = 0.5) -> list[Results]

Detect objects in source (path / PIL / array, or a list of them).

Returns one :class:~dfine.results.Results per image; boxes are in the original pixel scale. conf drops low-scoring detections. For a task="segment" model, each result also carries per-instance :class:~dfine.results.Masks (original scale), thresholded at mask_thresh.

Source code in dfine/model.py

@torch.no_grad()
def predict(
    self,
    source,
    conf: float = 0.25,
    imgsz: int | None = None,
    mask_thresh: float = 0.5,
) -> list[Results]:
    """Detect objects in ``source`` (path / PIL / array, or a list of them).

    Returns one :class:`~dfine.results.Results` per image; boxes are in the
    original pixel scale. ``conf`` drops low-scoring detections. For a
    ``task="segment"`` model, each result also carries per-instance
    :class:`~dfine.results.Masks` (original scale), thresholded at ``mask_thresh``.
    """
    images = _load_images(source)
    size = imgsz or self.config.imgsz
    if size != self.config.imgsz:
        raise ValueError(
            f"predict(imgsz={size}) must equal the model's imgsz ({self.config.imgsz}): the "
            "encoder's positional embeddings are precomputed for that resolution. Build the "
            f"model at this size instead — DFINE(size=..., imgsz={size})."
        )
    transform = T.Compose([T.Resize((size, size)), T.ToTensor()])

    batch = torch.stack([transform(im) for im in images]).to(self.device)
    orig_sizes = torch.tensor([[im.width, im.height] for im in images], device=self.device)

    outputs = self.model(batch)
    detections = self.postprocessor(outputs, orig_sizes)
    pred_masks = outputs.get("pred_masks")  # [B,Q,Hm,Wm] sigmoid probs, or None
    return [
        self._to_results(
            im, det, conf, None if pred_masks is None else pred_masks[b], mask_thresh
        )
        for b, (im, det) in enumerate(zip(images, detections))
    ]

train

train(train_loader=None, epochs: int | None = None, *, data: str | PathLike | None = None, batch_size: int = 4, num_workers: int = 4, augment: bool = True, remap_mscoco_category: bool = False, devices: int | None = None, val_loader=None, val_fn=None, output_dir: str = 'runs/train', use_wandb: bool = False, visualize: bool = True)

Fine-tune the model (Phase 4).

Provide the data one of two ways:

• data="path/to/coco" — a standard COCO dataset root (train2017/ + annotations/instances_train2017.json, optional val2017/). The train loader (full two-phase augmentation + multi-scale) and, if present, a val loader are built for you via :func:~dfine.train.dataset.build_coco_dataloaders. batch_size, num_workers, augment and remap_mscoco_category tune that build (set remap_mscoco_category=True for stock 80-class MS-COCO ids).
• train_loader=... — a ready dataloader yielding (samples, targets) batches: samples a float BCHW image tensor, each target a dict with labels (LongTensor) and boxes (cxcywh, normalized).

Multi-GPU: pass devices=N to train on N GPUs — this call becomes the launcher and spawns one DDP worker per GPU (no torchrun needed); it requires data= (in-memory loaders can't be shipped to workers). Alternatively launch the script yourself with torchrun --nproc_per_node=N and call train(...) without devices — each worker detects the distributed env and joins the group.

Optimizer groups, LR schedule, EMA, AMP and grad-clip all come from this model's :class:~dfine.config.DFINEConfig. Progress is visualized like upstream D-FINE: a live console readout plus TensorBoard scalars and a loss_curve.png under output_dir (and W&B if use_wandb); only rank 0 writes them. Returns self; the trained (EMA) weights replace self.model.

When a val_loader is available (passed, or auto-built from data) and no val_fn is given, COCO metrics are computed each epoch via :func:~dfine.train.evaluator.coco_val_fn and logged alongside the loss.

Source code in dfine/model.py

def train(
    self,
    train_loader=None,
    epochs: int | None = None,
    *,
    data: str | os.PathLike | None = None,
    batch_size: int = 4,
    num_workers: int = 4,
    augment: bool = True,
    remap_mscoco_category: bool = False,
    devices: int | None = None,
    val_loader=None,
    val_fn=None,
    output_dir: str = "runs/train",
    use_wandb: bool = False,
    visualize: bool = True,
):
    """Fine-tune the model (Phase 4).

    Provide the data one of two ways:

    * ``data="path/to/coco"`` — a standard COCO dataset root (``train2017/`` +
      ``annotations/instances_train2017.json``, optional ``val2017/``). The train
      loader (full two-phase augmentation + multi-scale) and, if present, a val
      loader are built for you via
      :func:`~dfine.train.dataset.build_coco_dataloaders`. ``batch_size``,
      ``num_workers``, ``augment`` and ``remap_mscoco_category`` tune that build
      (set ``remap_mscoco_category=True`` for stock 80-class MS-COCO ids).
    * ``train_loader=...`` — a ready dataloader yielding ``(samples, targets)``
      batches: ``samples`` a float ``BCHW`` image tensor, each ``target`` a dict
      with ``labels`` (``LongTensor``) and ``boxes`` (``cxcywh``, normalized).

    **Multi-GPU:** pass ``devices=N`` to train on ``N`` GPUs — this call becomes the
    launcher and spawns one DDP worker per GPU (no ``torchrun`` needed); it requires
    ``data=`` (in-memory loaders can't be shipped to workers). Alternatively launch
    the script yourself with ``torchrun --nproc_per_node=N`` and call ``train(...)``
    without ``devices`` — each worker detects the distributed env and joins the group.

    Optimizer groups, LR schedule, EMA, AMP and grad-clip all come from this
    model's :class:`~dfine.config.DFINEConfig`. Progress is visualized like upstream
    D-FINE: a live console readout plus TensorBoard scalars and a ``loss_curve.png``
    under ``output_dir`` (and W&B if ``use_wandb``); only rank 0 writes them. Returns
    ``self``; the trained (EMA) weights replace ``self.model``.

    When a ``val_loader`` is available (passed, or auto-built from ``data``) and no
    ``val_fn`` is given, COCO metrics are computed each epoch via
    :func:`~dfine.train.evaluator.coco_val_fn` and logged alongside the loss.
    """
    from .train.distributed import launched_via_torchrun, setup_distributed

    # This process is the launcher: spawn one worker per GPU and reload the result.
    if devices is not None and int(devices) > 1 and not launched_via_torchrun():
        return self._train_multigpu(
            int(devices),
            data=data,
            epochs=epochs,
            batch_size=batch_size,
            num_workers=num_workers,
            augment=augment,
            remap_mscoco_category=remap_mscoco_category,
            output_dir=output_dir,
            use_wandb=use_wandb,
            visualize=visualize,
        )

    # Launched under torchrun: join the group and bind this rank's GPU.
    if launched_via_torchrun():
        setup_distributed()
        self._bind_local_rank_device()

    self._fit(
        train_loader=train_loader,
        epochs=epochs,
        data=data,
        batch_size=batch_size,
        num_workers=num_workers,
        augment=augment,
        remap_mscoco_category=remap_mscoco_category,
        val_loader=val_loader,
        val_fn=val_fn,
        output_dir=output_dir,
        use_wandb=use_wandb,
        visualize=visualize,
    )
    return self

val

val(data: str | PathLike | None = None, *, val_loader=None, batch_size: int = 4, num_workers: int = 4, remap_mscoco_category: bool = False) -> dict[str, float]

Evaluate the model on a COCO val set and return the metrics dict.

Provide the data one of two ways (mutually exclusive):

• data="path/to/coco" — a COCO root; the val loader is built from val2017/ + annotations/instances_val2017.json for you.
• val_loader=... — a ready loader from build_coco_dataloader (its dataset must carry the ground-truth .coco).

Returns the 12 standard COCO metrics keyed by name (AP is the primary mAP@[.50:.95]); see :data:~dfine.train.evaluator.COCO_STAT_NAMES. For stock MS-COCO ground truth (sparse category ids), build the model with remap_mscoco_category=True so predicted labels match the annotations.

Source code in dfine/model.py

def val(
    self,
    data: str | os.PathLike | None = None,
    *,
    val_loader=None,
    batch_size: int = 4,
    num_workers: int = 4,
    remap_mscoco_category: bool = False,
) -> dict[str, float]:
    """Evaluate the model on a COCO val set and return the metrics dict.

    Provide the data one of two ways (mutually exclusive):

    * ``data="path/to/coco"`` — a COCO root; the val loader is built from
      ``val2017/`` + ``annotations/instances_val2017.json`` for you.
    * ``val_loader=...`` — a ready loader from ``build_coco_dataloader`` (its
      dataset must carry the ground-truth ``.coco``).

    Returns the 12 standard COCO metrics keyed by name (``AP`` is the primary
    mAP@[.50:.95]); see :data:`~dfine.train.evaluator.COCO_STAT_NAMES`. For stock
    MS-COCO ground truth (sparse category ids), build the model with
    ``remap_mscoco_category=True`` so predicted labels match the annotations.
    """
    if data is None and val_loader is None:
        raise ValueError("Provide validation data via `data=` or `val_loader=`.")
    if data is not None and val_loader is not None:
        raise ValueError("Pass either `data=` or `val_loader=`, not both.")
    if data is not None:
        from .train.dataset import build_coco_val_dataloader

        val_loader = build_coco_val_dataloader(
            data,
            cfg=self.config,
            batch_size=batch_size,
            num_workers=num_workers,
            remap_mscoco_category=remap_mscoco_category,
        )

    from .train.evaluator import evaluate

    return evaluate(self.model, self.postprocessor, val_loader, self.device)

export

export(format: str = 'onnx', file: str | PathLike | None = None, *, imgsz: int | None = None, batch: int = 1, dynamic: bool = True, simplify: bool = False, opset: int = 16) -> Path

Export the model to a deployable graph (Phase 3).

Currently format="onnx": writes a single ONNX graph with the two-input signature (images, orig_target_sizes) → (labels, boxes, scores) (boxes xyxy in original scale), batch dim dynamic by default. Returns the output :class:~pathlib.Path. Needs pip install pydfine[export].

file defaults to dfine-<size>.onnx. Use simplify=True for onnxsim, and :func:dfine.export.tensorrt_command for a downstream trtexec engine.

Source code in dfine/model.py

def export(
    self,
    format: str = "onnx",
    file: str | os.PathLike | None = None,
    *,
    imgsz: int | None = None,
    batch: int = 1,
    dynamic: bool = True,
    simplify: bool = False,
    opset: int = 16,
) -> Path:
    """Export the model to a deployable graph (Phase 3).

    Currently ``format="onnx"``: writes a single ONNX graph with the two-input
    signature ``(images, orig_target_sizes)`` → ``(labels, boxes, scores)`` (boxes
    ``xyxy`` in original scale), batch dim dynamic by default. Returns the output
    :class:`~pathlib.Path`. Needs ``pip install pydfine[export]``.

    ``file`` defaults to ``dfine-<size>.onnx``. Use ``simplify=True`` for ``onnxsim``,
    and :func:`dfine.export.tensorrt_command` for a downstream ``trtexec`` engine.
    """
    if format != "onnx":
        raise ValueError(f"Unsupported export format {format!r}; only 'onnx' is available.")
    from .export.onnx import export_onnx

    # The encoder precomputes positional embeddings sized to cfg.imgsz, so the export
    # resolution must match the model it was built with — otherwise the traced graph
    # crashes deep in the encoder with a cryptic shape mismatch.
    imgsz = imgsz or self.config.imgsz
    if imgsz != self.config.imgsz:
        raise ValueError(
            f"export imgsz={imgsz} must match the model's imgsz={self.config.imgsz}; "
            f"rebuild the model with DFINE(size=..., imgsz={imgsz}) to export at that size."
        )
    file = (
        Path(file) if file is not None else Path(f"dfine-{self.config.size or 'custom'}.onnx")
    )
    return export_onnx(
        self.model,
        self.postprocessor,
        file,
        imgsz=imgsz,
        batch=batch,
        opset=opset,
        dynamic=dynamic,
        simplify=simplify,
        device=self.device,
    )

predict_video

predict_video(source, output: str | PathLike = 'output.mp4', conf: float = 0.25, imgsz: int | None = None, stream: bool = False, track: bool = False)

Detect objects frame-by-frame in a video.

With stream=True returns a generator of per-frame :class:Results and writes nothing. Otherwise writes an annotated video to output (original resolution/fps) and returns its :class:~pathlib.Path.

With track=True each frame's detections are run through a ByteTrack tracker so boxes carry a persistent boxes.id across frames (rendered as #id and colored per track). Needs scipy (the [track] extra).

Source code in dfine/model.py

def predict_video(
    self,
    source,
    output: str | os.PathLike = "output.mp4",
    conf: float = 0.25,
    imgsz: int | None = None,
    stream: bool = False,
    track: bool = False,
):
    """Detect objects frame-by-frame in a video.

    With ``stream=True`` returns a generator of per-frame :class:`Results` and
    writes nothing. Otherwise writes an annotated video to ``output`` (original
    resolution/fps) and returns its :class:`~pathlib.Path`.

    With ``track=True`` each frame's detections are run through a ByteTrack tracker
    so boxes carry a persistent ``boxes.id`` across frames (rendered as ``#id`` and
    colored per track). Needs scipy (the ``[track]`` extra).
    """
    if stream:
        return self._iter_video(source, conf, imgsz, track)

    cv2 = _require_cv2()
    cap = cv2.VideoCapture(str(source))
    if not cap.isOpened():
        raise FileNotFoundError(f"Could not open video source: {source!r}")

    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    writer = cv2.VideoWriter(str(output), cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
    tracker = self._make_tracker(fps) if track else None
    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            result = self.predict(rgb, conf=conf, imgsz=imgsz)[0]
            if tracker is not None:
                result = tracker.update(result)
            writer.write(cv2.cvtColor(result.plot(), cv2.COLOR_RGB2BGR))
    finally:
        cap.release()
        writer.release()
    return Path(output)