JMT-YOLO：分割引导一致性学习，让 YOLO 在陌生域里也"看得清"

# jmt_yolo_core_improvements.py
# JMT-YOLO (https://github.com/25043221/JMT-YOLO) 核心改进整理
# 主要包含: DenseRFB (分割头核心) 与 Mask-Consistency (掩码一致性机制)
# 说明: 以下代码基于论文描述、官方仓库结构及公开技术资料整理，
#       旨在展示核心设计逻辑，具体实现请以官方仓库代码为准。

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


# ============================================================
# 1. DenseRFB: 密集感受野模块 (分割头核心)
# ============================================================
class DenseRFB(nn.Module):
    """
    Dense Receptive Field Block.
    作为辅助分割头的核心，通过密集连接和空洞卷积捕获多尺度上下文信息，
    缓解风格抽象导致的细粒度特征退化。
    参考: 论文指出 DenseRFB 使用密集连接 (dense connection) 和
          空洞卷积 (dilated convolution)，膨胀率分别为 d1=2, d2=3。
    """
    def __init__(self, in_channels, out_channels):
        super().__init__()
        # 主分支: 1x1 卷积 + 3x3 标准卷积，得到初始特征
        self.main_branch = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.SiLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.SiLU(inplace=True),
        )

        # 密集连接分支: 两个不同膨胀率的空洞卷积分支
        # 膨胀率 d1=2, d2=3 (来自论文描述)
        self.dilated_branch1 = nn.Sequential(
            nn.Conv2d(out_channels, out_channels, kernel_size=3,
                      padding=2, dilation=2, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.SiLU(inplace=True),
        )
        self.dilated_branch2 = nn.Sequential(
            nn.Conv2d(out_channels, out_channels, kernel_size=3,
                      padding=3, dilation=3, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.SiLU(inplace=True),
        )

        # 融合层: 将主分支、两个空洞分支以及可能的残差连接进行拼接
        # 密集连接: 后续分支的输入包含前序分支的输出
        # 这里简化演示融合逻辑，实际实现可能包含更多密集连接细节
        self.fuse = nn.Sequential(
            nn.Conv2d(out_channels * 3, out_channels, kernel_size=1, bias=False),
            nn.BatchNorm2d(out_channels),
            nn.SiLU(inplace=True),
        )

    def forward(self, x):
        # 主分支特征
        main_feat = self.main_branch(x)

        # 第一个空洞分支 (d1=2)，输入为主分支特征
        d1_feat = self.dilated_branch1(main_feat)

        # 第二个空洞分支 (d2=3)，输入为第一个分支的输出 (密集连接)
        d2_feat = self.dilated_branch2(d1_feat)

        # 拼接融合: 主分支 + 两个空洞分支
        fused = torch.cat([main_feat, d1_feat, d2_feat], dim=1)
        out = self.fuse(fused)
        return out


# ============================================================
# 2. 辅助分割头 (Segment Head)
# ============================================================
class SegmentHead(nn.Module):
    """
    辅助语义分割头。
    接收融合特征，通过 DenseRFB 模块进行多尺度上下文增强，
    然后输出像素级分割预测。
    """
    def __init__(self, in_channels, num_classes):
        super().__init__()
        # 特征融合后的通道压缩 (可选)
        self.reduce = nn.Sequential(
            nn.Conv2d(in_channels, in_channels // 2, kernel_size=1, bias=False),
            nn.BatchNorm2d(in_channels // 2),
            nn.SiLU(inplace=True),
        )
        # DenseRFB 核心模块
        self.dense_rfb = DenseRFB(in_channels // 2, in_channels // 2)
        # 最终分割预测层
        self.seg_conv = nn.Conv2d(in_channels // 2, num_classes, kernel_size=1)

    def forward(self, x):
        x = self.reduce(x)
        x = self.dense_rfb(x)
        seg_logits = self.seg_conv(x)
        return seg_logits


# ============================================================
# 3. 掩码一致性机制 (Mask-Consistency)
# ============================================================
def generate_random_mask(x, mask_ratio=0.5, patch_size=16):
    """
    对目标图像或特征图应用随机掩码。
    论文中描述为对目标图像应用随机掩码操作。
    这里提供一个基于 patch 的随机掩码生成器。
    """
    B, C, H, W = x.shape
    num_patches_h = H // patch_size
    num_patches_w = W // patch_size
    num_patches = num_patches_h * num_patches_w

    # 随机选择要掩码的 patch
    mask = torch.rand(B, num_patches, device=x.device) > mask_ratio
    mask = mask.float().reshape(B, 1, num_patches_h, num_patches_w)
    # 上采样到原始尺寸
    mask = F.interpolate(mask, size=(H, W), mode='nearest')
    return mask


class MaskConsistencyLoss(nn.Module):
    """
    掩码一致性损失。
    对学生模型在掩码输入上的预测和教师模型在原始输入上的预测
    计算交叉熵损失，引导网络学习结构上下文。
    参考: 论文指出通过随机掩码操作构建掩码一致性损失 (mask-consistency loss)。
    """
    def __init__(self, ignore_index=255):
        super().__init__()
        self.ce = nn.CrossEntropyLoss(ignore_index=ignore_index)

    def forward(self, student_logits, teacher_logits, mask):
        """
        student_logits: 学生模型对掩码图像的预测 (B, C, H, W)
        teacher_logits: 教师模型对原始图像的预测 (B, C, H, W)
        mask: 掩码图，1 表示保留，0 表示掩码 (B, 1, H, W)
        """
        # 教师预测作为伪标签 (stop-gradient)
        teacher_probs = teacher_logits.softmax(dim=1)
        # 在掩码区域计算一致性损失 (或全局一致性，视实现而定)
        # 这里演示在掩码区域进行加权的交叉熵
        # 将教师概率转换为硬标签或使用软标签
        loss = F.cross_entropy(student_logits, teacher_probs.argmax(dim=1),
                                reduction='none')
        # 仅在掩码区域 (mask == 0) 或非掩码区域加权，具体策略需参考原文
        # 示例: 在掩码区域给予更高权重
        weight = (1.0 - mask.squeeze(1))  # 掩码区域权重为 1
        loss = (loss * weight).sum() / (weight.sum() + 1e-6)
        return loss


# ============================================================
# 4. 整合示例: JMT-YOLO 核心组件
# ============================================================
if __name__ == "__main__":
    # 模拟设置
    B, C, H, W = 2, 256, 64, 64
    num_seg_classes = 19  # 例如 Cityscapes

    # 模拟融合特征 (来自检测主干和颈部)
    fused_feature = torch.randn(B, C, H, W)

    # 1. 辅助分割头
    seg_head = SegmentHead(in_channels=C, num_classes=num_seg_classes)
    seg_logits = seg_head(fused_feature)
    print(f"Segment Head 输入: {fused_feature.shape}, 输出: {seg_logits.shape}")

    # 2. 随机掩码
    mask = generate_random_mask(fused_feature, mask_ratio=0.5, patch_size=16)
    masked_feature = fused_feature * mask
    print(f"掩码形状: {mask.shape}, 保留比例: {mask.mean().item():.2f}")

    # 3. 掩码一致性损失 (模拟学生和教师输出)
    student_logits = seg_head(masked_feature)  # 学生模型对掩码输入的预测
    teacher_logits = seg_head(fused_feature).detach()  # 教师模型对原始输入的预测 (停止梯度)
    consistency_loss = MaskConsistencyLoss()
    loss = consistency_loss(student_logits, teacher_logits, mask)
    print(f"掩码一致性损失: {loss.item():.4f}")