CandyEye 0.1.0__py3-none-any.whl

candyeye/__init__.py

@@ -0,0 +1,57 @@
+"""CandyEye — a lightweight, CPU-first object detector.
+
+Public API::
+
+    from candyeye import CandyEye, train
+
+    model = CandyEye()                     # bundled yolo11n architecture
+    results = model.train(data="dataset.yaml", epochs=100, pretrained=True)
+
+    # ...or the function form:
+    results = train(data="dataset.yaml", epochs=100)
+
+Heavy submodules (torch, torchvision) are imported lazily so ``import
+candyeye`` stays cheap.
+"""
+from __future__ import annotations
+
+from candyeye.paths import (
+    available_configs,
+    cache_dir,
+    default_config_path,
+    default_weights_path,
+    download_default_weights,
+    resolve_config,
+)
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "CandyEye",
+    "CandyEyePredictor",
+    "train",
+    "predict",
+    "resolve_config",
+    "default_config_path",
+    "default_weights_path",
+    "download_default_weights",
+    "available_configs",
+    "cache_dir",
+    "__version__",
+]
+
+
+def __getattr__(name):
+    if name == "CandyEye":
+        from candyeye.core.candyeye import CandyEye
+        return CandyEye
+    if name == "train":
+        from candyeye.training.trainer import train
+        return train
+    if name == "predict":
+        from candyeye.inference.predict import predict
+        return predict
+    if name == "CandyEyePredictor":
+        from candyeye.inference.predict import CandyEyePredictor
+        return CandyEyePredictor
+    raise AttributeError(f"module 'candyeye' has no attribute {name!r}")

candyeye/configs/yolo11.yaml

@@ -0,0 +1,37 @@
+# CandyEye — YOLO11n transcription, object detection (nc=20)
+#
+# Entry format: [from, repeats, module, args]
+#   from < 0  -> offset into the built layer list (-1 = previous layer)
+#   from >= 0 -> absolute layer index
+#   from list -> gather several layers (Concat / Detect)
+# Channel args are in BASE units; the builder scales them by width_mult.
+nc: 20
+scale: n
+
+backbone:
+  - [-1, 1, Conv, [64, 3, 2]]          # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]]         # 1-P2/4
+  - [-1, 2, C3k2, [256, False, 0.25]]  # 2
+  - [-1, 1, Conv, [256, 3, 2]]         # 3-P3/8
+  - [-1, 2, C3k2, [512, False, 0.25]]  # 4
+  - [-1, 1, Conv, [512, 3, 2]]         # 5-P4/16
+  - [-1, 2, C3k2, [512, True]]         # 6
+  - [-1, 1, Conv, [1024, 3, 2]]        # 7-P5/32
+  - [-1, 2, C3k2, [1024, True]]        # 8
+  - [-1, 1, SPPF, [1024, 5]]           # 9
+  - [-1, 2, C2PSA, [1024]]             # 10
+
+head:
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]   # 11
+  - [[-1, 6], 1, Concat, [1]]                     # 12 cat backbone P4
+  - [-1, 1, C3k2, [512, False]]                   # 13
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]    # 14
+  - [[-1, 4], 1, Concat, [1]]                     # 15 cat backbone P3
+  - [-1, 1, C3k2, [256, False]]                   # 16 P3/8 (small)
+  - [-1, 1, Conv, [256, 3, 2]]                    # 17
+  - [[-1, 13], 1, Concat, [1]]                    # 18 cat head P4
+  - [-1, 1, C3k2, [512, False]]                   # 19 P4/16 (medium)
+  - [-1, 1, Conv, [512, 3, 2]]                    # 20
+  - [[-1, 10], 1, Concat, [1]]                    # 21 cat head P5
+  - [-1, 1, C3k2, [1024, True]]                   # 22 P5/32 (large)
+  - [[16, 19, 22], 1, Detect, [nc]]               # 23 box+cls, parallel

candyeye/configs/yolo11_exchange.yaml

