Add files via upload

Readme.md

@@ -0,0 +1,103 @@
+#### Voxelmorph无监督图像配准说明
+
+##### **简介**
+
+​ 基于Voxelmoph（https://github.com/voxelmorph/voxelmorph）的无监督SAR图片配准。对于同一场景不同波段或入射角度获取SAR图像进行变形配准。网络结构可选择voxelmorph代码包networks里面的模型，默认使用unet-core基础模型。原理参考论文：**VoxelMorph: A Learning Framework for Deformable Medical Image Registration**
+
+##### **安装依赖库**
+
+1. Tensoflow-1.14.0
+2. keras-2.1.5
+3. PIL
+4. Opencv-python
+5. Matplotlib
+6. Pickle
+
+##### **训练**
+
+1. 数据：无监督训练，只需要准备图像对，即基准图像和待配准图像，放到不同的文件夹。需要分别生成图片的路径txt用于训练，验证，测试。
+
+ train_fixed.txt, train_moving.txt 分别是用于训练的基准图像和待配准图像，数量相同
+
+ val_fixed.txt, val_moving.txt. 用于训练时验证的基准图和待配准图
+
+2. 运行mydata_train.py,
+
+ 可自定义：
+
+ --model_dir 模型保存路径
+
+ --gpu 自定义GPU
+
+ --lr 学习率
+
+ --epochs 训练集循环次数
+
+ --img_size 训练图片大小
+
+ --ambda_param 损失函数权重，MSE损失默认使用0.01
+
+ --steps_per_epoch 根据bats_szie和训练集大小调整，计算方式为：训练样本数量/batch_size
+
+ -- batch_size 批训练大小
+
+ 命令行运行：mydata_train.py /my/path/to/train_fixed.txt /my/path/to/train_moving.txt /my/path/to/val_fixed.txt /my/path/to/val_moving.txt --gpu 0 --model_dir /my/path/to/save/models 
+
+ --img_size image size -- batch_size 16 --steps_per_epoch 488
+
+3. 训练监控训练时loss，默认每五个epoch保存模型可在源码中调整
+
+##### **测试**
+
+运行my data_test.py
+
+1. --root_dir 测试结果保存路径
+2. --fixed_dir 基准图像放置文件夹地址，将要配准的基准图像放在该路径下面
+3. --moving_dir 待配准图像放置文件夹地址，将待配准图像放在该路径下面
+4. --model_dir 加载训练好的权重，传递模型路径
+5. --img_size 测试图片大小，可以自定义，不用和训练时一致
+
+配准图示例：
+
+![predict](data/predict.png)
+
+流场可视化图：
+
+<img src="data/flow.png" alt="flow" style="zoom:72%;" />
+
+##### **训练过程曲线**
+
+训练完自动保存训练losss数据SAR_model_hist.pickle，可直接运行plot_loss_curve，绘制训练总损失，验证损失，相似性损失，流场损失等。
+
+总Loss:
+
+<img src="data/loss.jpg" alt="loss" style="zoom:72%;" />
+
+总Val_loss:
+
+<img src="data/val_loss.jpg" alt="val_loss" style="zoom:72%;" />
+
+相似性loss：
+
+<img src="data/spatial_transformer_loss.jpg" alt="spatial_transformer_loss" style="zoom:72%;" />
+
+验证：
+
+<img src="data/val_disp_loss.jpg" alt="val_disp_loss" style="zoom:72%;" />
+
+流场损失：
+
+<img src="data/disp_loss.jpg" alt="disp_loss" style="zoom:72%;" />
+
+验证流场损失：
+
+<img src="data/val_spatial_transformer_loss.jpg" alt="val_spatial_transformer_loss" style="zoom:72%;" />
+
+##### **比较配准前后结构相似性**
+
+运行SSAIM.py ,传递用于测试的基准图和待配准图路径txt文件，模型将会比较每一个样本配准前后结构相似性。
+
+<img src="/Users/huangwenbin/Desktop/未命名文件夹/data/基准_vs_待配准.png" alt="基准_vs_待配准" style="zoom:36%;" />
+
+<img src="/Users/huangwenbin/Desktop/未命名文件夹/data/_基准vs_配准后.png" alt="_基准vs_配准后" style="zoom:36%;" />
+

