DINOv3 助力 YOLO26 攻小目标：召回 +6.2pt、参数零增加，收益从哪来、能用在哪

# -*- coding: utf-8 -*-
"""
DINOv3 × YOLO26 小目标检测方案 —— 伪代码
==========================================

对应文章：《DINOv3 助力 YOLO26 攻小目标：召回 +6.2pt、参数零增加》

本文伪代码覆盖两条技术路线：
  A. 骨干替换路线（NEU-DET 实测）：DINOv3 ViT-B/16 替换 YOLO26 原有骨干，
     冻结前 8 层，get_intermediate_layers(n=4) 取 4 层中间特征，
     经 FPN 重建 P3/P4/P5 三尺度，接 YOLO26 检测头（DFL 解码）。
  B. 并行融合路线（yaml 驱动）：DINOv3 作为冻结并行分支与原图并行，
     轻量投影后注入骨干 P4（Single）或 P3+P4（Dual）。

NEU-DET 实测锚点（640×640，300 epochs）：
  原始 YOLO26n  : mAP50 0.722 / Recall 0.643 / 2,376,006 参数 / 5.2 GFLOPs / 3.0 it/s
  + DINOv3      : mAP50 0.743 / Recall 0.705 / 2,376,006 参数 / 5.3 GFLOPs / 3.1 it/s
  crazing 召回  : 0.215 → 0.356（相对 +66%）
"""

from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn


# ============================================================
# 0. 配置区
# ============================================================

class FusionMode(Enum):
    """融合模式开关（--fusion 的四种取值）"""
    NONE = "none"    # 原始 YOLO26 基线
    P4 = "p4"        # 并行融合 · Single P4
    P34 = "p34"      # 并行融合 · Dual P3+P4
    AUTO = "auto"    # 原样使用 --model 指定的 yaml


@dataclass
class DINOv3Config:
    """DINOv3 基座配置（骨干替换 / 并行融合共用）"""
    arch: str = "vitb16"            # vits16 / vits16plus / vitb16 / vitl16 ...
    img_size: int = 640             # 输入尺寸（RoPE 支持任意分辨率）
    embed_dim: int = 768            # ViT-B/16 输出通道
    num_blocks: int = 12            # transformer block 数
    freeze_blocks: int = 8          # 冻结前 8 层（骨干替换路线）
    n_intermediate: int = 4         # get_intermediate_layers(n=4)
    weights: Optional[str] = None   # 官方预训练权重（需向 Meta 申请）
    lora_r: int = 0                 # attn.qkv 上的 LoRA 秩，0 = 关闭
    unfreeze_tail: int = 0          # 末尾 N 个 block 保持可训练，0 = 全冻结


@dataclass
class NeckConfig:
    """FPN 多尺度配置（NEU-DET 640 输入下的三尺度）"""
    p3_size: Tuple[int, int] = (40, 40)    # stride 16
    p3_ch: int = 256
    p4_size: Tuple[int, int] = (20, 20)    # stride 32
    p4_ch: int = 512
    p5_size: Tuple[int, int] = (10, 10)    # stride 64
    p5_ch: int = 512


@dataclass
class TrainConfig:
    """训练配置"""
    data: str = "data/NEU-DET.yaml"
    imgsz: int = 640
    epochs: int = 300
    batch: int = 16
    close_mosaic: int = 10
    workers: int = 2
    device: str = "0"
    lr0: float = 0.01               # 开启 unfreeze / LoRA 时建议降到 1e-3 量级
    # NEU-DET 六类缺陷
    nc: int = 6
    names: List[str] = field(default_factory=lambda: [
        "crazing", "inclusion", "patches",
        "pitted_surface", "rolled-in_scale", "scratches",
    ])


# ============================================================
# 1. DINOv3 基座（骨干替换路线）
# ============================================================

