AP50提升2.2%、精确率提升3.9%！CF-YOLO上下文感知特征细化，破解工业微缺陷伪装难题

"""
CF-YOLO: Context-Aware Feature Refinement for Camouflaged Industrial Micro-Defect Detection
PyTorch-style Pseudocode Implementation

This pseudocode outlines the core architectural components and training/inference flow
as described in the paper. For brevity, complex operations are represented with placeholder
functions that convey the intended logic rather than full CUDA-optimized implementations.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
import numpy as np

# ============================================================================
# 1. Configuration
# ============================================================================
class Config:
    # Dataset
    num_classes = 1                 # single-class defect detection
    img_size = 640
    batch_size = 16
    epochs = 200
    
    # Model architecture
    backbone_channels = [64, 128, 256, 512]   # typical YOLO stages
    cpam_kernel_large = 11
    cpam_kernel_small = 7
    farm_hidden_dim = 256
    
    # Loss weights
    lambda_box = 7.5
    lambda_cls = 0.5
    lambda_dfl = 1.5
    
    # Training
    lr = 0.01
    momentum = 0.937
    weight_decay = 0.0005
    nms_threshold = 0.45
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# ============================================================================
# 2. Context-Perception Aggregation Module (CPAM)
# ============================================================================
class CPAM(nn.Module):
    """
    CPAM decouples feature processing into large-kernel perception (LKP)
    and small-kernel aggregation (SKA) to capture macro-texture and refine local details.
    """
    def __init__(self, in_channels, out_channels, kernel_large=11, kernel_small=7, groups=4):
        super().__init__()
        # LKP branch
        self.pw1 = nn.Conv2d(in_channels, in_channels//2, 1)
        self.dw_large = nn.Conv2d(in_channels//2, in_channels//2, kernel_large, padding=kernel_large//2, groups=in_channels//2)
        self.pw2 = nn.Conv2d(in_channels//2, out_channels, 1)
        
        # SKA branch: group-wise dynamic convolution will be applied during forward
        self.groups = groups
        self.kernel_small = kernel_small
        self.out_channels = out_channels
        
    def forward(self, x):
        # Input: x [B, C, H, W]
        B, C, H, W = x.shape
        
        # ---- LKP stage ----
        # Pointwise reduce channels
        x_reduced = self.pw1(x)                     # [B, C/2, H, W]
        # Large-kernel depthwise conv to perceive macro-texture
        x_context = self.dw_large(x_reduced)        # [B, C/2, H, W]
        # Generate dynamic weight tensor W_agg
        W_agg = self.pw2(x_context)                 # [B, out_channels, H, W]
        
        # ---- SKA stage ----
        # Reshape W_agg into dynamic kernels of size (kernel_small x kernel_small)
        # For simplicity, we simulate group-wise dynamic convolution using unfold + matmul
        # In practice, this would be a custom CUDA kernel for efficiency.
        # Here we provide a high-level logical implementation.
        
        # Pad input to maintain spatial dimensions
        pad = self.kernel_small // 2
        x_pad = F.pad(x, (pad, pad, pad, pad))      # [B, C, H+pad*2, W+pad*2]
        
        # Extract local patches (unfold)
        patches = F.unfold(x_pad, kernel_size=self.kernel_small, stride=1)  # [B, C*K*K, H*W]
        patches = patches.view(B, -1, self.kernel_small*self.kernel_small, H*W)  # [B, C, K*K, N]
        
        # Group-wise split
        group_size = C // self.groups
        patches_groups = patches.chunk(self.groups, dim=1)  # list of [B, group_size, K*K, N]
        # W_agg also needs to be grouped: W_agg [B, out_channels, H, W] -> we use it as group-wise kernels
        # We approximate: use W_agg to weight the patches
        # For each group, compute weighted sum over spatial positions
        # Simplified: element-wise multiply patches with appropriate weights (omitted complex reshape)
        # Pseudocode:
        out = []
        for g in range(self.groups):
            patch_g = patches_groups[g]              # [B, gC, K*K, N]
            # Generate dynamic weights for this group
            # W_agg slice: take corresponding output channels (out_channels/g)
            w_g = W_agg[:, g*self.out_channels//self.groups : (g+1)*self.out_channels//self.groups, :, :]  # [B, out_g, H, W]
            # Reshape w_g to [B, out_g, 1, N] (N=H*W)
            w_g_flat = w_g.view(B, -1, 1, H*W)      # [B, out_g, 1, N]
            # Weighted aggregation: sum over kernel positions (weighted by w_g)
            # This is a simplified representation: actual implementation uses depthwise conv with dynamic kernels
            aggregated = torch.sum(patch_g * w_g_flat, dim=2, keepdim=True)  # [B, out_g, 1, N]
            out.append(aggregated)
        
        # Concatenate groups and reshape back
        out = torch.cat(out, dim=1)                  # [B, out_channels, 1, H*W]
        out = out.view(B, self.out_channels, H, W)   # [B, out_channels, H, W]
        
        # Residual connection (if in_channels == out_channels, else no residual)
        if in_channels == self.out_channels:
            out = out + x
        
        return out

# ============================================================================
# 3. Feature Additive Refinement Module (FARM)
# ============================================================================
class FARM(nn.Module):
    """
    FARM uses a linear-complexity additive token mixer to globally verify and refine features.
    It consists of spatial and channel interaction branches.
    """
    def __init__(self, dim, hidden_dim=None):
        super().__init__()
        self.dim = dim
        self.hidden_dim = hidden_dim or dim
        
        # Query, Key, Value projections
        self.q_proj = nn.Linear(dim, dim)
        self.k_proj = nn.Linear(dim, dim)
        self.v_proj = nn.Linear(dim, dim)
        
        # Spatial branch: 3x3 depthwise conv + sigmoid
        self.dw_spatial = nn.Conv2d(dim, dim, 3, padding=1, groups=dim)
        # Channel branch: global avg pool + pointwise conv + sigmoid
        self.pw_channel = nn.Conv2d(dim, dim, 1)
        
        # Output projection
        self.out_proj = nn.Linear(dim, dim)
        
    def forward(self, x):
        # x: [B, C, H, W] -> flatten to [B, N, C] where N=H*W
        B, C, H, W = x.shape
        N = H * W
        x_flat = x.flatten(2).transpose(1, 2)       # [B, N, C]
        
        # Project to Q, K, V
        Q = self.q_proj(x_flat)                     # [B, N, C]
        K = self.k_proj(x_flat)
        V = self.v_proj(x_flat)
        
        # ---- Context mapping Phi(T) ----
        # Spatial interaction: 3x3 depthwise conv on spatial grid (use 2D conv)
        # Reshape Q and K to 2D for spatial conv
        Q_spatial = Q.transpose(1,2).view(B, C, H, W)   # [B, C, H, W]
        K_spatial = K.transpose(1,2).view(B, C, H, W)
        
        # Spatial branch: depthwise conv + sigmoid
        Q_spat = torch.sigmoid(self.dw_spatial(Q_spatial))  # spatial attention from Q
        K_spat = torch.sigmoid(self.dw_spatial(K_spatial))
        # Combine spatial terms (addition)
        spat_term = Q_spat + K_spat
        
        # Channel branch: global avg pool + pointwise conv + sigmoid, then multiply with original
        Q_chan = torch.sigmoid(self.pw_channel(Q_spatial.mean(dim=(2,3), keepdim=True))) * Q_spatial
        K_chan = torch.sigmoid(self.pw_channel(K_spatial.mean(dim=(2,3), keepdim=True))) * K_spatial
        
        # Combine channel terms
        chan_term = Q_chan + K_chan
        
        # Overall Phi(Q) + Phi(K) = spat_term + chan_term
        phi_sum = spat_term + chan_term          # [B, C, H, W]
        
        # Flatten back to [B, N, C]
        phi_sum_flat = phi_sum.flatten(2).transpose(1,2)  # [B, N, C]
        
        # Modulate Value branch: element-wise multiply
        refined = phi_sum_flat * V                # [B, N, C]
        
        # Output projection
        refined = self.out_proj(refined)          # [B, N, C]
        
        # Reshape back to [B, C, H, W]
        refined = refined.transpose(1,2).view(B, C, H, W)
        
        # Residual connection
        return x + refined

# ============================================================================
# 4. CF-YOLO Backbone with CPAM integrated
# ============================================================================
class CFYOLOBackbone(nn.Module):
    """
    Backbone alternates C3k2 blocks (standard YOLO) and CPAM modules.
    """
    def __init__(self, in_channels=3, out_channels_list=[64, 128, 256, 512]):
        super().__init__()
        # This is a placeholder; actual YOLO backbone would be more complex.
        # We simulate stages.
        self.stage1 = nn.Sequential(
            nn.Conv2d(in_channels, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.SiLU(),
        )
        self.stage2 = self._make_stage(64, 128, 2)   # C3k2-like
        self.stage3 = self._make_stage(128, 256, 2)
        self.stage4 = self._make_stage(256, 512, 2)
        
        # CPAM modules inserted after certain stages (according to paper)
        self.cpam3 = CPAM(128, 128)   # after stage2
        self.cpam4 = CPAM(256, 256)   # after stage3
        self.cpam5 = CPAM(512, 512)   # after stage4
        
    def _make_stage(self, in_ch, out_ch, num_blocks):
        # Simplified C3k2 (just a few conv layers)
        layers = []
        layers.append(nn.Conv2d(in_ch, out_ch, 3, stride=2, padding=1))
        layers.append(nn.BatchNorm2d(out_ch))
        layers.append(nn.SiLU())
        for _ in range(num_blocks):
            layers.append(nn.Conv2d(out_ch, out_ch, 3, padding=1))
            layers.append(nn.BatchNorm2d(out_ch))
            layers.append(nn.SiLU())
        return nn.Sequential(*layers)
    
    def forward(self, x):
        # Feature pyramid outputs: F3, F4, F5
        x1 = self.stage1(x)          # stride 2
        x2 = self.stage2(x1)         # stride 4
        F3 = self.cpam3(x2)          # stride 8? (depending on design)
        x3 = self.stage3(F3)
        F4 = self.cpam4(x3)          # stride 16
        x4 = self.stage4(F4)
        F5 = self.cpam5(x4)          # stride 32
        return F3, F4, F5

# ============================================================================
# 5. FARM applied on F5 before neck
# ============================================================================
class CFYOLONeck(nn.Module):
    """
    Neck performs cross-scale feature fusion. FARM is applied to F5 before fusion.
    """
    def __init__(self, channels=[128, 256, 512]):
        super().__init__()
        self.farm = FARM(channels[-1])   # apply to highest-level feature (F5)
        # Then standard neck (e.g., PANet) would be here.
        # For brevity, we skip detailed implementation.
    
    def forward(self, F3, F4, F5):
        # Apply FARM on F5
        F5_refined = self.farm(F5)
        # Return refined pyramid (for simplicity, we just return them)
        return F3, F4, F5_refined

# ============================================================================
# 6. Decoupled Head (Classification + Regression)
# ============================================================================
class DecoupledHead(nn.Module):
    """
    Decoupled detection head with separate branches for classification and regression.
    """
    def __init__(self, in_channels, num_classes, num_anchors=1, reg_max=16):
        super().__init__()
        self.num_classes = num_classes
        self.reg_max = reg_max
        
        # Classification branch: DWConv + 1x1
        self.cls_conv = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, 3, padding=1, groups=in_channels),
            nn.Conv2d(in_channels, num_classes * num_anchors, 1)
        )
        # Regression branch: Conv3x3 + 1x1
        self.reg_conv = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, 3, padding=1),
            nn.Conv2d(in_channels, 4 * reg_max * num_anchors, 1)
        )
    
    def forward(self, x):
        cls_out = self.cls_conv(x)   # [B, num_classes, H, W]
        reg_out = self.reg_conv(x)   # [B, 4*reg_max, H, W]
        return cls_out, reg_out

# ============================================================================
# 7. Full CF-YOLO Model
# ============================================================================
class CFYOLO(nn.Module):
    def __init__(self, num_classes=1, config=Config):
        super().__init__()
        self.backbone = CFYOLOBackbone()
        self.neck = CFYOLONeck()
        # Detection heads for each scale
        self.heads = nn.ModuleList([
            DecoupledHead(128, num_classes),
            DecoupledHead(256, num_classes),
            DecoupledHead(512, num_classes)
        ])
        self.nms_thresh = config.nms_threshold
        
    def forward(self, x):
        # Extract features
        F3, F4, F5 = self.backbone(x)
        # Neck with FARM
        F3, F4, F5_refined = self.neck(F3, F4, F5)
        
        # Predictions
        cls3, reg3 = self.heads[0](F3)
        cls4, reg4 = self.heads[1](F4)
        cls5, reg5 = self.heads[2](F5_refined)
        
        # In training, return losses; in inference, decode and NMS
        if self.training:
            return cls3, reg3, cls4, reg4, cls5, reg5
        else:
            # Decode and NMS (pseudo-code)
            detections = self._decode_and_nms(cls3, reg3, cls4, reg4, cls5, reg5)
            return detections
    
    def _decode_and_nms(self, cls3, reg3, cls4, reg4, cls5, reg5):
        # Combine predictions from all scales, decode boxes, apply NMS
        # This is a placeholder for actual decoding logic.
        boxes = []
        scores = []
        # ... (omitted for brevity)
        return boxes, scores

# ============================================================================
# 8. Loss Function (CIoU + BCE + DFL)
# ============================================================================
def compute_loss(cls_pred, reg_pred, targets, config):
    """
    Compute total loss: L_CIoU + L_BCE + L_DFL
    """
    # Placeholder: actual implementation would match predictions to targets.
    loss_box = torch.tensor(0.0)
    loss_cls = torch.tensor(0.0)
    loss_dfl = torch.tensor(0.0)
    # ... (detailed matching and loss computation)
    total_loss = config.lambda_box * loss_box + config.lambda_cls * loss_cls + config.lambda_dfl * loss_dfl
    return total_loss

# ============================================================================
# 9. Training Loop (Pseudocode)
# ============================================================================
def train():
    config = Config()
    model = CFYOLO(num_classes=config.num_classes, config=config).to(config.device)
    optimizer = torch.optim.SGD(model.parameters(), lr=config.lr, momentum=config.momentum, weight_decay=config.weight_decay)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.epochs)
    
    # Dummy dataloader
    train_loader = DataLoader(...)
    
    for epoch in range(config.epochs):
        model.train()
        for batch_idx, (images, targets) in enumerate(train_loader):
            images = images.to(config.device)
            # targets: list of dicts with boxes, labels
            
            # Forward
            cls3, reg3, cls4, reg4, cls5, reg5 = model(images)
            
            # Compute loss
            loss = compute_loss((cls3, reg3, cls4, reg4, cls5, reg5), targets, config)
            
            # Backward
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        
        scheduler.step()
        print(f"Epoch {epoch+1}/{config.epochs}, Loss: {loss.item():.4f}")

# ============================================================================
# 10. Inference (Pseudocode)
# ============================================================================
def inference(model, image):
    model.eval()
    with torch.no_grad():
        detections = model(image.unsqueeze(0).to(Config.device))
        # detections contain boxes and scores after NMS
    return detections

# ============================================================================
# Entry point
# ============================================================================
if __name__ == '__main__':
    # This is just a placeholder to show the structure.
    print("CF-YOLO Pseudocode Structure")
    # train() would be called if dataset is available.