SAAIM.py

@@ -0,0 +1,183 @@
+# imports
+import os
+import sys
+import glob
+from argparse import ArgumentParser
+from PIL import Image
+from skimage.measure import compare_ssim as ssim
+import cv2
+
+# third party imports
+import numpy as np
+import keras.layers
+from keras.models import load_model
+from keras.optimizers import Adam
+import tensorflow as tf
+from keras.backend.tensorflow_backend import set_session
+from keras.utils import plot_model
+import matplotlib.pyplot as plt
+import keras.backend as K
+
+# 添加voxelmorph支持包路径
+sys.path.append('/Users/huangwenbin/Desktop/SAR-voxelmorph/ext/pynd-lib/')
+sys.path.append('/Users/huangwenbin/Desktop/SAR-voxelmorph/ext/pytools-lib/')
+sys.path.append('/Users/huangwenbin/Desktop/SAR-voxelmorph/ext/neuron/')
+sys.path.append('/usr/local/lib/python3.7/site-packages')
+import neuron
+
+# 添加本地函数
+import networks
+import losses
+import datagenerators
+
+os.environ['KMP_DUPLICATE_LIB_OK']='True' # 避免macos 报错OMP
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+
+
+def test(root_dir,fixed_dir,moving_dir,model_dir,gpu_id,img_size):
+ fixed_file=open(fixed_dir,'r')
+ moving_file=open(moving_dir,'r')
+
+
+ def str_strip(file):
+ list=[]
+ for f in file.readlines():
+ list.append(f.strip())
+ return list
+ fixed_vol_names = str_strip(fixed_file)
+ moving_vol_names = str_strip(moving_file) 
+ assert len(fixed_vol_names) > 0, "fixed路径中找不到训练数据"
+ assert len(moving_vol_names) > 0, "moving路径中找不到训练数据"
+
+
+
+ #GPU handling
+ gpu = '/gpu:%d' % 0 # gpu_id
+ os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id
+ config = tf.ConfigProto()
+ config.gpu_options.allow_growth = True
+ config.allow_soft_placement = True
+ set_session(tf.Session(config=config))
+
+ ndims = 2 
+ vol_shape = (img_size,img_size)
+ nb_enc_features = [32, 32, 32, 32] # 下采样卷积核个数
+ nb_dec_features = [32, 32, 32, 32, 32, 16] # 上采样卷积核个数
+
+
+ # 网络定义U-net
+ unet = networks.unet_core(vol_shape, nb_enc_features, nb_dec_features);
+
+ # 输入
+ print('numer of inputs', len(unet.inputs))
+ moving_input_tensor = unet.inputs[0]
+ fixed_input_tensor = unet.inputs[1]
+
+ # 输出
+ print('output:', unet.output)
+
+ # 转换为流场维度
+ disp_tensor = keras.layers.Conv2D(ndims, kernel_size=3, padding='same', name='disp')(unet.output)
+
+ # 显示流场维度
+ print('displacement tensor:', disp_tensor)
+
+ spatial_transformer = neuron.layers.SpatialTransformer(name='spatial_transformer')
+
+ # 扭转图像
+ moved_image_tensor = spatial_transformer([moving_input_tensor, disp_tensor])
+
+ inputs = [moving_input_tensor, fixed_input_tensor]
+ outputs = [moved_image_tensor, disp_tensor]
+ vxm_model = keras.models.Model(inputs, outputs)
+
+ # losses. Keras recognizes the string 'mse' as mean squared error, so we don't have to code it
+ loss = ['mse', losses.Grad('l2').loss]
+
+ # 损失函数
+ lambda_param = 0.01
+ loss_weights = [1, lambda_param]
+
+ #---------------加载模型权重-------------------------
+ vxm_model.compile(optimizer='Adam', loss=loss, loss_weights=loss_weights)
+ vxm_model.load_weights(model_dir)
+
+ #------------定义DICE函数-------------------------------
+ def dice_coef(y_true, y_pred):
+ y_true_f = y_true.flatten() # 将 y_true 拉伸为一维.
+ y_pred_f = y_pred.flatten()
+ intersection = sum(y_true_f * y_pred_f)
+ return (2. * intersection ) / (sum(y_true_f * y_true_f) + sum(y_pred_f * y_pred_f))
+
+ def compare_images(imageA, imageB, title):
+ # 分别计算输入图片的MSE和SSIM指标值的大小
+ s = ssim(imageA, imageB)
+ #return s
+
+ # 创建figure
+ fig = plt.figure(title)
+ plt.suptitle("SSIM: %.2f" % (s))
+
+ # 显示第一张图片
+ ax = fig.add_subplot(1, 2, 1)
+ plt.imshow(imageA, cmap = plt.cm.gray)
+ plt.axis("off")
+
+ # 显示第二张图片
+ ax = fig.add_subplot(1, 2, 2)
+ plt.imshow(imageB, cmap = plt.cm.gray)
+ plt.axis("off")
+ plt.tight_layout()
+ plt.show()
+
+
+ #--------------前向推理DICE计算------------------------------------
+ length=len(fixed_vol_names)
+ dice_before=0
+ dice_after=0
+ for i in range(300,length):
+ data=datagenerators.my_data_generator([fixed_vol_names[i]],[moving_vol_names[i]], batch_size=1,img_size=img_size)
+ sample, _ = next(data)
+ sample_pred = vxm_model.predict(sample)
+
+ fixed=sample[1].squeeze()
+ moving=sample[0].squeeze()
+ warped=sample_pred[0].squeeze()
+
+ # dice_before+=dice_coef(fixed,moving)
+ # dice_after+=dice_coef(fixed,warped)
+
+ #compare_images(fixed, fixed, "基准 vs 基准")
+ # dice_before+=compare_images(fixed, moving, "基准 vs 待配准")
+ # dice_after+=compare_images(fixed, warped, " 基准vs 配准后")
+ #compare_images(fixed, fixed, "基准 vs 待配准")
+ compare_images(fixed, moving, "基准 vs 待配准")
+ compare_images(fixed, warped, " 基准vs 配准后")
+
+ # Mean_Dice_befroe=dice_before/length
+ # Mean_Dice_after=dice_after/length
+ # print('配准前DICE：',Mean_Dice_befroe)
+ # print('配准后DICE：',Mean_Dice_after)
+
+if __name__ == "__main__":
+ parser = ArgumentParser()
+
+ parser.add_argument("--root_dir", type=str,default='data/test/',
+ help="结果输出目录")
+ parser.add_argument("--fixed_dir", type=str,default='/Users/huangwenbin/Desktop/SAR-voxelmorph/data/SAR_CUTS/SARcuts2/txt/test_fixed.txt',
+ help="固定图像目录")
+ parser.add_argument("--moving_dir", type=str,default='/Users/huangwenbin/Desktop/SAR-voxelmorph/data/SAR_CUTS/SARcuts2/txt/test_moving.txt',
+ help="待配准图像目录")
+
+ parser.add_argument("--model_dir", type=str,
+ dest="model_dir", default='models/TR1-MSE-300.h5',
+ help="models folder")
+ parser.add_argument("--gpu", type=str, default='0',
+ dest="gpu_id", help="gpu id number (or numbers separated by comma)")
+ parser.add_argument("--img_size", type=int,
+ default=512,
+ help="image size")
+
+ args = parser.parse_args()
+ print('测试参数：',args)
+ test(**vars(args))

compare_SSIAM.py

@@ -0,0 +1,26 @@
+
+from skimage.measure import compare_ssim as ssim
+import cv2
+import matplotlib.pyplot as plt
+
+
+def compare_images(imageA, imageB, title):
+ # 分别计算输入图片的MSE和SSIM指标值的大小
+ s = ssim(imageA, imageB)
+ #return s
+
+ # 创建figure
+ fig = plt.figure(title)
+ plt.suptitle("SSIM: %.2f" % (s))
+
+ # 显示第一张图片
+ ax = fig.add_subplot(1, 2, 1)
+ plt.imshow(imageA, cmap = plt.cm.gray)
+ plt.axis("off")
+
+ # 显示第二张图片
+ ax = fig.add_subplot(1, 2, 2)
+ plt.imshow(imageB, cmap = plt.cm.gray)
+ plt.axis("off")
+ plt.tight_layout()
+ plt.show()