DINOv3助力YOLO26提升NEU-DET缺陷检测,map提升2%

"""
DINOv3 + YOLO26 for NEU-DET Surface Defect Detection
=====================================================
Pseudo-code implementation.
Target: +2% mAP on NEU-DET (hot-rolled steel surface defects).
Key idea: DINOv3 as a frozen parallel feature branch, injected into YOLO26
backbone at P4 (Single) or P3+P4 (Dual). Only lightweight projection heads
and fusion layers are trainable.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


# =============================================================================
# 1. DINOv3 FROZEN FEATURE EXTRACTOR
# =============================================================================

class DINOv3FrozenExtractor(nn.Module):
    """
    Wraps a pretrained DINOv3 ViT model. Fully frozen.
    Extracts patch tokens from intermediate blocks for multi-scale injection.

    DINOv3 outputs:
      - cls_token:    [B, C]
      - patch_tokens: [B, N, C]  (N = (img_size / patch_size)^2)
      - register_tokens: [B, 4, C] (DINOv3 specific)
    """

    def __init__(self, model_name="dinov3_vits14", img_size=640):
        super().__init__()
        # Load official DINOv3 weights (ViT-Small, patch 14)
        self.dinov3 = torch.hub.load("facebookresearch/dinov3", model_name)
        self.patch_size = 14
        self.img_size = img_size

        # Freeze everything -- no gradients through DINOv3
        for param in self.dinov3.parameters():
            param.requires_grad = False

        # Spatial grid: 640 / 14 = 45 (floor) -> 45 x 45 patch tokens
        self.grid_h = img_size // self.patch_size   # 45
        self.grid_w = img_size // self.patch_size   # 45
        self.embed_dim = self.dinov3.embed_dim      # e.g., 384 for ViT-S/14

        # Layer indices for injection (ViT-S has 12 blocks)
        # Shallow block -> spatial detail (for P3)
        # Deep block   -> semantic richness (for P4)
        self.layer_for_p3 = 5    # block index
        self.layer_for_p4 = 11   # last block

    @torch.no_grad()
    def forward(self, x):
        """
        Args:
            x: [B, 3, H, W], H=W=img_size (640)
        Returns:
            feat_p3_raw: [B, embed_dim, grid_h, grid_w]  (shallow patch tokens)
            feat_p4_raw: [B, embed_dim, grid_h, grid_w]  (deep patch tokens)
        """
        B = x.shape[0]

        # DINOv3 forward with intermediate layer outputs
        # Using hook or get_intermediate_layers API
        outputs = self.dinov3.get_intermediate_layers(
            x,
            n=[self.layer_for_p3, self.layer_for_p4],
            reshape=True,           # return [B, C, H, W]
            return_class_token=False,
        )
        feat_p3_raw = outputs[0]   # [B, C, 45, 45]
        feat_p4_raw = outputs[1]   # [B, C, 45, 45]

        return feat_p3_raw, feat_p4_raw


# =============================================================================
# 2. LIGHTWEIGHT PROJECTION & FUSION MODULES
# =============================================================================

class DINOv3Projector(nn.Module):
    """
    Lightweight projection: maps DINOv3 channel dim -> YOLO channel dim,
    then resizes spatial resolution to match target P-level.
    """

    def __init__(self, in_dim, out_dim, target_stride):
        super().__init__()
        # 1x1 conv: channel alignment only
        self.proj = nn.Conv2d(in_dim, out_dim, kernel_size=1, bias=False)
        self.bn = nn.BatchNorm2d(out_dim)
        self.act = nn.SiLU()

        # Spatial target stride (P3=8, P4=16)
        self.target_stride = target_stride

    def forward(self, x, input_size=640):
        """
        Args:
            x: [B, in_dim, grid_h, grid_w]  (e.g., 45x45)
        Returns:
            [B, out_dim, target_h, target_w]  (e.g., 80x80 for P3)
        """
        x = self.act(self.bn(self.proj(x)))

        # Upsample to match YOLO feature map resolution
        target_h = input_size // self.target_stride
        target_w = input_size // self.target_stride
        if x.shape[-2:] != (target_h, target_w):
            x = F.interpolate(
                x, size=(target_h, target_w),
                mode="bilinear", align_corners=False
            )
        return x


class FusionBlock(nn.Module):
    """
    Fuses YOLO CNN feature map with projected DINOv3 features.
    Simple additive fusion with a learnable gate.
    """

    def __init__(self, channels):
        super().__init__()
        self.gate = nn.Sequential(
            nn.Conv2d(channels * 2, channels, 1, bias=False),
            nn.BatchNorm2d(channels),
            nn.Sigmoid()
        )

    def forward(self, yolo_feat, dino_feat):
        """
        Args:
            yolo_feat: [B, C, H, W]
            dino_feat: [B, C, H, W]  (already projected & resized)
        Returns:
            fused: [B, C, H, W]
        """
        g = self.gate(torch.cat([yolo_feat, dino_feat], dim=1))
        return yolo_feat + g * dino_feat


# =============================================================================
# 3. FULL MODEL: YOLO26 + DINOv3 PARALLEL BRANCH
# =============================================================================

class YOLO26DINOv3(nn.Module):
    """
    YOLO26 with a frozen DINOv3 parallel branch.

    Fusion modes:
      'p4'     : inject at P4 only (Single)
      'p3p4'   : inject at P3 and P4 (Dual)

    YOLO26 backbone channels (YOLO26n):
      P3 (stride 8):  64 channels
      P4 (stride 16): 128 channels
    """

    def __init__(self, yolo_model, fusion_mode="p3p4",
                 dinov3_name="dinov3_vits14", img_size=640):
        super().__init__()
        self.yolo = yolo_model
        self.fusion_mode = fusion_mode
        self.img_size = img_size

        # --- DINOv3 branch ---
        self.dinov3_extractor = DINOv3FrozenExtractor(dinov3_name, img_size)
        dino_dim = self.dinov3_extractor.embed_dim  # 384 for ViT-S/14

        # --- Projection heads (trainable) ---
        # YOLO26n backbone output channels at P3 and P4
        c_p3 = 64    # stride 8
        c_p4 = 128   # stride 16

        if fusion_mode in ("p4", "p3p4"):
            self.proj_p4 = DINOv3Projector(dino_dim, c_p4, target_stride=16)
            self.fuse_p4 = FusionBlock(c_p4)

        if fusion_mode == "p3p4":
            self.proj_p3 = DINOv3Projector(dino_dim, c_p3, target_stride=8)
            self.fuse_p3 = FusionBlock(c_p3)

    def forward(self, x):
        """
        Returns YOLO26 detection outputs with DINOv3-enhanced backbone features.
        """
        # --- 1. DINOv3 branch (no gradients) ---
        dino_feat_p3, dino_feat_p4 = self.dinov3_extractor(x)

        # --- 2. Project DINOv3 features to YOLO channel space & resolution ---
        if self.fusion_mode == "p3p4":
            dino_p3 = self.proj_p3(dino_feat_p3, self.img_size)  # [B, 64, 80, 80]
        dino_p4 = self.proj_p4(dino_feat_p4, self.img_size)      # [B, 128, 40, 40]

        # --- 3. YOLO26 backbone forward (up to the fusion points) ---
        # In practice, this is implemented via YOLO's forward hooks
        # or by splitting the backbone forward into stages.
        yolo_feats = self.yolo.backbone_forward(x)

        # --- 4. Fusion at P3 and/or P4 ---
        if self.fusion_mode == "p3p4":
            yolo_feats["P3"] = self.fuse_p3(yolo_feats["P3"], dino_p3)

        if self.fusion_mode in ("p4", "p3p4"):
            yolo_feats["P4"] = self.fuse_p4(yolo_feats["P4"], dino_p4)

        # --- 5. YOLO26 neck + head ---
        detections = self.yolo.neck_head_forward(yolo_feats)
        return detections


# =============================================================================
# 4. TRAINING UTILITIES
# =============================================================================

def build_optimizer(model):
    """
    Only train: projection heads + fusion blocks.
    DINOv3: frozen. YOLO26 backbone/neck/head: optional unfreeze or frozen.
    """
    trainable_params = []
    for name, param in model.named_parameters():
        if param.requires_grad:
            trainable_params.append(param)
        # Uncomment below to also fine-tune YOLO head
        # if "yolo.head" in name or "yolo.neck" in name:
        #     param.requires_grad = True
        #     trainable_params.append(param)

    optimizer = torch.optim.AdamW(
        trainable_params,
        lr=1e-3,
        weight_decay=1e-4
    )
    return optimizer


def train_one_epoch(model, dataloader, optimizer, device, epoch, scaler=None):
    model.train()
    total_loss = 0.0

    for batch_idx, (images, targets) in enumerate(dataloader):
        images = images.to(device)
        targets = [t.to(device) for t in targets]

        optimizer.zero_grad()

        # Mixed precision: DINOv3 runs in fp16, projections in fp32
        with torch.cuda.amp.autocast(enabled=(scaler is not None)):
            preds = model(images)
            loss, loss_items = compute_yolo_loss(preds, targets)

        if scaler is not None:
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            optimizer.step()

        total_loss += loss.item()

    return total_loss / len(dataloader)


def compute_yolo_loss(preds, targets):
    """
    Wrapper around YOLO26's native loss (e.g., TAL / STAL).
    """
    from ultralytics.utils.loss import v8DetectionLoss
    criterion = v8DetectionLoss(...)  # configured with model head
    return criterion(preds, targets)


# =============================================================================
# 5. EVALUATION (mAP ON NEU-DET)
# =============================================================================

@torch.no_grad()
def evaluate(model, dataloader, device, iou_thresholds=(0.5, 0.75)):
    """
    Compute mAP@0.5, mAP@0.5:0.95, per-class AP on NEU-DET.
    NEU-DET classes: crazing, inclusion, patches, pitted_surface,
                     rolled-in_scale, scratches
    """
    model.eval()

    all_preds = []
    all_targets = []

    for images, targets in dataloader:
        images = images.to(device)
        preds = model(images)
        all_preds.append(preds.cpu())
        all_targets.append(targets)

    # Use pycocotools or torchmetrics for mAP computation
    from torchmetrics.detection import MeanAveragePrecision
    metric = MeanAveragePrecision(
        iou_thresholds=[0.5, 0.75],
        class_metrics=True
    )
    metric.update(all_preds, all_targets)
    results = metric.compute()

    print(f"mAP@0.5:      {results['map_50']:.4f}")
    print(f"mAP@0.5:0.95: {results['map']:.4f}")
    print(f"mAP@0.75:     {results['map_75']:.4f}")
    return results


# =============================================================================
# 6. MAIN PIPELINE
# =============================================================================

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # --- Data ---
    # NEU-DET: 1800 images, 6 defect classes, 200x200 px originals
    # Resize to 640x640 with letterbox
    train_loader, val_loader = build_neu_det_dataloaders(
        data_yaml="data/NEU-DET.yaml",
        imgsz=640,
        batch_size=16
    )

    # --- Model ---
    yolo26 = load_yolo26("yolo26n.pt")           # pretrained YOLO26n
    model = YOLO26DINOv3(
        yolo_model=yolo26,
        fusion_mode="p3p4",                       # Dual P3+P4 injection
        dinov3_name="dinov3_vits14",
        img_size=640
    ).to(device)

    # --- Optimizer ---
    optimizer = build_optimizer(model)
    scaler = torch.cuda.amp.GradScaler()

    # --- Training loop ---
    best_map50 = 0.0
    NUM_EPOCHS = 200

    for epoch in range(NUM_EPOCHS):
        train_loss = train_one_epoch(
            model, train_loader, optimizer, device, epoch, scaler
        )
        results = evaluate(model, val_loader, device)
        current_map50 = results["map_50"].item()

        print(f"Epoch {epoch+1}/{NUM_EPOCHS} | Loss: {train_loss:.4f} | "
              f"mAP@0.5: {current_map50:.4f}")

        if current_map50 > best_map50:
            best_map50 = current_map50
            torch.save(model.state_dict(), "best_yolo26_dinov3_neu_det.pt")
            print(f"  -> Best model saved (mAP@0.5: {best_map50:.4f})")

    print(f"\nTraining complete. Best mAP@0.5: {best_map50:.4f}")

    # --- Expected improvement ---
    # YOLO26n baseline on NEU-DET: ~0.722 mAP@0.5
    # + DINOv3 Dual P3+P4:        ~0.737 mAP@0.5  (+2.1%)


if __name__ == "__main__":
    main()

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO


if __name__ == '__main__':
    model = YOLO('ultralytics/cfg/models/26/yolo26n.yaml')
    #model.load('yolo26n.pt') # loading pretrain weights
    model.train(data='data/NEU-DET.yaml',
                cache=False,
                imgsz=640,
                epochs=300,
                batch=16,
                close_mosaic=10,
                workers=2,
                device='0',
                #optimizer='SGD',   # using SGD,auto
                project='runs/train',
                name='exp',
                )

YOLO26n summary (fused): 122 layers, 2,376,006 parameters, 0 gradients, 5.2 GFLOPs
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100% ━━━━━━━━━━━━ 16/16 3.0it/s 5.4s
                   all        486       1069      0.745      0.643      0.722      0.407
               crazing         69        149      0.615      0.215      0.367      0.139
             inclusion         87        222      0.784      0.725      0.796      0.396
               patches         92        243      0.834      0.848      0.918      0.605
        pitted_surface         93        130      0.823      0.716      0.788      0.499
       rolled-in_scale         81        171      0.643      0.538      0.615       0.29
             scratches         90        154       0.77      0.818       0.85       0.51

YOLO26n-dinov3 summary (fused): 122 layers, 2,376,006 parameters, 0 gradients, 5.3 GFLOPs
                 Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100% ━━━━━━━━━━━━ 16/16 3.1it/s 5.1s
                   all        486       1069      0.719      0.705      0.743      0.413
               crazing         69        149       0.46      0.356      0.392      0.136
             inclusion         87        222      0.734      0.758       0.79      0.421
               patches         92        243      0.829       0.86      0.927       0.59
        pitted_surface         93        130      0.861      0.738      0.814      0.507
       rolled-in_scale         81        171      0.606      0.658      0.641      0.305
             scratches         90        154       0.82      0.861      0.896      0.518