DocLayout-YOLO：页面轴注意力 + 百万级合成预训练，让 YOLO 读懂文档版面

# doclayout_yolo_core_improvements.py
# DocLayout-YOLO (release-v2) 核心改进模块整理
# 主要包含: PALA (Page-Axis Layout Attention) 与 GL-CRM
# 说明: 以下代码基于论文描述、官方仓库结构及公开技术博客整理，
#       旨在展示核心设计逻辑，具体实现请以官方 release-v2 代码为准。

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


# ============================================================
# 1. Page-Axis Layout Attention (PALA)
# ============================================================
class PALA(nn.Module):
    """
    Page-Axis Layout Attention.
    论文描述: 沿文档页面的水平轴与垂直轴捕获长距离依赖，
    同时保持计算效率。保留了与 Area Attention 相同的特征投影
    与位置编码，但通过更高效的轴式聚合降低延迟。
    参考: release-v2 README 中 "Page-Axis Layout Attention (PALA),
    which captures long-range dependencies along the horizontal
    and vertical page axes"[reference:0]。
    """
    def __init__(self, dim, num_heads=8, qkv_bias=False, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        self.dim = dim
        head_dim = dim // num_heads
        self.scale = head_dim ** -0.5

        # 与 Area Attention 保持一致的特征投影
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        # 可学习的位置编码 (与 Area Attention 一致)
        self.pos_embed = nn.Parameter(torch.zeros(1, dim))

    def forward(self, x):
        """
        x: (B, N, C) 其中 N = H * W
        """
        B, N, C = x.shape
        # 假设输入是展平的二维特征图，尝试推断 H, W
        # 实际实现中通常由调用方传入 H, W 或通过 shape 维护
        # 这里使用简单的方形假设进行演示，实际请根据特征图尺寸调整
        H = W = int(N ** 0.5)
        if H * W != N:
            # 如果不是完美方形，回退到 Area Attention 的通用处理
            # 或需要外部传入 H, W
            raise ValueError("PALA 需要方形特征图或显式的 H, W 输入")

        qkv = self.qkv(x + self.pos_embed).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]  # (B, heads, N, head_dim)

        # 重塑为二维空间 (B, heads, H, W, head_dim)
        q = q.reshape(B, self.num_heads, H, W, -1)
        k = k.reshape(B, self.num_heads, H, W, -1)
        v = v.reshape(B, self.num_heads, H, W, -1)

        # --- 水平轴注意力 ---
        # 对每一行进行注意力聚合 (沿宽度轴)
        q_h = q.permute(0, 1, 2, 4, 3)  # (B, heads, H, head_dim, W)
        k_h = k.permute(0, 1, 2, 4, 3)
        v_h = v.permute(0, 1, 2, 4, 3)
        attn_h = (q_h @ k_h.transpose(-1, -2)) * self.scale
        attn_h = attn_h.softmax(dim=-1)
        attn_h = self.attn_drop(attn_h)
        out_h = (attn_h @ v_h).permute(0, 1, 2, 4, 3)  # 回到 (B, heads, H, W, head_dim)

        # --- 垂直轴注意力 ---
        # 对每一列进行注意力聚合 (沿高度轴)
        q_v = q.permute(0, 1, 3, 4, 2)  # (B, heads, W, head_dim, H)
        k_v = k.permute(0, 1, 3, 4, 2)
        v_v = v.permute(0, 1, 3, 4, 2)
        attn_v = (q_v @ k_v.transpose(-1, -2)) * self.scale
        attn_v = attn_v.softmax(dim=-1)
        attn_v = self.attn_drop(attn_v)
        out_v = (attn_v @ v_v).permute(0, 1, 4, 2, 3)  # 回到 (B, heads, H, W, head_dim)

        # 融合双轴信息
        out = out_h + out_v
        out = out.reshape(B, self.num_heads, N, -1).permute(0, 2, 1, 3).reshape(B, N, C)
        out = self.proj(out)
        out = self.proj_drop(out)
        return out


# ============================================================
# 2. GL-CRM (Global-to-Local Controllable Receptive Module)
# ============================================================
class CRM(nn.Module):
    """
    可控感受野模块 (Controllable Receptive Module).
    通过不同膨胀率的卷积捕获多尺度上下文。
    """
    def __init__(self, channels, dilation_rates=(1, 2, 3)):
        super().__init__()
        self.branches = nn.ModuleList()
        for rate in dilation_rates:
            self.branches.append(
                nn.Conv2d(
                    channels, channels, kernel_size=3,
                    padding=rate, dilation=rate, groups=channels, bias=False
                )
            )
        self.fuse = nn.Conv2d(channels * len(dilation_rates), channels, kernel_size=1)

    def forward(self, x):
        outs = [branch(x) for branch in self.branches]
        out = torch.cat(outs, dim=1)
        return self.fuse(out)


class GLCRM(nn.Module):
    """
    全局到局部可控感受野模块 (Global-to-Local Controllable Receptive Module).
    作为标准 CSP 瓶颈的替代，更好地处理文档元素的多尺度变化。
    参考: 官方博客指出 "GL-CRM 模块作为标准 CSP 瓶颈的替代"，
          "包括可控感受野模块 (CRM) 和全局到局部设计 (GL)"。
    """
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):
        super().__init__()
        c_ = int(c2 * e)  # 隐藏通道数
        self.cv1 = nn.Conv2d(c1, c_, 1, 1)
        self.cv2 = nn.Conv2d(c1, c_, 1, 1)
        self.cv3 = nn.Conv2d(c_, c_, 1, 1)
        self.cv4 = nn.Conv2d(c_, c2, 1, 1)

        # 局部分支: 使用 CRM 进行多尺度特征提取
        self.local_branch = nn.Sequential(
            *[CRM(c_) for _ in range(n)]
        )

        # 全局分支: 简单的全局上下文 (可通过池化或大核卷积实现)
        self.global_branch = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(c_, c_, 1, bias=False),
            nn.SiLU(),
            nn.Conv2d(c_, c_, 1, bias=False)
        )

        self.shortcut = shortcut and c1 == c2

    def forward(self, x):
        y1 = self.cv1(x)  # 分支1
        y2 = self.cv2(x)  # 分支2

        # 局部多尺度特征
        y_local = self.local_branch(y2)
        # 全局上下文
        y_global = self.global_branch(y2)

        # 全局到局部的融合 (例如通过门控或相加)
        # 这里采用简单的门控机制，实际实现可能更复杂
        gate = torch.sigmoid(y_global)
        y_combined = y_local * gate + y_local  # 残差式融合

        y_combined = self.cv3(y_combined)
        out = self.cv4(y_combined)

        if self.shortcut:
            out = out + x
        return out


# ============================================================
# 3. 整合示例: 在 YOLO 主干中替换标准 Bottleneck
# ============================================================
if __name__ == "__main__":
    # 示例: 构建一个包含 PALA 和 GL-CRM 的简化模块
    B, C, H, W = 2, 256, 32, 32
    x = torch.randn(B, C, H, W)

    # 使用 GL-CRM 替代 CSP 瓶颈
    glcrm = GLCRM(c1=C, c2=C, n=2, shortcut=True)
    out_glcrm = glcrm(x)
    print(f"GL-CRM 输入: {x.shape}, 输出: {out_glcrm.shape}")

    # PALA 通常作用于展平后的序列特征
    # 假设特征图被展平为序列
    x_seq = x.flatten(2).transpose(1, 2)  # (B, N, C)
    pala = PALA(dim=C, num_heads=8)
    out_pala = pala(x_seq)
    print(f"PALA 输入: {x_seq.shape}, 输出: {out_pala.shape}")