#萌新记录某些日子的学习经历。#
前言
以GAN为基础的生成式对抗网络从抽象的角度来说是结合一组图像的风格和另一组图像的内容,或者说实现将初始图像从一个域迁移到目标图像的域。
前段时间我尝试了包括pix2pix, CycleGAN以及3d CycleGAN在内的三种图像生成网络,用于脑部图像T1-weighted和T1-weighted的互相转换。
应用于什么情形呢?比如说我们手上的数据集只有T1或者T2加权模态,那么我们可以用图像生成网络生成另一种模态,(甚至说生成其它种类的模态)。作为一个数据补充,比如可以用于其它模型的数据增强环节。
(注,因为大部分步骤源码已经说的很清晰我就不赘述了,主要说明一下整个操作的流程和注意事项。
摘要
基于Linux,使用pix2pix, CycleGAN和3d CycleGAN实现T1-weighted和T2-weighted模态的互转。
目录
Part one: pix2pix
pix2pix有两个特点:
1.训练时输入和输出的图像是成对的。
2.模型单向的。只能实现T1转T2或者T2转T1。如果需要互转需要训练两个模型(一个模型实现一个方向)。
准备工作
1.下载模型源码:pix2pix和CycleGAN的源码
2.创建虚拟环境、下载requirements.txt的所有依赖
3.如果要实现自己的图像转换,还需要准备数据集。
关于数据集的整理,通常有两个部分:首先是将3d图像切片为一个系列的2d图像(因为这里使用的pix2pix模型是针对2d输入的),然后是将2d图像以合适的文件格式保存起来(便于pix2pix的读取)。
切片环节,如果是脑部MR图像,一般沿着z轴取每个横截面作为2d图像。
如图所示,比如原图像大小为100*150*150,切片后变为100张大小为150*150的2d图像。保存格式,比如原图像是.nii或.nii.gz,2d图像格式为.png。
需要注意的是,pix2pix的训练需要创建一个文件夹,文件夹下包含"A"和"B"两个文件夹。分别存放初始图象和目标图像。为了训练时图像一一对应,两个类别的图像名字需要一一对应。具体实施规则:tip
训练
1.首先进入options文件夹,看看base和train的options.py文件有没有需要修改的超参数。
2.如果要实现T1和T2的互转,每个方向都需要一个pix2pix模型。
我的训练指令如下:
python train.py --dataroot (训练数据集路径)
--name (模型名称)
--model pix2pix
--direction BtoA (选择方向)
--gpu_ids 0,1,3
--batch_size 32
--load_size 256
最终得到的模型结果保存在checkpoints文件夹中和模型名称同名的文件夹下。
G,generator,生成器;D,discriminator,判别器。
预测
1.使用checkpoints文件夹中和模型名称同名的文件夹下的"latest_net_G.pth"。
我的预测指令如下,注意最后两行指令很重要,是模型参数:
python test.py --dataroot (预测数据集路径)
--name (模型名称)
--direction BtoA (方向)
--model test
--results_dir (输出文件夹)
--num_test (待处理图像数量)
--dataset_mode single
--norm batch
--netG unet_256
最终在输出文件夹中得到预测的结果。
需要注意的是,得到的结果图像是2d图像,还需要resize并按原来的顺序组合成3d图像,并以nii或nii.gz的格式保存。
Part two: CycleGAN
CycleGAN相比于pix2pix有两个特点:
1.训练时输入和输出的图像不需要一一对应,只需要两个图象集合对应即可。
2.模型双向的。CycleGAN是一组镜像对称的GAN网络组成的环形循环网络。训练一个CycleGAN即可实现T1和T2加权模态的互转。
模型同样是处理2d图像的,预处理需要切片(3d转2d,nii/nii.gz转png),后处理需要resize和组合切片(png转nii/nii.gz)。
准备工作
1.模型下载、环境配置和pix2pix相同,参照Part one。
2.同样需要切片操作。但数据的保存方式略有不同,参照tip
训练
1.首先进入options文件夹,看看base和train的options.py文件有没有需要修改的超参数。
2.实现T1和T2的互转,只需要训练一个模型
我的训练指令如下:
python train.py --dataroot (训练数据集路径)
--name (模型名称)
--model cycle_gan
--gpu_ids 0,1,3
--batch_size 32
--load_size 256
最终得到的模型结果保存在checkpoints文件夹中和模型名称同名的文件夹下。
G,generator,生成器;D,discriminator,判别器。
需要注意的是这里有两组(每组表示一个方向)的G和D。
预测
1.使用checkpoints文件夹中和模型名称同名的文件夹下的"latest_net_G.pth"。
需要注意的是,
如训练所示,保存的模型是不能直接用的,需要手动设置。我的建议是:
在checkpoints文件夹中新建两个文件夹表示不同方向的GAN网络。在每个文件夹下保存各自方向的生成器参数。比如说新建文件夹"cycle_gan_a"和"cycle_gan_b"。将训练得到的"latest_net_G_A.pth"保存在"cycle_gan_a"中并重命名为"latest_net_G.pth"。将训练得到的"latest_net_G_B.pth"保存在"cycle_gan_b"中并重命名为"latest_net_G.pth"。以此分别存储两个方向的生成器。
我的预测指令如下,注意最后两行指令很重要,是模型参数:
python test.py --dataroot (预测数据集路径)
--name ***_a (模型名称,表示A方向)
--model test
--results_dir (输出文件夹)
--num_test (待处理图像数量)
--dataset_mode single
--norm instance
--netG resnet_9blocks --no_dropout
最终在输出文件夹中得到预测的结果。
需要注意的是,得到的结果图像是2d图像,还需要resize并按原来的顺序组合成3d图像,并以nii或nii.gz的格式保存。
Part three: 3d CycleGAN
Part one和Part two都是使用的2d模型。
3d模型从理论上来说会关注切片间的信息。也许生成的图像质量更高。
但实际上存在两个问题:
1.3d模型训练想要充分训练,达到很好的拟合需要更多的训练资源,在硬件受限的情况下很难。
2.3d模型可以说是魔改后的2d模型,具体效果如何谁也不知道。
模型的使用来源于:3d CycleGAN
因为作者大部分内容说的很清楚了所以我这里只是简单过一下,并补充一些可能遇到的问题。
1.下载模型,配置依赖
2.按照要求整理好数据集(自己调整),nii格式,并按照预处理指令处理
3.在options里修改自己想修改的参数后,开始训练
4.训练好的模型,在预测时会面临两个问题。
首先,详细操作参照Part two预测部份。因为得到的模型是两组G和D,为了便于使用需要放置于不同的文件夹下并重命名,调用的生成器是模型名称下的"latest_net_G.pth"。
其次,这个模型好像一次只能预测一张图像,如果需要批量处理的话需要改一下。
直接修改“test.py”,参考代码如下:
#if __name__ == '__main__':
# opt = TestOptions().parse()
# model = create_model(opt)
# model.setup(opt)
# inference(model, opt.image, opt.result, opt.resample, opt.new_resolution, opt.patch_size[0],
# opt.patch_size[1], opt.patch_size[2], opt.stride_inplane, opt.stride_layer, 1)
if __name__ == '__main__':
opt = TestOptions().parse()
model = create_model(opt)
model.setup(opt)
for filename in os.listdir(opt.image):
resultname = filename[0:-5] + "2" + filename[-4:]
inference(model, opt.image+'/'+filename, opt.result+'/'+resultname, opt.resample, opt.new_resolution, opt.patch_size[0],
opt.patch_size[1], opt.patch_size[2], opt.stride_inplane, opt.stride_layer, 1)
总结
使用了pix2pix,CycleGAN和3d CycleGAN进行T1-weighted和T2-weighted模态的互转。
我自己训练大概每个模型训练了接近10天。
从得到的结果来看性能效果:pix2pix>CycleGAN>3d CycleGAN。不过从衡量指标来看基本上生成的模态还是有着较高质量的,可以模拟真实模态作为其他模型的输入。这是一个很有潜力的数据增强策略。
未来可以考虑使用更多图像生成网络生成更多模态(比如StarGAN生成FLAIR, T1ce?)。
(PS 我应该当时一边做一边记录的,现在回想起来有些细节已经模糊了所以就没写太多。
感谢阅读,请多指教!
Appendix
最终衡量生成的虚假模态质量的指标可以用PSNR和SSIM。
计算生成的虚假模态和真实模态之间的偏差。
from skimage.metrics import peak_signal_noise_ratio
from skimage.metrics import structural_similarity
Reference
[1] Goodfellow I J, Pouget-Abadie J, Mirza M , et al. Generative Adversarial Networks[J]. Advances in Neural Information Processing Systems, 2014, 3:2672-2680.
[2] Isola P, Zhu J Y, Zhou T, et al. Image-to-image translation with conditional adversarial networks[C]//Proceedings of the IEEE conference on computer vision and pattern recognition. 2017: 1125-1134.
[3] Zhu J Y, Park T, Isola P, et al. Unpaired image-to-image translation using cycle-consistent adversarial networks[C]//Proceedings of the IEEE international conference on computer vision. 2017: 2223- 2232.
总之就是很感谢把模型分享出来的各位大佬,ml的学习之路完全不能缺少前人们智慧的结晶~泪目