def train(data_dir,
atlas_file,
model_dir,
model,
gpu_id,
lr,
nb_epochs,
prior_lambda,
image_sigma,
mean_lambda,
steps_per_epoch,
batch_size,
load_model_file,
bidir,
atlas_wt,
bias_mult,
smooth_pen_layer,
data_loss,
reg_param,
ncc_win,
initial_epoch=0):
"""
model training function
:param data_dir: folder with npz files for each subject.
:param atlas_file: atlas filename. So far we support npz file with a 'vol' variable
:param model_dir: model folder to save to
:param gpu_id: integer specifying the gpu to use
:param lr: learning rate
:param nb_epochs: number of training iterations
:param prior_lambda: the prior_lambda, the scalar in front of the smoothing laplacian, in MICCAI paper
:param image_sigma: the image sigma in MICCAI paper
:param steps_per_epoch: frequency with which to save models
:param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size
:param load_model_file: optional h5 model file to initialize with
:param bidir: logical whether to use bidirectional cost function
"""
# prepare data files
# we have data arranged in train/validate/test folders
# inside each folder is a /vols/ and a /asegs/ folder with the volumes
# and segmentations. All of our papers use npz formated data.
train_vol_names = glob.glob(data_dir)
train_vol_names = [f for f in train_vol_names if 'ADNI' not in f]
random.shuffle(train_vol_names) # shuffle volume list
assert len(train_vol_names) > 0, "Could not find any training data"
# data generator
train_example_gen = datagenerators.example_gen(train_vol_names, batch_size=batch_size)
# prepare the initial weights for the atlas "layer"
if atlas_file is None or atlas_file == "":
nb_atl_creation = 100
print('creating "atlas" by averaging %d subjects' % nb_atl_creation)
x_avg = 0
for _ in range(nb_atl_creation):
x_avg += next(train_example_gen)[0][0,...,0]
x_avg /= nb_atl_creation
x_avg = x_avg[np.newaxis,...,np.newaxis]
atlas_vol = x_avg
else:
atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis]
vol_size = atlas_vol.shape[1:-1]
# Diffeomorphic network architecture used in MICCAI 2018 paper
nf_enc = [16,32,32,32]
nf_dec = [32,32,32,32,16,3]
if model == 'm1':
pass
elif model == 'm1double':
nf_enc = [f*2 for f in nf_enc]
nf_dec = [f*2 for f in nf_dec]
# prepare model folder
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
assert data_loss in ['mse', 'cc', 'ncc'], 'Loss should be one of mse or cc, found %s' % data_loss
if data_loss in ['ncc', 'cc']:
data_loss = losses.NCC(win=[ncc_win]*3).loss
else:
data_loss = lambda y_t, y_p: K.mean(K.square(y_t-y_p))
# gpu handling
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# prepare the model
with tf.device(gpu):
# the MICCAI201 model takes in [image_1, image_2] and outputs [warped_image_1, velocity_stats]
# in these experiments, we use image_2 as atlas
model = networks.img_atlas_diff_model(vol_size, nf_enc, nf_dec,
atl_mult=1, bidir=bidir,
smooth_pen_layer=smooth_pen_layer)
# compile
mean_layer_loss = lambda _, y_pred: mean_lambda * K.mean(K.square(y_pred))
flow_vol_shape = model.outputs[-2].shape[1:-1]
loss_class = losses.Miccai2018(image_sigma, prior_lambda, flow_vol_shape=flow_vol_shape)
if bidir:
model_losses = [data_loss,
lambda _,y_p: data_loss(model.get_layer('atlas').output, y_p),
mean_layer_loss,
losses.Grad('l2').loss]
loss_weights = [atlas_wt, 1-atlas_wt, 1, reg_param]
else:
model_losses = [loss_class.recon_loss, loss_class.kl_loss, mean_layer_loss]
loss_weights = [1, 1, 1]
model.compile(optimizer=Adam(lr=lr), loss=model_losses, loss_weights=loss_weights)
# set initial weights in model
model.get_layer('atlas').set_weights([atlas_vol[0,...]])
# load initial weights. # note this overloads the img_param weights
if load_model_file is not None and len(load_model_file) > 0:
model.load_weights(load_model_file, by_name=True)
# save first iteration
model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch))
# atlas_generator specific to this model. Once we're convinced of this, move to datagenerators
def atl_gen(gen):
zero_flow = np.zeros([batch_size, *vol_size, len(vol_size)])
zero_flow_half = np.zeros([batch_size] + [f//2 for f in vol_size] + [len(vol_size)])
while 1:
x2 = next(train_example_gen)[0]
# TODO: note this is the opposite of train_miccai and it might be confusing.
yield ([atlas_vol, x2], [x2, atlas_vol, zero_flow, zero_flow])
atlas_gen = atl_gen(train_example_gen)
# prepare callbacks
save_file_name = os.path.join(model_dir, '{epoch:02d}.h5')
save_callback = ModelCheckpoint(save_file_name)
# fit generator
with tf.device(gpu):
model.fit_generator(atlas_gen,
initial_epoch=initial_epoch,
epochs=nb_epochs,
callbacks=[save_callback],
steps_per_epoch=steps_per_epoch,
verbose=1)
def train(data_dir,
atlas_file,
model_dir,
model,
gpu_id,
lr,
nb_epochs,
prior_lambda,
image_sigma,
mean_lambda,
steps_per_epoch,
batch_size,
load_model_file,
bidir,
atlas_wt,
bias_mult,
smooth_pen_layer,
data_loss,
reg_param,
ncc_win,
initial_epoch=0):
"""
model training function
:param data_dir: folder with npz files for each subject.
:param atlas_file: atlas filename. So far we support npz file with a 'vol' variable
:param model_dir: model folder to save to
:param gpu_id: integer specifying the gpu to use
:param lr: learning rate
:param nb_epochs: number of training iterations
:param prior_lambda: the prior_lambda, the scalar in front of the smoothing laplacian, in MICCAI paper
:param image_sigma: the image sigma in MICCAI paper
:param steps_per_epoch: frequency with which to save models
:param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size
:param load_model_file: optional h5 model file to initialize with
:param bidir: logical whether to use bidirectional cost function
"""
# prepare data files
# we have data arranged in train/validate/test folders
# inside each folder is a /vols/ and a /asegs/ folder with the volumes
# and segmentations. All of our papers use npz formated data.
train_vol_names = glob.glob(data_dir)
train_vol_names = [f for f in train_vol_names if 'ADNI' not in f]
random.shuffle(train_vol_names) # shuffle volume list
assert len(train_vol_names) > 0, "Could not find any training data"
# prepare data generation specific to conditional templates
train_vols_dir = os.path.join(data_dir, 'vols')
train_vol_basenames = [
f for f in os.listdir(train_vols_dir) if ('ADNI' in f or 'ABIDE' in f)
]
train_vol_names = [
os.path.join(train_vols_dir, f) for f in train_vol_basenames
]
# csv pruning
# our csv is a file of the form:
# file,age,sex
# ADNI_ADNI-1.5T-FS-5.3-Long_63196.long.094_S_1314_base_mri_talairach_norm.npz,81.0,2
csv_file_path = os.path.join(data_dir, 'combined_pheno.csv')
train_atr_dct = load_pheno_csv(csv_file_path)
train_vol_basenames = [
f for f in train_vol_basenames if f in list(train_atr_dct.keys())
]
train_vol_names = [
os.path.join(train_vols_dir, f) for f in train_vol_basenames
]
# replace keys with full path
for key in list(train_atr_dct.keys()):
if key in train_vol_basenames:
train_atr_dct[os.path.join(data_dir, 'vols',
key)] = train_atr_dct[key]
train_atr_dct.pop(key, None)
# prepare the initial weights for the atlas "layer"
if atlas_file is None or atlas_file == "":
nb_atl_creation = 100
print('creating "atlas" by averaging %d subjects' % nb_atl_creation)
x_avg = 0
for _ in range(nb_atl_creation):
x_avg += next(train_example_gen)[0][0, ..., 0]
x_avg /= nb_atl_creation
x_avg = x_avg[np.newaxis, ..., np.newaxis]
atlas_vol = x_avg
else:
atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis]
vol_size = atlas_vol.shape[1:-1]
genobj = Generator(train_vol_names, atlas_vol, y_k=train_atr_dct)
# Diffeomorphic network architecture used in MICCAI 2018 paper
nf_enc = [16, 32, 32, 32]
nf_dec = [32, 32, 32, 32, 16, 3]
if model == 'm1':
pass
elif model == 'm1double':
nf_enc = [f * 2 for f in nf_enc]
nf_dec = [f * 2 for f in nf_dec]
# prepare model folder
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
assert data_loss in [
'mse', 'cc', 'ncc'
], 'Loss should be one of mse or cc, found %s' % data_loss
if data_loss in ['ncc', 'cc']:
data_loss = losses.NCC(win=[ncc_win] * 3).loss
else:
data_loss = lambda y_t, y_p: K.mean(K.square(y_t - y_p))
# gpu handling
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# prepare the model
with tf.device(gpu):
# parameters for atlas construction.
nb_conv_features = 4
cond_im_input_shape = [160, 192, 224, 4] # note the 8 features
cond_nb_levels = 0
cond_conv_size = [3, 3, 3]
extra_conv_layers = 3
pheno_input_shape = [2]
bidir = True
smooth_pen_layer = 'diffflow'
model, mn = networks.cond_img_atlas_diff_model(
vol_size,
nf_enc,
nf_dec,
atl_mult=1,
bidir=bidir,
smooth_pen_layer=smooth_pen_layer,
nb_conv_features=nb_conv_features,
cond_im_input_shape=cond_im_input_shape,
cond_nb_levels=cond_nb_levels,
cond_conv_size=cond_conv_size,
pheno_input_shape=pheno_input_shape,
extra_conv_layers=extra_conv_layers,
ret_vm=True,
)
outputs = [model.outputs[f]
for f in [0, 2, 3, 3]] # latest model used in paper
model = keras.models.Model(model.inputs, outputs)
# compile
mean_layer_loss = lambda _, y_pred: mean_lambda * K.mean(
K.square(y_pred))
model_losses = [
data_loss, # could be mse or ncc or a combination, etc
mean_layer_loss,
lambda _, yp: losses.Grad('l2').loss(_, yp),
lambda _, yp: K.mean(K.square(yp))
]
# parameters used in paper
msmag_param = 0.01
reg_param = 1
centroid_reg = 1
loss_weights = [atlas_wt, centroid_reg, reg_param, msmag_param]
model.compile(optimizer=keras.optimizers.Adam(lr=lr),
loss=model_losses,
loss_weights=loss_weights)
# load initial weights. # note this overloads the img_param weights
if load_model_file is not None and len(load_model_file) > 0:
model.load_weights(load_model_file, by_name=True)
# save first iteration
model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch))
# prepare callbacks
save_file_name = os.path.join(model_dir, '{epoch:02d}.h5')
save_callback = ModelCheckpoint(save_file_name)
# fit generator
with tf.device(gpu):
model.fit_generator(genobj.cond_mean_flow_x2(batch_size=batch_size),
initial_epoch=initial_epoch,
epochs=nb_epochs,
callbacks=[save_callback],
steps_per_epoch=steps_per_epoch,
verbose=1)
def train(data_dir,
atlas_file,
model,
model_name,
gpu_id,
lr,
nb_epochs,
reg_param,
steps_per_epoch,
batch_size,
load_model_file,
data_loss,
initial_epoch=0):
"""
model training function
:param data_dir: folder with npz files for each subject.
:param atlas_file: atlas filename. So far we support npz file with a 'vol' variable
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param model_dir: the model directory to save to
:param gpu_id: integer specifying the gpu to use
:param lr: learning rate
:param n_iterations: number of training iterations
:param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper)
:param steps_per_epoch: frequency with which to save models
:param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size
:param load_model_file: optional h5 model file to initialize with
:param data_loss: data_loss: 'mse' or 'ncc
"""
# load atlas from provided files. The atlas we used is 160x192x224.
# atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis]
atlas_vol = nib.load(atlas_file).get_data()[np.newaxis, ..., np.newaxis]
vol_size = atlas_vol.shape[1:-1]
# prepare data files
# for the CVPR and MICCAI papers, we have data arranged in train/validate/test folders
# inside each folder is a /vols/ and a /asegs/ folder with the volumes
# and segmentations. All of our papers use npz formated data.
train_vol_names = glob.glob(os.path.join(data_dir, '*.npz'))
random.shuffle(train_vol_names) # shuffle volume list
assert len(train_vol_names) > 0, "Could not find any training data"
# UNET filters for voxelmorph-1 and voxelmorph-2,
# these are architectures presented in CVPR 2018
nf_enc = [16, 32, 32, 32]
if model == 'vm1':
nf_dec = [32, 32, 32, 32, 8, 8]
elif model == 'vm2':
nf_dec = [32, 32, 32, 32, 32, 16, 16]
else: # 'vm2double':
nf_enc = [f * 2 for f in nf_enc]
nf_dec = [f * 2 for f in [32, 32, 32, 32, 32, 16, 16]]
assert data_loss in [
'mse', 'cc', 'ncc'
], 'Loss should be one of mse or cc, found %s' % data_loss
if data_loss in ['ncc', 'cc']:
data_loss = losses.NCC().loss
model_dir = "../models/" + model_name
# prepare model folder
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
# GPU handling
gpu = '/gpu:%d' % gpu_id
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# prepare the model
with tf.device(gpu):
# prepare the model
# in the CVPR layout, the model takes in [image_1, image_2] and outputs [warped_image_1, flow]
# in the experiments, we use image_2 as atlas
model = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
# load initial weights
if load_model_file is not None and load_model_file != '':
print('loading', load_model_file)
model.load_weights(load_model_file)
# save first iteration
model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch))
# data generator
# nb_gpus = len(gpu_id.split(','))
# assert np.mod(batch_size, nb_gpus) == 0, \
# 'batch_size should be a multiple of the nr. of gpus. ' + \
# 'Got batch_size %d, %d gpus' % (batch_size, nb_gpus)
nb_gpus = 1
train_example_gen = datagenerators.example_gen(train_vol_names,
batch_size=batch_size)
atlas_vol_bs = np.repeat(atlas_vol, batch_size, axis=0)
cvpr2018_gen = datagenerators.cvpr2018_gen(train_example_gen,
atlas_vol_bs,
batch_size=batch_size)
# prepare callbacks
save_file_name = os.path.join(model_dir, '{epoch:02d}.h5')
# fit generator
with tf.device(gpu):
# multi-gpu support
if nb_gpus > 1:
save_callback = nrn_gen.ModelCheckpointParallel(save_file_name)
mg_model = multi_gpu_model(model, gpus=nb_gpus)
# single-gpu
else:
save_callback = ModelCheckpoint(save_file_name)
mg_model = model
# compile
mg_model.compile(optimizer=Adam(lr=lr),
loss=[data_loss, losses.Grad('l2').loss],
loss_weights=[1.0, reg_param])
# fit
mg_model.fit_generator(cvpr2018_gen,
initial_epoch=initial_epoch,
epochs=nb_epochs,
callbacks=[save_callback],
steps_per_epoch=steps_per_epoch,
verbose=1)
def train(data_dir,
atlas_file,
model_dir,
gpu_id,
lr,
nb_epochs,
prior_lambda,
image_sigma,
steps_per_epoch,
batch_size,
load_model_file,
bidir,
bool_cc,
initial_epoch=0):
"""
model training function
:param data_dir: folder with npz files for each subject.
:param atlas_file: atlas filename. So far we support npz file with a 'vol' variable
:param model_dir: model folder to save to
:param gpu_id: integer specifying the gpu to use
:param lr: learning rate
:param nb_epochs: number of training iterations
:param prior_lambda: the prior_lambda, the scalar in front of the smoothing laplacian, in MICCAI paper
:param image_sigma: the image sigma in MICCAI paper
:param steps_per_epoch: frequency with which to save models
:param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size
:param load_model_file: optional h5 model file to initialize with
:param bidir: logical whether to use bidirectional cost function
:param bool_cc: Train CC or MICCAI version
"""
# load atlas from provided files. The atlas we used is 160x192x224.
#atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis]
vm_dir = '/home/jdram/voxelmorph/'
base = np.load(
os.path.join(vm_dir, "data",
"ts12_dan_a88_fin_o_trim_adpc_002661_256.npy"))
monitor = np.load(
os.path.join(vm_dir, "data",
"ts12_dan_a05_fin_o_trim_adpc_002682_256.npy"))
#base = np.load(os.path.join(vm_dir, "data","ts12_dan_a88_fin_o_trim_adpc_002661_abs.npy"))
#monitor = np.load(os.path.join(vm_dir, "data","ts12_dan_a05_fin_o_trim_adpc_002682_abs.npy"))
#vol_size = (64, 64, 64)
vol_size = (64, 64, 256 - 64)
#vol_size = (128, 128, 256)
# prepare data files
# for the CVPR and MICCAI papers, we have data arranged in train/validate/test folders
# inside each folder is a /vols/ and a /asegs/ folder with the volumes
# and segmentations. All of our papers use npz formated data.
#train_vol_names = glob.glob(os.path.join(data_dir, '*.npy'))
#random.shuffle(train_vol_names) # shuffle volume list
#assert len(train_vol_names) > 0, "Could not find any training data"
# Diffeomorphic network architecture used in MICCAI 2018 paper
nf_enc = [32, 64, 64, 64]
nf_dec = [64, 64, 64, 64, 32, 3]
# prepare model folder
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
tf.reset_default_graph()
if bool_cc:
pre_net = "cc_"
else:
if bidir:
pre_net = "miccai_bidir_"
else:
pre_net = "miccai_"
# gpu handling
gpu = '/device:GPU:%d' % int(gpu_id) # 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))
# prepare the model
with tf.device(gpu):
# prepare the model
# in the CVPR layout, the model takes in [image_1, image_2] and outputs [warped_image_1, flow]
# in the experiments, we use image_2 as atlas
if bool_cc:
model = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
else:
model = networks.miccai2018_net(vol_size,
nf_enc,
nf_dec,
bidir=bidir,
vel_resize=.5)
# load initial weights
if load_model_file is not None and load_model_file != "":
print('loading', load_model_file)
model.load_weights(load_model_file)
# save first iteration
model.save(os.path.join(model_dir, f'{pre_net}{initial_epoch:02d}.h5'))
model.summary()
if bool_cc:
model_losses = [losses.NCC().loss, losses.Grad('l2').loss]
loss_weights = [1.0, 0.01] # recommend 1.0 for ncc, 0.01 for mse
else:
flow_vol_shape = model.outputs[-1].shape[1:-1]
loss_class = losses.Miccai2018(image_sigma,
prior_lambda,
flow_vol_shape=flow_vol_shape)
if bidir:
model_losses = [
loss_class.recon_loss, loss_class.recon_loss,
loss_class.kl_loss
]
loss_weights = [0.5, 0.5, 1]
else:
model_losses = [loss_class.recon_loss, loss_class.kl_loss]
loss_weights = [1, 1]
segy_gen = datagenerators.segy_gen(base, monitor, batch_size=batch_size)
# prepare callbacks
save_file_name = os.path.join(model_dir, pre_net + '{epoch:02d}.h5')
with tf.device(gpu):
# fit generator
save_callback = ModelCheckpoint(save_file_name, period=5)
csv_cb = CSVLogger(f'{pre_net}log.csv')
nan_cb = TerminateOnNaN()
rlr_cb = ReduceLROnPlateau(monitor='loss', verbose=1)
els_cb = EarlyStopping(monitor='loss',
patience=15,
verbose=1,
restore_best_weights=True)
cbs = [save_callback, csv_cb, nan_cb, rlr_cb, els_cb]
mg_model = model
# compile
mg_model.compile(optimizer=Adam(lr=lr),
loss=model_losses,
loss_weights=loss_weights)
mg_model.fit(
[base, monitor],
[monitor, np.zeros_like(base)],
initial_epoch=initial_epoch,
batch_size=8,
epochs=nb_epochs,
callbacks=cbs,
#steps_per_epoch=steps_per_epoch,
verbose=1)
def test(root_dir, fixed_dir, moving_dir, model_dir, gpu_id, img_size):
#
# 定义模型，加载权重
# 输入维度
#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]
loss = [losses.NCC().loss, 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)
#--------------前向推理------------------------------------
fixed_vol_names = glob.glob(os.path.join(fixed_dir, '*.png'))
moving_vol_names = glob.glob(os.path.join(moving_dir, '*.png'))
print(fixed_vol_names, moving_vol_names)
# fixed_vol_names.sort()
# moving_vol_names.sort()
data = datagenerators.my_data_generator(fixed_vol_names,
moving_vol_names,
batch_size=1,
img_size=img_size)
sample, _ = next(data)
print('输入维度：', sample[0].shape)
sample_pred = vxm_model.predict(sample)
#---------------保存流场维度数据------------------------------
print('流场输出维度：', sample_pred[1].shape)
#a=sample_pred[0].squeeze()
slices_in = [sample_pred[1].squeeze()]
np.save(root_dir + 'flow.npy', slices_in)
u, v = slices_in[0][..., 0], slices_in[0][..., 1]
#u, v = flow_smooth.smooth(u,v,1)
# np.savetxt(root_dir+'a.csv',a,delimiter=",")
np.savetxt(root_dir + 'u.csv', u, delimiter=",")
np.savetxt(root_dir + 'v.csv', v, delimiter=",") # 保存偏移值
#------------输出配准点的列表-------txt---------------------------
txt = open(root_dir + 'match.txt', 'w')
moving_img = Image.open(moving_vol_names[0])
# moving_pixeles=str(moving_vol_names[0]).split('_')
# fixed_pixeles=str(fixed_vol_names[0]).split('_')
# print(moving_pixeles)
# print(fixed_pixeles)
# x_moving=float(moving_pixeles[2])
# y_moving=float(moving_pixeles[3].strip('.png')) # 读取文件名中左上角初始坐标
# x_fixed=float(fixed_pixeles[2])
# y_fixed=float(fixed_pixeles[3].strip('.png'))
# width,height=moving_img.size
# print('待配准图片长宽：',width,height)
# w_scale=width/img_size
# h_scale=height/img_size
# for i in range(img_size):
# for j in range(img_size):
# x1=str(i*w_scale+x_moving) # 坐标转换
# y1=str(j*h_scale+y_moving)
# x2=str(i+u[i][j]+x_fixed)
# y2=str(j-v[i][j]+y_fixed)
# txt.write(x1+' '+y1+' '+x2+' '+y2)
# txt.write('\n')
width, height = moving_img.size
print('待配准图片长宽：', width, height)
w_scale = width / img_size
h_scale = height / img_size
for i in range(img_size):
for j in range(img_size):
# if u[i][j]!=0:
x1 = str(i * w_scale) # 坐标转换
y1 = str(j * h_scale)
x2 = str(i + u[i][j])
y2 = str(j - v[i][j])
txt.write(x1 + ' ' + y1 + ' ' + x2 + ' ' + y2)
txt.write('\n')
#---------------可视化配准图像和流场-----------------------------
slices_2d = [f[0, ..., 0] for f in sample + sample_pred]
# fixed=sample[1].squeeze()
# moving=sample[0].squeeze()
warped = sample_pred[0].squeeze()
#print(warped.shape)
plt.imshow(warped, cmap='Greys_r')
plt.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0, wspace=0)
plt.margins(0, 0)
plt.axis('off')
plt.savefig(root_dir + 'warped_picture.png',
transparent=True,
dpi=300,
pad_inches=0.0)
# plt.show()
titles = [
'input_moving', 'input_fixed', 'predicted_moved', 'deformation_x'
]
neuron.plot.slices(slices_2d,
titles=titles,
cmaps=['gray'],
do_colorbars=True,
path=root_dir + 'predict.png')
neuron.plot.flow([sample_pred[1].squeeze()],
width=5,
path=root_dir + 'flow.png')
def train(
data_dir,
val_data_dir,
atlas_file,
val_atlas_file,
model,
model_dir,
gpu_id,
lr,
nb_epochs,
reg_param,
gama_param,
steps_per_epoch,
batch_size,
load_model_file,
data_loss,
seg_dir=None, # one file
val_seg_dir=None,
Sf_file=None, # one file
val_Sf_file=None,
auxi_label=None,
initial_epoch=0):
"""
model training function
:param data_dir: folder with npz files for each subject.
:param atlas_file: atlas filename. So far we support npz file with a 'vol' variable
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param model_dir: the model directory to save to
:param gpu_id: integer specifying the gpu to use
:param lr: learning rate
:param n_iterations: number of training iterations
:param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper)
:param steps_per_epoch: frequency with which to save models
:param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size
:param load_model_file: optional h5 model file to initialize with
:param data_loss: 'mse' or 'ncc
:param auxi_label: whether to use auxiliary informmation during the training
"""
# load atlas from provided files. The atlas we used is 160x192x224.
# atlas_file = 'D:/voxel/data/t064.tif'
atlas = Image.open(atlas_file) # is a TiffImageFile _size is (628, 690)
atlas_vol = np.array(atlas)[
np.newaxis, ..., np.newaxis] # is a ndarray, shape is (1, 690, 628, 1)
# new = Image.fromarray(X) new.size is (628, 690)
vol_size = atlas_vol.shape[1:-1] # (690, 628)
print(vol_size)
val_atlas = Image.open(
val_atlas_file) # is a TiffImageFile _size is (628, 690)
val_atlas_vol = np.array(val_atlas)[
np.newaxis, ..., np.newaxis] # is a ndarray, shape is (1, 690, 628, 1)
# new = Image.fromarray(X) new.size is (628, 690)
val_vol_size = val_atlas_vol.shape[1:-1] # (690, 628)
print(val_vol_size)
Sm = Image.open(seg_dir) # is a TiffImageFile _size is (628, 690)
Sm_ = np.array(Sm)[np.newaxis, ..., np.newaxis]
val_Sm = Image.open(val_seg_dir) # is a TiffImageFile _size is (628, 690)
val_Sm_ = np.array(val_Sm)[np.newaxis, ..., np.newaxis]
# prepare data files
# for the CVPR and MICCAI papers, we have data arranged in train/validate/test folders
# inside each folder is a /vols/ and a /asegs/ folder with the volumes
# and segmentations. All of our papers use npz formated data.
# data_dir = D:/voxel/data/01
train_vol_names = data_dir # glob.glob(os.path.join(data_dir, '*.tif')) # is a list contain file path(name)
# random.shuffle(train_vol_names) # shuffle volume list tif
assert len(train_vol_names) > 0, "Could not find any training data"
val_vol_names = val_data_dir # glob.glob(os.path.join(data_dir, '*.tif')) # is a list contain file path(name)
# random.shuffle(train_vol_names) # shuffle volume list tif
assert len(val_vol_names) > 0, "Could not find any training data"
# UNET filters for voxelmorph-1 and voxelmorph-2,
# these are architectures presented in CVPR 2018
nf_enc = [16, 32, 32, 32]
if model == 'vm1':
nf_dec = [32, 32, 32, 32, 8, 8]
elif model == 'vm2':
nf_dec = [32, 32, 32, 32, 32, 16, 16]
else: # 'vm2double':
nf_enc = [f * 2 for f in nf_enc]
nf_dec = [f * 2 for f in [32, 32, 32, 32, 32, 16, 16]]
assert data_loss in [
'mse', 'cc', 'ncc'
], 'Loss should be one of mse or cc, found %s' % data_loss
if data_loss in ['ncc', 'cc']:
data_loss = losses.NCC().loss
if Sf_file is not None:
Sf = Image.open(Sf_file)
Sf_ = np.array(Sf)[np.newaxis, ..., np.newaxis]
if val_Sf_file is not None:
val_Sf = Image.open(val_Sf_file)
val_Sf_ = np.array(val_Sf)[np.newaxis, ..., np.newaxis]
# prepare model folder
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
# 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))
#gpu = gpu_id
# data generator
nb_gpus = len(gpu_id.split(',')) # 1
assert np.mod(batch_size, nb_gpus) == 0, \
'batch_size should be a multiple of the nr. of gpus. ' + \
'Got batch_size %d, %d gpus' % (batch_size, nb_gpus)
train_example_gen = datagenerators.example_gen(
train_vol_names,
batch_size=batch_size) # it is a list contain a ndarray
atlas_vol_bs = np.repeat(
atlas_vol, batch_size,
axis=0) # is a ndarray, if batch_size is 2, shape is (2, 690, 628, 1)
cvpr2018_gen = datagenerators.cvpr2018_gen(train_example_gen,
atlas_vol_bs,
batch_size=batch_size)
val_example_gen = datagenerators.example_gen(
val_vol_names, batch_size=batch_size) # it is a list contain a ndarray
val_atlas_vol_bs = np.repeat(
val_atlas_vol, batch_size,
axis=0) # is a ndarray, if batch_size is 2, shape is (2, 690, 628, 1)
val_cvpr2018_gen = datagenerators.cvpr2018_gen(val_example_gen,
val_atlas_vol_bs,
batch_size=batch_size)
# prepare the model
with tf.device(gpu):
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# prepare the model
# in the CVPR layout, the model takes in [image_1, image_2] and outputs [warped_image_1, flow]
# in the experiments, we use image_2 as atlas
model = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
# load initial weights
if load_model_file is not None:
print('loading', load_model_file)
model.load_weights(load_model_file)
# save first iteration
model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch))
# if auxi_label is not None:
# print('yes')
# loss_model= [data_loss, losses.Grad('l2').loss, losses.Lseg()._lseg(Sf_) ] ##########################
# loss_weight= [1.0, reg_param, gama_param]
# else:
loss_model = [
data_loss,
losses.Grad(gama_param, Sf_, Sm_, penalty='l2').loss
] # real gama: reg_param*gama_param
loss_weight = [1.0, reg_param]
# reg_param_tensor = tf.constant(5, dtype=tf.float32)
metrics_2 = losses.Grad(gama_param,
val_Sf_,
val_Sm_,
penalty='l2',
flag_vali=True).loss # reg_param
# prepare callbacks
save_file_name = os.path.join(model_dir, '{epoch:02d}.h5')
# fit generator
with tf.device(gpu):
# multi-gpu support
if nb_gpus > 1:
save_callback = nrn_gen.ModelCheckpointParallel(save_file_name)
mg_model = multi_gpu_model(model, gpus=nb_gpus)
# single-gpu
else:
save_callback = ModelCheckpoint(save_file_name)
mg_model = model
# compile
mg_model.compile(optimizer=Adam(lr=lr),
loss=loss_model,
loss_weights=loss_weight,
metrics={'flow': metrics_2})
# fit
history = mg_model.fit_generator(cvpr2018_gen,
initial_epoch=initial_epoch,
epochs=nb_epochs,
callbacks=[save_callback],
steps_per_epoch=steps_per_epoch,
validation_data=val_cvpr2018_gen,
validation_steps=1,
verbose=2)
# plot
print('model', mg_model.metrics_names)
print('keys()', history.history.keys())
# print(metrics.name)
plt.plot(history.history['loss'])
# plt.plot(history.history['val_spatial_transformer_1_loss'])
plt.title('cvpr_auxi_loss')
plt.ylabel('loss')
plt.xlabel('Epoch')
plt.legend(['Train', 'Validation'])
plt.show()
def train(src_dir,
tgt_dir,
model_dir,
model_lr_dir,
lr,
nb_epochs,
reg_param,
steps_per_epoch,
batch_size,
load_model_file=None,
data_loss='ncc',
initial_epoch=0):
"""
model training function
:param data_dir: folder with npz files for each subject.
:param atlas_file: atlas filename. So far we support npz file with a 'vol' variable
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param model_dir: the model directory to save to
:param lr: learning rate
:param n_iterations: number of training iterations
:param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper)
:param steps_per_epoch: frequency with which to save models
:param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size
:param load_model_file: optional h5 model file to initialize with
:param data_loss: data_loss: 'mse' or 'ncc
"""
# prepare data files
# for the CVPR and MICCAI papers, we have data arranged in train/validate/test folders
# inside each folder is a /vols/ and a /asegs/ folder with the volumes
# and segmentations. All of our papers use npz formated data.
src_vol_names = glob.glob(os.path.join(src_dir, '*.npz'))
tgt_vol_names = glob.glob(os.path.join(tgt_dir, '*.npz'))
random.shuffle(src_vol_names) # shuffle volume list
random.shuffle(tgt_vol_names) # shuffle volume list
assert len(src_vol_names) > 0, "Could not find any training data"
assert data_loss in [
'mse', 'ncc'
], 'Loss should be one of mse or cc, found %s' % data_loss
if data_loss == 'ncc':
data_loss = losses.NCC().loss
# GPU handling
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
# set_session(tf.Session(config=config))
vol_size = (56, 56, 56)
# prepare the model
src_lr = tf.placeholder(tf.float32, [None, *vol_size, 1],
name='input_src_lr')
tgt_lr = tf.placeholder(tf.float32, [None, *vol_size, 1],
name='input_tgt_lr')
srm_lr = tf.placeholder(tf.float32, [None, *vol_size, 1],
name='mask_src_lr')
attn_lr = tf.placeholder(tf.float32, [None, *vol_size, 1], name='attn_lr')
src_mr = tf.placeholder(tf.float32, [None, *vol_size, 1],
name='input_src_mr')
tgt_mr = tf.placeholder(tf.float32, [None, *vol_size, 1],
name='input_tgt_mr')
srm_mr = tf.placeholder(tf.float32, [None, *vol_size, 1],
name='mask_src_mr')
df_lr2mr = tf.placeholder(tf.float32, [None, *vol_size, 3],
name='df_lr2mr')
attn_mr = tf.placeholder(tf.float32, [None, *vol_size, 1], name='attn_mr')
src_hr = tf.placeholder(tf.float32, [None, *vol_size, 1],
name='input_src_hr')
tgt_hr = tf.placeholder(tf.float32, [None, *vol_size, 1],
name='input_tgt_hr')
srm_hr = tf.placeholder(tf.float32, [None, *vol_size, 1],
name='mask_src_hr')
df_mr2hr = tf.placeholder(tf.float32, [None, *vol_size, 3],
name='df_mr2hr')
attn_hr = tf.placeholder(tf.float32, [None, *vol_size, 1], name='attn_hr')
model_lr = networks.net_lr(src_lr, tgt_lr, srm_lr)
model_mr = networks.net_mr(src_mr, tgt_mr, srm_mr, df_lr2mr)
model_hr = networks.net_hr(src_hr, tgt_hr, srm_hr, df_mr2hr)
# the loss functions
lr_ncc = data_loss(model_lr[0].outputs, tgt_lr)
#lr_grd = losses.Grad('l2').loss(model_lr[0].outputs, model_lr[2].outputs)
lr_grd = losses.Anti_Folding('l2').loss(model_lr[0].outputs,
model_lr[2].outputs)
cost_lr = lr_ncc + reg_param * lr_grd # + lr_attn
mr_ncc = data_loss(model_mr[0].outputs, tgt_mr)
#mr_grd = losses.Grad('l2').loss(model_mr[0].outputs, model_mr[2].outputs)
mr_grd = losses.Anti_Folding('l2').loss(model_mr[0].outputs,
model_mr[2].outputs)
cost_mr = mr_ncc + reg_param * mr_grd
hr_ncc = data_loss(model_hr[0].outputs, tgt_hr)
#hr_grd = losses.Grad('l2').loss(model_hr[0].outputs, model_hr[2].outputs)
hr_grd = losses.Anti_Folding('l2').loss(model_hr[0].outputs,
model_hr[2].outputs)
cost_hr = hr_ncc + reg_param * hr_grd
# the training operations
def get_v(name):
t_vars = tf.trainable_variables()
d_vars = [var for var in t_vars if name in var.name]
return d_vars
#attn_vars = tl.layers.get_variables_with_name('cbam_1', True, True)
attn_vars = get_v('cbam_1')
for a_v in attn_vars:
print(a_v)
train_op_lr = tf.train.AdamOptimizer(lr).minimize(cost_lr)
train_op_mr = tf.train.AdamOptimizer(lr).minimize(cost_mr)
train_op_hr = tf.train.AdamOptimizer(lr).minimize(cost_hr)
# data generator
src_example_gen = datagenerators.example_gen(src_vol_names,
batch_size=batch_size)
tgt_example_gen = datagenerators.example_gen(tgt_vol_names,
batch_size=batch_size)
data_gen = datagenerators.gen_with_mask(src_example_gen,
tgt_example_gen,
batch_size=batch_size)
variables_to_restore = tf.contrib.framework.get_variables_to_restore(
exclude=['net_hr'])
saver = tf.train.Saver(variables_to_restore)
#saver = tf.train.Saver(max_to_keep=3)
# fit generator
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
# load initial weights
try:
if load_model_file is not None:
model_file = tf.train.latest_checkpoint(load_model_file) #
saver.restore(sess, model_file)
except:
print('No files in', load_model_file)
saver.save(sess, model_dir + 'dfnet', global_step=0)
def resize_df(df, zoom):
df1 = nd.interpolation.zoom(
df[0, :, :, :, 0], zoom=zoom, mode='nearest', order=3) * zoom[
0] # Cubic: order=3; Bilinear: order=1; Nearest: order=0
df2 = nd.interpolation.zoom(
df[0, :, :, :,
1], zoom=zoom, mode='nearest', order=3) * zoom[1]
df3 = nd.interpolation.zoom(
df[0, :, :, :,
2], zoom=zoom, mode='nearest', order=3) * zoom[2]
dfs = np.stack((df1, df2, df3), axis=3)
return dfs[np.newaxis, :, :, :]
class logPrinter(object):
def __init__(self):
self.n_batch = 0
self.total_dice = []
self.cost = []
self.ncc = []
self.grd = []
def addLog(self, dice, cost, ncc, grd):
self.n_batch += 1
self.dice.append(dice)
self.cost.append(cost)
self.ncc.append(ncc)
self.grd.append(grd)
def output(self):
dice = np.array(self.dice).mean(axis=0).round(3).tolist()
cost = np.array(self.cost).mean()
ncc = np.array(self.ncc).mean()
grd = np.array(self.grd).mean()
return dice, cost, ncc, grd, self.n_batch
def clear(self):
self.n_batch = 0
self.dice = []
self.cost = []
self.ncc = []
self.grd = []
lr_log = logPrinter()
mr_log = logPrinter()
hr_log = logPrinter()
# train low resolution
# load initial weights
saver = tf.train.Saver(max_to_keep=1)
#if model_lr_dir is not None:
# model_lr_dir = tf.train.latest_checkpoint(model_lr_dir) #
# print(model_lr_dir)
# saver.restore(sess, model_lr_dir)
nb_epochs = 20 #20#10
steps_per_epoch = 30 * 29
for epoch in range(nb_epochs):
tbar = trange(steps_per_epoch, unit='batch', ncols=100)
lr_log.clear()
for i in tbar:
image, mask = data_gen.__next__()
global_X, global_atlas = image
global_X_mask, global_atlas_mask = mask
global_diff = global_X[0, :, :, :,
0] - global_atlas[0, :, :, :, 0]
# low resolution
global_X_64 = nd.interpolation.zoom(global_X[0, :, :, :, 0],
zoom=(0.25, 0.25, 0.25),
mode='nearest')
global_A_64 = nd.interpolation.zoom(global_atlas[0, :, :, :,
0],
zoom=(0.25, 0.25, 0.25),
mode='nearest')
global_XM_64 = nd.interpolation.zoom(global_X_mask[0, :, :, :,
0],
zoom=(0.25, 0.25, 0.25),
mode='nearest',
order=0)
global_AM_64 = nd.interpolation.zoom(
global_atlas_mask[0, :, :, :, 0],
zoom=(0.25, 0.25, 0.25),
mode='nearest',
order=0)
global_diff_16 = nd.interpolation.zoom(global_diff,
zoom=(0.25, 0.25, 0.25),
mode='nearest')
global_X_64 = global_X_64[np.newaxis, :, :, :, np.newaxis]
global_A_64 = global_A_64[np.newaxis, :, :, :, np.newaxis]
global_XM_64 = global_XM_64[np.newaxis, :, :, :, np.newaxis]
global_AM_64 = global_AM_64[np.newaxis, :, :, :, np.newaxis]
global_diff_16 = global_diff_16[np.newaxis, :, :, :,
np.newaxis]
feed_dict = {
src_lr: global_X_64,
tgt_lr: global_A_64,
srm_lr: global_XM_64,
attn_lr: global_diff_16
}
err_lr, _ = sess.run([cost_lr, train_op_lr],
feed_dict=feed_dict)
df_lr, warp_seg, elr_ncc, elr_grad, lr_attn_map, lr_attn_feature = sess.run(
[
model_lr[2].outputs, model_lr[1].outputs, lr_ncc,
lr_grd, model_lr[3], model_lr[4]
],
feed_dict=feed_dict)
# print(df_lr.shape)
lr_dice, _ = dice(warp_seg[0, :, :, :, 0],
global_AM_64[0, :, :, :, 0],
labels=[0, 10, 150, 250],
nargout=2)
lr_log.addLog(lr_dice, err_lr, elr_ncc, elr_grad)
lr_out = lr_log.output()
tbar.set_description('Epoch %d/%d ### step %i' %
(epoch + 1, nb_epochs, i))
tbar.set_postfix(lr_dice=lr_out[0],
lr_cost=lr_out[1],
lr_ncc=lr_out[2],
lr_grd=lr_out[3])
saver.save(sess, model_lr_dir + 'dfnet', global_step=0)
# train middle resolution
nb_epochs = 1 #1
steps_per_epoch = 30 * 29
for epoch in range(nb_epochs):
lr_log.clear()
for lr_step in range(steps_per_epoch):
image, mask = data_gen.__next__()
global_X, global_atlas = image
global_X_mask, global_atlas_mask = mask
global_diff = global_X[0, :, :, :,
0] - global_atlas[0, :, :, :, 0]
# low resolution
global_X_64 = nd.interpolation.zoom(global_X[0, :, :, :, 0],
zoom=(0.25, 0.25, 0.25),
mode='nearest')
global_A_64 = nd.interpolation.zoom(global_atlas[0, :, :, :,
0],
zoom=(0.25, 0.25, 0.25),
mode='nearest')
global_XM_64 = nd.interpolation.zoom(global_X_mask[0, :, :, :,
0],
zoom=(0.25, 0.25, 0.25),
mode='nearest',
order=0)
global_AM_64 = nd.interpolation.zoom(
global_atlas_mask[0, :, :, :, 0],
zoom=(0.25, 0.25, 0.25),
mode='nearest',
order=0)
global_diff_16 = nd.interpolation.zoom(global_diff,
zoom=(0.25, 0.25, 0.25),
mode='nearest')
global_X_64 = global_X_64[np.newaxis, :, :, :, np.newaxis]
global_A_64 = global_A_64[np.newaxis, :, :, :, np.newaxis]
global_XM_64 = global_XM_64[np.newaxis, :, :, :, np.newaxis]
global_AM_64 = global_AM_64[np.newaxis, :, :, :, np.newaxis]
global_diff_16 = global_diff_16[np.newaxis, :, :, :,
np.newaxis]
feed_dict = {
src_lr: global_X_64,
tgt_lr: global_A_64,
srm_lr: global_XM_64,
attn_lr: global_diff_16
}
err_lr, _ = sess.run([cost_lr, train_op_lr],
feed_dict=feed_dict)
df_lr, warp_seg, elr_ncc, elr_grad, lr_attn_map, lr_attn_feature = sess.run(
[
model_lr[2].outputs, model_lr[1].outputs, lr_ncc,
lr_grd, model_lr[3], model_lr[4]
],
feed_dict=feed_dict)
lr_dice, _ = dice(warp_seg[0, :, :, :, 0],
global_AM_64[0, :, :, :, 0],
labels=[0, 10, 150, 250],
nargout=2)
lr_log.addLog(lr_dice, err_lr, elr_ncc, elr_grad)
lr_out = lr_log.output()
print('\nEpoch %d/%d ### step %i' %
(epoch + 1, nb_epochs, lr_out[-1]))
print(
'[lr] lr_dice={}, lr_cost={:.3f}, lr_ncc={:.3f}, lr_grd={:.3f}'
.format(lr_out[0], lr_out[1], lr_out[2], lr_out[3]))
# middle part
df_lr_res2mr = resize_df(df_lr, zoom=(2, 2, 2))
select_points_lr = patch_selection_attn(lr_attn_map,
zoom_scales=[8, 8, 8],
kernel=7,
mi=10,
ma=18)
print(select_points_lr)
mr_log.clear()
for sp in select_points_lr:
mov_img_112 = global_X[0, sp[0] - 56:sp[0] + 56,
sp[1] - 56:sp[1] + 56,
sp[2] - 56:sp[2] + 56, 0]
fix_img_112 = global_atlas[0, sp[0] - 56:sp[0] + 56,
sp[1] - 56:sp[1] + 56,
sp[2] - 56:sp[2] + 56, 0]
mov_seg_112 = global_X_mask[0, sp[0] - 56:sp[0] + 56,
sp[1] - 56:sp[1] + 56,
sp[2] - 56:sp[2] + 56, 0]
fix_seg_112 = global_atlas_mask[0, sp[0] - 56:sp[0] + 56,
sp[1] - 56:sp[1] + 56,
sp[2] - 56:sp[2] + 56, 0]
dif_img_112 = global_diff[sp[0] - 56:sp[0] + 56,
sp[1] - 56:sp[1] + 56,
sp[2] - 56:sp[2] + 56]
#print(mov_img_112.shape)
if fix_img_112.shape != (112, 112, 112):
print(mov_img_112.shape)
continue
fix_112_56 = nd.interpolation.zoom(fix_img_112,
zoom=(0.5, 0.5, 0.5),
mode='nearest')
mov_112_56 = nd.interpolation.zoom(mov_img_112,
zoom=(0.5, 0.5, 0.5),
mode='nearest')
fix_112_56m = nd.interpolation.zoom(fix_seg_112,
zoom=(0.5, 0.5, 0.5),
mode='nearest',
order=0)
mov_112_56m = nd.interpolation.zoom(mov_seg_112,
zoom=(0.5, 0.5, 0.5),
mode='nearest',
order=0)
dif_112_56 = nd.interpolation.zoom(dif_img_112,
zoom=(0.5, 0.5, 0.5),
mode='nearest')
mid_fix_img = fix_112_56[np.newaxis, :, :, :, np.newaxis]
mid_mov_img = mov_112_56[np.newaxis, :, :, :, np.newaxis]
mid_fix_seg = fix_112_56m[np.newaxis, :, :, :, np.newaxis]
mid_mov_seg = mov_112_56m[np.newaxis, :, :, :, np.newaxis]
mid_dif_img = dif_112_56[np.newaxis, :, :, :, np.newaxis]
df_mr_feed = df_lr_res2mr[:,
sp[0] // 2 - 28:sp[0] // 2 + 28,
sp[1] // 2 - 28:sp[1] // 2 + 28,
sp[2] // 2 - 28:sp[2] // 2 +
28, :]
feed_dict = {
src_mr: mid_mov_img,
tgt_mr: mid_fix_img,
srm_mr: mid_mov_seg,
df_lr2mr: df_mr_feed,
attn_mr: mid_dif_img
}
err_mr, _ = sess.run([cost_mr, train_op_mr],
feed_dict=feed_dict)
df_mr, warp_seg, emr_ncc, emr_grad = sess.run(
[
model_mr[2].outputs, model_mr[1].outputs, mr_ncc,
mr_grd
],
feed_dict=feed_dict)
mr_dice, _ = dice(warp_seg[0, :, :, :, 0],
mid_fix_seg[0, :, :, :, 0],
labels=[0, 10, 150, 250],
nargout=2)
mr_log.addLog(mr_dice, err_mr, emr_ncc, emr_grad)
mr_out = mr_log.output()
# print(' Epoch %d/%d ### step %i' % (epoch+1, nb_epochs, mr_out[-1]))
print(
' [mr] {}/{} mr_dice={}, mr_cost={:.3f}, mr_ncc={:.3f}, mr_grd={:.3f}'
.format(mr_out[-1], len(select_points_lr), mr_out[0],
mr_out[1], mr_out[2], mr_out[3]))
saver.save(sess, model_dir + 'dfnet', global_step=0)
# train high resolution
nb_epochs = 1
steps_per_epoch = 300
for epoch in range(nb_epochs):
lr_log.clear()
for lr_step in range(steps_per_epoch):
image, mask = data_gen.__next__()
global_X, global_atlas = image
global_X_mask, global_atlas_mask = mask
global_diff = global_X[0, :, :, :,
0] - global_atlas[0, :, :, :, 0]
# low resolution
global_X_64 = nd.interpolation.zoom(global_X[0, :, :, :, 0],
zoom=(0.25, 0.25, 0.25),
mode='nearest')
global_A_64 = nd.interpolation.zoom(global_atlas[0, :, :, :,
0],
zoom=(0.25, 0.25, 0.25),
mode='nearest')
global_XM_64 = nd.interpolation.zoom(global_X_mask[0, :, :, :,
0],
zoom=(0.25, 0.25, 0.25),
mode='nearest',
order=0)
global_AM_64 = nd.interpolation.zoom(
global_atlas_mask[0, :, :, :, 0],
zoom=(0.25, 0.25, 0.25),
mode='nearest',
order=0)
global_diff_16 = nd.interpolation.zoom(global_diff,
zoom=(0.25, 0.25, 0.25),
mode='nearest')
global_X_64 = global_X_64[np.newaxis, :, :, :, np.newaxis]
global_A_64 = global_A_64[np.newaxis, :, :, :, np.newaxis]
global_XM_64 = global_XM_64[np.newaxis, :, :, :, np.newaxis]
global_AM_64 = global_AM_64[np.newaxis, :, :, :, np.newaxis]
global_diff_16 = global_diff_16[np.newaxis, :, :, :,
np.newaxis]
feed_dict = {
src_lr: global_X_64,
tgt_lr: global_A_64,
srm_lr: global_XM_64,
attn_lr: global_diff_16
}
err_lr, _ = sess.run([cost_lr, train_op_lr],
feed_dict=feed_dict)
df_lr, warp_seg, elr_ncc, elr_grad, lr_attn_map, lr_attn_feature = sess.run(
[
model_lr[2].outputs, model_lr[1].outputs, lr_ncc,
lr_grd, model_lr[3], model_lr[4]
],
feed_dict=feed_dict)
lr_dice, _ = dice(warp_seg[0, :, :, :, 0],
global_AM_64[0, :, :, :, 0],
labels=[0, 10, 150, 250],
nargout=2)
lr_log.addLog(lr_dice, err_lr, elr_ncc, elr_grad)
lr_out = lr_log.output()
print('\nEpoch %d/%d ### step %i' %
(epoch + 1, nb_epochs, lr_out[-1]))
print(
'[lr] lr_dice={}, lr_cost={:.3f}, lr_ncc={:.3f}, lr_grd={:.3f}'
.format(lr_out[0], lr_out[1], lr_out[2], lr_out[3]))
# middle part
df_lr_res2mr = resize_df(df_lr, zoom=(2, 2, 2))
select_points_lr = patch_selection_attn(lr_attn_map,
zoom_scales=[8, 8, 8],
kernel=7,
mi=10,
ma=18)
print(select_points_lr)
mr_log.clear()
for sp in select_points_lr:
mov_img_112 = global_X[0, sp[0] - 56:sp[0] + 56,
sp[1] - 56:sp[1] + 56,
sp[2] - 56:sp[2] + 56, 0]
fix_img_112 = global_atlas[0, sp[0] - 56:sp[0] + 56,
sp[1] - 56:sp[1] + 56,
sp[2] - 56:sp[2] + 56, 0]
mov_seg_112 = global_X_mask[0, sp[0] - 56:sp[0] + 56,
sp[1] - 56:sp[1] + 56,
sp[2] - 56:sp[2] + 56, 0]
fix_seg_112 = global_atlas_mask[0, sp[0] - 56:sp[0] + 56,
sp[1] - 56:sp[1] + 56,
sp[2] - 56:sp[2] + 56, 0]
dif_img_112 = global_diff[sp[0] - 56:sp[0] + 56,
sp[1] - 56:sp[1] + 56,
sp[2] - 56:sp[2] + 56]
#print(mov_img_112.shape)
if fix_img_112.shape != (112, 112, 112):
print(mov_img_112.shape)
continue
fix_112_56 = nd.interpolation.zoom(fix_img_112,
zoom=(0.5, 0.5, 0.5),
mode='nearest')
mov_112_56 = nd.interpolation.zoom(mov_img_112,
zoom=(0.5, 0.5, 0.5),
mode='nearest')
fix_112_56m = nd.interpolation.zoom(fix_seg_112,
zoom=(0.5, 0.5, 0.5),
mode='nearest',
order=0)
mov_112_56m = nd.interpolation.zoom(mov_seg_112,
zoom=(0.5, 0.5, 0.5),
mode='nearest',
order=0)
dif_112_56 = nd.interpolation.zoom(dif_img_112,
zoom=(0.5, 0.5, 0.5),
mode='nearest')
mid_fix_img = fix_112_56[np.newaxis, :, :, :, np.newaxis]
mid_mov_img = mov_112_56[np.newaxis, :, :, :, np.newaxis]
mid_fix_seg = fix_112_56m[np.newaxis, :, :, :, np.newaxis]
mid_mov_seg = mov_112_56m[np.newaxis, :, :, :, np.newaxis]
mid_dif_img = dif_112_56[np.newaxis, :, :, :, np.newaxis]
df_mr_feed = df_lr_res2mr[:,
sp[0] // 2 - 28:sp[0] // 2 + 28,
sp[1] // 2 - 28:sp[1] // 2 + 28,
sp[2] // 2 - 28:sp[2] // 2 +
28, :]
feed_dict = {
src_mr: mid_mov_img,
tgt_mr: mid_fix_img,
srm_mr: mid_mov_seg,
df_lr2mr: df_mr_feed,
attn_mr: mid_dif_img
}
err_mr, _ = sess.run([cost_mr, train_op_mr],
feed_dict=feed_dict)
df_mr, warp_seg, emr_ncc, emr_grad, mr_attn_map, mr_attn_feature = sess.run(
[
model_mr[2].outputs, model_mr[1].outputs, mr_ncc,
mr_grd, model_mr[3], model_mr[4]
],
feed_dict=feed_dict)
mr_dice, _ = dice(warp_seg[0, :, :, :, 0],
mid_fix_seg[0, :, :, :, 0],
labels=[0, 10, 150, 250],
nargout=2)
mr_log.addLog(mr_dice, err_mr, emr_ncc, emr_grad)
mr_out = mr_log.output()
# print(' Epoch %d/%d ### step %i' % (epoch+1, nb_epochs, mr_out[-1]))
print(
' [mr] {}/{} mr_dice={}, mr_cost={:.3f}, mr_ncc={:.3f}, mr_grd={:.3f}'
.format(mr_out[-1], len(select_points_lr), mr_out[0],
mr_out[1], mr_out[2], mr_out[3]))
# high part
df_mr_res2hr = resize_df(df_mr, zoom=(2, 2, 2))
hr_log.clear()
select_points_mr = patch_selection_attn(
mr_attn_map,
zoom_scales=[4, 4, 4],
kernel=7,
mi=8,
ma=20)
print(30 * '-')
print('High Part')
print(select_points_mr)
for spm in select_points_mr:
fix_img_56 = fix_img_112[spm[0] - 28:spm[0] + 28,
spm[1] - 28:spm[1] + 28,
spm[2] - 28:spm[2] + 28]
mov_img_56 = mov_img_112[spm[0] - 28:spm[0] + 28,
spm[1] - 28:spm[1] + 28,
spm[2] - 28:spm[2] + 28]
fix_seg_56 = fix_seg_112[spm[0] - 28:spm[0] + 28,
spm[1] - 28:spm[1] + 28,
spm[2] - 28:spm[2] + 28]
mov_seg_56 = mov_seg_112[spm[0] - 28:spm[0] + 28,
spm[1] - 28:spm[1] + 28,
spm[2] - 28:spm[2] + 28]
dif_img_56 = dif_img_112[spm[0] - 28:spm[0] + 28,
spm[1] - 28:spm[1] + 28,
spm[2] - 28:spm[2] + 28]
if fix_img_56.shape != (56, 56, 56):
continue
hig_fix_img = fix_img_56[np.newaxis, :, :, :,
np.newaxis]
hig_mov_img = mov_img_56[np.newaxis, :, :, :,
np.newaxis]
hig_fix_seg = fix_seg_56[np.newaxis, :, :, :,
np.newaxis]
hig_mov_seg = mov_seg_56[np.newaxis, :, :, :,
np.newaxis]
hig_dif_img = dif_img_56[np.newaxis, :, :, :,
np.newaxis]
df_hr_feed = df_mr_res2hr[:, spm[0] - 28:spm[0] + 28,
spm[1] - 28:spm[1] + 28,
spm[2] - 28:spm[2] + 28, :]
feed_dict = {
src_hr: hig_mov_img,
tgt_hr: hig_fix_img,
srm_hr: hig_mov_seg,
df_mr2hr: df_hr_feed,
attn_hr: hig_dif_img
}
err_hr, _ = sess.run([cost_hr, train_op_hr],
feed_dict=feed_dict)
df_hr, warp_seg, ehr_ncc, ehr_grad = sess.run(
[
model_hr[2].outputs, model_hr[1].outputs,
hr_ncc, hr_grd
],
feed_dict=feed_dict)
hr_dice, _ = dice(warp_seg[0, :, :, :, 0],
hig_fix_seg[0, :, :, :, 0],
labels=[0, 10, 150, 250],
nargout=2)
hr_log.addLog(hr_dice, err_hr, ehr_ncc, ehr_grad)
hr_out = hr_log.output()
# print(' Epoch %d/%d ### step %i' % (epoch+1, nb_epochs, mr_out[-1]))
print(
' [hr] {}/{} hr_dice={}, hr_cost={:.3f}, hr_ncc={:.3f}, hr_grd={:.3f}'
.format(hr_out[-1], len(select_points_mr),
hr_out[0], hr_out[1], hr_out[2],
hr_out[3]))
saver.save(sess, model_dir + 'dfnet', global_step=lr_step)
def train(data_dir,
depth_size,
model,
model_dir,
gpu_id,
lr,
nb_epochs,
reg_param,
steps_per_epoch,
train_mode,
load_model_file,
data_loss,
batch_size,
initial_epoch=0):
"""
model training function
:param data_dir: folder with npz files for each subject.
:param atlas_file: atlas filename. So far we support npz file with a 'vol' variable
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param model_dir: the model directory to save to
:param gpu_id: integer specifying the gpu to use
:param lr: learning rate
:param n_iterations: number of training iterations
:param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper)
:param steps_per_epoch: frequency with which to save models
:param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size
:param load_model_file: optional h5 model file to initialize with
:param data_loss: data_loss: 'mse' or 'ncc'
"""
# load data
train_x, train_y = load_data(data_dir=data_dir,
depth_size=depth_size,
mode='train',
fixed='first')
# test_x, test_y = load_data(data_dir=data_dir, depth_size=depth_size, mode='test', fixed='first')
vol_size = train_x[0].shape[1:-1] # (width, height, depth)
# set encoder, decoder feature number
nf_enc = [16, 32, 32, 32]
if model == 'vm1':
nf_dec = [32, 32, 32, 32, 8, 8]
elif model == 'vm2':
nf_dec = [32, 32, 32, 32, 32, 16, 16]
else: # 'vm2double':
nf_enc = [f * 2 for f in nf_enc]
nf_dec = [f * 2 for f in [32, 32, 32, 32, 32, 16, 16]]
# set loss function
# Mean Squared Error, Cross-Correlation, Negative Cross-Correlation
assert data_loss in [
'mse', 'cc', 'ncc'
], 'Loss should be one of mse or cc, found %s' % data_loss
if data_loss in ['ncc', 'cc']:
data_loss = losses.NCC().loss
# prepare model folder
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
# 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))
# prepare the model
with tf.device(gpu):
# in the CVPR layout, the model takes in [moving image, fixed image] and outputs [warped image, flow]
model = networks.cvpr2018_net(tuple(vol_size), nf_enc, nf_dec)
model.load_weights(load_model_file)
# save first iteration
model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch))
# prepare callbacks
save_file_name = os.path.join(model_dir, '{epoch:02d}.h5')
# fit
with tf.device(gpu):
save_callback = ModelCheckpoint(save_file_name)
mg_model = model
# compile
mg_model.compile(optimizer=Adam(lr=lr),
loss=[data_loss, losses.Grad('l2').loss],
loss_weights=[1.0, reg_param])
# fit
mg_model.fit(x=train_x,
y=train_y,
batch_size=None,
epochs=nb_epochs,
verbose=1,
callbacks=[save_callback],
steps_per_epoch=steps_per_epoch)
def train(data_dir, model, model_dir, gpu_id, lr, nb_epochs, reg_param,
steps_per_epoch, load_model_file, data_loss, window_size,
batch_size):
"""
model training function
:param data_dir: folder with npz files for each subject.
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param model_dir: the model directory to save to
:param gpu_id: integer specifying the gpu to use
:param lr: learning rate
:param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper)
:param steps_per_epoch: frequency with which to save models
:param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size
:param load_model_file: optional h5 model file to initialize with
:param data_loss: data_loss: 'mse' or 'ncc'
"""
vol_size = [256, 256, 144] # (width, height, depth)
# set encoder, decoder feature number
nf_enc = [16, 32, 32, 32]
if model == 'vm1':
nf_dec = [32, 32, 32, 32, 8, 8]
elif model == 'vm2':
nf_dec = [32, 32, 32, 32, 32, 16, 16]
else: # 'vm2double':
nf_enc = [f * 2 for f in nf_enc]
nf_dec = [f * 2 for f in [32, 32, 32, 32, 32, 16, 16]]
# set loss function
# Mean Squared Error, Cross-Correlation, Negative Cross-Correlation
assert data_loss in [
'mse', 'cc', 'ncc'
], 'Loss should be one of mse or cc, found %s' % data_loss
if data_loss in ['ncc', 'cc']:
NCC = losses.NCC(win=window_size)
data_loss = NCC.loss
# prepare model folder
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
# 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))
# prepare the model
with tf.device(gpu):
# in the CVPR layout, the model takes in [moving image, fixed image] and outputs [warped image, flow]
model = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
if load_model_file is not None:
model.load_weights(load_model_file)
# save first iteration
# model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch))
# load data
# path = "../../dataset/urinary"
# vol_names = [filename for filename in os.listdir(data_dir) if (int(filename.split("_")[-1].split('.')[0]) < 206) and
# (int(filename.split("_")[-1].split('.')[0]) not in except_list)]
vol_names = [
filename for filename in os.listdir(data_dir)
if int(filename.split("_")[-1].split(".")[0]) in normal
]
# vol_names = [filename for filename in os.listdir(data_dir) if int(filename.split("_")[-1].split(".")[0]) in (9, 130, 128)]
vol_names.sort()
uro_gen = uro_generator(vol_names, data_dir, fixed='joyoungje')
# test_path = os.path.join(data_dir, 'test')
# test_vol_names = [filename for filename in os.listdir(test_path) if '.npz']
# test_gen = uro_generator(test_vol_names, test_path)
# fit
with tf.device(gpu):
mg_model = model
# compile
mg_model.compile(optimizer=Adam(lr=lr),
loss=[data_loss, losses.Grad('l2').loss],
loss_weights=[1.0, reg_param])
# fit
save_file_name = os.path.join(model_dir, '{epoch:04d}.h5')
save_callback = ModelCheckpoint(save_file_name)
mg_model.fit_generator(uro_gen,
epochs=nb_epochs,
verbose=1,
callbacks=[save_callback],
steps_per_epoch=steps_per_epoch)
