nnUNet V2代码——图像增强（四）

本文阅读的nnU-Net V2图像增强有伽马校正、镜像

各个类内的各个函数的调用关系见前文nnUNet V2代码——图像增强（一）的BasicTransform类

安装batchgeneratorsv2，nnU-Net V2关于图像增强的代码都在这个库中，点击链接，将其clone到本地后，在命令行进入文件夹内，pip install -e . 即可（注意-e后有个点）。

一 GammaTransform

该类负责伽马校正，继承自ImageOnlyTransform类，只对image施加，seg不施加。

代码在batchgeneratorsv2 \ transforms \ intensity \ intensity.py文件中

1. __init__函数

变量名称	含义
self.gamma	伽马调整范围
self.p_invert_image	反转图像
self.synchronize_channels	各通道是否施加相同的低分辨率处理
self.p_per_channel	各通道施加伽马调整的概率
self.p_retain_stats	伽马校正前后图像的均值、方差是否保持一致

2. get_parameters函数

def get_parameters(self, **data_dict) -> dict:
	## 获取图像大小
    shape = data_dict['image'].shape
    ## 确定哪些通道需要伽马校正
    apply_to_channel = torch.where(torch.rand(shape[0]) < self.p_per_channel)[0]
    ## 是否保持伽马校正前后图像的均值、方差是否保持一致，nnU-Net V2设置为一致
    retain_stats = torch.rand(len(apply_to_channel)) < self.p_retain_stats
    ## 是否选择反转图像进行伽马校正，nnU-Net V2会进行两次伽马校正，一次反转，一次不反转
    invert_image = torch.rand(len(apply_to_channel)) < self.p_invert_image
	
	## 各通道同步应用相同的伽马校正
    if self.synchronize_channels:
        gamma = torch.Tensor([sample_scalar(self.gamma, image=data_dict['image'], channel=None)] * len(apply_to_channel))
    ## 各通道应用各自的伽马校正
    else:
        gamma = torch.Tensor([sample_scalar(self.gamma, image=data_dict['image'], channel=c) for c in apply_to_channel])
    return {
        'apply_to_channel': apply_to_channel,
        'retain_stats': retain_stats,
        'invert_image': invert_image,
        'gamma': gamma
    }

sample_scalar函数见nnUNet V2代码——图像增强（一）的其余函数

3. _apply_to_image函数

def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
	## 遍历所有需要应用伽马矫正的通道
    for c, r, i, g in zip(params['apply_to_channel'], params['retain_stats'], params['invert_image'], params['gamma']):
        ## 如果本次反转图像
        if i:
            img[c] *= -1
        ## 如果本次保持校正前后图像均值、方差一致
        if r:
            mean = torch.mean(img[c])
            std = torch.std(img[c])
        ## 记录最值
        minm = torch.min(img[c])
        rnge = torch.max(img[c]) - minm
        ## 伽马校正
        img[c] = torch.pow(((img[c] - minm) / torch.clamp(rnge, min=1e-7)), g) * rnge + minm
        if r:
            mn_here = torch.mean(img[c])
            std_here = torch.std(img[c])
            img[c] -= mn_here
            img[c] *= (std / torch.clamp(std_here, min=1e-7))
            img[c] += mean
		## 反转回来
        if i:
            img[c] *= -1
    return img

二 MirrorTransform

该类负责图像镜像

代码在batchgeneratorsv2 \ transforms \ spatial\ mirroring.py文件中

1. __init__函数

## 允许镜像的轴
self.allowed_axes = allowed_axes

2. get_parameters函数

def get_parameters(self, **data_dict) -> dict:
	## 判断允许镜像的轴是否真的应用镜像，概率固定为0.5
    axes = [i for i in self.allowed_axes if torch.rand(1) < 0.5]
    return {
        'axes': axes
    }

3. _apply_to_image函数

def _apply_to_image(self, img: torch.Tensor, **params) -> torch.Tensor:
	## 没有镜像的轴，直接返回
    if len(params['axes']) == 0:
        return img
    ## 这里+1是为了匹配具有通道维度的img
    ## 计算允许镜像的轴时，维度是从0开始的
    ## 而应用镜像的img的维度是从1开始的
    ## img的0维度是通道维度
    axes = [i + 1 for i in params['axes']]
    ## 调用flip镜像
    return torch.flip(img, axes)

_apply_to_segmentation函数、_apply_to_regr_target函数和_apply_to_image函数流程一致