class DINOv3Backbone(nn.Module):
    """
    DINOv3 ViT-B/16 骨干。

    要点：
      - 加载官方预训练权重，冻结前 freeze_blocks 层（no_grad，0 梯度）；
      - get_intermediate_layers(n=4) 提取 4 层中间特征，供 FPN 多尺度融合；
      - ViT 原生输出为单尺度（640 输入下 40×40），多尺度由 Neck 负责；
      - RoPE 无绝对位置编码，640/1280 任意尺寸直接推理。
    """

    def __init__(self, cfg: DINOv3Config):
        super().__init__()
        self.cfg = cfg
        self.blocks = nn.ModuleList([
            self._build_transformer_block(i) for i in range(cfg.num_blocks)
        ])
        self.cls_token = nn.Parameter(torch.zeros(1, 1, cfg.embed_dim))
        self.register_tokens = nn.Parameter(
            torch.zeros(1, 4, cfg.embed_dim)   # 4 个 register token
        )
        self.rope = self._build_rope()          # 轴向旋转位置编码
        self._load_pretrained(cfg.weights)
        self._freeze_first_n_blocks(cfg.freeze_blocks)

    # ---------- 冻结策略：冻结层走 no_grad，不建反向图 ----------
    def _freeze_first_n_blocks(self, n: int) -> None:
        for i, blk in enumerate(self.blocks):
            requires_grad = (i >= n)
            for p in blk.parameters():
                p.requires_grad = requires_grad
            if not requires_grad:
                blk.eval()  # 冻结层永远处于 eval（无 BN/DO 也保持习惯一致）

    @torch.no_grad()
    def _frozen_forward(self, x: torch.Tensor) -> torch.Tensor:
        """冻结 block 的前向（no_grad，显存几乎不增）"""
        for blk in self.blocks[: self.cfg.freeze_blocks]:
            x = blk(x)
        return x

    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
        """
        输入 x: [B, 3, 640, 640]
        输出  : 4 层中间 patch token 特征 [B, N, C]（供 FPN 使用）
        """
        x = self._patch_embed(x)                       # [B, N=1600, 768]
        x = self._prepend_special_tokens(x)            # [cls, reg×4, patches]
        x = self._frozen_forward(x)                    # 前 8 层：冻结
        feats: List[torch.Tensor] = []
        for blk in self.blocks[self.cfg.freeze_blocks:]:
            x = blk(x)
            feats.append(x)                            # 收集后 4 层
        return self._get_intermediate_layers(feats, n=self.cfg.n_intermediate)


# ============================================================
# 2. Multi-Scale Neck（FPN）：单尺度 → P3/P4/P5 三尺度
# ============================================================

class MultiScaleFPN(nn.Module):
    """
    将 DINOv3 的 4 层中间特征做 C2/C3/C4 多层融合，
    重建 YOLO 检测头需要的三尺度金字塔。

      P3: 40×40, 256 通道（stride 16，小目标主战场）
      P4: 20×20, 512 通道（stride 32）
      P5: 10×10, 512 通道（stride 64）
    """

    def __init__(self, cfg: NeckConfig, in_dim: int = 768):
        super().__init__()
        self.cfg = cfg
        # C2/C3/C4 融合的轻量卷积组（示意）
        self.lateral_c2 = nn.Conv2d(in_dim, cfg.p3_ch, 1)
        self.lateral_c3 = nn.Conv2d(in_dim, cfg.p4_ch, 1)
        self.lateral_c4 = nn.Conv2d(in_dim, cfg.p5_ch, 1)
        self.fuse_p3 = nn.Sequential(nn.Conv2d(cfg.p3_ch, cfg.p3_ch, 3, padding=1), nn.BatchNorm2d(cfg.p3_ch), nn.SiLU())
        self.fuse_p4 = nn.Sequential(nn.Conv2d(cfg.p4_ch, cfg.p4_ch, 3, padding=1), nn.BatchNorm2d(cfg.p4_ch), nn.SiLU())
        self.fuse_p5 = nn.Sequential(nn.Conv2d(cfg.p5_ch, cfg.p5_ch, 3, padding=1), nn.BatchNorm2d(cfg.p5_ch), nn.SiLU())

    def forward(self, intermediate_feats: List[torch.Tensor]) -> Dict[str, torch.Tensor]:
        """
        intermediate_feats: DINOv3 的 4 层中间 token 特征
        返回: {"P3": ..., "P4": ..., "P5": ...} 三个尺度的特征图
        """
        c2, c3, c4, _ = [self._tokens_to_featmap(f) for f in intermediate_feats]
        p5 = self.fuse_p5(self.lateral_c4(c4))          # 10×10, 512
        p4 = self.fuse_p4(self.lateral_c3(c3) + self._up(p5))
        p3 = self.fuse_p3(self.lateral_c2(c2) + self._up(p4))
        return {"P3": p3, "P4": p4, "P5": p5}


