pydfine 0.0.1__py3-none-any.whl

dfine/__init__.py

@@ -0,0 +1,48 @@
+"""dfine — a config-first, ultralytics-style wrapper around D-FINE.
+
+The whole model is configured by typed params on one class. Phase 0/1 ships the
+config surface (``DFINEConfig``, presets); ``DFINE`` and ``Results`` land with the
+inference backend (Phase 2) and are exposed lazily so importing this package never
+requires torch until you actually build a model.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from .config import SIZE_PRESETS, SIZES, DFINEConfig, list_presets
+from .convert import yolo_to_coco
+from .registry import list_checkpoints
+
+__version__ = "0.0.1"
+
+__all__ = [
+    "DFINEConfig",
+    "SIZE_PRESETS",
+    "SIZES",
+    "list_presets",
+    "list_checkpoints",
+    "yolo_to_coco",
+    "__version__",
+]
+
+# Inference symbols live in modules that import torch; expose them lazily so a bare
+# `import dfine` (config/CLI only) never requires the torch extra.
+_LAZY = {
+    "DFINE": "model",
+    "Results": "results",
+    "Boxes": "results",
+}
+
+
+def __getattr__(name: str) -> Any:
+    if name in _LAZY:
+        try:
+            mod = __import__(f"dfine.{_LAZY[name]}", fromlist=[name])
+        except ImportError as exc:  # torch/torchvision/pillow not installed
+            raise AttributeError(
+                f"dfine.{name} needs the inference deps — install with "
+                f"`pip install pydfine[torch]` (missing: {exc.name})."
+            ) from exc
+        return getattr(mod, name)
+    raise AttributeError(f"module 'dfine' has no attribute {name!r}")

dfine/backends/__init__.py

@@ -0,0 +1,6 @@
+"""Model backends.
+
+The native port (Path A) lives in ``dfine.backends.native`` — upstream D-FINE
+``nn.Module``s copied here with the YAML/registry layer stripped and wired directly
+from :class:`dfine.config.DFINEConfig`.
+"""

dfine/backends/native/__init__.py

@@ -0,0 +1,29 @@
+"""Native D-FINE modules (Path A port).
+
+Each module is copied from upstream ``D-FINE/src`` with the registry/YAML layer
+removed and a ``from_config`` constructor added. Layer/parameter names match
+upstream so released checkpoints load without a remap.
+"""
+
+from __future__ import annotations
+
+from .criterion import DFINECriterion
+from .dfine import DFINE
+from .dfine_decoder import DFINETransformer
+from .hgnetv2 import HGNetv2
+from .hybrid_encoder import HybridEncoder
+from .loader import extract_state_dict, load_checkpoint
+from .matcher import HungarianMatcher
+from .postprocessor import DFINEPostProcessor
+
+__all__ = [
+    "DFINE",
+    "DFINECriterion",
+    "DFINEPostProcessor",
+    "DFINETransformer",
+    "HGNetv2",
+    "HungarianMatcher",
+    "HybridEncoder",
+    "extract_state_dict",
+    "load_checkpoint",
+]

dfine/backends/native/box_ops.py

@@ -0,0 +1,86 @@
+"""Bounding-box helpers.
+
+Ported from ``D-FINE/src/zoo/dfine/box_ops.py``, which originates from DETR
+(© Facebook, Inc.; https://github.com/facebookresearch/detr). Kept verbatim.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+from torchvision.ops.boxes import box_area
+
+
+def box_cxcywh_to_xyxy(x: Tensor) -> Tensor:
+    x_c, y_c, w, h = x.unbind(-1)
+    b = [
+        (x_c - 0.5 * w.clamp(min=0.0)),
+        (y_c - 0.5 * h.clamp(min=0.0)),
+        (x_c + 0.5 * w.clamp(min=0.0)),
+        (y_c + 0.5 * h.clamp(min=0.0)),
+    ]
+    return torch.stack(b, dim=-1)
+
+
+def box_xyxy_to_cxcywh(x: Tensor) -> Tensor:
+    x0, y0, x1, y1 = x.unbind(-1)
+    b = [(x0 + x1) / 2, (y0 + y1) / 2, (x1 - x0), (y1 - y0)]
+    return torch.stack(b, dim=-1)
+
+
+def box_iou(boxes1: Tensor, boxes2: Tensor) -> tuple[Tensor, Tensor]:
+    """IoU plus the union area (modified from torchvision to also return union)."""
+    area1 = box_area(boxes1)
+    area2 = box_area(boxes2)
+
+    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]
+    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]
+
+    wh = (rb - lt).clamp(min=0)  # [N,M,2]
+    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]
+
+    union = area1[:, None] + area2 - inter
+
+    iou = inter / union
+    return iou, union
+
+
+def generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
+    """Generalized IoU (https://giou.stanford.edu/). Boxes in ``xyxy``.
+
+    Returns an ``[N, M]`` pairwise matrix.
+    """
+    # degenerate boxes give inf/nan, so check first
+    assert (boxes1[:, 2:] >= boxes1[:, :2]).all()
+    assert (boxes2[:, 2:] >= boxes2[:, :2]).all()
+    iou, union = box_iou(boxes1, boxes2)
+
+    lt = torch.min(boxes1[:, None, :2], boxes2[:, :2])
+    rb = torch.max(boxes1[:, None, 2:], boxes2[:, 2:])
+
+    wh = (rb - lt).clamp(min=0)  # [N,M,2]
+    area = wh[:, :, 0] * wh[:, :, 1]
+
+    return iou - (area - union) / area
+
+
+def masks_to_boxes(masks: Tensor) -> Tensor:
+    """Bounding boxes (``xyxy``) around ``[N, H, W]`` masks."""
+    if masks.numel() == 0:
+        return torch.zeros((0, 4), device=masks.device)
+
+    h, w = masks.shape[-2:]
+
+    y = torch.arange(0, h, dtype=torch.float)
+    x = torch.arange(0, w, dtype=torch.float)
+    y, x = torch.meshgrid(y, x, indexing="ij")
+
+    x_mask = masks * x.unsqueeze(0)
+    x_max = x_mask.flatten(1).max(-1)[0]
+    x_min = x_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]
+
+    y_mask = masks * y.unsqueeze(0)
+    y_max = y_mask.flatten(1).max(-1)[0]
+    y_min = y_mask.masked_fill(~(masks.bool()), 1e8).flatten(1).min(-1)[0]
+
+    return torch.stack([x_min, y_min, x_max, y_max], 1)