@@ -0,0 +1,38 @@
+# CandyEye — YOLO11n transcription with the adaptive cross-scale exchange neck.
+#
+# Same backbone/neck as configs/yolo11.yaml; layer 23 is a ScaleExchange that
+# lets P3/P4/P5 exchange gated messages before the Detect head. The gate mode is
+# the last arg: "none" | "static" | "dynamic" (optional second arg = iterations).
+#
+# Entry format: [from, repeats, module, args]
+nc: 20
+scale: n
+
+backbone:
+  - [-1, 1, Conv, [64, 3, 2]]          # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]]         # 1-P2/4
+  - [-1, 2, C3k2, [256, False, 0.25]]  # 2
+  - [-1, 1, Conv, [256, 3, 2]]         # 3-P3/8
+  - [-1, 2, C3k2, [512, False, 0.25]]  # 4
+  - [-1, 1, Conv, [512, 3, 2]]         # 5-P4/16
+  - [-1, 2, C3k2, [512, True]]         # 6
+  - [-1, 1, Conv, [1024, 3, 2]]        # 7-P5/32
+  - [-1, 2, C3k2, [1024, True]]        # 8
+  - [-1, 1, SPPF, [1024, 5]]           # 9
+  - [-1, 2, C2PSA, [1024]]             # 10
+
+head:
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]   # 11
+  - [[-1, 6], 1, Concat, [1]]                     # 12 cat backbone P4
+  - [-1, 1, C3k2, [512, False]]                   # 13
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]    # 14
+  - [[-1, 4], 1, Concat, [1]]                     # 15 cat backbone P3
+  - [-1, 1, C3k2, [256, False]]                   # 16 P3/8 (small)
+  - [-1, 1, Conv, [256, 3, 2]]                    # 17
+  - [[-1, 13], 1, Concat, [1]]                    # 18 cat head P4
+  - [-1, 1, C3k2, [512, False]]                   # 19 P4/16 (medium)
+  - [-1, 1, Conv, [512, 3, 2]]                    # 20
+  - [[-1, 10], 1, Concat, [1]]                    # 21 cat head P5
+  - [-1, 1, C3k2, [1024, True]]                   # 22 P5/32 (large)
+  - [[16, 19, 22], 1, ScaleExchange, ["dynamic"]] # 23 gated cross-scale exchange
+  - [-1, 1, Detect, [nc]]                         # 24 box+cls, parallel

candyeye/core/__init__.py

@@ -0,0 +1,13 @@
+"""Public CandyEye model and training entry points, loaded on demand."""
+
+__all__ = ["CandyEye", "train"]
+
+
+def __getattr__(name):
+    if name == "CandyEye":
+        from candyeye.core.candyeye import CandyEye
+        return CandyEye
+    if name == "train":
+        from candyeye.training.trainer import train
+        return train
+    raise AttributeError(name)

candyeye/core/backbone_mobilenet.py

