病斑只有几像素、叶脉纹理全是干扰？把 DINOv3 接进 YOLO，农业视觉的“老大难“有解了

# =============================================================================
#  dino_yolo_agri.py
# -----------------------------------------------------------------------------
#  DINOv3 + YOLO 农业病害检测 伪代码骨架
#
#  解决的问题：
#    1) 病斑只有几像素  → ViT 的高分辨率 patch token + P2 检测头 + 滑窗切片推理
#    2) 叶脉纹理全是干扰 → DINOv3 自监督语义先验 + CLS-attention 显著性门控
#    3) CNN 主干感受野固定，对"语义异常"不敏感 → 换 DINOv3 做 backbone
#
#  本文件是"伪代码"：所有 nn.Module 的结构/形状/数据流都是真实可跑的，
#  但 DINOv3 权重加载、YOLO 的 loss/decode/NMS 用占位函数表示，
#  实际落地时替换为 timm/transformers + ultralytics 即可。
# =============================================================================

import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass


# =============================================================================
# 0. 配置
# =============================================================================
@dataclass
class Config:
    img_size: int = 640          # 训练输入尺寸
    patch: int = 16              # DINOv3 patch 大小
    vit_dim: int = 1024          # DINOv3 特征维度 (ViT-L 量级)
    num_classes: int = 1         # 病害类别数
    reg_max: int = 16            # DFL 分箱数
    lambda_sal: float = 0.1      # 显著性辅助损失权重
    small_ratio_thr: float = 0.005   # "病斑像素占比 < 0.5%" 视为小目标样本
    tile: int = 640              # 推理滑窗尺寸
    overlap: int = 64            # 滑窗重叠
    conf_thr: float = 0.25
    iou_thr: float = 0.6


# =============================================================================
# 1. DINOv3 主干 (占位)
# -----------------------------------------------------------------------------
#  真实实现：
#     from transformers import AutoModel
#     self.vit = AutoModel.from_pretrained("facebook/dinov3-vitl16-pretrain-lvd1689m")
#  这里用随机张量模拟其输出契约：
#     tokens: (B, N, C)   N = (H/P) * (W/P)   最后一层 patch token
#     attn:   (B, N)     CLS → 每个 patch 的注意力 (已 softmax)
# =============================================================================
class DINOv3Backbone(nn.Module):
    def __init__(self, dim: int = 1024, patch: int = 16, depth: int = 24):
        super().__init__()
        self.patch = patch
        self.dim = dim
        self.depth = depth
        # 真实实现中这里是 24 层 ViT Block
        self.blocks = nn.ModuleList([nn.Identity() for _ in range(depth)])

    def forward(self, x: torch.Tensor, return_attn: bool = True):
        B, _, H, W = x.shape
        hp, wp = H // self.patch, W // self.patch
        N = hp * wp
        # —— 占位输出 ——
        tokens = torch.randn(B, N, self.dim, device=x.device)
        attn = torch.rand(B, N, device=x.device)
        attn = attn / attn.sum(dim=-1, keepdim=True)
        return tokens, (attn if return_attn else None)


# =============================================================================
# 2. ViT → 多尺度适配器 (ViTDet / SimpleFPN 风格)
# -----------------------------------------------------------------------------
#  ViT 只有单尺度 (stride = patch = 16)，YOLO 需要 P2~P5 多尺度金字塔。
#  这里用 1x1 投影 + 反卷积/步长卷积生成四个尺度：
#     P2: stride 4   ← 小病斑的关键层
#     P3: stride 8
#     P4: stride 16  ← 与 ViT 原生尺度一致
#     P5: stride 32
# =============================================================================
class ViTAdapter(nn.Module):
    def __init__(self, in_dim: int = 1024,
                 out_channels=(96, 192, 384, 768),
                 patch: int = 16):
        super().__init__()
        self.patch = patch
        self.out_channels = out_channels

        c2, c3, c4, c5 = out_channels
        # 原生尺度 P4
        self.proj = nn.Conv2d(in_dim, c4, kernel_size=1)
        # 上采到 P3 / P2
        self.up_p4_to_p3 = nn.ConvTranspose2d(c4, c3, 2, 2)
        self.up_p3_to_p2 = nn.ConvTranspose2d(c3, c2, 2, 2)
        # 下采到 P5
        self.down_p4_to_p5 = nn.Conv2d(c4, c5, kernel_size=3, stride=2, padding=1)

        # 每层加一个 3x3 平滑，减轻反卷积棋盘效应
        self.smooth = nn.ModuleList([
            nn.Conv2d(c, c, 3, 1, 1) for c in out_channels
        ])

    def forward(self, tokens: torch.Tensor, H: int, W: int):
        B, N, C = tokens.shape
        hp, wp = H // self.patch, W // self.patch
        x = tokens.transpose(1, 2).reshape(B, C, hp, wp)   # (B,C,hp,wp)

        p4 = self.proj(x)
        p3 = self.up_p4_to_p3(p4)
        p2 = self.up_p3_to_p2(p3)
        p5 = self.down_p4_to_p5(p4)

        feats = [p2, p3, p4, p5]
        feats = [self.smooth[i](f) for i, f in enumerate(feats)]
        return feats


