0.08M参数吊打YOLO！MicroCharNet车牌字符检测mAP@50达0.85，推理速度1ms，轻量王者

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


class C2fBlock(nn.Module):
    """
    C2f block: split-transform-concatenate paradigm.
    Paper Eq. (1)-(3): X1,X2 = split(Conv1x1(X)); X̂2 = X2 + Conv3x3(Conv3x3(X2));
    Y = Conv1x1(Concat(X1, X̂2))
    """
    def __init__(self, in_channels, out_channels, hidden_ratio=0.5):
        super().__init__()
        hidden_channels = int(in_channels * hidden_ratio)
        
        self.proj = nn.Conv2d(in_channels, hidden_channels * 2, kernel_size=1)
        self.conv1 = nn.Conv2d(hidden_channels, hidden_channels, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(hidden_channels, hidden_channels, kernel_size=3, padding=1)
        self.fuse = nn.Conv2d(hidden_channels * 2, out_channels, kernel_size=1)
        
    def forward(self, x):
        # Project and split
        x_proj = self.proj(x)                    # [B, 2*C', H, W]
        x1, x2 = x_proj.chunk(2, dim=1)          # [B, C', H, W], [B, C', H, W]
        
        # Transform branch
        x2_hat = self.conv2(self.conv1(x2))      # [B, C', H, W]
        x2_hat = x2 + x2_hat                     # Residual connection (Eq. 2)
        
        # Concatenate and fuse (Eq. 3)
        out = torch.cat([x1, x2_hat], dim=1)     # [B, 2*C', H, W]
        out = self.fuse(out)                     # [B, out_channels, H, W]
        return out


class CoordAtt(nn.Module):
    """
    Coordinate Attention module.
    Decomposes contextual encoding into two separate spatial directions.
    """
    def __init__(self, in_channels, reduction=32):
        super().__init__()
        self.reduction = reduction
        reduced_channels = max(1, in_channels // reduction)
        
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))
        
        self.shared_conv = nn.Sequential(
            nn.Conv2d(in_channels, reduced_channels, kernel_size=1),
            nn.BatchNorm2d(reduced_channels),
            nn.Hardswish()
        )
        self.conv_h = nn.Conv2d(reduced_channels, in_channels, kernel_size=1)
        self.conv_w = nn.Conv2d(reduced_channels, in_channels, kernel_size=1)
        
    def forward(self, x):
        B, C, H, W = x.shape
        
        # Height-aware and width-aware pooling
        x_h = self.pool_h(x)                     # [B, C, H, 1]
        x_w = self.pool_w(x).permute(0, 1, 3, 2) # [B, C, W, 1] -> [B, C, 1, W]
        
        # Shared transformation
        y = torch.cat([x_h, x_w], dim=2)         # [B, C, H+W, 1]? Actually [B,C,H+W,1]? 
        # Correct implementation: concatenate along spatial dimension
        # x_h: [B,C,H,1], x_w: [B,C,1,W] -> need to permute x_w to [B,C,W,1]
        x_w = x_w.permute(0, 1, 3, 2)            # [B, C, W, 1]
        y = torch.cat([x_h, x_w], dim=2)         # [B, C, H+W, 1]
        y = self.shared_conv(y)                  # [B, C//r, H+W, 1]
        
        # Split back
        y_h, y_w = torch.split(y, [H, W], dim=2)
        y_h = y_h.unsqueeze(3)                   # [B, C//r, H, 1]
        y_w = y_w.unsqueeze(3)                   # [B, C//r, W, 1]
        
        # Generate attention maps
        a_h = torch.sigmoid(self.conv_h(y_h))    # [B, C, H, 1]
        a_w = torch.sigmoid(self.conv_w(y_w))    # [B, C, 1, W]
        
        # Apply attention (Eq. 4)
        out = x * a_h * a_w
        return out