@@ -0,0 +1,82 @@
+"""MobileNetV3-Small backbone with a feature pyramid and detection head."""
+from __future__ import annotations
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torchvision.models import MobileNet_V3_Small_Weights, mobilenet_v3_small
+
+from candyeye.core.modules.blocks import C3k2
+from candyeye.core.modules.conv import Conv
+from candyeye.core.modules.detect import Detect
+from candyeye.core.modules.exchange import ScaleExchange
+
+
+class MobileNetV3SmallDetector(nn.Module):
+    """Small VOC detector using MobileNetV3-Small features at strides 8/16/32.
+
+    ImageNet initialization is optional because torchvision may need to
+    download the checkpoint when it is not in the local torch cache.
+
+    ``neck="light"`` (default) keeps the compact custom FPN. ``neck="exchange"``
+    appends the gated cross-scale exchange neck before the Detect head.
+    """
+    def __init__(self, nc: int = 20, img_size: int = 128,
+                 pretrained: bool = False, neck: str = "light",
+                 exchange_gate: str = "none", exchange_iters: int = 1):
+        super().__init__()
+        if img_size % 32:
+            raise ValueError(f"img_size must be divisible by 32, got {img_size}")
+        if neck not in ("light", "exchange"):
+            raise ValueError(f"unknown neck {neck!r} (have 'light', 'exchange')")
+        weights = MobileNet_V3_Small_Weights.DEFAULT if pretrained else None
+        backbone = mobilenet_v3_small(weights=weights)
+        self.backbone = backbone.features
+
+        # MobileNet feature indices 3, 8, and 12 are at /8, /16, and /32.
+        self.proj3 = Conv(24, 64, 1)
+        self.proj4 = Conv(48, 128, 1)
+        self.proj5 = Conv(576, 256, 1)
+        self.fuse4 = C3k2(384, 128, n=1, c3k=False)
+        self.fuse3 = C3k2(192, 64, n=1, c3k=False)
+        self.down4 = Conv(64, 64, 3, 2)
+        self.pan4 = C3k2(192, 128, n=1, c3k=False)
+        self.down5 = Conv(128, 128, 3, 2)
+        self.pan5 = C3k2(384, 256, n=1, c3k=False)
+
+        self.neck = neck
+        self.exchange = (
+            ScaleExchange(64, 128, 256, gate=exchange_gate, iters=exchange_iters)
+            if neck == "exchange" else None
+        )
+
+        detect = Detect(nc=nc, ch=(64, 128, 256))
+        detect.stride = torch.tensor([8., 16., 32.])
+        detect.bias_init()
+        self.model = nn.ModuleList([detect])
+        self.nc = nc
+        self.img_size = img_size
+        self.register_buffer("stride", torch.tensor([8., 16., 32.]))
+
+    def forward(self, x: torch.Tensor, *, decode: bool | None = None):
+        p3 = p4 = None
+        for i, layer in enumerate(self.backbone):
+            x = layer(x)
+            if i == 3:
+                p3 = x
+            elif i == 8:
+                p4 = x
+            elif i == 12:
+                p5 = x
+
+        p3, p4, p5 = self.proj3(p3), self.proj4(p4), self.proj5(p5)
+        p4_td = self.fuse4(torch.cat((F.interpolate(p5, scale_factor=2,
+                                                   mode="nearest"), p4), 1))
+        p3_out = self.fuse3(torch.cat((F.interpolate(p4_td, scale_factor=2,
+                                                    mode="nearest"), p3), 1))
+        p4_out = self.pan4(torch.cat((self.down4(p3_out), p4_td), 1))
+        p5_out = self.pan5(torch.cat((self.down5(p4_out), p5), 1))
+        features = [p3_out, p4_out, p5_out]
+        if self.exchange is not None:
+            features = self.exchange(features)
+        return self.model[-1](features, decode=decode)

candyeye/core/candyeye.py