# =============================================================================
# 3. 显著性门控 (核心创新点)
# -----------------------------------------------------------------------------
#  动机：叶脉纹理在 DINOv3 的 CLS-attention 上是"低显著性"的，
#        病斑是"语义异常" → 高显著性。
#  做法：把 patch 级 attention 上采到特征图分辨率，
#        作为空间门控乘到颈部特征上；再叠加一个通道注意力。
#  效果：叶脉响应被压低，病斑响应被保留/放大。
# =============================================================================
class SaliencyGate(nn.Module):
    def __init__(self, channels: int, patch: int = 16):
        super().__init__()
        self.patch = patch
        self.channel_reduce = nn.Conv2d(channels, 1, kernel_size=1)

    def forward(self, feat: torch.Tensor, dino_attn: torch.Tensor,
                H: int, W: int) -> torch.Tensor:
        B, C, h, w = feat.shape
        hp, wp = H // self.patch, W // self.patch

        # patch attention → 特征图尺寸的显著性图
        attn = dino_attn.reshape(B, 1, hp, wp)
        attn = F.interpolate(attn, size=(h, w), mode="bilinear",
                             align_corners=False)
        # 归一化到 [0,1]，避免 attention 尺度漂移
        attn = attn / (attn.amax(dim=(2, 3), keepdim=True) + 1e-6)
        attn = torch.sigmoid(attn * 4.0 - 2.0)     # 锐化：>均值放大，<均值压低

        # 通道注意力
        ca = torch.sigmoid(self.channel_reduce(feat))

        # 残差式门控：即使门控为 0 也保留 50% 原始响应，防止梯度消失
        return feat * (0.5 + attn) * (0.5 + ca)


# =============================================================================
# 4. YOLO 检测头 (P2~P5 四尺度)
# =============================================================================
class YOLOHead(nn.Module):
    def __init__(self, in_channels=(96, 192, 384, 768),
                 num_classes: int = 1, reg_max: int = 16):
        super().__init__()
        self.num_classes = num_classes
        self.reg_max = reg_max
        self.stems = nn.ModuleList(
            [nn.Conv2d(c, 128, 3, 1, 1) for c in in_channels]
        )
        self.cls_branch = nn.ModuleList(
            [nn.Conv2d(128, num_classes, 1) for _ in in_channels]
        )
        self.reg_branch = nn.ModuleList(
            [nn.Conv2d(128, 4 * reg_max, 1) for _ in in_channels]
        )

    def forward(self, feats):
        outs = []
        for i, f in enumerate(feats):
            s = self.stems[i](f)
            outs.append((self.cls_branch[i](s), self.reg_branch[i](s)))
        return outs


# =============================================================================
# 5. 完整模型
# =============================================================================
class DINOv3YOLO(nn.Module):
    def __init__(self, cfg: Config):
        super().__init__()
        self.cfg = cfg
        self.backbone = DINOv3Backbone(dim=cfg.vit_dim, patch=cfg.patch)
        self.adapter = ViTAdapter(in_dim=cfg.vit_dim, patch=cfg.patch)
        # 只对 P2 做门控：小病斑最需要压制叶脉干扰
        self.gate = SaliencyGate(channels=96, patch=cfg.patch)
        self.head = YOLOHead(num_classes=cfg.num_classes, reg_max=cfg.reg_max)

    def forward(self, x: torch.Tensor):
        B, _, H, W = x.shape

        # 1) DINOv3 提取语义 token + CLS-attention
        tokens, attn = self.backbone(x, return_attn=True)

        # 2) 单尺度 → 多尺度金字塔
        feats = self.adapter(tokens, H, W)         # [P2,P3,P4,P5]

        # 3) 显著性门控（只作用于 P2）
        feats[0] = self.gate(feats[0], attn, H, W)

        # 4) 检测头
        preds = self.head(feats)
        return preds, feats, attn


