Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/cafe_interior.jpg filter=lfs diff=lfs merge=lfs -text
+examples/girl_with_dog.jpg filter=lfs diff=lfs merge=lfs -text
+examples/living_room_blue_couch.jpg filter=lfs diff=lfs merge=lfs -text
+examples/motorcycle_street.jpg filter=lfs diff=lfs merge=lfs -text
+examples/pizza_board.jpg filter=lfs diff=lfs merge=lfs -text
+examples/skateboarder_rail.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,47 @@
 ---
-title: Objectmodel V1 Detection
-emoji: 🌍
-colorFrom: yellow
-colorTo: red
+title: ObjectModel-v1 Detection
+emoji: 🔍
+colorFrom: red
+colorTo: green
 sdk: gradio
 sdk_version: 6.24.0
-python_version: '3.12'
 app_file: app.py
-pinned: false
+short_description: NMS-free COCO object detection with ObjectModel-v1
+python_version: "3.12"
+startup_duration_timeout: 30m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ObjectModel-v1 — compact NMS-free object detection
+
+Demo of [`bench-labs/objectmodel-v1`](https://huggingface.co/bench-labs/objectmodel-v1), a
+40.8M-parameter clean-room object detector trained from scratch on COCO 2017
+(AP 0.358 / AP50 0.544 at epoch 95).
+
+The architecture compresses multi-scale pyramid features into a fixed 64-slot latent memory
+for global semantics, then recovers geometry with query-conditioned local sampling whose
+radius scales with each query's current box. Prediction is a fixed set of 300 object
+queries trained with Hungarian matching — no anchors and no NMS.
+
+## What this Space does
+
+- Letterboxes the uploaded image to 640×640 (grey padding, ImageNet normalisation), exactly
+  as the upstream evaluation pipeline does.
+- Runs a single forward pass of the released EMA checkpoint (`objectmodel_v1_best.pt`) on a
+  ZeroGPU worker.
+- Takes the per-query max-sigmoid score, keeps the top-k above the confidence threshold, and
+  maps boxes back into original image coordinates.
+- Renders boxes plus a table of `label / score / x0 y0 x1 y1`.
+
+## Notes
+
+- Small objects are the model's known weak spot (AP_small 0.188 vs AP_large 0.493), and
+  out-of-domain footage degrades faster than COCO-style photography.
+- Because there is no NMS, overlapping duplicate boxes are suppressed by the set-prediction
+  training rather than post-processing; raising the confidence threshold is the intended way
+  to clean up marginal detections.
+
+## Credits
+
+Model and inference code © Bench Labs, Apache-2.0. The `objectmodel_v1/` package in this
+Space is vendored unchanged from the model repository. Example images come from
+[`linoyts/repo-to-space-example-inputs`](https://huggingface.co/datasets/linoyts/repo-to-space-example-inputs).

app.py

@@ -0,0 +1,270 @@
+"""Gradio demo for ObjectModel-v1 — a compact, NMS-free COCO object detector."""
+
+from __future__ import annotations
+
+import os
+import time
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # noqa: E402  — must precede torch / any CUDA-touching import
+
+import numpy as np  # noqa: E402
+import torch  # noqa: E402
+import gradio as gr  # noqa: E402
+from PIL import Image, ImageDraw, ImageFont  # noqa: E402
+from huggingface_hub import hf_hub_download  # noqa: E402
+
+from objectmodel_v1.boxes import box_cxcywh_to_xyxy  # noqa: E402
+from objectmodel_v1.config import load_config  # noqa: E402
+from objectmodel_v1.model import build_model  # noqa: E402
+
+MODEL_ID = "bench-labs/objectmodel-v1"
+CHECKPOINT_FILE = "objectmodel_v1_best.pt"
+CONFIG_FILE = "objectmodel_v1.yaml"
+
+# The training dataloader maps COCO category ids -> contiguous labels by ascending
+# category id, so label i is the i-th COCO category in id order.
+COCO_CLASSES = [
+    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck",
+    "boat", "traffic light", "fire hydrant", "stop sign", "parking meter", "bench",
+    "bird", "cat", "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra",
+    "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
+    "skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove",
+    "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
+    "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange",
+    "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
+    "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
+    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink",
+    "refrigerator", "book", "clock", "vase", "scissors", "teddy bear", "hair drier",
+    "toothbrush",
+]
+
+# ImageNet statistics, matching the training-time preprocessing.
+MEAN = torch.tensor([0.485, 0.456, 0.406], dtype=torch.float32)[:, None, None]
+STD = torch.tensor([0.229, 0.224, 0.225], dtype=torch.float32)[:, None, None]
+
+# ---------------------------------------------------------------------------
+# Model — loaded once at module scope and moved to CUDA eagerly (ZeroGPU rule 2).
+# ---------------------------------------------------------------------------
+_config = load_config(CONFIG_FILE)
+_checkpoint_path = hf_hub_download(MODEL_ID, CHECKPOINT_FILE)
+
+model = build_model(_config)
+_state = torch.load(_checkpoint_path, map_location="cpu", weights_only=False)
+model.load_state_dict(_state.get("ema", _state.get("model", _state)))
+del _state
+model.eval().to("cuda")
+
+INPUT_SIZE = int(model.spec.input_size)
+NUM_PARAMS = sum(p.numel() for p in model.parameters())
+print(f"ObjectModel-v1 loaded: {NUM_PARAMS/1e6:.1f}M params, input {INPUT_SIZE}px", flush=True)
+
+
+def letterbox(image: Image.Image, size: int) -> tuple[torch.Tensor, float, int, int]:
+    """Resize keeping aspect ratio and pad to a square canvas (eval-time transform)."""
+    width, height = image.size
+    ratio = min(size / width, size / height)
+    resized_width = max(1, round(width * ratio))
+    resized_height = max(1, round(height * ratio))
+    resized = image.resize((resized_width, resized_height), Image.Resampling.BILINEAR)
+    offset_x = (size - resized_width) // 2
+    offset_y = (size - resized_height) // 2
+    canvas = Image.new("RGB", (size, size), (114, 114, 114))
+    canvas.paste(resized, (offset_x, offset_y))
+    array = np.asarray(canvas, dtype=np.float32).copy() / 255.0
+    tensor = torch.from_numpy(array).permute(2, 0, 1)
+    tensor = (tensor - MEAN) / STD
+    return tensor, ratio, offset_x, offset_y
+
+
+def palette(index: int) -> tuple[int, int, int]:
+    """Stable, well-spread colour per class index."""
+    hue = (index * 0.6180339887) % 1.0
+    i = int(hue * 6)
+    f = hue * 6 - i
+    q, t = 1 - f, f
+    table = [(1, t, 0), (q, 1, 0), (0, 1, t), (0, q, 1), (t, 0, 1), (1, 0, q)]
+    r, g, b = table[i % 6]
+    return (int(70 + 185 * r), int(70 + 185 * g), int(70 + 185 * b))
+
+
+def _font(size: int):
+    for name in ("DejaVuSans-Bold.ttf", "DejaVuSans.ttf"):
+        try:
+            return ImageFont.truetype(name, size)
+        except Exception:
+            continue
+    return ImageFont.load_default()
+
+
+def draw_detections(image: Image.Image, detections: list[dict]) -> Image.Image:
+    canvas = image.convert("RGB").copy()
+    draw = ImageDraw.Draw(canvas)
+    scale = max(canvas.width, canvas.height) / 640.0
+    thickness = max(2, round(3 * scale))
+    font = _font(max(13, round(16 * scale)))
+    for det in detections:
+        x0, y0, x1, y1 = det["box"]
+        colour = palette(det["label_index"])
+        draw.rectangle([x0, y0, x1, y1], outline=colour, width=thickness)
+        caption = f"{det['label']} {det['score']:.2f}"
+        left, top, right, bottom = draw.textbbox((0, 0), caption, font=font)
+        text_w, text_h = right - left, bottom - top
+        pad = max(2, round(3 * scale))
+        box_h = text_h + 2 * pad
+        text_y = y0 - box_h if y0 - box_h >= 0 else y0
+        draw.rectangle([x0, text_y, x0 + text_w + 2 * pad, text_y + box_h], fill=colour)
+        draw.text((x0 + pad - left, text_y + pad - top), caption, fill=(20, 20, 20), font=font)
+    return canvas
+
+
+@spaces.GPU(duration=45)
+def detect(
+    image: Image.Image,
+    confidence: float = 0.35,
+    max_detections: int = 100,
+) -> tuple[Image.Image, list[list], str]:
+    """Detect COCO objects in an image with ObjectModel-v1.
+
+    Args:
+        image: input photograph to run detection on.
+        confidence: minimum score (0-1) a detection must reach to be kept.
+        max_detections: hard cap on how many boxes are returned.
+
+    Returns:
+        The image with boxes drawn, a table of detections, and a short summary.
+    """
+    if image is None:
+        raise gr.Error("Please upload an image first.")
+
+    image = image.convert("RGB")
+    original_width, original_height = image.size
+    tensor, ratio, offset_x, offset_y = letterbox(image, INPUT_SIZE)
+
+    started = time.perf_counter()
+    with torch.inference_mode():
+        outputs = model(tensor[None].to("cuda"))
+        probabilities = outputs["pred_logits"].sigmoid()[0]
+        boxes = box_cxcywh_to_xyxy(outputs["pred_boxes"][0]).clamp(0.0, 1.0) * INPUT_SIZE
+        scores, labels = probabilities.max(dim=-1)
+        count = min(int(max_detections), scores.numel())
+        scores, indices = scores.topk(count)
+        keep = scores >= float(confidence)
+        scores, indices = scores[keep], indices[keep]
+        labels = labels[indices]
+        boxes = boxes[indices]
+        # Undo the letterbox transform back into original image coordinates.
+        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - offset_x) / ratio
+        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - offset_y) / ratio
+        boxes[:, [0, 2]] = boxes[:, [0, 2]].clamp(0, original_width)
+        boxes[:, [1, 3]] = boxes[:, [1, 3]].clamp(0, original_height)
+        scores = scores.float().cpu().tolist()
+        labels = labels.cpu().tolist()
+        boxes = boxes.float().cpu().tolist()
+    elapsed = time.perf_counter() - started
+
+    detections = [
+        {
+            "label": COCO_CLASSES[label] if label < len(COCO_CLASSES) else str(label),
+            "label_index": int(label),
+            "score": float(score),
+            "box": [float(v) for v in box],
+        }
+        for score, label, box in zip(scores, labels, boxes)
+    ]
+
+    annotated = draw_detections(image, detections)
+    table = [
+        [
+            det["label"],
+            round(det["score"], 3),
+            round(det["box"][0]),
+            round(det["box"][1]),
+            round(det["box"][2]),
+            round(det["box"][3]),
+        ]
+        for det in detections
+    ]
+    if not table:
+        table = [["(nothing above threshold)", 0.0, 0, 0, 0, 0]]
+    summary = (
+        f"**{len(detections)} object(s)** above {float(confidence):.2f} confidence · "
+        f"forward pass {elapsed * 1000:.0f} ms · NMS-free (300 queries)"
+    )
+    return annotated, table, summary
+
+
+CSS = """
+#col-container { max-width: 1200px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="ObjectModel-v1 Detection") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """
+            # ObjectModel-v1 — compact NMS-free object detection
+
+            A 40.8M-parameter clean-room detector from Bench Labs: compressed global latent
+            memory + box-scaled local pyramid sampling, trained from scratch on COCO
+            (AP 0.358). No anchors, no NMS — a fixed set of 300 queries.
+
+            [Model card](https://huggingface.co/bench-labs/objectmodel-v1)
+            """
+        )
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_input = gr.Image(label="Input image", type="pil", height=380)
+                confidence = gr.Slider(
+                    label="Confidence threshold",
+                    minimum=0.05,
+                    maximum=0.95,
+                    step=0.01,
+                    value=0.35,
+                )
+                run_button = gr.Button("Detect objects", variant="primary")
+                with gr.Accordion("Advanced settings", open=False):
+                    max_detections = gr.Slider(
+                        label="Max detections",
+                        minimum=1,
+                        maximum=300,
+                        step=1,
+                        value=100,
+                        info="Top-k queries kept before the confidence filter.",
+                    )
+            with gr.Column(scale=1):
+                image_output = gr.Image(label="Detections", type="pil", height=380)
+                summary_output = gr.Markdown()
+                detections_output = gr.Dataframe(
+                    headers=["label", "score", "x0", "y0", "x1", "y1"],
+                    label="Detections",
+                    wrap=True,
+                    row_count=(1, "dynamic"),
+                )
+
+        gr.Examples(
+            examples=[
+                ["examples/skateboarder_rail.jpg"],
+                ["examples/motorcycle_street.jpg"],
+                ["examples/cafe_interior.jpg"],
+                ["examples/living_room_blue_couch.jpg"],
+                ["examples/girl_with_dog.jpg"],
+                ["examples/pizza_board.jpg"],
+            ],
+            inputs=[image_input],
+            outputs=[image_output, detections_output, summary_output],
+            fn=detect,
+            cache_examples=True,
+            cache_mode="lazy",
+        )
+
+    run_button.click(
+        fn=detect,
+        inputs=[image_input, confidence, max_detections],
+        outputs=[image_output, detections_output, summary_output],
+        api_name="detect",
+    )
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

objectmodel_v1.yaml

@@ -0,0 +1,61 @@
+model:
+  num_classes: 80
+  input_size: 640
+  stem_channels: 64
+  backbone_channels: [96, 192, 384, 512]
+  backbone_depths: [2, 3, 6, 3]
+  hidden_dim: 384
+  fpn_depth: 2
+  latent_count: 64
+  latent_pool_sizes: [12, 6, 3]
+  latent_layers: 2
+  decoder_layers: 8
+  num_queries: 300
+  num_heads: 8
+  local_points: 4
+  dropout: 0.0
+  dense_aux: true
+
+loss:
+  cost_class: 2.0
+  cost_bbox: 5.0
+  cost_giou: 2.0
+  weight_class: 2.0
+  weight_bbox: 5.0
+  weight_giou: 2.0
+  weight_dense: 1.0
+  focal_alpha: 0.25
+  focal_gamma: 2.0
+  aux_weight: 1.0
+  dense_topk: 5
+
+train:
+  epochs: 150
+  batch_size: 16
+  eval_batch_size: 8
+  workers: 8
+  prefetch_factor: 4
+  lr: 0.0002
+  backbone_lr: 0.0001
+  min_lr_ratio: 0.05
+  weight_decay: 0.05
+  warmup_steps: 1500
+  clip_grad_norm: 0.1
+  amp: true
+  amp_dtype: bfloat16
+  channels_last: false
+  compile: false
+  ema_decay: 0.9998
+  seed: 42
+  eval_every: 1
+  print_freq: 50
+
+data:
+  train_image_dir: train2017
+  train_annotations: annotations/instances_train2017.json
+  val_image_dir: val2017
+  val_annotations: annotations/instances_val2017.json
+  hflip_prob: 0.5
+  scale_range: [0.65, 1.0]
+  mean: [0.485, 0.456, 0.406]
+  std: [0.229, 0.224, 0.225]

objectmodel_v1/__init__.py

@@ -0,0 +1,6 @@
+"""ObjectModel-v1 compact object detection research package."""
+
+from .model import ObjectModelV1, ObjectModelV1Spec, build_model
+
+__all__ = ["ObjectModelV1", "ObjectModelV1Spec", "build_model"]
+__version__ = "0.1.0"

objectmodel_v1/boxes.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+
+def box_cxcywh_to_xyxy(boxes: Tensor) -> Tensor:
+    cx, cy, width, height = boxes.unbind(-1)
+    return torch.stack(
+        (cx - 0.5 * width, cy - 0.5 * height, cx + 0.5 * width, cy + 0.5 * height),
+        dim=-1,
+    )
+
+
+def box_xyxy_to_cxcywh(boxes: Tensor) -> Tensor:
+    x0, y0, x1, y1 = boxes.unbind(-1)
+    return torch.stack(
+        ((x0 + x1) * 0.5, (y0 + y1) * 0.5, x1 - x0, y1 - y0), 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]:
+    area1 = box_area(boxes1)
+    area2 = box_area(boxes2)
+    top_left = torch.maximum(boxes1[:, None, :2], boxes2[:, :2])
+    bottom_right = torch.minimum(boxes1[:, None, 2:], boxes2[:, 2:])
+    intersection = (bottom_right - top_left).clamp(min=0).prod(dim=-1)
+    union = area1[:, None] + area2 - intersection
+    return intersection / union.clamp(min=1e-7), union
+
+
+def generalized_box_iou(boxes1: Tensor, boxes2: Tensor) -> Tensor:
+    """Pairwise generalized IoU for boxes in x0, y0, x1, y1 format."""
+    iou, union = box_iou(boxes1, boxes2)
+    top_left = torch.minimum(boxes1[:, None, :2], boxes2[:, :2])
+    bottom_right = torch.maximum(boxes1[:, None, 2:], boxes2[:, 2:])
+    enclosing = (bottom_right - top_left).clamp(min=0).prod(dim=-1)
+    return iou - (enclosing - union) / enclosing.clamp(min=1e-7)
+
+
+def generalized_box_iou_batched(boxes1: Tensor, boxes2: Tensor) -> Tensor:
+    """Batched pairwise generalized IoU, boxes in x0, y0, x1, y1 format.
+
+    boxes1: [B, N, 4], boxes2: [B, M, 4] -> [B, N, M]
+    """
+    area1 = box_area(boxes1)
+    area2 = box_area(boxes2)
+    top_left = torch.maximum(boxes1[:, :, None, :2], boxes2[:, None, :, :2])
+    bottom_right = torch.minimum(boxes1[:, :, None, 2:], boxes2[:, None, :, 2:])
+    intersection = (bottom_right - top_left).clamp(min=0).prod(dim=-1)
+    union = area1[:, :, None] + area2[:, None, :] - intersection
+    iou = intersection / union.clamp(min=1e-7)
+    enc_top_left = torch.minimum(boxes1[:, :, None, :2], boxes2[:, None, :, :2])
+    enc_bottom_right = torch.maximum(boxes1[:, :, None, 2:], boxes2[:, None, :, 2:])
+    enclosing = (enc_bottom_right - enc_top_left).clamp(min=0).prod(dim=-1)
+    return iou - (enclosing - union) / enclosing.clamp(min=1e-7)
+
+
+def inverse_sigmoid(value: Tensor, eps: float = 1e-5) -> Tensor:
+    value = value.clamp(min=0.0, max=1.0)
+    return torch.log(value.clamp(min=eps) / (1.0 - value).clamp(min=eps))

objectmodel_v1/config.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+from copy import deepcopy
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+def load_config(path: str | Path) -> dict[str, Any]:
+    with Path(path).open("r", encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    if not isinstance(config, dict):
+        raise ValueError(f"Configuration must be a mapping: {path}")
+    return config
+
+
+def apply_overrides(config: dict[str, Any], overrides: list[str]) -> dict[str, Any]:
+    result = deepcopy(config)
+    for item in overrides:
+        if "=" not in item:
+            raise ValueError(f"Override must have key=value form: {item}")
+        dotted_key, raw_value = item.split("=", 1)
+        keys = dotted_key.split(".")
+        node = result
+        for key in keys[:-1]:
+            if key not in node or not isinstance(node[key], dict):
+                raise KeyError(f"Unknown configuration path: {dotted_key}")
+            node = node[key]
+        if keys[-1] not in node:
+            raise KeyError(f"Unknown configuration key: {dotted_key}")
+        node[keys[-1]] = yaml.safe_load(raw_value)
+    return result
+
+
+def save_config(config: dict[str, Any], path: str | Path) -> None:
+    with Path(path).open("w", encoding="utf-8") as handle:
+        yaml.safe_dump(config, handle, sort_keys=False)

objectmodel_v1/model.py

@@ -0,0 +1,413 @@
+from __future__ import annotations
+
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+from .boxes import inverse_sigmoid
+
+
+class ConvNormAct(nn.Sequential):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int = 1,
+        stride: int = 1,
+        groups: int = 1,
+        activation: bool = True,
+    ) -> None:
+        padding = kernel_size // 2
+        layers: list[nn.Module] = [
+            nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                padding,
+                groups=groups,
+                bias=False,
+            ),
+            nn.BatchNorm2d(out_channels),
+        ]
+        if activation:
+            layers.append(nn.SiLU(inplace=True))
+        super().__init__(*layers)
+
+
+class GatedConvBlock(nn.Module):
+    """Inverted residual block with a cheap learned residual gate."""
+
+    def __init__(self, channels: int, expansion: float = 2.0) -> None:
+        super().__init__()
+        hidden = int(channels * expansion)
+        self.expand = ConvNormAct(channels, hidden)
+        self.depthwise = ConvNormAct(hidden, hidden, 3, groups=hidden)
+        self.project = ConvNormAct(hidden, channels, activation=False)
+        self.gate = nn.Parameter(torch.zeros(1))
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        return inputs + torch.tanh(self.gate) * self.project(self.depthwise(self.expand(inputs)))
+
+
+class BackboneStage(nn.Sequential):
+    def __init__(self, in_channels: int, out_channels: int, depth: int, stride: int) -> None:
+        super().__init__(
+            ConvNormAct(in_channels, out_channels, 3, stride=stride),
+            *(GatedConvBlock(out_channels) for _ in range(depth)),
+        )
+
+
+class CompactBackbone(nn.Module):
+    def __init__(
+        self, stem_channels: int, channels: list[int], depths: list[int]
+    ) -> None:
+        super().__init__()
+        if len(channels) != 4 or len(depths) != 4:
+            raise ValueError("Backbone requires four channel and depth values")
+        self.stem = nn.Sequential(
+            ConvNormAct(3, stem_channels, 3, stride=2),
+            ConvNormAct(stem_channels, stem_channels, 3, stride=2),
+        )
+        stages: list[nn.Module] = []
+        in_channels = stem_channels
+        for index, (out_channels, depth) in enumerate(zip(channels, depths, strict=True)):
+            stages.append(
+                BackboneStage(in_channels, out_channels, depth, stride=1 if index == 0 else 2)
+            )
+            in_channels = out_channels
+        self.stages = nn.ModuleList(stages)
+        self.out_channels = channels[1:]
+
+    def forward(self, images: Tensor) -> list[Tensor]:
+        features = self.stem(images)
+        outputs = []
+        for index, stage in enumerate(self.stages):
+            features = stage(features)
+            if index > 0:
+                outputs.append(features)
+        return outputs
+
+
+class PyramidFusion(nn.Module):
+    def __init__(self, in_channels: list[int], hidden_dim: int, depth: int) -> None:
+        super().__init__()
+        self.lateral = nn.ModuleList(ConvNormAct(c, hidden_dim) for c in in_channels)
+        self.refine = nn.ModuleList(
+            nn.Sequential(*(GatedConvBlock(hidden_dim, expansion=1.5) for _ in range(depth)))
+            for _ in in_channels
+        )
+
+    def forward(self, inputs: list[Tensor]) -> list[Tensor]:
+        projected = [layer(x) for layer, x in zip(self.lateral, inputs, strict=True)]
+        outputs = list(projected)
+        for index in range(len(outputs) - 2, -1, -1):
+            outputs[index] = outputs[index] + F.interpolate(
+                outputs[index + 1], size=outputs[index].shape[-2:], mode="nearest"
+            )
+        return [block(x) for block, x in zip(self.refine, outputs, strict=True)]
+
+
+def sine_position_encoding(
+    height: int, width: int, dim: int, device: torch.device, dtype: torch.dtype
+) -> Tensor:
+    if dim % 4 != 0:
+        raise ValueError("Position encoding dimension must be divisible by four")
+    y, x = torch.meshgrid(
+        torch.linspace(0, 1, height, device=device, dtype=dtype),
+        torch.linspace(0, 1, width, device=device, dtype=dtype),
+        indexing="ij",
+    )
+    frequencies = torch.arange(dim // 4, device=device, dtype=dtype)
+    frequencies = 2.0 * torch.pi * (10000.0 ** (-frequencies / max(dim // 4, 1)))
+    x = x.flatten()[:, None] * frequencies[None]
+    y = y.flatten()[:, None] * frequencies[None]
+    return torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=-1)
+
+
+class FeedForward(nn.Sequential):
+    def __init__(self, dim: int, expansion: int = 4, dropout: float = 0.0) -> None:
+        super().__init__(
+            nn.Linear(dim, dim * expansion),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(dim * expansion, dim),
+            nn.Dropout(dropout),
+        )
+
+
+class LatentLayer(nn.Module):
+    def __init__(self, dim: int, num_heads: int, dropout: float) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attention = nn.MultiheadAttention(dim, num_heads, dropout, batch_first=True)
+        self.norm2 = nn.LayerNorm(dim)
+        self.ffn = FeedForward(dim, dropout=dropout)
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        normalized = self.norm1(inputs)
+        inputs = inputs + self.attention(normalized, normalized, normalized, need_weights=False)[0]
+        return inputs + self.ffn(self.norm2(inputs))
+
+
+class LatentMemory(nn.Module):
+    """Compresses multi-scale maps into a fixed-size global reasoning memory."""
+
+    def __init__(
+        self,
+        dim: int,
+        latent_count: int,
+        pool_sizes: list[int],
+        layers: int,
+        num_heads: int,
+        dropout: float,
+    ) -> None:
+        super().__init__()
+        if len(pool_sizes) != 3:
+            raise ValueError("One latent pool size is required for each pyramid level")
+        self.pool_sizes = pool_sizes
+        self.latents = nn.Parameter(torch.empty(latent_count, dim))
+        self.level_embedding = nn.Parameter(torch.empty(len(pool_sizes), dim))
+        self.query_norm = nn.LayerNorm(dim)
+        self.token_norm = nn.LayerNorm(dim)
+        self.compress = nn.MultiheadAttention(dim, num_heads, dropout, batch_first=True)
+        self.layers = nn.ModuleList(LatentLayer(dim, num_heads, dropout) for _ in range(layers))
+        nn.init.normal_(self.latents, std=0.02)
+        nn.init.normal_(self.level_embedding, std=0.02)
+
+    def forward(self, features: list[Tensor]) -> Tensor:
+        tokens = []
+        for level, (feature, size) in enumerate(zip(features, self.pool_sizes, strict=True)):
+            pooled = F.adaptive_avg_pool2d(feature, (size, size)).flatten(2).transpose(1, 2)
+            position = sine_position_encoding(
+                size, size, feature.shape[1], feature.device, feature.dtype
+            )
+            tokens.append(pooled + position[None] + self.level_embedding[level][None, None])
+        token_memory = self.token_norm(torch.cat(tokens, dim=1))
+        latents = self.latents[None].expand(features[0].shape[0], -1, -1)
+        latents = latents + self.compress(
+            self.query_norm(latents), token_memory, token_memory, need_weights=False
+        )[0]
+        for layer in self.layers:
+            latents = layer(latents)
+        return latents
+
+
+class QueryLocalSampler(nn.Module):
+    """Samples high-resolution pyramid evidence around each evolving query box."""
+
+    def __init__(self, dim: int, num_levels: int, points: int) -> None:
+        super().__init__()
+        self.num_levels = num_levels
+        self.points = points
+        self.offsets = nn.Linear(dim, num_levels * points * 2)
+        self.weights = nn.Linear(dim, num_levels * points)
+        self.output = nn.Linear(dim, dim)
+        nn.init.zeros_(self.offsets.weight)
+        nn.init.zeros_(self.offsets.bias)
+        nn.init.zeros_(self.weights.weight)
+        nn.init.zeros_(self.weights.bias)
+
+    def forward(self, queries: Tensor, boxes: Tensor, features: list[Tensor]) -> Tensor:
+        batch, query_count, _ = queries.shape
+        offsets = self.offsets(queries).view(
+            batch, query_count, self.num_levels, self.points, 2
+        )
+        offsets = offsets.tanh() * boxes[..., None, None, 2:] * 0.5
+        centers = boxes[..., None, None, :2]
+        sample_points = (centers + offsets).clamp(0.0, 1.0)
+        weights = self.weights(queries).view(
+            batch, query_count, self.num_levels * self.points
+        )
+        weights = weights.softmax(dim=-1).view(
+            batch, query_count, self.num_levels, self.points
+        )
+
+        sampled_levels = []
+        for level, feature in enumerate(features):
+            grid = sample_points[:, :, level] * 2.0 - 1.0
+            sampled = F.grid_sample(
+                feature,
+                grid,
+                mode="bilinear",
+                padding_mode="zeros",
+                align_corners=False,
+            )
+            sampled = sampled.permute(0, 2, 3, 1)
+            sampled_levels.append(sampled)
+        sampled_features = torch.stack(sampled_levels, dim=2)
+        fused = (sampled_features * weights[..., None]).sum(dim=(2, 3))
+        return self.output(fused)
+
+
+class DecoderLayer(nn.Module):
+    def __init__(
+        self, dim: int, num_heads: int, num_levels: int, local_points: int, dropout: float
+    ) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.self_attention = nn.MultiheadAttention(dim, num_heads, dropout, batch_first=True)
+        self.norm2 = nn.LayerNorm(dim)
+        self.global_attention = nn.MultiheadAttention(dim, num_heads, dropout, batch_first=True)
+        self.norm3 = nn.LayerNorm(dim)
+        self.local_sampler = QueryLocalSampler(dim, num_levels, local_points)
+        self.norm4 = nn.LayerNorm(dim)
+        self.ffn = FeedForward(dim, dropout=dropout)
+
+    def forward(
+        self, queries: Tensor, memory: Tensor, boxes: Tensor, features: list[Tensor]
+    ) -> Tensor:
+        normalized = self.norm1(queries)
+        queries = queries + self.self_attention(
+            normalized, normalized, normalized, need_weights=False
+        )[0]
+        queries = queries + self.global_attention(
+            self.norm2(queries), memory, memory, need_weights=False
+        )[0]
+        queries = queries + self.local_sampler(self.norm3(queries), boxes, features)
+        return queries + self.ffn(self.norm4(queries))
+
+
+class MLP(nn.Sequential):
+    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, layers: int) -> None:
+        modules: list[nn.Module] = []
+        for index in range(layers):
+            in_dim = input_dim if index == 0 else hidden_dim
+            out_dim = output_dim if index == layers - 1 else hidden_dim
+            modules.append(nn.Linear(in_dim, out_dim))
+            if index < layers - 1:
+                modules.append(nn.ReLU(inplace=True))
+        super().__init__(*modules)
+
+
+class DenseAuxiliaryHead(nn.Module):
+    def __init__(self, dim: int, num_classes: int) -> None:
+        super().__init__()
+        self.shared = nn.ModuleList(
+            nn.Sequential(ConvNormAct(dim, dim, 3, groups=dim), ConvNormAct(dim, dim))
+            for _ in range(3)
+        )
+        self.classification = nn.Conv2d(dim, num_classes, 1)
+        self.regression = nn.Conv2d(dim, 4, 1)
+
+    def forward(self, features: list[Tensor]) -> list[dict[str, Tensor]]:
+        outputs = []
+        for feature, tower in zip(features, self.shared, strict=True):
+            hidden = tower(feature)
+            outputs.append(
+                {
+                    "logits": self.classification(hidden),
+                    "distances": F.softplus(self.regression(hidden)),
+                }
+            )
+        return outputs
+
+
+@dataclass(frozen=True)
+class ObjectModelV1Spec:
+    num_classes: int = 80
+    input_size: int = 640
+    stem_channels: int = 48
+    backbone_channels: tuple[int, int, int, int] = (64, 128, 256, 384)
+    backbone_depths: tuple[int, int, int, int] = (2, 3, 6, 3)
+    hidden_dim: int = 256
+    fpn_depth: int = 2
+    latent_count: int = 64
+    latent_pool_sizes: tuple[int, int, int] = (12, 6, 3)
+    latent_layers: int = 2
+    decoder_layers: int = 6
+    num_queries: int = 300
+    num_heads: int = 8
+    local_points: int = 4
+    dropout: float = 0.0
+    dense_aux: bool = True
+
+
+class ObjectModelV1(nn.Module):
+    """NMS-free detector with compressed global memory and local geometric sampling."""
+
+    def __init__(self, spec: ObjectModelV1Spec) -> None:
+        super().__init__()
+        self.spec = spec
+        self.backbone = CompactBackbone(
+            spec.stem_channels, list(spec.backbone_channels), list(spec.backbone_depths)
+        )
+        self.neck = PyramidFusion(self.backbone.out_channels, spec.hidden_dim, spec.fpn_depth)
+        self.memory = LatentMemory(
+            spec.hidden_dim,
+            spec.latent_count,
+            list(spec.latent_pool_sizes),
+            spec.latent_layers,
+            spec.num_heads,
+            spec.dropout,
+        )
+        decoder_template = DecoderLayer(
+            spec.hidden_dim, spec.num_heads, 3, spec.local_points, spec.dropout
+        )
+        self.decoder = nn.ModuleList(deepcopy(decoder_template) for _ in range(spec.decoder_layers))
+        self.query_embedding = nn.Embedding(spec.num_queries, spec.hidden_dim)
+        self.reference_points = nn.Embedding(spec.num_queries, 4)
+        self.class_heads = nn.ModuleList(
+            nn.Linear(spec.hidden_dim, spec.num_classes) for _ in range(spec.decoder_layers)
+        )
+        self.box_heads = nn.ModuleList(
+            MLP(spec.hidden_dim, spec.hidden_dim, 4, 3) for _ in range(spec.decoder_layers)
+        )
+        self.dense_head = (
+            DenseAuxiliaryHead(spec.hidden_dim, spec.num_classes) if spec.dense_aux else None
+        )
+        self._reset_parameters()
+
+    def _reset_parameters(self) -> None:
+        prior_probability = 0.01
+        class_bias = -torch.log(torch.tensor((1.0 - prior_probability) / prior_probability))
+        for head in self.class_heads:
+            nn.init.constant_(head.bias, class_bias)
+        nn.init.zeros_(self.reference_points.weight)
+        with torch.no_grad():
+            self.reference_points.weight[:, 2:] = -2.0
+        for head in self.box_heads:
+            nn.init.zeros_(head[-1].weight)
+            nn.init.zeros_(head[-1].bias)
+        if self.dense_head is not None:
+            nn.init.constant_(self.dense_head.classification.bias, class_bias)
+            nn.init.zeros_(self.dense_head.regression.weight)
+            nn.init.constant_(self.dense_head.regression.bias, 1.0)
+
+    def forward(self, images: Tensor) -> dict[str, Any]:
+        features = self.neck(self.backbone(images))
+        memory = self.memory(features)
+        batch = images.shape[0]
+        queries = self.query_embedding.weight[None].expand(batch, -1, -1)
+        boxes = self.reference_points.weight.sigmoid()[None].expand(batch, -1, -1)
+        layer_outputs: list[dict[str, Tensor]] = []
+        for layer, class_head, box_head in zip(
+            self.decoder, self.class_heads, self.box_heads, strict=True
+        ):
+            queries = layer(queries, memory, boxes, features)
+            boxes = (inverse_sigmoid(boxes) + box_head(queries)).sigmoid()
+            layer_outputs.append({"pred_logits": class_head(queries), "pred_boxes": boxes})
+            boxes = boxes.detach() if self.training else boxes
+
+        output: dict[str, Any] = dict(layer_outputs[-1])
+        output["aux_outputs"] = layer_outputs[:-1]
+        if self.training and self.dense_head is not None:
+            output["dense_outputs"] = self.dense_head(features)
+        return output
+
+
+def build_model(config: dict[str, Any]) -> ObjectModelV1:
+    model_config = config.get("model", config)
+    fields = ObjectModelV1Spec.__dataclass_fields__
+    unknown = set(model_config) - set(fields)
+    if unknown:
+        raise ValueError(f"Unknown model configuration keys: {sorted(unknown)}")
+    values = dict(model_config)
+    for key in ("backbone_channels", "backbone_depths", "latent_pool_sizes"):
+        if key in values:
+            values[key] = tuple(values[key])
+    return ObjectModelV1(ObjectModelV1Spec(**values))

objectmodel_v1/postprocess.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+from .boxes import box_cxcywh_to_xyxy
+
+
+@torch.no_grad()
+def decode_predictions(
+    outputs: dict[str, Tensor],
+    image_sizes: list[tuple[int, int]],
+    confidence: float = 0.25,
+    top_k: int = 300,
+) -> list[dict[str, Tensor]]:
+    logits = outputs["pred_logits"].sigmoid()
+    boxes = box_cxcywh_to_xyxy(outputs["pred_boxes"]).clamp(0.0, 1.0)
+    results = []
+    for index, (height, width) in enumerate(image_sizes):
+        scores, labels = logits[index].max(dim=-1)
+        keep = scores >= confidence
+        if keep.sum() > top_k:
+            selected = scores.masked_fill(~keep, -1).topk(top_k).indices
+        else:
+            selected = torch.where(keep)[0]
+        selected_boxes = boxes[index, selected].clone()
+        selected_boxes[:, [0, 2]] *= width
+        selected_boxes[:, [1, 3]] *= height
+        results.append(
+            {"scores": scores[selected], "labels": labels[selected], "boxes": selected_boxes}
+        )
+    return results

requirements.txt

@@ -0,0 +1,3 @@
+numpy
+pillow
+pyyaml