@@ -0,0 +1,293 @@
+"""YAML-driven model builder + forward graph for the YOLO11n detector.
+
+The YAML in ``configs/yolo11.yaml`` is the single source of truth for the
+architecture. The builder walks it exactly like the official implementation:
+
+- each entry is ``[from, repeats, module, args]``
+- ``from`` < 0  -> offset into the built list (-1 = previous layer)
+- ``from`` >= 0 -> absolute layer index; a list = several layers (Concat / Detect)
+- output channels are tracked alongside the built layers
+- width : ``c2 = make_divisible(min(c2, max_channels) * width_mult, 8)``
+- depth : ``repeats = max(round(repeats * depth_mult), 1)``
+
+The result is an ``nn.Sequential`` indexed exactly like the official yolo11n
+(0..23), so state_dict keys land at ``model.0.conv.weight``,
+``model.23.cv3.1.0.1.conv.weight``, ... and Phase 2 can load the official
+``yolo11n.pt`` 1:1 after re-initialising the head's class branch for nc=20.
+"""
+from __future__ import annotations
+
+import ast
+import copy
+from pathlib import Path
+
+import torch
+from torch import nn
+
+from candyeye.core.functions.layer_utils import make_divisible
+from candyeye.core.modules.blocks import Attention, Bottleneck, C2PSA, C3k, C3k2, PSABlock, SPPF
+from candyeye.core.modules.conv import Conv
+from candyeye.core.modules.detect import Detect
+from candyeye.core.modules.exchange import ScaleExchange
+from candyeye.paths import resolve_config
+
+
+class Concat(nn.Module):
+    """Concatenate several tensors along a dimension (neck merges)."""
+
+    def __init__(self, dimension: int = 1):
+        super().__init__()
+        self.d = dimension
+
+    def forward(self, x: list[torch.Tensor]) -> torch.Tensor:
+        return torch.cat(x, self.d)
+
+
+_MODULE_MAP = {
+    "Conv": Conv,
+    "Bottleneck": Bottleneck,
+    "C3k": C3k,
+    "C3k2": C3k2,
+    "SPPF": SPPF,
+    "Attention": Attention,
+    "PSABlock": PSABlock,
+    "C2PSA": C2PSA,
+    "Concat": Concat,
+    "Detect": Detect,
+    "ScaleExchange": ScaleExchange,
+    "nn.Upsample": nn.Upsample,
+    "nn.Conv2d": nn.Conv2d,
+}
+
+_CHANNEL_MODULES = {"Conv", "C3k2", "C2PSA", "SPPF", "C3k", "Bottleneck", "Attention", "PSABlock"}
+
+# depth_mult, width_mult, max_channels
+_SCALES = {"n": (0.50, 0.25, 1024)}
+
+
+def _source_channels(src: list[int], cur: int | list[int],
+                     ch: dict[int, int | list[int]]) -> list[int]:
+    """Flatten the input-channel spec of a multi-input layer.
+
+    A source entry may itself be a channel *list* (a previous multi-output
+    layer such as ``ScaleExchange``), in which case its channels are spliced in.
+    """
+    channels: list[int] = []
+    for index in src:
+        entry = cur if index == -1 else ch[index]
+        if isinstance(entry, (list, tuple)):
+            channels.extend(entry)
+        else:
+            channels.append(entry)
+    return channels
+
+
+
+def parse_model(data: dict) -> nn.Sequential:
+    """Expand the yaml layer table into an indexed nn.Sequential."""
+    nc = data["nc"]
+    scale = data.get("scale", "n")
+    if scale not in _SCALES:
+        raise ValueError(f"scale {scale!r} not supported (have {sorted(_SCALES)})")
+    depth, width, max_channels = _SCALES[scale]
+
+    x = data["backbone"] + data["head"]
+    layers, ch, cur = [], {}, 3  # ch[i] = out channels of layer i; cur = last output
+
+    for i, (f, n, m, args) in enumerate(x):
+        args = [nc if a == "nc" else a for a in args]
+        for j, a in enumerate(args):
+            if isinstance(a, str) and a != "nc":
+                try:
+                    args[j] = ast.literal_eval(a)  # "None" -> None, keep "nearest"
+                except (ValueError, SyntaxError):
+                    pass
+        n = max(round(n * depth), 1)
+
+        if m in _CHANNEL_MODULES:
+            c1, c2 = cur, args[0]
+            if c2 != nc:
+                c2 = make_divisible(min(c2, max_channels) * width, 8)
+            args = [c1, c2, *args[1:]]
+            if m == "C3k2":
+                args.insert(2, n)  # repeats slot inside the block
+                n = 1
+        elif m == "nn.Upsample":
+            c2 = cur
+        elif m == "Concat":
+            c2 = sum(cur if x == -1 else ch[x] for x in f)
+        elif m == "Detect":
+            src = f if isinstance(f, (list, tuple)) else [f]
+            args.append(_source_channels(src, cur, ch))
+            c2 = nc  # not referenced downstream; placeholder
+        elif m == "ScaleExchange":
+            src = f if isinstance(f, (list, tuple)) else [f]
+            in_ch = _source_channels(src, cur, ch)
+            gate = str(args[0]) if args else "none"
+            iters = int(args[1]) if len(args) > 1 else 1
+            args = [*in_ch, gate, iters]
+            c2 = list(in_ch)  # this layer emits a channel *list*
+        else:
+            raise ValueError(f"unhandled module {m!r}")
+
+        module_cls = _MODULE_MAP[m]
+        if m == "Detect":
+            module = module_cls(nc=args[0], ch=args[1])
+        elif n > 1:
+            module = nn.Sequential(*(copy.deepcopy(module_cls(*args)) for _ in range(n)))
+        else:
+            module = module_cls(*args)
+        module.type = m
+        module.i = i  # absolute index
+        module.f = f  # from-list (int or list of ints)
+        module.rep = n  # repeats used (rep avoids clobbering SPPF/C2PSA self.n)
+
+        layers.append(module)
+        ch[i] = c2
+        cur = c2
+    return nn.Sequential(*layers)
+
+
+class CandyEye(nn.Module):
+    """CandyEye detector: YAML-built graph and anchor-free Detect head."""
+
+    def __init__(self, cfg=None, nc: int | None = None, img_size: int = 128):
+        super().__init__()
+        resolved = cfg if isinstance(cfg, dict) else resolve_config(cfg)
+        data = copy.deepcopy(resolved) if isinstance(resolved, dict) else self._load_yaml(resolved)
+        if nc is not None:
+            data["nc"] = nc
+        self.model = parse_model(data)
+        self.nc = self.model[-1].nc
+        self.yaml = resolved
+        self.img_size = img_size
+        if img_size % 32:
+            raise ValueError(f"img_size must be divisible by 32, got {img_size}")
+        self.stride = self._detect_stride(img_size)
+        self.model[-1].stride = self.stride
+        self.model[-1].bias_init()
+
+    def set_classes(self, nc: int):
+        """Resize the detection class head, preserving all compatible weights."""
+        nc = int(nc)
+        if nc <= 0:
+            raise ValueError(f"nc must be positive, got {nc}")
+        old = self.model[-1]
+        if nc == old.nc:
+            self.nc = nc
+            return self
+
+        channels = tuple(branch[0].conv.in_channels for branch in old.cv2)
+        new = Detect(nc=nc, reg_max=old.reg_max, ch=channels).to(
+            device=old.stride.device, dtype=next(old.parameters()).dtype)
+        new.stride = old.stride.clone()
+        new.bias_init()
+        for attribute in ("type", "i", "f", "rep"):
+            if hasattr(old, attribute):
+                setattr(new, attribute, getattr(old, attribute))
+        old_state, new_state = old.state_dict(), new.state_dict()
+        with torch.no_grad():
+            for key, value in new_state.items():
+                if key in old_state and old_state[key].shape == value.shape:
+                    value.copy_(old_state[key])
+        self.model[-1] = new
+        self.nc = nc
+        return self
+
+    @staticmethod
+    def _load_yaml(cfg: str) -> dict:
+        import yaml
+
+        with open(cfg) as fh:
+            data = yaml.safe_load(fh)
+        CandyEye._validate_yaml(data)
+        return data
+
+    @staticmethod
+    def _validate_yaml(data: dict) -> None:
+        path = Path(data.get("path", "?"))
+        for section in ("backbone", "head"):
+            if section not in data:
+                raise ValueError(f"yaml {path} missing section {section!r}")
+            for entry in data[section]:
+                if not (isinstance(entry, list) and len(entry) == 4):
+                    raise ValueError(f"bad {section} entry {entry!r} (want [from, repeats, module, args])")
+        if "nc" not in data:
+            raise ValueError(f"yaml {path} missing nc")
+
+    def _detect_stride(self, img_size: int) -> torch.Tensor:
+        """Feed a blank image and infer P3/P4/P5 strides from output sizes."""
+        was_training = self.training
+        self.eval()  # BatchNorm train mode rejects 1x1 feature maps (img_size=32)
+        try:
+            with torch.no_grad():
+                # decode=False: raw per-level maps regardless of train/eval.
+                feats = self(torch.zeros(1, 3, img_size, img_size), decode=False)
+        finally:
+            if was_training:
+                self.train()
+        if not isinstance(feats, (list, tuple)):
+            return torch.ones(1)
+        return torch.tensor([img_size / f.shape[-2] for f in feats])
+
+    def forward(self, x: torch.Tensor, *, decode: bool | None = None) -> torch.Tensor | list[torch.Tensor]:
+        """Run the graph; ``decode=False`` returns raw maps even in eval mode."""
+        y = []  # layer-output history, index == absolute layer index
+        for m in self.model:
+            f = m.f
+            if isinstance(f, int):
+                xi = x if f == -1 else y[f]
+            else:
+                xi = [x if j == -1 else y[j] for j in f]
+            x = m(xi, decode=decode) if isinstance(m, Detect) else m(xi)
+            y.append(x)
+        return x
+
+    def train(self, mode: bool = True, *, data=None, epochs: int = 100,
+              imgsz: int | None = None, batch: int = 16, patience: int = 50,
+              workers: int = 0, device: str = "cpu",
+              project: str | Path = "runs/train", name: str = "exp",
+              resume: bool | str | Path = False, optimizer: str = "AdamW",
+              lr0: float = 2e-4, weight_decay: float = 5e-4,
+              warmup_epochs: float = 3, mosaic: float = .5,
+              hsv: bool = True, fliplr: float = .5,
+              pretrained: bool | str | Path = False, seed: int = 23,
+              exist_ok: bool = False, max_batches: int | None = None,
+              threads: int = 4):
+        """Set PyTorch mode or launch CandyEye training when `data` is set.
+
+        Example: ``model.train(data="configs/default.yaml", epochs=100,
+        imgsz=128, batch=16, patience=20)``. When called without `data`, this
+        retains the standard ``nn.Module.train(mode)`` behavior.
+        """
+        if data is None:
+            return super().train(mode)
+        from candyeye.training.trainer import train_model
+
+        return train_model(
+            self, data=data, epochs=epochs, imgsz=imgsz or self.img_size,
+            batch=batch, patience=patience, workers=workers, device=device,
+            project=project, name=name, resume=resume, optimizer=optimizer,
+            lr0=lr0, weight_decay=weight_decay, warmup_epochs=warmup_epochs,
+            mosaic=mosaic, hsv=hsv, fliplr=fliplr, pretrained=pretrained,
+            seed=seed, exist_ok=exist_ok, max_batches=max_batches,
+            threads=threads,
+        )
+
+    def fuse(self):
+        """Fold every BatchNorm into its conv (in place).
+
+        Puts the model in eval mode first (BN running stats are the whole
+        point of fusing).  After this, state_dict keys match the ONNX
+        versions of the model (``model.0.conv.weight`` +
+        ``model.0.conv.bias``, no ``.bn.*``), and ``load_fused_from_onnx``
+        can copy the official fp32 weights 1:1.  Numerical parity with the
+        official export is then exact, not "fp16 checkpoint drift" close.
+        """
+        from candyeye.core.modules.conv import Conv
+
+        self.eval()
+        for m in self.modules():
+            if isinstance(m, Conv):
+                m.fuse()
+        return self

