Upload 2 files

utils/common_util.py

@@ -0,0 +1,301 @@
+from torch.nn import functional as F
+import torch
+import matplotlib.pyplot as plt
+
+from torchvision import transforms
+import numpy as np
+from PIL import Image
+import cv2
+import torch.nn as nn
+import scipy.stats as st
+
+
+# 从提供的关键点和对应分数中移除那些过于靠近图像边缘的关键点
+# 这个函数是在推理的时候用的，训练的时候没有使用
+def remove_borders(keypoints, scores, border: int, height: int, width: int):
+    """ Removes keypoints too close to the border """
+    '''
+    keypoints: 关键点坐标的二维数组，形状为 (N, 2)，其中每行表示一个关键点的 (y, x) 坐标。
+    scores: 每个关键点对应的分数，形状为 (N,)。
+    border: 表示需要移除的边界宽度。    推理时预设的是4像素
+    height: 图像的高度。
+    width: 图像的宽度。
+    '''
+    # 创建高度方向掩码: 关键点必须在 [border, height-border) 范围内
+    mask_h = (keypoints[:, 0] >= border) & (keypoints[:, 0] < (height - border))
+    # 创建宽度方向掩码: 关键点必须在 [border, width-border) 范围内
+    mask_w = (keypoints[:, 1] >= border) & (keypoints[:, 1] < (width - border))
+    # 组合掩码 (必须同时满足高度和宽度条件)
+    mask = mask_h & mask_w  # 所以这个mask是判定条件
+    # 返回过滤后的关键点和分数
+    return keypoints[mask], scores[mask]
+
+
+def simple_nms(scores, nms_radius: int):
+    """
+    快速非极大值抑制 (NMS) 算法，用于移除相邻关键点
+    
+    参数:
+        scores: 关键点分数图 (B, H, W) 或 (H, W)
+        nms_radius: NMS邻域半径
+    
+    返回:
+        NMS处理后的分数图
+    """
+    assert (nms_radius >= 0)
+    # 计算NMS窗口大小 (2*半径+1)
+    size = nms_radius * 2 + 1
+    avg_size = 2
+    # 定义最大池化函数 (使用固定步长1和适当填充)
+    def max_pool(x):
+        return torch.nn.functional.max_pool2d(x, kernel_size=size, stride=1, padding=nms_radius)
+    
+    # 创建与输入相同形状的零张量
+    zeros = torch.zeros_like(scores)
+    # max_map = max_pool(scores)
+
+    # 步骤1: 识别局部最大值点
+    # 比较每个点与其邻域内的最大值
+    max_mask = scores == max_pool(scores)   # max_pool(scores):每个像素点被替换为其局部窗口内的最大值。
+    # 步骤2: 添加微小随机扰动 (避免多个相同最大值)
+    # 生成 [0, 0.1) 范围内的随机数
+    max_mask_ = torch.rand(max_mask.shape).to(max_mask.device) / 10 
+    # 生成与 max_mask 相同形状的随机数（范围在 [0, 0.1)），作为微小扰动。
+    # 非局部最大值点置零
+    max_mask_[~max_mask] = 0
+
+    # 步骤3: 对扰动后的图再次应用NMS
+    # 识别扰动后仍然是局部最大值的点
+    mask = ((max_mask_ == max_pool(max_mask_)) & (max_mask_ > 0))   # mask:布尔掩码，仅保留扰动后仍然是局部最大值的点。
+
+    # 步骤4: 保留局部最大值点，其他点置零
+    return torch.where(mask, scores, zeros) # 如果 mask 为 True，保留原始分数。否则，将得分设置为零。
+
+
+def pre_processing(data):
+    """ Enhance retinal images """
+    train_imgs = datasets_normalized(data)
+    train_imgs = clahe_equalized(train_imgs)
+    train_imgs = adjust_gamma(train_imgs, 1.2)
+
+    train_imgs = train_imgs / 255.
+
+    return train_imgs.astype(np.float32)
+
+
+def rgb2gray(rgb):
+    """ Convert RGB image to gray image """
+    r, g, b = rgb.split()
+    return g
+
+
+# 对输入图像的 CLAHE（Contrast Limited Adaptive Histogram Equalization）增强处理
+# 用于提高图像的对比度，特别是在光照不均或细节难以分辨的情况下。
+def clahe_equalized(images):
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    # clipLimit=2.0: 限制对比度的参数，值越低，增强的对比度越小，避免高对比度导致的过曝光区域。
+    # tileGridSize=(8, 8): 将图像分成大小为 8x8 的网格，每个网格单独进行直方图均衡化，减少全局对比度增强引入的伪影。
+    images_equalized = np.empty(images.shape)
+    images_equalized[:, :] = clahe.apply(np.array(images[:, :],
+                                                  dtype=np.uint8))
+
+    return images_equalized
+
+
+def datasets_normalized(images):
+    # 归一化之后还需要把值映射到0到255
+    # images_normalized = np.empty(images.shape)
+    images_std = np.std(images)
+    images_mean = np.mean(images)
+    images_normalized = (images - images_mean) / (images_std + 1e-6)
+    minv = np.min(images_normalized)
+    images_normalized = ((images_normalized - minv) /
+                         (np.max(images_normalized) - minv)) * 255
+
+    return images_normalized
+
+
+def adjust_gamma(images, gamma=1.0):
+    invGamma = 1.0 / gamma  # invGamma: 伽马值的倒数，用于生成查找表。
+    table = np.array([((i / 255.0) ** invGamma) * 255
+                      for i in np.arange(0, 256)]).astype("uint8")  
+    # 预计算伽��校正的转换值，用于快速查找。
+    # 每个输入像素值（0-255）都映射到一个经过伽马变换的输出值。
+    # 生成过程：
+    #     i / 255.0: 将像素值归一化到 [0, 1] 范围。
+    #    (i / 255.0) ** invGamma: 应用伽马校正公式。
+    #    * 255: 将归一化后的值还原到 [0, 255] 范围。
+    #   astype("uint8"): 转换为 uint8 数据类型，适配图像格式。
+    new_images = np.empty(images.shape)
+    new_images[:, :] = cv2.LUT(np.array(images[:, :],
+                                        dtype=np.uint8), table)
+    # cv2.LUT: OpenCV 的快速像素值映射函数。
+    # 输入图像的每个像素值通过查找表 table 进行伽马校正。
+    # 大幅提高效率，避免逐像素计算。
+
+    return new_images
+
+
+def nms(detector_pred, nms_thresh=0.1, nms_size=10, detector_label=None, mask=False):
+    """
+    在检测器预测上应用非极大值抑制 (NMS)
+    
+    参数:
+        detector_pred: 检测器预测 (B, 1, H, W)
+        nms_thresh: NMS阈值
+        nms_size: NMS邻域大小
+        detector_label: 检测器标签 (可选)
+        mask: 是否使用标签掩码 (当前未实现)
+    
+    返回:
+        关键点位置列表 (每个元素是 (N, 2) 的数组)
+    """
+    # 创建预测副本 (避免修改原始数据)
+    detector_pred = detector_pred.clone().detach()
+    # 获取批次大小和图像尺寸
+    B, _, h, w = detector_pred.shape
+
+    # if mask:
+    #     assert detector_label is not None
+    #     detector_pred[detector_pred < nms_thresh] = 0
+    #     label_mask = detector_label
+    #
+    #     # more area
+    #
+    #     detector_label = detector_label.long().cpu().numpy()
+    #     detector_label = detector_label.astype(np.uint8)
+    #     kernel = np.ones((3, 3), np.uint8)
+    #     label_mask = np.array([cv2.dilate(detector_label[s, 0], kernel, iterations=1)
+    #                            for s in range(len(detector_label))])
+    #     label_mask = torch.from_numpy(label_mask).unsqueeze(1)
+    #     detector_pred[label_mask > 1e-6] = 0
+    
+    # 应用快速NMS算法
+    scores = simple_nms(detector_pred, nms_size)
+    # 重塑分数图形状 (B, H, W)
+    scores = scores.reshape(B, h, w)
+    # 找出分数高于阈值的点
+    points = [
+        torch.nonzero(s > nms_thresh)
+        for s in scores]
+    # 提取这些点的分数值
+    scores = [s[tuple(k.t())] for s, k in zip(scores, points)]
+    # 移除靠近边界的点
+    points, scores = list(zip(*[
+        remove_borders(k, s, 8, h, w)
+        for k, s in zip(points, scores)]))
+    # 翻转坐标顺序: [y, x] -> [x, y]
+    points = [torch.flip(k, [1]).long() for k in points]
+
+    return points
+
+
+# 实际上模型生成的描述符是1/8尺寸*1/8尺寸的描述符特征，这里是上采样到原尺寸的
+# 这个是在PKE算损失的时候用的
+def sample_keypoint_desc(keypoints, descriptors, s: int = 8):
+    """
+    在关键点位置采样描述符
+    
+    参数:
+        keypoints: 关键点坐标 (B, N, 2) 格式 [x, y]
+        descriptors: 描述符图 (B, C, H, W)
+        s: 描述符图相对于原始图像的下采样比例
+    
+    返回:
+        采样后的描述符 (B, C, N)
+    """
+    # 获取描述符张量的形状信息，用于后续处理
+    b, c, h, w = descriptors.shape  # 原始输入 descriptors: (b, c, h, w)
+    
+    # 克隆关键点并将其转换为浮点类型，以便进行坐标计算
+    keypoints = keypoints.clone().float()
+
+    # 将关键点坐标归一化到范围 (0, 1)
+    keypoints /= torch.tensor([(w * s - 1), (h * s - 1)]).to(keypoints)[None]
+    
+    # 将关键点坐标缩放到范围 (-1, 1)，以适应 grid_sample 函数的要求
+    keypoints = keypoints * 2 - 1
+
+    # 根据 PyTorch 版本准备 grid_sample 函数的参数，确保兼容性
+    args = {'align_corners': True} if int(torch.__version__[2]) > 2 else {}
+    
+    # 使用 grid_sample 函数在关键点位置插值描述符
+    descriptors = torch.nn.functional.grid_sample(
+        descriptors, keypoints.view(b, 1, -1, 2), mode='bilinear', **args)  # 经过 grid_sample:     (b, c, 1, n) n是关键点数量
+
+    # 对描述符进行 L2 归一化，使其长度为 1（1个像素），以便后续处理
+    descriptors = torch.nn.functional.normalize(
+        descriptors.reshape(b, c, -1), p=2, dim=1)  # reshape 后: (b, c, n)  channel=256
+    
+    # 返回处理后的描述符
+    return descriptors
+
+
+# 这个是在模型算损失的时候用的，可以同时处理关键点和被映射后的关键点损失
+def sample_descriptors(detector_pred, descriptor_pred, affine_descriptor_pred, grid_inverse,
+                        nms_size=10, nms_thresh=0.1, scale=8, affine_detector_pred=None):
+    """
+    基于关键点采样描述符
+    
+    参数:
+        detector_pred: 原始图像的检测器预测 (B, 1, H, W)
+        descriptor_pred: 原始图像的描述符预测 (B, C, H, W)
+        affine_descriptor_pred: 仿射图像的���述符预测 (B, C, H, W)
+        grid_inverse: 逆变换网格 (B, H, W, 2)
+        nms_size: NMS邻域大小
+        nms_thresh: NMS阈值
+        scale: 描述符图相对于原始图像的下采样比例
+        affine_detector_pred: 仿射图像的检测器预测 (可选)
+    
+    返回:
+        descriptors: 原始图像关键点的描述符列表
+        affine_descriptors: 仿射图像对应关键点的描述符列表
+        keypoints: 原始图像的关键点位置列表
+    """
+    # 获取批次大小和图像尺寸
+    B, _, h, w = detector_pred.shape
+    
+    # 应用NMS获取关键点位置
+    keypoints = nms(detector_pred, nms_size=nms_size, nms_thresh=nms_thresh)
+    
+    # 使用逆变换网格将关键点映射到仿射空间
+    affine_keypoints = [
+        grid_inverse[s, k[:, 1].long(), k[:, 0].long()]  # 使用网格插值
+        for s, k in enumerate(keypoints)
+    ]
+    
+    # 初始化存储列表
+    kp = []        # 过滤后的原始关键点
+    affine_kp = []  # 过滤后的仿射关键点
+    
+    # 处理每个样本
+    for s, k in enumerate(affine_keypoints):
+        # 过滤超出仿射图像边界的点
+        idx = (k[:, 0] < 1) & (k[:, 0] > -1) & (k[:, 1] < 1) & (k[:, 1] > -1)
+        
+        # 存储过滤后的原始关键点
+        kp.append(keypoints[s][idx])
+        
+        # 获取过滤后的仿射关键点
+        ak = k[idx]
+        
+        # 将归一化坐标转换回像素坐标
+        ak[:, 0] = (ak[:, 0] + 1) / 2 * (w - 1)  # x坐标
+        ak[:, 1] = (ak[:, 1] + 1) / 2 * (h - 1)  # y坐标
+        
+        # 存储转换后的仿射关键点
+        affine_kp.append(ak)
+    
+    # 在原始图像关键点位置采样描述符
+    descriptors = [
+        sample_keypoint_desc(k[None], d[None], s=scale)[0]
+        for k, d in zip(kp, descriptor_pred)
+    ]
+    
+    # 在仿射图像关键点位置采样描述符
+    affine_descriptors = [
+        sample_keypoint_desc(k[None], d[None], s=scale)[0]
+        for k, d in zip(affine_kp, affine_descriptor_pred)
+    ]
+    
+    return descriptors, affine_descriptors, keypoints