# ============================================================
# 3. YOLO26 检测头（DFL 解码）
# ============================================================

class YOLO26DetectHead(nn.Module):
    """
    三个尺度分别预测边界框与类别；
    采用 DFL（Distribution Focal Loss）解码策略优化框回归精度。
    """

    def __init__(self, nc: int, ch: Tuple[int, int, int] = (256, 512, 512), reg_max: int = 16):
        super().__init__()
        self.nc = nc
        self.reg_max = reg_max
        self.box_branches = nn.ModuleList([nn.Conv2d(c, 4 * reg_max, 1) for c in ch])
        self.cls_branches = nn.ModuleList([nn.Conv2d(c, nc, 1) for c in ch])

    def forward(
        self, pyramid: Dict[str, torch.Tensor], training: bool = True
    ):
        """
        训练时输出字典 dict(boxes, scores, feats)；
        推理时输出已解码的边界框张量 [B, anchors, 4 + nc]。
        """
        feats = [pyramid["P3"], pyramid["P4"], pyramid["P5"]]
        boxes, scores, raw_feats = [], [], []
        for i, f in enumerate(feats):
            box_dist = self.box_branches[i](f)     # DFL 分布
            score = self.cls_branches[i](f).sigmoid()
            raw_feats.append(f)
            if training:
                boxes.append(box_dist)
                scores.append(score)
            else:
                boxes.append(self._dfl_decode(box_dist))   # 分布期望解码
                scores.append(score.flatten(2).permute(0, 2, 1))
        if training:
            return {"boxes": boxes, "scores": scores, "feats": raw_feats}
        return torch.cat([torch.cat([b, s], dim=-1) for b, s in zip(boxes, scores)], dim=1)

    def _dfl_decode(self, box_dist: torch.Tensor) -> torch.Tensor:
        """DFL：对 reg_max 个 bin 的分布取期望，得到 4 条边距离"""
        B, C, H, W = box_dist.shape
        prob = box_dist.view(B, 4, self.reg_max, H, W).softmax(dim=2)
        bins = torch.arange(self.reg_max, device=box_dist.device, dtype=prob.dtype)
        return (prob * bins).sum(dim=2)             # [B, 4, H, W]


# ============================================================
# 4. 路线 A：骨干替换版模型（NEU-DET 实测路线）
# ============================================================

class YOLO26_DINOv3_BackboneReplaced(nn.Module):
    """
    DINOv3 ViT-B/16 整体替换 YOLO26 原有骨干 + FPN 重建三尺度 + YOLO26 检测头。

    收益锚点（NEU-DET）：
      mAP50   0.722 → 0.743（+2.1pt）
      Recall  0.643 → 0.705（+6.2pt）
      参数量 零增加 / GFLOPs +0.1 / 推理速度 3.0 → 3.1 it/s
    """

    def __init__(self, dcfg: DINOv3Config, ncfg: NeckConfig, train_cfg: TrainConfig):
        super().__init__()
        self.backbone = DINOv3Backbone(dcfg)          # 冻结前 8 层
        self.neck = MultiScaleFPN(ncfg, in_dim=dcfg.embed_dim)
        self.head = YOLO26DetectHead(nc=train_cfg.nc)

    def forward(self, img: torch.Tensor, training: bool = True):
        intermediate = self.backbone(img)             # 4 层中间特征
        pyramid = self.neck(intermediate)             # P3/P4/P5
        return self.head(pyramid, training=training)