candyeye/core/convert_yolo11.py

@@ -0,0 +1,157 @@
+"""Load the official YOLO11n weights (nc=80) into our model.
+
+Reads a *clean* state_dict (plain tensors only — no ultralytics dependency).
+By default it uses the wheel's bundled ``assets/yolo11n.pth``; pass an explicit
+path to override.  The loader is shape-checked: any tensor whose shape doesn't
+match is skipped and reported, everything else is copied 1:1 (parameters *and*
+buffers such as BatchNorm running statistics, so eval-mode inference is
+faithful).
+
+Two load scenarios:
+- ``CandyEye(yaml, nc=80)``   — exact full load: 0 missing / 0 unexpected.
+- ``CandyEye(yaml, nc=20)``   — our VOC training target: the whole class branch
+  ``model.23.cv3.*`` is nc/c3-dependent, so exactly ``EXPECTED_NC20_SKIPPED``
+  (51 tensors, generated below from the known structure) are skipped and stay
+  random-initialised.  The box branch (``cv2`` + ``dfl``) still loads 1:1, so
+  box proposals remain the official ones.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+
+from candyeye.paths import default_weights_path
+
+ONNX_PATH = "tmp/yolo11n.onnx"  # official fp32 export (BN fused)
+
+# The official checkpoint was trained with nc=80 (COCO), our training model
+# uses nc=20 (VOC).  In Detect, c3 = max(ch[0], min(nc, 100)) pins the class
+# branch's hidden width (80 -> 64) and its final Conv2d in/out (-> 20).
+# Per scale that unmatchable part is:
+#   cv3.<s>.0.1  Conv(x -> c3)      bn.weight/bias/mean/var + conv.weight
+#   cv3.<s>.1.0  DWConv(c3 -> c3)    bn.*                          + conv.weight
+#   cv3.<s>.1.1  Conv(c3 -> c3)      bn.*                          + conv.weight
+#   cv3.<s>.2    Conv2d(c3 -> nc)    weight + bias
+# = 17 tensors/scale x 3 scales = 51.  Everything else loads unchanged.
+EXPECTED_NC20_SKIPPED = frozenset(
+    [
+        *(
+            f"model.23.cv3.{s}.{a}.{b}.{comp}"
+            for s in range(3)
+            for (a, b) in ((0, 1), (1, 0), (1, 1))
+            for comp in ("bn.bias", "bn.running_mean", "bn.running_var",
+                         "bn.weight", "conv.weight")
+        ),
+        *(
+            f"model.23.cv3.{s}.2.{comp}"
+            for s in range(3)
+            for comp in ("weight", "bias")
+        ),
+    ]
+)
+
+
+def load_official_state_dict(path: str | Path | None = None) -> dict[str, torch.Tensor]:
+    """Read the clean .pth (tensors only -> safe with torch.load defaults).
+
+    Defaults to the bundled ``assets/yolo11n.pth`` so callers can load the
+    default weights without knowing the package location.
+    """
+    if path is None:
+        path = default_weights_path()
+    sd = torch.load(path, map_location="cpu")
+    if not (isinstance(sd, dict) and all(isinstance(v, torch.Tensor) for v in sd.values())):
+        raise TypeError(f"{path} is not a clean state_dict of tensors "
+                        f"(re-run scripts/bootstrap_weights.py)")
+    return sd
+
+
+def load_weights(
+    model: torch.nn.Module,
+    official: dict[str, torch.Tensor],
+    verbose: bool = False,
+) -> dict[str, list]:
+    """Shape-checked copy of *official* into *model*; nothing in-place original.
+
+    Returns {"loaded": [...], "skipped": [(key, our_shape, off_shape)],
+             "unexpected": [(key, off_shape)]}.  Buffers are copied too.
+    """
+    ours = model.state_dict()
+    loaded, skipped, unexpected = [], [], []
+
+    for key, value in official.items():  # first pass: classify
+        if key not in ours:
+            unexpected.append((key, tuple(value.shape)))
+        elif ours[key].shape == value.shape:
+            loaded.append(key)
+        else:
+            skipped.append((key, tuple(ours[key].shape), tuple(value.shape)))
+
+    with torch.no_grad():  # in-place copy onto the model's existing tensors
+        for key in loaded:
+            ours[key].copy_(official[key])
+
+    if verbose:
+        print(f"loaded {len(loaded)} tensors, skipped {len(skipped)}, "
+              f"unexpected {len(unexpected)}")
+        for key, a, b in skipped:
+            print(f"  skip {key}: ours {a} vs official {b}")
+    return {"loaded": loaded, "skipped": skipped, "unexpected": unexpected}
+
+
+# With BN fused away, the nc=20 class branch (c3 = min(nc, 100) -> 64 vs the
+# official 80) differs in 4 tensors per scale x 3 scales = 24.  Per scale:
+#   cv3.<s>.0.0.1  Conv(c3 -> c3 fixed 64<->80)     weight + bias
+#   cv3.<s>.0.1.0  DWConv(c3 -> c3, depthwise)      weight + bias
+#   cv3.<s>.0.1.1  Conv(c3 -> c3)                   weight + bias
+#   cv3.<s>.2      Conv2d(c3 -> nc)                 weight + bias
+EXPECTED_FUSED_NC20_SKIPPED = frozenset(
+    [
+        *(
+            f"model.23.cv3.{s}.{a}.{b}.{comp}"
+            for s in range(3)
+            for (a, b) in ((0, 1), (1, 0), (1, 1))
+            for comp in ("conv.weight", "conv.bias")
+        ),
+        *(
+            f"model.23.cv3.{s}.2.{comp}"
+            for s in range(3)
+            for comp in ("weight", "bias")
+        ),
+    ]
+)
+
+
+def remap_prefix(official: dict[str, torch.Tensor], source_index: int,
+                 target_index: int) -> dict[str, torch.Tensor]:
+    """Shift ``model.<source_index>.*`` keys onto another layer index.
+
+    Neck variants that add layers after the original Detect (e.g. the
+    ``ScaleExchange`` neck) move Detect from layer 23 to a new index, so the
+    official head keys must be re-keyed before loading.
+    """
+    if source_index == target_index:
+        return official
+    source = f"model.{source_index}."
+    target = f"model.{target_index}."
+    return {
+        (target + key[len(source):] if key.startswith(source) else key): value
+        for key, value in official.items()
+    }
+
+
+def load_fused_from_onnx(path: str = ONNX_PATH) -> dict[str, torch.Tensor]:
+    """Read the official ONNX graph initializers as a name -> fp32 Tensor dict.
+
+    The graph's weights are BN-fused, so this only makes sense for a model
+    that has been ``CandyEye(...).fuse()``d — the key set then matches 1:1
+    (``model.0.conv.weight`` + ``model.0.conv.bias``, no ``.bn.*``).
+    """
+    import onnx
+
+    model = onnx.load(path)
+    return {
+        init.name: torch.tensor(onnx.numpy_helper.to_array(init), dtype=torch.float32)
+        for init in model.graph.initializer
+    }