class C3k2Block(nn.Module):
    """
    C3k2: lightweight variant derived from C2f structure.
    Enhanced feature aggregation with short information paths.
    """
    def __init__(self, in_channels, out_channels, hidden_ratio=0.5):
        super().__init__()
        hidden_channels = int(in_channels * hidden_ratio)
        
        self.proj = nn.Conv2d(in_channels, hidden_channels, kernel_size=1)
        self.conv1 = nn.Conv2d(hidden_channels, hidden_channels, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(hidden_channels, hidden_channels, kernel_size=3, padding=1)
        self.fuse = nn.Conv2d(hidden_channels + hidden_channels, out_channels, kernel_size=1)
        
    def forward(self, x):
        # Short path + local transformations
        x_proj = self.proj(x)
        x1 = self.conv1(x_proj)
        x2 = self.conv2(x1)
        
        # Concatenate short path and transformed features
        out = torch.cat([x1, x2], dim=1)
        out = self.fuse(out)
        return out


class DetectionHead(nn.Module):
    """
    Single-level anchor-free detection head.
    reg_max=1 for direct bounding-box regression (end-to-end).
    """
    def __init__(self, in_channels, num_classes=36, reg_max=1):
        super().__init__()
        self.num_classes = num_classes
        self.reg_max = reg_max
        
        # Shared convolution
        self.shared_conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
        
        # Classification branch
        self.cls_conv = nn.Conv2d(in_channels, num_classes, kernel_size=1)
        
        # Regression branch (4 coordinates * reg_max)
        self.reg_conv = nn.Conv2d(in_channels, 4 * reg_max, kernel_size=1)
        
    def forward(self, x):
        x = self.shared_conv(x)
        
        # Class predictions: [B, num_classes, H, W]
        cls_pred = self.cls_conv(x)
        
        # Box predictions: [B, 4*reg_max, H, W]
        reg_pred = self.reg_conv(x)
        
        # Reshape to [B, H, W, 4] for bounding boxes
        B, _, H, W = reg_pred.shape
        reg_pred = reg_pred.view(B, 4, self.reg_max, H, W).sum(dim=2)  # reg_max=1 -> direct
        reg_pred = reg_pred.permute(0, 2, 3, 1)  # [B, H, W, 4]
        
        # cls_pred: [B, num_classes, H, W] -> [B, H, W, num_classes]
        cls_pred = cls_pred.permute(0, 2, 3, 1)
        
        return cls_pred, reg_pred


class MicroCharNet(nn.Module):
    """
    Ultra-lightweight model for license plate character detection.
    Total: ~0.08M parameters, 0.096 GFLOPs.
    Input: 128x128 RGB image.
    """
    def __init__(self, num_classes=36, input_channels=3):
        super().__init__()
        
        # ---- Backbone ----
        # Stem
        self.stem = nn.Conv2d(input_channels, 16, kernel_size=3, stride=2, padding=1)
        
        # Stage 1
        self.conv1 = nn.Conv2d(16, 32, kernel_size=3, stride=2, padding=1)
        self.c2f1 = C2fBlock(32, 32)
        
        # Stage 2 (deepest)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
        self.c2f2 = C2fBlock(64, 64)
        
        # Coordinate Attention
        self.coord_att = CoordAtt(64)
        
        # ---- Neck ----
        # Upsample and fuse (deep -> shallow)
        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
        self.fusion_conv = nn.Conv2d(64 + 32, 64, kernel_size=1)
        self.c3k2 = C3k2Block(64, 64)
        
        # ---- Detection Head ----
        self.head = DetectionHead(64, num_classes=num_classes, reg_max=1)
        
    def forward(self, x):
        # ---- Backbone ----
        # Stem
        x = self.stem(x)                        # [B, 16, H/2, W/2]
        
        # Stage 1 (shallow feature for fusion)
        x1 = self.conv1(x)                      # [B, 32, H/4, W/4]
        x1 = self.c2f1(x1)                     # [B, 32, H/4, W/4]
        
        # Stage 2 (deep feature)
        x2 = self.conv2(x1)                     # [B, 64, H/8, W/8]
        x2 = self.c2f2(x2)                     # [B, 64, H/8, W/8]
        
        # Coordinate Attention
        x2 = self.coord_att(x2)                 # [B, 64, H/8, W/8]
        
        # ---- Neck ----
        # Upsample deep feature and fuse with shallow
        x_up = self.upsample(x2)                # [B, 64, H/4, W/4]
        x_fused = torch.cat([x_up, x1], dim=1)  # [B, 96, H/4, W/4]
        x_fused = self.fusion_conv(x_fused)     # [B, 64, H/4, W/4]
        x_fused = self.c3k2(x_fused)            # [B, 64, H/4, W/4]
        
        # ---- Detection Head ----
        cls_pred, reg_pred = self.head(x_fused)  # [B, H/4, W/4, C], [B, H/4, W/4, 4]
        
        return cls_pred, reg_pred


def decode_predictions(cls_pred, reg_pred, threshold=0.5):
    """
    Decode model outputs to bounding boxes and class labels.
    
    Args:
        cls_pred: [B, H, W, num_classes] class logits
        reg_pred: [B, H, W, 4] bounding box offsets (cx, cy, w, h) in grid coordinates
        threshold: confidence threshold
    
    Returns:
        List of detection results per image: [{'boxes': [...], 'labels': [...], 'scores': [...]}]
    """
    batch_size = cls_pred.shape[0]
    H, W = cls_pred.shape[1], cls_pred.shape[2]
    results = []
    
    for b in range(batch_size):
        boxes, labels, scores = [], [], []
        
        # Get class probabilities
        probs = torch.softmax(cls_pred[b], dim=-1)  # [H, W, num_classes]
        max_probs, max_labels = probs.max(dim=-1)    # [H, W]
        
        # Find detections above threshold
        mask = max_probs > threshold
        indices = torch.nonzero(mask)                # [N, 2] (h, w)
        
        for h, w in indices:
            # Grid coordinates (center of cell)
            cx = (w + 0.5) / W
            cy = (h + 0.5) / H
            
            # Decode bounding box (relative to image)
            dx, dy, dw, dh = reg_pred[b, h, w]
            cx_pred = cx + dx
            cy_pred = cy + dy
            w_pred = dw
            h_pred = dh
            
            # Clamp to valid range
            x1 = max(0, cx_pred - w_pred / 2)
            y1 = max(0, cy_pred - h_pred / 2)
            x2 = min(1, cx_pred + w_pred / 2)
            y2 = min(1, cy_pred + h_pred / 2)
            
            boxes.append([x1, y1, x2, y2])
            labels.append(max_labels[h, w].item())
            scores.append(max_probs[h, w].item())
        
        results.append({
            'boxes': torch.tensor(boxes) if boxes else torch.empty(0, 4),
            'labels': torch.tensor(labels) if labels else torch.empty(0, dtype=torch.long),
            'scores': torch.tensor(scores) if scores else torch.empty(0)
        })
    
    return results


# ============ Usage Example ============
if __name__ == "__main__":
    # Create model
    model = MicroCharNet(num_classes=36)
    
    # Count parameters
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Total parameters: {total_params:,} (~{total_params/1e6:.2f}M)")
    # Expected: ~0.08M
    
    # Forward pass
    dummy_input = torch.randn(1, 3, 128, 128)
    cls_pred, reg_pred = model(dummy_input)
    
    print(f"Class predictions shape: {cls_pred.shape}")  # [1, 32, 32, 36]
    print(f"Box predictions shape: {reg_pred.shape}")    # [1, 32, 32, 4]
    
    # Decode predictions
    results = decode_predictions(cls_pred, reg_pred, threshold=0.5)
    print(f"Detected {len(results[0]['boxes'])} characters")