utils/image_process.py

@@ -0,0 +1,403 @@
+import math
+import os
+import random
+import shutil
+import time
+
+import torch
+from torchvision import transforms as T
+from torch.nn import functional as F
+import numpy as np
+import cv2
+import scipy.stats as st
+# import config
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+
+
+def get_gaussian_kernel(kernlen=21, nsig=5):
+    """
+    生成二维高斯核并返回一个不可训练的核权重，用于创建高斯热图
+    
+    参数:
+        kernlen (int): 核的大小（边长），必须是奇数
+        nsig (float): 高斯分布的标准差（控制分布的宽度）
+    
+    返回:
+        torch.Tensor: 形状为(1, 1, kernlen, kernlen)的高斯核张量
+    """
+    # 1. 计算采样间隔
+    # 在[-nsig, nsig]范围内均匀采样kernlen个点
+    interval = (2 * nsig + 1.) / kernlen
+    # 2. 创建一维坐标点数组
+    # 从-nsig-interval/2到nsig+interval/2，共kernlen+1个点
+    x = np.linspace(-nsig - interval / 2., nsig + interval / 2., kernlen + 1)
+    # 3. 计算一维高斯核
+    # 使用标准正态分布的累积分布函数(CDF)的差分
+    kern1d = np.diff(st.norm.cdf(x))
+    # 4. 创建二维高斯核
+    # 通过两个一维高斯核的外积(outer product)
+    kernel_raw = np.sqrt(np.outer(kern1d, kern1d))
+    # 5. 归一化核
+    # 使所有元素之和为1，确保核的总"能量"为1
+    kernel = kernel_raw / kernel_raw.sum()
+    # 6. 转换为PyTorch张量并调整维度
+    # 添加批次和通道维度 (1, 1, H, W)
+    kernel = torch.FloatTensor(kernel).unsqueeze(0).unsqueeze(0)
+    # 7. 创建不可训练的高斯核权重参数
+    weight = torch.nn.Parameter(data=kernel, requires_grad=False)
+    # 8. 将值归一化到[0, 1]范围
+    # 使最小值为0，最大值为1
+    weight = (weight - weight.min()) / (weight.max() - weight.min())
+
+    return weight
+
+
+# 对用于监督的点标签进行高斯模糊
+def label_gaussian_blur(label_point_positions, kernel, stride=1):
+    '''
+    对用于监督的点标签进行高斯模糊
+    该函数通过应用高斯卷积核来模糊标签点，以生成更平滑的热图
+    参数:
+        label_point_positions: 标签点的位置，形状为(B, C, H, W)的张量
+        kernel: 高斯卷积核，形状为(1, 1, Hk, Wk)的张量
+        stride: 卷积步长，默认为1
+    返回值:
+        blurred_label: 模糊后的标签热图，形状与输入相同
+    '''
+    # 应用高斯卷积进行模糊
+    blurred_label = F.conv2d(label_point_positions, kernel, stride=stride, padding=(kernel.shape[-1] - 1) // 2)
+    # 裁剪值范围：限制热图最大值为1
+    blurred_label[blurred_label > 1] = 1
+
+    return blurred_label
+
+
+def affine_images(images, used_for='detector'):
+    '''
+    根据detector和descriptor两种计算loss和训练的情况分别生成两种不同的仿射变及其逆变换，支持数据增强
+    参数:
+        images: 输入图像张量，形状为 (B, C, H, W)
+        used_for: 指定变换用途 - 'detector'用于检测器，'descriptor'用于描述器
+    返回值:
+        CPU上的仿射变换后的图像 output.detach().clone()
+        仿射变换网格 grid
+        逆仿射变换网格 grid_inverse
+    '''
+
+    """
+    Perform affine transformation on images
+    :param images: (B, C, H, W)
+    :param keypoint_labels: corresponding labels
+    :param value_map: value maps, used to record history learned geo_points
+    :return: results of affine images, affine labels, affine value maps, affine transformed grid_inverse, inverse transformed grid_inverse
+    """
+
+    h, w = images.shape[2:] # 获取图像的尺寸
+    theta = None
+    thetaI = None
+
+    # 对一个batch中的每个图像进行随机仿射变换参数的生成
+    for i in range(len(images)):
+        if used_for == 'detector':
+            affine_params = T.RandomAffine(20).get_params(degrees=[-15, 15],                # 旋转角度范围：±15度
+                                                        translate=[0.2, 0.2],               # 平移范围：±20%
+                                                        scale_ranges=[0.9, 1.35],           # 缩放范围：0.9-1.35倍
+                                                        shears=None,                        # 不使用剪切变换
+                                                        img_size=[h, w])
+        else:
+            affine_params = T.RandomAffine(20).get_params(degrees=[-3, 3],                  # 旋转角度范围：±3度
+                                                            translate=[0.1, 0.1],           # 平移范围：宽高的10%
+                                                            scale_ranges=[0.9, 1.1],        # 缩放范围：0.9到1.1倍
+                                                            shears=None,                    # 不使用剪切变换
+                                                            img_size=[h, w])
+        # 根据仿射变换参数计算变换矩阵和逆变换矩阵
+        angle = -affine_params[0] * math.pi / 180
+        theta_ = torch.tensor([
+            [1 / affine_params[2] * math.cos(angle), math.sin(-angle), -affine_params[1][0] / images.shape[2]],
+            [math.sin(angle), 1 / affine_params[2] * math.cos(angle), -affine_params[1][1] / images.shape[3]],
+            [0, 0, 1]
+        ], dtype=torch.float).to(images)
+        thetaI_ = theta_.inverse()
+        theta_ = theta_[:2]
+        thetaI_ = thetaI_[:2]
+
+        # 将变换矩阵和逆变换矩阵叠成一个batch
+        theta_ = theta_.unsqueeze(0)
+        thetaI_ = thetaI_.unsqueeze(0)
+        theta = theta_ if theta is None else torch.cat((theta, theta_)) # 如果theta为None，意味着这是第一个，否则则说明前面已经有参数值了
+        thetaI = thetaI_ if thetaI is None else torch.cat((thetaI, thetaI_))
+
+    # 根据变换矩阵生成网格
+    # 变换网格（Transformation Grid）是在进行图像的仿射变换时生成的一个规则的坐标网格。它定义了原始图像中每个像素在变换后的目标位置。
+    grid = F.affine_grid(theta, images.size(), align_corners=True)
+    grid = grid.to(images)  # 将参数移动到GPU上
+    grid_inverse = F.affine_grid(thetaI, images.size(), align_corners=True)
+    grid_inverse = grid_inverse.to(images)  # 将参数移动到GPU上
+    output = F.grid_sample(images, grid, align_corners=True)     # 对图像进行采样得到仿射变换后的图像
+
+    # 对于用于描述符的情况进行一些额外的随机增强操作
+    if used_for == 'descriptor':
+        if random.random() >= 0.4:
+            # 颜色抖动（其实是灰度变化）
+            output = output.repeat(1, 3, 1, 1)   # 将单通道图像复制为三通道（颜色变换需要RGB）
+            output = T.ColorJitter(brightness=0.4, contrast=0.3, saturation=0.3, hue=0.2)(output)
+            output = T.Grayscale()(output)      # 灰度化
+
+    return output.detach().clone(), grid, grid_inverse
+
+
+def affine_images_with_mask(images, used_for='detector'):
+    '''
+    仿射变换图像后，额外返回一个公共区域mask
+    根据detector和descriptor两种计算loss和训练的情况分别生成两种不同的仿射变及其逆变换，支持数据增强
+    参数:
+        images: 输入图像张量，形状为 (B, C, H, W)
+        used_for: 指定变换用途 - 'detector'用于检测器，'descriptor'用于描述器
+    返回值:
+        CPU上的仿射变换后的图像 output.detach().clone()
+        仿射变换网格 grid
+        逆仿射变换网格 grid_inverse
+    '''
+
+    """
+    Perform affine transformation on images
+    :param images: (B, C, H, W)
+    :param keypoint_labels: corresponding labels
+    :param value_map: value maps, used to record history learned geo_points
+    :return: results of affine images, affine labels, affine value maps, affine transformed grid_inverse, inverse transformed grid_inverse
+    """
+
+    h, w = images.shape[2:] # 获取图像的尺寸
+    theta = None
+    thetaI = None
+
+    # ---------------------------------------------------------------------
+    # 新增: 创建原始图像坐标网格
+    base_grid = F.affine_grid(torch.eye(2, 3).unsqueeze(0), (1, 1, h, w), align_corners=True)
+    base_grid = base_grid.expand(images.size(0), *base_grid.shape[1:])
+    # ---------------------------------------------------------------------
+
+    # 对一个batch中的每个图像进行随机仿射变换参数的生成
+    for i in range(len(images)):
+        if used_for == 'detector':
+            affine_params = T.RandomAffine(20).get_params(degrees=[-15, 15],                # 旋转角度范围：±15度
+                                                        translate=[0.2, 0.2],               # 平移范围：±20%
+                                                        scale_ranges=[0.9, 1.35],           # 缩放范围：0.9-1.35倍
+                                                        shears=None,                        # 不使用剪切变换
+                                                        img_size=[h, w])
+        else:
+            affine_params = T.RandomAffine(20).get_params(degrees=[-3, 3],                  # 旋转角度范围：±3度
+                                                            translate=[0.1, 0.1],           # 平移范围：宽高的10%
+                                                            scale_ranges=[0.9, 1.1],        # 缩放范围：0.9到1.1倍
+                                                            shears=None,                    # 不使用剪切变换
+                                                            img_size=[h, w])
+        # 根据仿射变换参数计算变换矩阵和逆变换矩阵
+        angle = -affine_params[0] * math.pi / 180
+        theta_ = torch.tensor([
+            [1 / affine_params[2] * math.cos(angle), math.sin(-angle), -affine_params[1][0] / images.shape[2]],
+            [math.sin(angle), 1 / affine_params[2] * math.cos(angle), -affine_params[1][1] / images.shape[3]],
+            [0, 0, 1]
+        ], dtype=torch.float).to(images)
+        thetaI_ = theta_.inverse()
+        theta_ = theta_[:2]
+        thetaI_ = thetaI_[:2]
+
+        # 将变换矩阵和逆变换矩阵叠成一个batch
+        theta_ = theta_.unsqueeze(0)
+        thetaI_ = thetaI_.unsqueeze(0)
+        theta = theta_ if theta is None else torch.cat((theta, theta_)) # 如果theta为None，意味着这是第一个，否则则说明前面已经有参数值了
+        thetaI = thetaI_ if thetaI is None else torch.cat((thetaI, thetaI_))
+
+    # 根据变换矩阵生成网格
+    # 变换网格（Transformation Grid）是在进行图像的仿射变换时生成的一个规则的坐标网格。它定义了原始图像中每个像素在变换后的目标位置。
+    grid = F.affine_grid(theta, images.size(), align_corners=True)
+    grid = grid.to(images)  # 将参数移动到GPU上
+    grid_inverse = F.affine_grid(thetaI, images.size(), align_corners=True)
+    grid_inverse = grid_inverse.to(images)  # 将参数移动到GPU上
+    output = F.grid_sample(images, grid, align_corners=True)     # 对图像进行采样得到仿射变换后的图像
+
+    # ---------------------------------------------------------------------
+    # 新增: 计算有效区域mask
+    # 1. 创建全1mask (表示所有像素初始有效)
+    valid_mask = torch.ones_like(images[:, :1])  # (B, 1, H, W)
+    # 2. 应用相同的变换到mask上
+    transformed_mask = F.grid_sample(valid_mask, grid, align_corners=True)
+    # 3. 反向映射回原始图像坐标
+    valid_mask = F.grid_sample(transformed_mask, grid_inverse, align_corners=True)
+    # 4. 二值化 (大于0.5视为有效)
+    valid_mask = (valid_mask > 0.5).float()
+    # ---------------------------------------------------------------------
+
+    # 对于用于描述符的情况进行一些额外的随机增强操作
+    if used_for == 'descriptor':
+        if random.random() >= 0.4:
+            # 颜色抖动（其实是灰度变化）
+            output = output.repeat(1, 3, 1, 1)   # 将单通道图像复制为三通道（颜色变换需要RGB）
+            output = T.ColorJitter(brightness=0.4, contrast=0.3, saturation=0.3, hue=0.2)(output)
+            output = T.Grayscale()(output)      # 灰度化
+
+    return output.detach().clone(), grid, grid_inverse, valid_mask.detach().clone()
+
+
+def pre_processing(data: np.ndarray) -> np.ndarray:
+    """
+    视网膜图像预处理流水线
+    包含标准化、CLAHE对比度增强和伽马校正
+    
+    参数:
+        image: 输入视网膜图像 (单通道，numpy数组)
+    
+    返回:
+        预处理后的图像 (float32, 范围[0,1])
+    """
+    train_imgs = datasets_normalized(data)
+    train_imgs = clahe_equalized(train_imgs)
+    train_imgs = adjust_gamma(train_imgs, 1.2)
+    train_imgs = train_imgs / 255.  # 最终归一化到[0,1]范围
+    return train_imgs.astype(np.float32)
+
+
+def datasets_normalized(images: np.ndarray) -> np.ndarray:
+    # 归一化之后还需要把值映射到0到255
+    # images_normalized = np.empty(images.shape)
+    # 计算全局统计量
+    images_std = np.std(images)
+    images_mean = np.mean(images)
+    # 应用标准化: (x - mean) / std
+    # 添加微小值避免除以零
+    images_normalized = (images - images_mean) / (images_std + 1e-6)
+    # 线性映射到[0,255]范围
+    minv = np.min(images_normalized)
+    images_normalized = ((images_normalized - minv) / (np.max(images_normalized) - minv)) * 255
+
+    return images_normalized
+
+
+def clahe_equalized(images):
+    # 对输入图像的 CLAHE（Contrast Limited Adaptive Histogram Equalization）增强处理
+    # 用于提高图像的对比度，特别是在光照不均或细节难以分辨的情况下。
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    # clipLimit=2.0: 限制对比度的参数，值越低，增强的对比度越小，避免高对比度导致的过曝光区域。
+    # tileGridSize=(8, 8): 将图像分成大小为 8x8 的网格，每个网格单独进行直方图均衡化，减少全局对比度增强引入的伪影。
+    images_equalized = np.empty(images.shape)
+    images_equalized[:, :] = clahe.apply(np.array(images[:, :], dtype=np.uint8))
+
+    return images_equalized
+
+
+def adjust_gamma(images, gamma=1.0):
+    # γ>1: 压缩高光区域，增强暗部细节 (适合视网膜图像)
+    # 创建伽马校正查找表
+    invGamma = 1.0 / gamma  # invGamma: 伽马值的倒数，用于生成查找表。
+    table = np.array([((i / 255.0) ** invGamma) * 255 for i in np.arange(0, 256)]).astype("uint8")  
+    # 预计算伽马校正的转换值，用于快速查找。
+    # 每个输入像素值（0-255）都映射到一个经过伽马变换的输出值。
+    # 生成过程：
+    #     i / 255.0: 将像素值归一化到 [0, 1] 范围。
+    #    (i / 255.0) ** invGamma: 应用伽马校正公式。
+    #    * 255: 将归一化后的值还原到 [0, 255] 范围。
+    #   astype("uint8"): 转换为 uint8 数据类型，适配图像格式。
+    new_images = np.empty(images.shape)
+    # 应用伽马校正
+    new_images[:, :] = cv2.LUT(np.array(images[:, :], dtype=np.uint8), table)
+    # cv2.LUT: OpenCV 的快速像素值映射函数。
+    # 输入图像的每个像素值通过查找表 table 进行伽马校正。
+    # 大幅提高效率，避免逐像素计算。
+
+    return new_images
+
+
+def simple_nms(scores, nms_radius: int):
+    """
+    快速非极大值抑制 (NMS) 算法，用于移除相邻关键点
+    参数:
+        scores: 关键点分数图 (B, H, W) 或 (H, W)
+        nms_radius: NMS邻域半径
+    返回:
+        NMS处理后的分数图
+    """
+    assert (nms_radius >= 0)
+    # 计算NMS窗口大小 (2*半径+1)
+    size = nms_radius * 2 + 1
+    avg_size = 2
+    # 定义最大池化函数 (使用固定步长1和适当填充)
+    def max_pool(x):
+        return torch.nn.functional.max_pool2d(x, kernel_size=size, stride=1, padding=nms_radius)
+    
+    # 创建与输入相同形状的零张量
+    zeros = torch.zeros_like(scores)
+    # max_map = max_pool(scores)
+
+    # 步骤1: 识别局部最大值点
+    # 比较每个点与其邻域内的最大值
+    max_mask = scores == max_pool(scores)   # max_pool(scores):每个像素点被替换为其局部窗口内的最大值。
+    # 步骤2: 添加微小随机扰动 (避免多个相同最大值)
+    # 生成 [0, 0.1) 范围内的随机数
+    max_mask_ = torch.rand(max_mask.shape).to(max_mask.device) / 10 
+    # 生成与 max_mask 相同形状的随机数（范围在 [0, 0.1)），作为微小扰动。
+    # 非局部最大值点置零
+    max_mask_[~max_mask] = 0
+    # 步骤3: 对扰动后的图再次应用NMS
+    # 识别扰动后仍然是局部最大值的点
+    mask = ((max_mask_ == max_pool(max_mask_)) & (max_mask_ > 0))   # mask:布尔掩码，仅保留扰动后仍然是局部最大值的点。
+    # 步骤4: 保留局部最大值点，其他点置零
+    return torch.where(mask, scores, zeros) # 如果 mask 为 True，保留原始分数。否则，将得分设置为零。
+
+
+def remove_borders(keypoints, scores, border: int, height: int, width: int):
+    '''
+    从提供的关键点和对应分数中移除那些过于靠近图像边缘的关键点
+    这个函数是在推理的时候用的，训练的时候没有使用
+    keypoints: 关键点坐标的二维数组，形状为 (N, 2)，其中每行表示一个关键点的 (y, x) 坐标。
+    scores: 每个关键点对应的分数，形状为 (N,)。
+    border: 表示需要移除的边界宽度。    推理时预设的是4像素
+    height: 图像的高度。
+    width: 图像的宽度。
+    '''
+    mask_h = (keypoints[:, 0] >= border) & (keypoints[:, 0] < (height - border))    # 创建高度方向掩码: 关键点必须在 [border, height-border) 范围内
+    mask_w = (keypoints[:, 1] >= border) & (keypoints[:, 1] < (width - border))     # 创建宽度方向掩码: 关键点必须在 [border, width-border) 范围内
+    mask = mask_h & mask_w  # 组合掩码 (必须同时满足高度和宽度条件)
+    return keypoints[mask], scores[mask]    # 返回过滤后的关键点和分数
+
+
+def nms(detector_pred, nms_thresh=0.1, nms_size=10, detector_label=None, mask=False):
+    """
+    在检测器预测上应用非极大值抑制 (NMS)
+    
+    参数:
+        detector_pred: 检测器预测 (B, 1, H, W)
+        nms_thresh: NMS阈值
+        nms_size: NMS邻域大小
+        detector_label: 检测器标签 (可选)
+        mask: 是否使用标签掩码 (当前未实现)
+    
+    返回:
+        关键点位置列表 (每个元素是 (N, 2) 的数组)
+    """
+    detector_pred = detector_pred.clone().detach()  # 创建预测副本 (避免修改原始数据)
+    B, _, h, w = detector_pred.shape    # 获取批次大小和图像尺寸
+    # if mask:
+    #     assert detector_label is not None
+    #     detector_pred[detector_pred < nms_thresh] = 0
+    #     label_mask = detector_label
+    #
+    #     # more area
+    #
+    #     detector_label = detector_label.long().cpu().numpy()
+    #     detector_label = detector_label.astype(np.uint8)
+    #     kernel = np.ones((3, 3), np.uint8)
+    #     label_mask = np.array([cv2.dilate(detector_label[s, 0], kernel, iterations=1)
+    #                            for s in range(len(detector_label))])
+    #     label_mask = torch.from_numpy(label_mask).unsqueeze(1)
+    #     detector_pred[label_mask > 1e-6] = 0
+    scores = simple_nms(detector_pred, nms_size)    # 应用快速NMS算法
+    scores = scores.reshape(B, h, w)    # 重塑分数图形状 (B, H, W)
+    # print(f"scores before thresh {nms_thresh}", scores)
+    points = [torch.nonzero(s > nms_thresh) for s in scores]    # 找出分数高于阈值的点
+    # print("points after thresh", points)
+    scores = [s[tuple(k.t())] for s, k in zip(scores, points)]  # 提取这些点的分数值
+    points, scores = list(zip(*[ remove_borders(k, s, 8, h, w) for k, s in zip(points, scores)]))   # 移除靠近边界的点
+    points = [torch.flip(k, [1]).long() for k in points]    # 翻转坐标顺序: [y, x] -> [x, y]
+
+    return points