# ============================================================
# 5. 路线 B：并行融合模块（yaml 驱动，Single P4 / Dual P3+P4）
# ============================================================

class LightProjection(nn.Module):
    """轻量投影：1×1 卷积降通道 + 可选上采样 + 3×3 深度可分离卷积"""

    def __init__(self, in_ch: int, out_ch: int, upsample: bool = False):
        super().__init__()
        self.conv1x1 = nn.Conv2d(in_ch, out_ch, 1)
        self.up = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False) if upsample else nn.Identity()
        self.dwconv = nn.Conv2d(out_ch, out_ch, 3, padding=1, groups=out_ch)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dwconv(self.up(self.conv1x1(x)))


class DINOv3(nn.Module):
    """
    并行融合分支：DINOv3 完全冻结，与原图并行前向。
    按 levels 取指定 transformer block 的 patch token
    （丢掉 cls / register token），reshape 成 stride 16 特征图。

    yaml 一行挂载（对应文章第 02 节）：
      - [0, 1, DINOv3, [256, "vits16", ["P3", "P4"], None, 12, 1.0]]
      #             c2    arch      levels      w   depth scale
    """

    def __init__(self, c2: int, arch: str = "vits16", levels: List[str] = ["P3", "P4"],
                 weights: Optional[str] = None, depth: int = 12, scale: float = 1.0):
        super().__init__()
        self.levels = levels
        self.depth = depth                          # 只保留前 N 个 block（降算力开关）
        self.scale = scale                          # ViT 输入降采样倍率（0.5 = 半分辨率）
        self.vit = self._load_frozen_vit(arch, weights, depth)
        for p in self.vit.parameters():
            p.requires_grad = False                 # 全冻结：no_grad 前向，显存几乎不变
        # 每个 level 一个轻量投影头（P3 需上采样 ×2 对齐 stride 8）
        self.projs = nn.ModuleDict({
            lv: LightProjection(self.vit.embed_dim, c2, upsample=(lv == "P3"))
            for lv in levels
        })

    @torch.no_grad()
    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]:
        tokens = self.vit.forward_tokens(x, n=self.depth)   # 冻结前向
        # P4 取末层 block，P3 取中间 block（Simple Feature Pyramid 思路，代价为零）
        feats = self.vit.get_intermediate_layers(tokens, blocks=self._block_map())
        return {lv: self.projs[lv](feats[lv]) for lv in self.levels}

    def _block_map(self) -> Dict[str, int]:
        """depth=12 时：P3 取 block 5（细节），P4 取 block 11（语义）"""
        return {"P3": self.depth // 2, "P4": self.depth - 1}


class YOLO26_DINOv3_Parallel(nn.Module):
    """
    并行融合版：保留 YOLO26 原生骨干，DINOv3 作为冻结并行分支，
    在骨干 P4（Single）或 P3+P4（Dual）处 Concat 融合，1×1 卷积降回原通道。

    可训练参数仅 +0.05M（Single）/ +0.09M（Dual）；
    neck / head 完全不动；检测 / 分割 / 分类共用同一套接法。
    """

    def __init__(self, base_yolo26: nn.Module, fusion: FusionMode, c2: int = 256):
        super().__init__()
        assert fusion in (FusionMode.P4, FusionMode.P34)
        self.base = base_yolo26                     # 原生 YOLO26（CNN 骨干）
        self.fusion = fusion
        levels = ["P4"] if fusion == FusionMode.P4 else ["P3", "P4"]
        self.dinov3 = DINOv3(c2=c2, levels=levels)  # 冻结并行分支
        # 融合层：Concat 后 1×1 卷积降回原通道
        self.fuse_p4 = nn.Conv2d(c2 + 512, 512, 1)
        self.fuse_p3 = nn.Conv2d(c2 + 256, 256, 1) if fusion == FusionMode.P34 else None

    def forward(self, img: torch.Tensor, training: bool = True):
        dino_feats = self.dinov3(img)               # {"P3": ..., "P4": ...}
        cnn_p4 = self.base.backbone_p4(img)         # 骨干 P4（stride 16）
        p4 = self.fuse_p4(torch.cat([cnn_p4, dino_feats["P4"]], dim=1))
        if self.fusion == FusionMode.P34:
            cnn_p3 = self.base.backbone_p3(img)     # 骨干 P3（stride 8）
            p3 = self.fuse_p3(torch.cat([cnn_p3, dino_feats["P3"]], dim=1))
        else:
            p3 = self.base.backbone_p3(img)
        return self.base.neck_and_head(p3, p4, training=training)


# ============================================================
# 6. 训练：--fusion 一个开关跑三组对照实验
# ============================================================

def build_model(fusion: FusionMode, tcfg: TrainConfig) -> nn.Module:
    """按融合模式构建模型（--fusion none / p4 / p34 / auto）"""
    if fusion == FusionMode.NONE:
        return YOLO26_Base(nc=tcfg.nc)                       # 原始 YOLO26 基线
    if fusion == FusionMode.AUTO:
        return YOLO26_Base(nc=tcfg.nc)                       # 原样使用 --model 的 yaml
    if fusion in (FusionMode.P4, FusionMode.P34):
        return YOLO26_DINOv3_Parallel(YOLO26_Base(nc=tcfg.nc), fusion)
    raise ValueError(f"unknown fusion mode: {fusion}")


def train(fusion: FusionMode, tcfg: TrainConfig) -> Dict[str, float]:
    """
    训练入口。断点续训 / 多卡 DDP 略（多卡时 DINOv3 仓库路径
    需用环境变量 DINOV3_REPO 传给 spawn 子进程）。
    """
    model = build_model(fusion, tcfg)
    # 学习率策略：开启 unfreeze / LoRA 时 lr0 降到 1e-3 量级；
    # 全冻结时默认 lr0 即可（DINOv3 侧不产生梯度）
    optimizer = torch.optim.AdamW(
        filter(lambda p: p.requires_grad, model.parameters()), lr=tcfg.lr0
    )
    for epoch in range(tcfg.epochs):
        for img, target in dataloader(tcfg.data, tcfg.imgsz, tcfg.batch):
            out = model(img, training=True)
            loss = detection_loss(out, target)      # DFL + cls loss
            optimizer.zero_grad(); loss.backward(); optimizer.step()
        if epoch >= tcfg.epochs - tcfg.close_mosaic:
            close_mosaic()                          # 最后 10 轮关闭 mosaic
    return validate(model)                          # 返回 mAP50 / Recall 等指标


def run_ablation(tcfg: TrainConfig) -> Dict[str, Dict[str, float]]:
    """
    三组对照实验（文章第 03 节的出数命令）：
      python scripts/dinov3/train.py --model yolo26n --fusion none --data data/NEU-DET.yaml --name base
      python scripts/dinov3/train.py --model yolo26n --fusion p4   --data data/NEU-DET.yaml --name p4
      python scripts/dinov3/train.py --model yolo26n --fusion p34  --data data/NEU-DET.yaml --name p34
    """
    results = {}
    for mode in (FusionMode.NONE, FusionMode.P4, FusionMode.P34):
        results[mode.value] = train(mode, tcfg)
    return results
    # 期望锚点（NEU-DET）：
    #   base : mAP50 0.722 / Recall 0.643
    #   p4   : 小步验证 DINOv3 特征有效性
    #   p34  : mAP50 0.743 / Recall 0.705（crazing 召回 0.215 → 0.356）


# ============================================================
# 7. 可视化三件套：验证"融合是否真的起作用"
# ============================================================

def visualize(weights: str, source: str, mode: str, out: Optional[str] = None) -> str:
    """
    --mode det  : 带框结果图，与基线做直观对比
    --mode feat : 原图 + DINOv3 P3/P4 响应热力图 + 融合后骨干特征热力图
                 （典型现象：缺陷区在 CNN 特征里几乎无响应，DINOv3 上有明显激活）
    --mode attn : CLS token 对 patch 的注意力热力图（ViT 已"看到"前景/缺陷）
    """
    model = load_weights(weights)
    img = load_image(source)
    if mode == "det":
        return draw_boxes(img, model.predict(img))
    if mode == "feat":
        heatmaps = model.dinov3_feature_maps(img)      # 需融合权重，原始 YOLO 会报"找不到 DINOv3 层"
        return concat_horizontally([img, *heatmaps, fused_feature_map(model, img)])
    if mode == "attn":
        return draw_heatmap(img, model.dinov3_cls_attention(img))
    raise ValueError(mode)


# ============================================================
# 8. 部署导出：ONNX + 边缘 checklist
# ============================================================

def export(weights: str, imgsz: int = 640, opset: int = 17) -> str:
    """
    边缘部署 checklist（文章第 07 节）：
      1. 固定 imgsz，关闭 dynamic（多数 NPU 不支持动态 shape）；
      2. ONNX opset ≥ 17（RoPE 的 sin/cos、LayerNorm、SDPA 需要较新算子）；
      3. TensorRT FP16 建议 TRT ≥ 10.6，上线前用真实图对数值；
      4. 导出后跑一遍 predict 确认端到端一致；
      5. 延迟不达标：yaml 改 depth: 8 + scale: 0.5，ViT 侧算力约降到 1/3。
    """
    model = load_weights(weights)
    return export_onnx(model, imgsz=imgsz, opset=opset, dynamic=False)


def benchmark(models: List[str], imgsz: int = 640) -> str:
    """参数量 / GFLOPs / 延迟对比，一次出一张 Markdown 表"""
    rows = []
    for name in models:
        m = load_model(name)
        rows.append({
            "model": name,
            "params_M": count_params(m) / 1e6,
            "trainable_M": count_trainable(m) / 1e6,
            "GFLOPs": calc_flops(m, imgsz),
            "latency_ms": measure_latency(m, imgsz),
        })
    return to_markdown_table(rows)


# ============================================================
# 9. 主流程：NEU-DET 复现实验
# ============================================================

if __name__ == "__main__":
    tcfg = TrainConfig()   # NEU-DET, 640, 300 epochs, batch 16

    # Step 1: 三组对照实验（none / p4 / p34）
    results = run_ablation(tcfg)

    # Step 2: 打印收益结构（文章第 03 节的核心解读）
    #   召回 +6.2pt，约为 mAP50 涨幅（+2.1pt）的三倍
    #   → 收益来自"原本检不出的现在检出来了"
    print(results)

    # Step 3: 可视化验证（det / feat / attn 三件套）
    visualize("runs/dinov3/train/weights/best.pt", "neu_det_sample.jpg", mode="feat")

    # Step 4: 导出部署
    export("runs/dinov3/train/weights/best.pt", imgsz=640, opset=17)

    # Step 5: 场景衍生：更换 data 与检测头即可切入新领域
    #   光伏 EL / PCB / 医学微小病灶 / 遥感小目标 / 农业视觉
    for scenario in ("pv_el", "pcb_defect", "pathology", "remote_sensing", "agri_vision"):
        cfg = TrainConfig(data=f"data/{scenario}.yaml")
        train(FusionMode.P34, cfg)