# =============================================================================
# 6. 损失函数
# =============================================================================
def yolo_loss(preds, targets, cfg: Config):
    """占位：真实实现用 ultralytics 的 v8DetectionLoss(分类 BCE + CIoU + DFL)。"""
    loss = 0.0
    for cls_logit, reg_pred in preds:
        loss = loss + cls_logit.mean() + reg_pred.mean()
    return loss * 0.0 + torch.tensor(0.0, requires_grad=True)


def saliency_loss(attn: torch.Tensor, lesion_mask: torch.Tensor) -> torch.Tensor:
    """
    显著性辅助损失：
      让 DINOv3 的 CLS-attention 与 GT 病斑 mask 对齐。
      直觉：模型必须学会"病斑在哪"，叶脉自然被推到低注意力区。
    """
    B, N = attn.shape
    hp = wp = int(N ** 0.5)          # 假设方形 patch 网格
    pred = attn.reshape(B, 1, hp, wp)
    gt = F.interpolate(lesion_mask.float(), size=(hp, wp), mode="area")
    pred = pred.clamp(1e-4, 1 - 1e-4)
    return F.binary_cross_entropy(pred, gt)


def compute_loss(preds, targets, attn, cfg: Config):
    loss_det = yolo_loss(preds, targets, cfg)
    loss_sal = saliency_loss(attn, targets["lesion_mask"])
    return loss_det + cfg.lambda_sal * loss_sal


# =============================================================================
# 7. 训练
# =============================================================================
def train_one_epoch(model, loader, optimizer, device, cfg: Config):
    model.train()
    total = 0.0

    for imgs, targets in loader:
        imgs = imgs.to(device)
        targets = {k: (v.to(device) if torch.is_tensor(v) else v)
                   for k, v in targets.items()}

        # --- 小目标增强：病斑像素占比过低时，放大输入，等效于"放大病斑" ---
        if targets["lesion_ratio"].mean() < cfg.small_ratio_thr:
            imgs = F.interpolate(imgs, scale_factor=1.5,
                                 mode="bilinear", align_corners=False)
            targets["lesion_mask"] = F.interpolate(
                targets["lesion_mask"], scale_factor=1.5,
                mode="nearest"
            )
            # 注意：boxes 也要同步缩放（此处略）

        preds, feats, attn = model(imgs)
        loss = compute_loss(preds, targets, attn, cfg)

        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 10.0)
        optimizer.step()

        total += loss.item()
    return total / max(len(loader), 1)


def build_optimizer(model, lr_backbone=1e-5, lr_head=1e-3):
    """ViT 主干用极小 lr（避免破坏自监督语义），检测头用大 lr。"""
    backbone_params, other_params = [], []
    for n, p in model.named_parameters():
        (backbone_params if n.startswith("backbone.") else other_params).append(p)
    return torch.optim.AdamW([
        {"params": backbone_params, "lr": lr_backbone},
        {"params": other_params,    "lr": lr_head},
    ], weight_decay=0.05)


# =============================================================================
# 8. 推理：滑窗 + 跨窗 NMS
# -----------------------------------------------------------------------------
#  为什么必须滑窗？
#    4000x3000 的原图直接 resize 到 640，几像素病斑会缩到 0.1 像素，
#    信息被彻底抹掉。滑窗保持原始像素尺度，让病斑"活着"进入网络。
# =============================================================================
def decode_yolo(preds, conf_thr: float):
    """占位：把 head 输出解码成 (boxes[N,4], scores[N])。"""
    boxes = torch.tensor([[10., 10., 30., 30.]])
    scores = torch.tensor([0.9])
    return boxes, scores