dfine/backends/native/coco.py

@@ -0,0 +1,103 @@
+"""MS-COCO category maps — native D-FINE port.
+
+Copied from ``D-FINE/src/data/dataset/coco_dataset.py`` (Apache-2.0, © 2024 The
+D-FINE Authors). COCO's 80 object classes carry non-contiguous ids in the 1..90
+range; models predict a contiguous 0..79 label. ``mscoco_label2category`` maps a
+predicted label back to the original COCO category id (used by the postprocessor
+when ``remap_mscoco_category=True``); ``mscoco_category2name`` gives the display
+name for each COCO id.
+"""
+
+from __future__ import annotations
+
+__all__ = [
+    "mscoco_category2name",
+    "mscoco_category2label",
+    "mscoco_label2category",
+]
+
+mscoco_category2name = {
+    1: "person",
+    2: "bicycle",
+    3: "car",
+    4: "motorcycle",
+    5: "airplane",
+    6: "bus",
+    7: "train",
+    8: "truck",
+    9: "boat",
+    10: "traffic light",
+    11: "fire hydrant",
+    13: "stop sign",
+    14: "parking meter",
+    15: "bench",
+    16: "bird",
+    17: "cat",
+    18: "dog",
+    19: "horse",
+    20: "sheep",
+    21: "cow",
+    22: "elephant",
+    23: "bear",
+    24: "zebra",
+    25: "giraffe",
+    27: "backpack",
+    28: "umbrella",
+    31: "handbag",
+    32: "tie",
+    33: "suitcase",
+    34: "frisbee",
+    35: "skis",
+    36: "snowboard",
+    37: "sports ball",
+    38: "kite",
+    39: "baseball bat",
+    40: "baseball glove",
+    41: "skateboard",
+    42: "surfboard",
+    43: "tennis racket",
+    44: "bottle",
+    46: "wine glass",
+    47: "cup",
+    48: "fork",
+    49: "knife",
+    50: "spoon",
+    51: "bowl",
+    52: "banana",
+    53: "apple",
+    54: "sandwich",
+    55: "orange",
+    56: "broccoli",
+    57: "carrot",
+    58: "hot dog",
+    59: "pizza",
+    60: "donut",
+    61: "cake",
+    62: "chair",
+    63: "couch",
+    64: "potted plant",
+    65: "bed",
+    67: "dining table",
+    70: "toilet",
+    72: "tv",
+    73: "laptop",
+    74: "mouse",
+    75: "remote",
+    76: "keyboard",
+    77: "cell phone",
+    78: "microwave",
+    79: "oven",
+    80: "toaster",
+    81: "sink",
+    82: "refrigerator",
+    84: "book",
+    85: "clock",
+    86: "vase",
+    87: "scissors",
+    88: "teddy bear",
+    89: "hair drier",
+    90: "toothbrush",
+}
+
+mscoco_category2label = {k: i for i, k in enumerate(mscoco_category2name.keys())}
+mscoco_label2category = {v: k for k, v in mscoco_category2label.items()}

