attr-eomt 0.1.0__py3-none-any.whl

eomt/api.py

@@ -0,0 +1,253 @@
+"""High-level :class:`EoMT` interface — init from a checkpoint or a size, then
+``train`` / ``val`` / ``predict``.
+
+    from eomt import EoMT
+
+    model = EoMT("l")                       # fresh large model (DINOv2 backbone)
+    model.train(data="coco", epochs=50)     # COCO auto-downloads if missing
+
+    model = EoMT("runs/train/eomt-l")       # reload a run (everything auto-detected)
+    model.val(data="coco")
+    results = model.predict("images/", plot=True)
+
+The class is a thin orchestration layer over the existing engine functions; the
+trainable network itself is :class:`~eomt.model.EoMTModel`, reachable via
+``model.model``.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+from .config import SIZES
+from .data import CocoValImages, load_data_config
+from .engine import evaluate as _evaluate
+from .engine import evaluate_detection as _evaluate_detection
+from .engine import predict as _predict
+from .engine import train as _train
+from .model import build_model, load_dinov2_backbone
+from .serialization import is_hf_ref, load_model, save_checkpoint, wrap_checkpoint
+
+#: Named dataset aliases that resolve to a bundled YAML config.
+_DATASET_ALIASES = {"coco": "configs/coco.yaml"}
+
+
+def _looks_like_checkpoint(spec: str | Path) -> bool:
+    """True if ``spec`` points at a checkpoint file, a run/weights folder, or a Hub ref."""
+    if is_hf_ref(spec):
+        return True
+    p = Path(spec)
+    if p.suffix == ".pt":
+        return True
+    if p.is_file():
+        return True
+    if p.is_dir():
+        # A run folder if it (or its weights/ subdir) holds a best/last checkpoint.
+        return any((p / sub).is_file() for sub in ("best.pt", "last.pt", "weights/best.pt", "weights/last.pt"))
+    return False
+
+
+def _resolve_data(data: str | Path) -> dict:
+    """Resolve a dataset spec (alias like ``"coco"`` or a YAML path) to absolute paths."""
+    yaml_path = _DATASET_ALIASES.get(str(data), str(data))
+    if not Path(yaml_path).is_file():
+        raise FileNotFoundError(
+            f"dataset config not found: {yaml_path!r}. Pass a dataset YAML path or one "
+            f"of {sorted(_DATASET_ALIASES)} (run from the repo root for the bundled configs)."
+        )
+    return load_data_config(yaml_path)
+
+
+class EoMT:
+    """An EoMT instance-segmentation model with train/val/predict.
+
+    Construct from either a model size (``"s"`` / ``"b"`` / ``"l"`` — a fresh model
+    with a pretrained DINOv2 backbone and a randomly initialized head) or a
+    checkpoint (a ``.pt`` file or a run/weights folder — size, classes, image size
+    and any secondary heads are all auto-detected from the checkpoint).
+    """
+
+    def __init__(self, model: str | Path = "l", *, device: str = "auto", pretrained: bool = True, **build_kwargs):
+        self.device = device
+        self._ckpt: str | None = None
+        self._pretrained = pretrained
+        self._build_kwargs = build_kwargs
+
+        if isinstance(model, (str, Path)) and _looks_like_checkpoint(model):
+            self._model = load_model(model, device=device)
+            self._ckpt = str(model)
+            self.size = self._model.size
+        elif str(model) in SIZES:
+            self.size = str(model)
+            self._model = None  # built lazily (avoids a wasted DINOv2 load before train())
+        else:
+            raise ValueError(
+                f"{model!r} is neither a known size {tuple(SIZES)} nor an existing checkpoint."
+            )
+
+    # ----------------------------------------------------------- huggingface
+    @classmethod
+    def from_pretrained(
+        cls,
+        repo_id: str,
+        *,
+        filename: str = "model.pt",
+        revision: str | None = None,
+        device: str = "auto",
+    ) -> "EoMT":
+        """Load a model from a Hugging Face Hub repo (downloaded once, then cached).
+
+            model = EoMT.from_pretrained("imagra93/eomt-l-coco")
+
+        Equivalent to ``EoMT("hf://<repo_id>/<filename>")``. ``revision`` pins a
+        branch, tag or commit; ``filename`` selects the checkpoint inside the repo.
+        """
+        ref = f"hf://{repo_id}/{filename}"
+        self = cls.__new__(cls)
+        self.device = device
+        self._pretrained = True
+        self._build_kwargs = {}
+        # Resolve via the Hub (cached) and load, recording the ref for repr.
+        from .serialization import download_from_hub
+        local = download_from_hub(ref, revision=revision)
+        self._model = load_model(local, device=device)
+        self._ckpt = ref
+        self.size = self._model.size
+        return self
+
+    def push_to_hub(
+        self,
+        repo_id: str,
+        *,
+        filename: str = "model.pt",
+        private: bool = True,
+        commit_message: str = "Upload EoMT checkpoint",
+    ) -> str:
+        """Upload this model's weights to a Hugging Face Hub repo (created if absent).
+
+        Writes a self-describing checkpoint (same format as :meth:`save`) and
+        uploads it as ``filename``. Returns the repo URL. Requires a Hub token
+        (``huggingface-cli login`` or ``HF_TOKEN``). Defaults to a **private** repo.
+        """
+        try:
+            from huggingface_hub import HfApi
+        except ImportError as e:  # pragma: no cover
+            raise ImportError("push_to_hub needs 'huggingface_hub' (pip install huggingface_hub).") from e
+
+        import tempfile
+
+        api = HfApi()
+        api.create_repo(repo_id, repo_type="model", private=private, exist_ok=True)
+        with tempfile.TemporaryDirectory() as tmp:
+            local = Path(tmp) / filename
+            self.save(local)
+            api.upload_file(
+                path_or_fileobj=str(local),
+                path_in_repo=filename,
+                repo_id=repo_id,
+                repo_type="model",
+                commit_message=commit_message,
+            )
+        return f"https://huggingface.co/{repo_id}"
+
+    # ------------------------------------------------------------------ model
+    @property
+    def model(self):
+        """The underlying trainable :class:`~eomt.model.EoMTModel` (built on demand)."""
+        if self._model is None:
+            self._model = build_model(self.size, **self._build_kwargs)
+            if self._pretrained:
+                load_dinov2_backbone(self._model)
+            dev = "cuda" if (self.device in ("", "auto") and torch.cuda.is_available()) else \
+                ("cpu" if self.device in ("", "auto") else self.device)
+            self._model = self._model.to(dev).eval()
+        return self._model
+
+    # ------------------------------------------------------------------ train
+    def train(self, data: str | Path = "coco", *, resume: bool = False, **hp) -> dict:
+        """Train on a COCO-format dataset.
+
+        ``data`` is a dataset YAML path or a known alias (``"coco"`` auto-downloads).
+        Extra keyword args (``epochs``, ``batch``, ``lr0``, ``aux_w``, …) are passed
+        straight through to the training engine. For a checkpoint-initialized model,
+        ``resume=True`` continues the original run; otherwise the checkpoint warm-starts
+        a fresh run (fine-tune). Returns the engine's result dict and reloads the best
+        weights into this object.
+        """
+        cfg = _resolve_data(data)
+        if not (cfg["train_images"] and cfg["train_json"]):
+            raise ValueError(f"dataset {data!r} has no train split (train_images/train_json).")
+
+        if self._ckpt is not None:
+            hp["resume" if resume else "init_weights"] = self._ckpt
+
+        result = _train(
+            train_images=cfg["train_images"],
+            train_json=cfg["train_json"],
+            val_images=cfg["val_images"],
+            val_json=cfg["val_json"],
+            size=self.size,
+            device=self.device,
+            **hp,
+        )
+        best = result.get("best") or result.get("last")
+        if best:
+            self._model = load_model(best, device=self.device)
+            self._ckpt = best
+        return result
+
+    # -------------------------------------------------------------------- val
+    def val(self, data: str | Path = "coco", *, batch: int = 4, workers: int = 4,
+            conf_thres: float = 0.0, max_det: int = 100, letterbox: bool | None = None, **kw) -> dict:
+        """Evaluate on a dataset's val split, returning COCO mAP metrics.
+
+        Segmentation models report ``segm/*`` (+ ``bbox/*``); detection
+        (``family="detect"``) models report ``bbox/*`` only.
+        """
+        cfg = _resolve_data(data)
+        if not (cfg["val_images"] and cfg["val_json"]):
+            raise ValueError(f"dataset {data!r} has no val split (val_images/val_json).")
+
+        model = self.model
+        dev = next(model.parameters()).device
+        lb = letterbox if letterbox is not None else bool(getattr(model, "preprocess_letterbox", False))
+        val_ds = CocoValImages(cfg["val_images"], cfg["val_json"], imgsz=int(model.image_size), letterbox=lb)
+        eval_fn = _evaluate_detection if getattr(model, "family", "instance") == "detect" else _evaluate
+        return eval_fn(
+            model, val_ds, device=dev, batch_size=batch, num_workers=workers,
+            conf_thres=conf_thres, max_det=max_det, **kw,
+        )
+
+    # ---------------------------------------------------------------- predict
+    def predict(self, source: str | Path, *, plot: bool = False, save: str | None = "runs/predict",
+                conf_thres: float = 0.3, max_det: int = 100, mask_thresh: float = 0.5, **kw) -> list[dict]:
+        """Run inference on an image or a directory.
+
+        Returns one result dict per image (``boxes`` / ``scores`` / ``classes`` /
+        ``masks`` and, for models with secondary heads, ``aux``). With ``plot=True``
+        each image is rendered with masks/boxes/labels and saved under ``save``.
+        """
+        return _predict(
+            self.model, str(source), plot=plot, save=save,
+            conf_thres=conf_thres, max_det=max_det, mask_thresh=mask_thresh, **kw,
+        )
+
+    # ------------------------------------------------------------------- save
+    def save(self, path: str | Path) -> None:
+        """Write a self-describing checkpoint (reloadable with ``EoMT(path)``)."""
+        m = self.model
+        ckpt = wrap_checkpoint(
+            m.state_dict(),
+            size=m.size, nc=m.nc, imgsz=m.image_size, names=m.names,
+            task=getattr(m, "family", "instance"),
+            aux_heads=m.aux_specs, aux_head_arch=m.aux_head_arch,
+            letterbox=bool(getattr(m, "preprocess_letterbox", False)),
+            loss_weights=m.loss_weights, num_upscale_blocks=m.num_upscale_blocks,
+        )
+        save_checkpoint(ckpt, path)
+
+    def __repr__(self) -> str:
+        src = f"ckpt={self._ckpt!r}" if self._ckpt else f"size={self.size!r}"
+        return f"EoMT({src})"

eomt/aux_cls.py

@@ -0,0 +1,225 @@
+"""Secondary per-instance classification heads ("attributes") for EoMT.
+
+The primary task (instance segmentation over ``nc`` classes) is unchanged. Each
+aux head predicts an extra attribute per *detected instance* — typology,
+laterality, severity, … — read from the matched query's embedding.
+
+The supervision reuses EoMT's own Hungarian matcher
+(``model.eomt.criterion.matcher``) so each attribute is trained on the **same**
+query→GT assignment the detection loss used. Several specs ⇒ several independent
+heads, summed (optionally weighted) into one scalar.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+
+
+@torch.no_grad()
+def match_queries(model, out: dict, geom_labels, class_labels):
+    """Hungarian query→GT indices from EoMT's matcher (list of ``(src, tgt)``).
+
+    ``geom_labels`` is the GT geometry for the model's family — instance masks for the
+    ``"instance"`` family, normalized ``cxcywh`` boxes for ``"detect"``. The right
+    matcher (mask-cost or box-cost) is selected by what the model emitted in ``out``.
+    """
+    geom = out["pred_boxes"] if "pred_boxes" in out else out["masks_queries_logits"]
+    return model.eomt.criterion.matcher(
+        geom, out["class_queries_logits"], geom_labels, class_labels
+    )
+
+
+def _gather_matched(out: dict, indices):
+    """Return ``(matched_feats [N, hidden], [(b, tgt_idx)...])`` for matched queries."""
+    q = out["query_embed"]  # [B, Q, hidden]
+    feats, batch_tgt = [], []
+    for b, (src, tgt) in enumerate(indices):
+        if src.numel() == 0:
+            continue
+        feats.append(q[b, src])
+        batch_tgt.append((b, tgt))
+    if not feats:
+        return None, batch_tgt
+    return torch.cat(feats), batch_tgt
+
+
+@torch.no_grad()
+def gate_indices(
+    out: dict,
+    indices,
+    geom_labels,
+    class_labels,
+    *,
+    iou_thr: float = 0.5,
+    require_class: bool = False,
+):
+    """Keep only well-localized (and optionally correctly-classified) matched pairs.
+
+    The Hungarian matcher assigns *every* GT a query, even one whose prediction barely
+    overlaps it (common early in training). For attribute supervision we want queries
+    that actually localize the instance, so we drop matched ``(src, tgt)`` pairs whose
+    predicted↔GT IoU is ``< iou_thr`` (mask IoU for the instance family, box IoU for
+    detect). With ``require_class`` we also drop pairs whose predicted primary class ≠
+    the GT class. Returns the same ``[(src, tgt), ...]`` structure with each pair
+    filtered (possibly to empty).
+
+    A no-op (returns ``indices`` unchanged) when ``iou_thr <= 0`` and not
+    ``require_class`` — i.e. the pre-gate behaviour.
+    """
+    if iou_thr <= 0 and not require_class:
+        return indices
+    is_detect = "pred_boxes" in out
+    cls = out["class_queries_logits"]  # [B, Q, C+1]
+    dev = cls.device
+    gated = []
+    for b, (src, tgt) in enumerate(indices):
+        src = src.to(dev)
+        tgt = tgt.to(dev)
+        if src.numel() == 0:
+            gated.append((src, tgt))
+            continue
+        keep = torch.ones(src.numel(), dtype=torch.bool, device=dev)
+        if iou_thr > 0:
+            keep &= _localization_iou(out, b, src, tgt, geom_labels, is_detect, dev) >= iou_thr
+        if require_class:
+            pred_cls = cls[b, src].argmax(-1)
+            keep &= pred_cls == class_labels[b][tgt].to(dev)
+        gated.append((src[keep], tgt[keep]))
+    return gated
+
+
+def _localization_iou(out, b, src, tgt, geom_labels, is_detect, dev):
+    """Per-pair predicted↔GT IoU for the matched queries (box IoU or mask IoU)."""
+    if is_detect:
+        from .box_loss import box_cxcywh_to_xyxy, box_iou
+
+        pred = box_cxcywh_to_xyxy(out["pred_boxes"][b, src].to(dev))  # [n, 4]
+        gt = box_cxcywh_to_xyxy(geom_labels[b][tgt].to(dev).float())  # [n, 4]
+        return box_iou(pred, gt)[0].diagonal()  # per-matched-pair IoU
+    masks = out["masks_queries_logits"]  # [B, Q, h, w]
+    pm = masks[b, src].sigmoid() > 0.5  # [n, h, w]
+    gm = geom_labels[b][tgt].to(dev).float()  # [n, H, W] in {0, 1}
+    if gm.shape[-2:] != pm.shape[-2:]:
+        gm = F.interpolate(gm.unsqueeze(1), size=pm.shape[-2:], mode="nearest").squeeze(1)
+    gm = gm > 0.5
+    inter = (pm & gm).flatten(1).sum(1).float()
+    union = (pm | gm).flatten(1).sum(1).float().clamp(min=1.0)
+    return inter / union
+
+
+def aux_loss(
+    model,
+    out: dict,
+    mask_labels,
+    class_labels,
+    aux_labels: dict[str, list[torch.Tensor]],
+    weights: dict[str, float] | None = None,
+    *,
+    indices=None,
+    class_weights: dict[str, torch.Tensor] | None = None,
+    ignore_index: int = -100,
+):
+    """Weighted sum of per-attribute CE over matched queries.
+
+    ``aux_labels`` maps head name → per-image label tensors (line-aligned with
+    ``class_labels``). Labels equal to ``ignore_index`` (missing / out-of-vocab
+    attributes) are skipped — they contribute no loss instead of being trained as
+    class 0. ``weights`` scales each head; ``class_weights`` optionally re-weights
+    classes within a head (imbalance). Pass ``indices`` to reuse a matching already
+    computed this step (avoids re-running the Hungarian matcher). Returns
+    ``(total_loss, {name: per_head_loss})``.
+    """
+    if indices is None:
+        indices = match_queries(model, out, mask_labels, class_labels)
+    matched, batch_tgt = _gather_matched(out, indices)
+    if matched is None:  # no query matched in this batch — keep the graph alive
+        zero = out["query_embed"].sum() * 0.0
+        return zero, {name: zero.detach() for name in model.aux_heads}
+
+    weights = weights or {}
+    class_weights = class_weights or {}
+    total: torch.Tensor | None = None
+    per_head: dict[str, torch.Tensor] = {}
+    for name, head in model.aux_heads.items():
+        logits = head(matched)  # [N, ns]
+        gt = torch.cat([aux_labels[name][b][tgt] for (b, tgt) in batch_tgt]).to(logits.device)
+        cw = class_weights.get(name)
+        if cw is not None:
+            cw = cw.to(logits.device, logits.dtype)
+        if (gt != ignore_index).any():
+            loss = F.cross_entropy(logits, gt, weight=cw, ignore_index=ignore_index)
+        else:  # every matched label ignored — keep the graph alive with a zero
+            loss = logits.sum() * 0.0
+        per_head[name] = loss
+        w = float(weights.get(name, 1.0))
+        total = w * loss if total is None else total + w * loss
+    return total, per_head
+
+
+@torch.no_grad()
+def aux_accuracy(
+    model,
+    out: dict,
+    mask_labels,
+    class_labels,
+    aux_labels: dict[str, list[torch.Tensor]],
+    *,
+    indices=None,
+    ignore_index: int = -100,
+) -> dict[str, tuple[int, int]]:
+    """Top-1 ``{name: (correct, total)}`` on matched queries (the ``typ_acc`` analogue).
+
+    Ignored labels (``ignore_index``) are excluded from both correct and total.
+    Pass ``indices`` to reuse a matching already computed this step.
+    """
+    if indices is None:
+        indices = match_queries(model, out, mask_labels, class_labels)
+    matched, batch_tgt = _gather_matched(out, indices)
+    if matched is None:
+        return {name: (0, 0) for name in model.aux_heads}
+    res: dict[str, tuple[int, int]] = {}
+    for name, head in model.aux_heads.items():
+        pred = head(matched).argmax(-1)
+        gt = torch.cat([aux_labels[name][b][tgt] for (b, tgt) in batch_tgt]).to(pred.device)
+        valid = gt != ignore_index
+        res[name] = (int(((pred == gt) & valid).sum()), int(valid.sum()))
+    return res
+
+
+@torch.no_grad()
+def aux_accuracy_by_primary(
+    model,
+    out: dict,
+    mask_labels,
+    class_labels,
+    aux_labels: dict[str, list[torch.Tensor]],
+    *,
+    indices=None,
+    ignore_index: int = -100,
+) -> dict[str, dict[int, tuple[int, int]]]:
+    """Per-head accuracy bucketed by GT **primary** class.
+
+    Returns ``{head: {primary_class_id: (correct, total)}}`` over matched queries —
+    a diagnostic for *which primary classes* the attribute is (in)accurate on. Pass
+    ``indices`` (e.g. IoU-gated, but not class-gated, so weak primary classes still
+    appear) to control the population. Ignored labels are excluded.
+    """
+    if indices is None:
+        indices = match_queries(model, out, mask_labels, class_labels)
+    matched, batch_tgt = _gather_matched(out, indices)
+    res: dict[str, dict[int, tuple[int, int]]] = {name: {} for name in model.aux_heads}
+    if matched is None:
+        return res
+    prim = torch.cat([class_labels[b][tgt] for (b, tgt) in batch_tgt]).to(matched.device)
+    for name, head in model.aux_heads.items():
+        pred = head(matched).argmax(-1)
+        gt = torch.cat([aux_labels[name][b][tgt] for (b, tgt) in batch_tgt]).to(pred.device)
+        valid = gt != ignore_index
+        correct = (pred == gt) & valid
+        d: dict[int, tuple[int, int]] = {}
+        for c in prim[valid].unique().tolist():
+            sel = (prim == c) & valid
+            d[int(c)] = (int(correct[sel].sum()), int(sel.sum()))
+        res[name] = d
+    return res

eomt/box_loss.py

@@ -0,0 +1,196 @@
+"""Pure-PyTorch DETR-style detection loss: Hungarian matcher + L1/GIoU + class CE.
+
+The detection counterpart to :mod:`eomt.loss` (the mask-classification stack). It
+keeps the **same** query→GT assignment machinery — a Hungarian matcher exposed as
+``DetectionLoss.matcher`` with a ``(geometry, class_queries_logits, geom_labels,
+class_labels)`` signature and the same list-of-``(src, tgt)`` return — so
+:func:`eomt.aux_cls.match_queries` can drive attribute supervision identically for
+the box head. The only difference from :mod:`eomt.loss` is the geometry term: per-query
+mask BCE + dice is replaced by per-query box **L1 + GIoU** (DETR's box loss).
+
+All boxes here are normalized ``cxcywh`` in ``[0, 1]`` (the box head emits sigmoid
+``cxcywh``; the detection dataset stores GT the same way), so no image size is needed.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+from scipy.optimize import linear_sum_assignment
+from torch import Tensor, nn
+
+
+def box_cxcywh_to_xyxy(boxes: Tensor) -> Tensor:
+    """``(..., 4)`` boxes from center form ``(cx, cy, w, h)`` to corner form ``(x1, y1, x2, y2)``."""
+    cx, cy, w, h = boxes.unbind(-1)
+    return torch.stack([cx - 0.5 * w, cy - 0.5 * h, cx + 0.5 * w, cy + 0.5 * h], dim=-1)
+
+
+def _box_area(boxes: Tensor) -> Tensor:
+    return (boxes[:, 2] - boxes[:, 0]).clamp(min=0) * (boxes[:, 3] - boxes[:, 1]).clamp(min=0)
+
+
+def box_iou(boxes1: Tensor, boxes2: Tensor) -> tuple[Tensor, Tensor]:
+    """Pairwise IoU and union between two sets of ``xyxy`` boxes ``([N,4], [M,4])``."""
+    area1 = _box_area(boxes1)
+    area2 = _box_area(boxes2)
+    lt = torch.max(boxes1[:, None, :2], boxes2[None, :, :2])
+    rb = torch.min(boxes1[:, None, 2:], boxes2[None, :, 2:])
+    wh = (rb - lt).clamp(min=0)
+    inter = wh[..., 0] * wh[..., 1]
+    union = area1[:, None] + area2[None, :] - inter
+    iou = inter / union.clamp(min=1e-6)
+    return iou, union
+
+
+def generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
+    """Pairwise GIoU between two sets of ``xyxy`` boxes (DETR's enclosing-box variant)."""
+    iou, union = box_iou(boxes1, boxes2)
+    lt = torch.min(boxes1[:, None, :2], boxes2[None, :, :2])
+    rb = torch.max(boxes1[:, None, 2:], boxes2[None, :, 2:])
+    wh = (rb - lt).clamp(min=0)
+    area = wh[..., 0] * wh[..., 1]
+    return iou - (area - union) / area.clamp(min=1e-6)
+
+
+class DetectionHungarianMatcher(nn.Module):
+    """1-to-1 assignment between queries and GT boxes via class + L1 + GIoU cost."""
+
+    def __init__(self, cost_class: float = 1.0, cost_bbox: float = 1.0, cost_giou: float = 1.0):
+        super().__init__()
+        if cost_class == 0 and cost_bbox == 0 and cost_giou == 0:
+            raise ValueError("All costs can't be 0")
+        self.cost_class = cost_class
+        self.cost_bbox = cost_bbox
+        self.cost_giou = cost_giou
+
+    @torch.no_grad()
+    def forward(
+        self,
+        pred_boxes: Tensor,
+        class_queries_logits: Tensor,
+        box_labels: list[Tensor],
+        class_labels: list[Tensor],
+    ) -> list[tuple[Tensor, Tensor]]:
+        indices: list = []
+        batch_size = pred_boxes.shape[0]
+        for i in range(batch_size):
+            # Cost math in fp32: cdist/GIoU are not implemented for fp16 on CUDA, and
+            # the matcher is called both inside autocast (training) and outside it
+            # (eval/aux), where the tensors stay half.
+            tgt_boxes = box_labels[i].to(device=pred_boxes.device, dtype=torch.float32)
+            tgt_cls = class_labels[i]
+            if tgt_boxes.numel() == 0:
+                indices.append(
+                    (
+                        torch.as_tensor([], dtype=torch.int64),
+                        torch.as_tensor([], dtype=torch.int64),
+                    )
+                )
+                continue
+            pred_probs = class_queries_logits[i].float().softmax(-1)  # [Q, C+1]
+            out_boxes = pred_boxes[i].float()  # [Q, 4]
+
+            cost_class = -pred_probs[:, tgt_cls]  # [Q, num_tgt]
+            cost_bbox = torch.cdist(out_boxes, tgt_boxes, p=1)  # [Q, num_tgt]
+            cost_giou = -generalized_box_iou(
+                box_cxcywh_to_xyxy(out_boxes), box_cxcywh_to_xyxy(tgt_boxes)
+            )
+            cost_matrix = (
+                self.cost_bbox * cost_bbox
+                + self.cost_class * cost_class
+                + self.cost_giou * cost_giou
+            )
+            cost_matrix = torch.minimum(cost_matrix, torch.tensor(1e10))
+            cost_matrix = torch.maximum(cost_matrix, torch.tensor(-1e10))
+            cost_matrix = torch.nan_to_num(cost_matrix, 0)
+            indices.append(linear_sum_assignment(cost_matrix.cpu()))
+
+        return [
+            (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64))
+            for i, j in indices
+        ]
+
+
+class DetectionLoss(nn.Module):
+    """DETR-style detection loss (class CE + box L1 + GIoU), matching EoMTLoss's API.
+
+    ``weight_dict`` carries ``{"loss_cross_entropy", "loss_bbox", "loss_giou"}``; the
+    matcher reuses ``class_weight`` (from the config) and the L1/GIoU weights from
+    ``weight_dict``. ``forward`` returns the three *unweighted* loss terms — the caller
+    (``EoMTEncoder.get_loss_dict``) applies ``weight_dict``, exactly as for the mask loss.
+    """
+
+    def __init__(self, config, weight_dict: dict[str, float]):
+        super().__init__()
+        self.num_labels = config.num_labels
+        self.weight_dict = weight_dict
+
+        self.eos_coef = config.no_object_weight
+        empty_weight = torch.ones(self.num_labels + 1)
+        empty_weight[-1] = self.eos_coef
+        self.register_buffer("empty_weight", empty_weight)
+
+        self.matcher = DetectionHungarianMatcher(
+            cost_class=config.class_weight,
+            cost_bbox=weight_dict["loss_bbox"],
+            cost_giou=weight_dict["loss_giou"],
+        )
+
+    def _get_predictions_permutation_indices(self, indices):
+        batch_indices = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
+        predictions_indices = torch.cat([src for (src, _) in indices])
+        return batch_indices, predictions_indices
+
+    def loss_labels(self, class_queries_logits, class_labels, indices) -> dict[str, Tensor]:
+        pred_logits = class_queries_logits
+        batch_size, num_queries, _ = pred_logits.shape
+        criterion = nn.CrossEntropyLoss(weight=self.empty_weight)
+        idx = self._get_predictions_permutation_indices(indices)
+        target_classes_o = torch.cat(
+            [target[j] for target, (_, j) in zip(class_labels, indices)]
+        ).to(pred_logits.device)
+        target_classes = torch.full(
+            (batch_size, num_queries), fill_value=self.num_labels, dtype=torch.int64,
+            device=pred_logits.device,
+        )
+        target_classes[idx] = target_classes_o
+        loss_ce = criterion(pred_logits.transpose(1, 2), target_classes)
+        return {"loss_cross_entropy": loss_ce}
+
+    def loss_boxes(self, pred_boxes, box_labels, indices, num_boxes) -> dict[str, Tensor]:
+        idx = self._get_predictions_permutation_indices(indices)
+        src_boxes = pred_boxes[idx]  # [N, 4] cxcywh
+        tgt_boxes = torch.cat(
+            [t[j] for t, (_, j) in zip(box_labels, indices)]
+        ).to(src_boxes)  # [N, 4] cxcywh
+        if src_boxes.numel() == 0:  # no matched query this batch — keep graph alive
+            zero = pred_boxes.sum() * 0.0
+            return {"loss_bbox": zero, "loss_giou": zero}
+
+        loss_bbox = F.l1_loss(src_boxes, tgt_boxes, reduction="sum") / num_boxes
+        giou = generalized_box_iou(
+            box_cxcywh_to_xyxy(src_boxes), box_cxcywh_to_xyxy(tgt_boxes)
+        ).diagonal()
+        loss_giou = (1 - giou).sum() / num_boxes
+        return {"loss_bbox": loss_bbox, "loss_giou": loss_giou}
+
+    def get_num_boxes(self, class_labels, device) -> Tensor:
+        num = sum(len(c) for c in class_labels)
+        return torch.clamp(torch.as_tensor(num, dtype=torch.float, device=device), min=1)
+
+    def forward(
+        self,
+        masks_queries_logits: Tensor,  # unused; kept for signature symmetry with EoMTLoss
+        class_queries_logits: Tensor,
+        mask_labels: list[Tensor],  # actually box_labels (cxcywh) for the detect family
+        class_labels: list[Tensor],
+        auxiliary_predictions=None,
+    ) -> dict[str, Tensor]:
+        pred_boxes, box_labels = masks_queries_logits, mask_labels
+        indices = self.matcher(pred_boxes, class_queries_logits, box_labels, class_labels)
+        num_boxes = self.get_num_boxes(class_labels, device=class_queries_logits.device)
+        return {
+            **self.loss_boxes(pred_boxes, box_labels, indices, num_boxes),
+            **self.loss_labels(class_queries_logits, class_labels, indices),
+        }