XNet/lib/tools/loss.py

from typing import Optional

import torch
import torch.nn as nn
import torch.fft as fft
from monai.losses import DiceLoss, DiceCELoss, HausdorffDTLoss


class FocalFrequencyLoss(nn.Module):
    """
    焦点频域损失函数 (Focal Frequency Loss)
    
    核心思想：
        传统的空间域损失（如 MSE、Dice）主要关注像素级别的差异，而频域损失通过傅里叶变换
        将图像转换到频率域，从频率角度衡量预测图像与真实图像的差异。
        
        该损失的创新点在于"焦点"机制：
        1. 自动计算频谱权重矩阵，对不同频率成分赋予不同的重要性
        2. 对难以重建的频率成分给予更高权重（类似 Focal Loss 的思想）
        3. 可以捕捉图像的全局结构和纹理细节，弥补空间域损失的不足
    
    适用场景：
        - 医学图像分割：增强边缘和纹理的恢复
        - 图像超分辨率：重建高频细节
        - 图像去噪/去模糊：平衡低频和高频信息
        
    参数说明：
        loss_weight: 损失权重系数，用于平衡该损失与其他损失的重要性
        alpha: 频谱权重的幂次参数，控制权重分布的陡峭程度
               alpha 越大，困难频率成分的权重越突出
        patch_factor: 图像分块因子，将图像分成多个小块分别进行 FFT
                     值为 1 表示不分块，对整个图像做 FFT
                     值大于 1 时，将图像分成 patch_factor×patch_factor 个小块
        ave_spectrum: 是否对 batch 内的频谱进行平均，用于减少 batch 内差异
        log_matrix: 是否对频谱差异取对数，用于压缩动态范围
        batch_matrix: 权重归一化方式
                     True: 在整个 batch 范围内归一化到 [0,1]
                     False: 对每张图像单独归一化
    """

    def __init__(self, loss_weight=1.0, alpha=1.0, patch_factor=1, ave_spectrum=False, log_matrix=False,
                 batch_matrix=False):
        """
        初始化焦点频域损失函数的所有超参数
        
        Args:
            loss_weight (float): 损失权重，默认 1.0
            alpha (float): 频谱权重指数，默认 1.0
            patch_factor (int): 图像分块因子，默认 1（不分块）
            ave_spectrum (bool): 是否对 batch 频谱平均，默认 False
            log_matrix (bool): 是否使用对数矩阵，默认 False
            batch_matrix (bool): 是否使用 batch 级矩阵，默认 False
        """
        super(FocalFrequencyLoss, self).__init__()
        self.loss_weight = loss_weight
        self.alpha = alpha
        self.patch_factor = patch_factor
        self.ave_spectrum = ave_spectrum
        self.log_matrix = log_matrix
        self.batch_matrix = batch_matrix

    def tensor2freq(self, x):
        """
        将空间域图像张量转换为频域表示
        
        工作原理：
            1. 如果 patch_factor > 1，先将图像分割成多个小块
            2. 对每个小块执行 2D 快速傅里叶变换 (FFT)
            3. 将复数形式的 FFT 结果分解为实部和虚部
        
        傅里叶变换的物理意义：
            - 低频成分：对应图像的平滑区域和整体轮廓
            - 高频成分：对应图像的边缘、纹理和噪声
            - 通过分析频谱，可以分离和处理不同频率的特征
        
        Args:
            x (torch.Tensor): 输入图像张量，形状为 (N, C, H, W)
                             N=batch_size, C=channels, H=height, W=width
        
        Returns:
            freq (torch.Tensor): 频域表示，形状为 (N, patch_factor², C, H/pf, W/pf, 2)
                                最后一维的 2 个通道分别是 [实部，虚部]
                                patch_factor² 表示分成了多少个小块
        
        Example:
            输入：x.shape = (4, 1, 256, 256), patch_factor=4
            输出：freq.shape = (4, 16, 1, 64, 64, 2)
                - 16 = 4×4 个小块
                - 64×64 = 每个小块的尺寸
                - 2 = [实部，虚部]
        """
        # 获取分块因子
        patch_factor = self.patch_factor

        # 获取输入图像的尺寸信息
        _, _, h, w = x.shape

        # 断言检查：确保图像尺寸可以被 patch_factor 整除
        # 这是为了保证分块时每个小块大小一致，避免边界问题
        assert h % patch_factor == 0 and w % patch_factor == 0, (
            'Patch factor should be divisible by image height and width')

        # 初始化列表用于存储所有小块的频域表示
        patch_list = []

        # 计算每个小块的高度和宽度
        # 例如：原图 256×256, patch_factor=4 → 每个小块 64×64
        patch_h = h // patch_factor
        patch_w = w // patch_factor

        # 双重循环遍历所有小块
        # i 控制垂直方向的索引，j 控制水平方向的索引
        for i in range(patch_factor):
            for j in range(patch_factor):
                # 切片操作：提取第 (i,j) 个小块
                # 垂直方向：从 i*patch_h 到 (i+1)*patch_h
                # 水平方向：从 j*patch_w 到 (j+1)*patch_w
                # [:, :, ...] 保持 batch 和 channel 维度不变
                patch_list.append(x[:, :, i * patch_h:(i + 1) * patch_h, j * patch_w:(j + 1) * patch_w])

        # 将所有小块堆叠成一个新的维度
        # 原始形状：(N, C, patch_h, patch_w) 的列表，长度为 patch_factor²
        # 堆叠后形状：(N, patch_factor², C, patch_h, patch_w)
        # dim=1 表示在第 1 个维度（channel 之后）插入新的分块维度
        y = torch.stack(patch_list, 1)

        # 对每个小块执行 2D 快速傅里叶变换
        # torch.fft.fft2: 计算二维离散傅里叶变换
        # norm='ortho': 使用正交归一化，保证变换前后能量守恒
        # 变换结果是一个复数张量，包含每个频率成分的振幅和相位信息
        freq = torch.fft.fft2(y, norm='ortho')

        # 将复数形式转换为实数表示，方便后续神经网络处理
        # torch.stack([freq.real, freq.imag], -1): 将实部和虚部堆叠到最后一个维度
        # freq.real: 复数的实部，代表余弦分量的幅度
        # freq.imag: 复数的虚部，代表正弦分量的幅度
        # 最终形状：(N, patch_factor², C, patch_h, patch_w, 2)
        freq = torch.stack([freq.real, freq.imag], -1)

        return freq

    def loss_formulation(self, recon_freq, real_freq, matrix=None):
        """
        构建并计算焦点频域损失的核心公式
        
        核心思想：
            传统频域损失直接计算预测频谱与真实频谱的距离（如 MSE）。
            本方法引入"焦点"机制：根据每个频率成分的重建难度动态调整权重。
            
            权重计算逻辑：
                1. 如果未提供预定义权重矩阵，则在线计算动态权重
                2. 重建误差大的频率 → 高权重（重点关注）
                3. 重建误差小的频率 → 低权重（减少关注）
            这类似于 Focal Loss 中"关注难例"的思想
        
        Args:
            recon_freq (torch.Tensor): 重建（预测）图像的频域表示
                                      形状：(N, P, C, H, W, 2)，P=patch_factor²
            real_freq (torch.Tensor): 真实（目标）图像的频域表示
                                    形状：(N, P, C, H, W, 2)
            matrix (torch.Tensor, optional): 预定义的频谱权重矩阵
                                           如果为 None，则动态计算权重
        
        Returns:
            loss (torch.Tensor): 标量损失值
        
        工作流程：
            Step 1: 确定权重矩阵（预定义或动态计算）
            Step 2: 计算频谱距离（复数空间中的欧氏距离）
            Step 3: 加权求和得到最终损失
        """

        # ==================== Step 1: 确定权重矩阵 ====================
        if matrix is not None:
            # 情况 A: 使用预定义的权重矩阵
            # 这种模式允许外部指定固定的频率权重，适用于某些先验知识已知的场景
            # 例如：人工指定某些频率更重要，或者使用其他算法计算的权重
            weight_matrix = matrix.detach()
            # .detach() 确保权重矩阵不参与梯度反向传播
            # 这样权重是固定的，不会影响梯度的流动
        else:
            # 情况 B: 动态计算自适应权重矩阵（推荐模式）

            # --- 子步骤 1: 计算初步的频谱差异 ---
            # 逐元素计算预测频谱与真实频谱的差值的平方
            # 这是一个复数差的平方，需要分别处理实部和虚部
            # 形状：(N, P, C, H, W, 2)
            matrix_tmp = (recon_freq - real_freq) ** 2

            # --- 子步骤 2: 计算频谱幅度差异 ---
            # 复数的模长公式：|a + bi| = sqrt(a² + b²)
            # 这里计算的是每个频率成分的预测误差的幅度
            # [..., 0]: 实部的平方, [..., 1]: 虚部的平方
            # torch.sqrt(...): 开平方得到欧几里得距离
            # ** self.alpha: 应用幂次变换，alpha 控制权重的分布特性
            #   - alpha > 1: 放大差异，使大误差更突出
            #   - alpha < 1: 缩小差异，使权重分布更均匀
            #   - alpha = 1: 线性关系（默认情况）
            matrix_tmp = torch.sqrt(matrix_tmp[..., 0] + matrix_tmp[..., 1]) ** self.alpha

            # --- 子步骤 3: 可选的对数变换 ---
            if self.log_matrix:
                # 对数变换的作用：压缩动态范围
                # 当频谱差异的范围很大时（几个数量级），对数变换可以防止某些频率主导损失
                # log(x + 1.0): 加 1 是为了避免 log(0) 的数值不稳定
                matrix_tmp = torch.log(matrix_tmp + 1.0)

            # --- 子步骤 4: 归一化权重到 [0, 1] 范围 ---
            # 归一化的目的是确保权重有统一的尺度，便于控制和解释
            if self.batch_matrix:
                # 模式 A: Batch 级归一化
                # 在整个 batch 的所有像素、所有频率上找最大值，然后统一归一化
                # 优点：batch 内所有样本的权重在同一尺度
                # 缺点：可能掩盖样本间的差异
                matrix_tmp = matrix_tmp / matrix_tmp.max()
            else:
                # 模式 B: 样本级归一化（推荐）
                # 对每个样本单独归一化，保持样本间的相对差异
                # .max(-1).values: 沿最后一个维度（实部/虚部维度）取最大值
                # .max(-1).values[:, :, :, None, None]: 再沿空间维度取最大值
                # [:, :, :, None, None]: 添加维度以保持广播兼容性
                # 最终每张图片的权重矩阵独立归一化到 [0,1]
                matrix_tmp = matrix_tmp / matrix_tmp.max(-1).values.max(-1).values[:, :, :, None, None]

            # --- 子步骤 5: 数值稳定性处理 ---
            # 处理可能出现的 NaN 值（例如 0/0 的情况）
            # 将 NaN 替换为 0，表示这些位置不参与加权
            matrix_tmp[torch.isnan(matrix_tmp)] = 0.0

            # --- 子步骤 6: 截断到合法范围 ---
            # 确保所有权值都在 [0, 1] 区间内
            # torch.clamp: 将小于 0 的值设为 0，大于 1 的值设为 1
            # 这是防御性编程，防止数值溢出导致训练不稳定
            matrix_tmp = torch.clamp(matrix_tmp, min=0.0, max=1.0)

            # --- 子步骤 7: 创建最终的权重矩阵 ---
            # .clone().detach() 创建副本并断开梯度连接
            # 这样权重矩阵在本次前向传播中是固定的，不会自我影响
            # 这是关键设计：权重基于当前误差计算，但不参与本次的梯度回传
            weight_matrix = matrix_tmp.clone().detach()

        # ==================== 权重矩阵有效性验证 ====================
        # 断言检查：确保权重矩阵的所有值都在 [0, 1] 范围内
        # 这是一个安全检查，帮助调试时发现问题
        # .item() 将标量张量转换为 Python 浮点数，方便打印
        assert weight_matrix.min().item() >= 0 and weight_matrix.max().item() <= 1, (
                'The values of spectrum weight matrix should be in the range [0, 1], '
                'but got Min: %.10f Max: %.10f' % (weight_matrix.min().item(), weight_matrix.max().item()))

        # ==================== Step 2: 计算频谱距离 ====================
        # 计算预测频谱与真实频谱的逐元素差异的平方
        # 与前面计算 matrix_tmp 的第一步相同
        # 形状：(N, P, C, H, W, 2)
        tmp = (recon_freq - real_freq) ** 2

        # 将实部和虚部的平方和相加，得到复数空间中的欧几里得距离平方
        # 这就是每个频率成分的重建误差
        # 形状：(N, P, C, H, W)
        freq_distance = tmp[..., 0] + tmp[..., 1]

        # ==================== Step 3: 加权求和得到损失 ====================
        # 逐元素相乘：权重矩阵 × 频谱距离
        # 效果：
        #   - 权重高的频率（重建困难）→ 损失贡献大 → 梯度大 → 模型重点关注
        #   - 权重低的频率（重建简单）→ 损失贡献小 → 梯度小 → 模型次要关注
        # 形状：(N, P, C, H, W)
        loss = weight_matrix * freq_distance

        # 对所有维度取平均，得到标量损失值
        # torch.mean(): 将整个张量压缩成一个标量
        # 这样得到的损失可以直接用于反向传播
        return torch.mean(loss)

    def forward(self, pred, target, matrix=None):
        """
        焦点频域损失的前向传播计算
        
        这是损失函数的主要入口，当调用 loss_function(pred, target) 时执行此方法
        
        Args:
            pred (torch.Tensor): 预测的图像张量
                                形状：(N, C, H, W)
                                N: batch size（批次大小）
                                C: channels（通道数，对于灰度图 C=1，RGB 图 C=3）
                                H: height（图像高度）
                                W: width（图像宽度）
            
            target (torch.Tensor): 目标的图像张量（真实标签）
                                  形状：(N, C, H, W)，必须与 pred 完全相同
            
            matrix (torch.Tensor, optional): 预定义的频谱权重矩阵
                                           如果提供，则使用该固定权重而非动态计算
                                           默认：None（动态计算自适应权重）
        
        Returns:
            torch.Tensor: 标量损失值（0 维张量），可以直接用于反向传播
        
        完整计算流程:
            ┌─────────────────┐
            │ 输入：pred, target │
            └────────┬──────────┘
                     │
                     ▼
            ┌─────────────────┐
            │ Step 1: 傅里叶变换 │  tensor2freq()
            │ pred → pred_freq  │  空间域 → 频域
            │ target → target_freq│
            └────────┬──────────┘
                     │
                     ▼
            ┌─────────────────┐
            │ Step 2: 可选的平均 │  if ave_spectrum
            │ 对 batch 维度平均   │  减少样本间差异
            └────────┬──────────┘
                     │
                     ▼
            ┌─────────────────┐
            │ Step 3: 计算损失  │  loss_formulation()
            │ 动态权重 × 频谱距离 │  焦点机制核心
            └────────┬──────────┘
                     │
                     ▼
            ┌─────────────────┐
            │ Step 4: 应用权重  │  × loss_weight
            │ 返回最终损失值    │
            └─────────────────┘
        
        物理意义解释:
            1. 傅里叶变换：将图像从"像素空间"转换到"频率空间"
               - 像素空间：关注每个点的亮度值
               - 频率空间：关注图像的周期性模式（纹理、边缘、轮廓）
            
            2. 频谱比较：在频率空间中衡量预测与真实的差异
               - 低频误差：反映整体结构的偏差
               - 高频误差：反映细节纹理的偏差
            
            3. 焦点权重：自动识别并强调难以重建的频率成分
               - 这是与传统频域损失（如简单的频谱 MSE）的关键区别
               - 类似于注意力机制，让模型"聚焦"于困难频率
        
        """

        # ==================== Step 1: 将预测图像转换为频域 ====================
        # 调用 tensor2freq 方法对预测图像进行傅里叶变换
        # 将空间域的像素表示转换为频域的频谱表示
        # pred_freq 包含了预测图像在各个频率上的振幅和相位信息
        pred_freq = self.tensor2freq(pred)

        # ==================== Step 2: 将目标图像转换为频域 ====================
        # 同样对真实标签图像进行傅里叶变换
        # 这样我们就可以在频率空间中比较预测与真实的差异
        target_freq = self.tensor2freq(target)

        # ==================== Step 3: 可选的 Batch 频谱平均 ====================
        if self.ave_spectrum:
            # 如果启用了 ave_spectrum 选项，对 batch 维度（第 0 维）取平均
            # keepdim=True 保持维度数量不变，只是将第 0 维的大小设为 1
            # 
            # 这个操作的效果：
            #   - 原始形状：(N, P, C, H, W, 2)
            #   - 平均后：(1, P, C, H, W, 2)
            #   
            # 为什么要这样做？
            #   1. 减少 batch 内样本间的随机波动
            #   2. 计算一个"平均频谱"作为代表
            #   3. 在某些任务中可以提高训练稳定性
            #   
            # 注意：这会改变损失的语义，从"逐个样本的损失"变成"batch 级别的损失"
            pred_freq = torch.mean(pred_freq, 0, keepdim=True)
            target_freq = torch.mean(target_freq, 0, keepdim=True)

        # ==================== Step 4: 计算最终的焦点频域损失 ====================
        # 调用 loss_formulation 方法计算加权后的频域损失
        # 该方法会：
        #   1. 动态计算频谱权重矩阵（如果 matrix=None）
        #   2. 计算预测频谱与真实频谱的距离
        #   3. 用权重矩阵对距离加权，得到最终损失
        #
        # 返回值是一个标量张量，表示整个 batch 的平均损失
        loss_value = self.loss_formulation(pred_freq, target_freq, matrix)

        # ==================== Step 5: 应用损失权重系数 ====================
        # 将计算得到的损失乘以预设的权重系数 loss_weight
        # 
        # 这个参数的作用：
        #   - 在多损失联合训练时，平衡不同损失的重要性
        #   - 例如：总损失 = 1.0 * DiceLoss + 0.1 * FocalFrequencyLoss
        #   - 这样可以让频域损失作为辅助损失，不会主导训练过程
        #
        # 为什么需要这样做？
        #   不同的损失函数量级可能差异很大
        #   DiceLoss 可能在 0-1 之间，而频域损失可能在 0-100 之间
        #   通过调整 loss_weight，可以确保各个损失在同一数量级

        final_loss = loss_value * self.loss_weight

        return final_loss