dfine/backends/native/common.py

@@ -0,0 +1,75 @@
+"""Shared low-level layers for the native D-FINE port.
+
+Ported from ``D-FINE/src/nn/backbone/common.py`` (Apache-2.0, © 2024 The D-FINE
+Authors). ``FrozenBatchNorm2d`` itself originates from DETR
+(facebookresearch/detr). Kept verbatim in behaviour and parameter/buffer names so
+upstream checkpoints load unchanged.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class FrozenBatchNorm2d(nn.Module):
+    """BatchNorm2d with fixed batch statistics and affine params.
+
+    Adds ``eps`` before ``rsqrt`` (as in torchvision's variant) to avoid NaNs on
+    non-resnet backbones.
+    """
+
+    def __init__(self, num_features: int, eps: float = 1e-5):
+        super().__init__()
+        n = num_features
+        self.register_buffer("weight", torch.ones(n))
+        self.register_buffer("bias", torch.zeros(n))
+        self.register_buffer("running_mean", torch.zeros(n))
+        self.register_buffer("running_var", torch.ones(n))
+        self.eps = eps
+        self.num_features = n
+
+    def _load_from_state_dict(
+        self,
+        state_dict,
+        prefix,
+        local_metadata,
+        strict,
+        missing_keys,
+        unexpected_keys,
+        error_msgs,
+    ):
+        num_batches_tracked_key = prefix + "num_batches_tracked"
+        if num_batches_tracked_key in state_dict:
+            del state_dict[num_batches_tracked_key]
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Reshapes up front to stay fuser-friendly.
+        w = self.weight.reshape(1, -1, 1, 1)
+        b = self.bias.reshape(1, -1, 1, 1)
+        rv = self.running_var.reshape(1, -1, 1, 1)
+        rm = self.running_mean.reshape(1, -1, 1, 1)
+        scale = w * (rv + self.eps).rsqrt()
+        bias = b - rm * scale
+        return x * scale + bias
+
+    def extra_repr(self) -> str:
+        return f"{self.num_features}, eps={self.eps}"
+
+
+def freeze_norm(module: nn.Module) -> nn.Module:
+    """Recursively replace every ``nn.BatchNorm2d`` with a ``FrozenBatchNorm2d``.
+
+    Returns the (possibly replaced) module so callers can reassign the root.
+    """
+    if isinstance(module, nn.BatchNorm2d):
+        frozen = FrozenBatchNorm2d(module.num_features)
+        return frozen
+    for name, child in module.named_children():
+        new_child = freeze_norm(child)
+        if new_child is not child:
+            setattr(module, name, new_child)
+    return module