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