def nms(boxes: torch.Tensor, scores: torch.Tensor, iou_thr: float = 0.6):
    """占位：真实实现用 torchvision.ops.nms。"""
    if boxes.numel() == 0:
        return torch.empty(0, dtype=torch.long)
    return torch.arange(boxes.shape[0])


@torch.no_grad()
def predict(model, img: torch.Tensor, cfg: Config):
    model.eval()
    B, C, H, W = img.shape
    step = cfg.tile - cfg.overlap

    all_boxes, all_scores = [], []

    for y in range(0, H, step):
        for x in range(0, W, step):
            patch = img[:, :, y:y + cfg.tile, x:x + cfg.tile]
            ph, pw = patch.shape[-2:]
            if ph < cfg.tile or pw < cfg.tile:
                patch = F.pad(patch, (0, cfg.tile - pw, 0, cfg.tile - ph))

            preds, _, _ = model(patch)
            boxes, scores = decode_yolo(preds, cfg.conf_thr)

            # 坐标从"切片坐标系"平移回"原图坐标系"
            boxes[:, [0, 2]] += x
            boxes[:, [1, 3]] += y
            all_boxes.append(boxes)
            all_scores.append(scores)

    boxes = torch.cat(all_boxes, dim=0)
    scores = torch.cat(all_scores, dim=0)
    keep = nms(boxes, scores, cfg.iou_thr)
    return boxes[keep], scores[keep]


# =============================================================================
# 9. main —— 最小可跑 demo
# =============================================================================
def main():
    cfg = Config()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = DINOv3YOLO(cfg).to(device)

    # ---- 伪造一个 batch，验证前向 shape 是否通 ----
    imgs = torch.randn(2, 3, cfg.img_size, cfg.img_size).to(device)
    targets = {
        "lesion_mask": (torch.rand(2, 1, cfg.img_size, cfg.img_size) > 0.98).float(),
        "lesion_ratio": torch.tensor([0.002, 0.003]),   # 典型的"几像素病斑"
    }

    preds, feats, attn = model(imgs)
    print("[OK] P2 shape:", feats[0].shape)   # 期望 (2, 96, 160, 160)
    print("[OK] P3 shape:", feats[1].shape)   # 期望 (2, 192, 80, 80)
    print("[OK] P4 shape:", feats[2].shape)   # 期望 (2, 384, 40, 40)
    print("[OK] P5 shape:", feats[3].shape)   # 期望 (2, 768, 20, 20)
    print("[OK] attn shape:", attn.shape)     # 期望 (2, 1600)

    loss = compute_loss(preds, targets, attn, cfg)
    print("[OK] loss:", float(loss))

    # ---- 推理 demo ----
    boxes, scores = predict(model, imgs[:1], cfg)
    print("[OK] predict boxes:", boxes.shape, "scores:", scores.shape)


if __name__ == "__main__":
    main()


# =============================================================================
# 设计要点速记 (放正文里可当"为什么这样做"的答案)
# -----------------------------------------------------------------------------
# 1. 为什么 DINOv3 能扛叶脉干扰？
#    DINOv3 是自监督训练的 ViT，"叶脉纹理"是低层重复纹理，在高层语义
#    token 里被自然压缩；而"病斑"是与健康组织不同的语义异常 → 更突出。
#
# 2. 为什么还要 SaliencyGate？
#    ViT 的 attention 是全局的，直接接到 CNN 头上，病斑信号会被大量
#    叶脉 patch 稀释。用 CLS→patch attention 当"空间先验"做门控，
#    等于强行告诉网络"先看哪块"，把小病斑的信噪比拉起来。
#
# 3. 为什么必须有 P2？
#    640 输入下 P3=stride8，一个 4px 病斑在 P3 上只剩 0.5 格；
#    stride4 的 P2 是它唯一能被 anchor-free head 采到的机会。
#
# 4. 为什么推理要滑窗？
#    resize 是"像素级信息湮灭"，滑窗是"原始像素尺度采样"。
#    几像素病斑这个任务上，滑窗 + 跨窗 NMS 是硬性要求。
#
# 5. 为什么显著性损失权重只有 0.1？
#    它是"软先验"不是硬约束。权重太大 → 网络只学 attention 不学检测；
#    0.05~0.2 之间是典型的"引导而不绑架"区间。
# =============================================================================