def train(model,
train_generator,
callbacks,
learning_rate,
number_of_epochs,
regularization_parameter,
steps_per_epoch,
loss_function='mse'):
gpu_options = tf.GPUOptions(allow_growth=True)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options,
log_device_placement=False))
K.set_session(sess)
model.compile(optimizer=Adam(lr=learning_rate),
loss=[loss_function, losses.Grad('l2').loss],
loss_weights=[1.0, regularization_parameter])
model.fit_generator(train_generator,
initial_epoch=0,
epochs=number_of_epochs,
callbacks=callbacks,
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 FAIM_loss(y_true, y_pred):
return losses.Grad('l2').loss(
y_true, y_pred) + 1e-5 * losses.NJ_loss(y_true, y_pred)
def train(data_dir, fixed_image, model_dir, device, lr, nb_epochs, AAN_param,
steps_per_epoch, batch_size, load_model_file, initial_epoch,
DLR_model):
# prepare data files
# inside the folder are npz files with the 'vol' and 'label'.
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"
vol_size = [144, 192, 160]
# load atlas from provided files, if atlas-based registration
if fixed_image != './':
fixed_vol = np.load(fixed_image)['vol'][np.newaxis, ..., np.newaxis]
def FAIM_loss(y_true, y_pred):
return losses.Grad('l2').loss(
y_true, y_pred) + 1e-5 * losses.NJ_loss(y_true, y_pred)
assert DLR_model in [
'VM', 'FAIM'
], 'DLR_model should be one of VM or FAIM, found %s' % LBR_model
if DLR_model == 'FAIM':
reg_loss = FAIM_loss
else:
reg_loss = losses.Grad('l2').loss
# prepare model folder
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
# device handling
if 'gpu' in device:
if '0' in device:
device = '/gpu:0'
if '1' in device:
device = '/gpu:1'
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
else:
device = '/cpu:0'
# prepare the model
with tf.device(device):
model = networks.AAN(vol_size, DLR_model)
# load initial weights
if 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
train_example_gen = datagenerators.example_gen(train_vol_names,
batch_size=batch_size,
return_boundary=True)
if fixed_image != './':
fixed_vol_bs = np.repeat(fixed_vol, batch_size, axis=0)
data_gen = datagenerators.gen_atlas(train_example_gen,
fixing_vol_bs,
batch_size=batch_size)
else:
data_gen = datagenerators.gen_s2s(train_example_gen,
batch_size=batch_size)
# prepare callbacks
save_file_name = os.path.join(model_dir, '{epoch:02d}.h5')
# fit generator
with tf.device(device):
save_callback = ModelCheckpoint(save_file_name)
# compile
model.compile(
optimizer=Adam(lr=lr),
loss=['mse', reg_loss,
losses.Grad('l1').loss_with_boundary],
loss_weights=[1.0, 0.01, AAN_param])
# fit
model.fit_generator(data_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"
# 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,
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 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 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_unsupervised_segmentation(data_dir,
atlas_file,
mapping_file,
model,
model_dir,
gpu_id,
lr,
nb_epochs,
init_stats,
reg_param,
steps_per_epoch,
batch_size,
stat_post_warp,
warp_method,
load_model_file,
initial_epoch=0):
"""
model training function
:param data_dir: folder with npz files (coregistered, intensity normalized)
:param atlas_file: file with probabilistic atlas (coregistered to images)
:param mapping_file: file with mapping from labels to tissue types
:param model: registration (voxelmorph) model: vm1, vm2, or vm2double
:param model_dir: the model directory to save to
:param gpu_id: integer specifying the gpu to use
:param lr: learning rate
:param nb_epochs: number of epochs
:param init_stats: file with guesses for means and log-variances (vectors init_mu, init_sigma)
:param reg_param: smoothness/reconstruction tradeoff parameter (lambda in the paper)
:param steps_per_epoch: frequency with which to save models
:param batch_size: default of 1. can be larger, depends on GPU memory and volume size
:param stat_post_warp: set to 1 to use warped atlas to estimate Gaussian parameters
:param warp_method: set to 'WARP' if you want to warp the atlas
:param load_model_file: optional h5 model file to initialize with
:param initial_epoch: initial epoch
"""
# load reference soft edge and corresponding mask from provided files
# (we used size 160x192x224).
# Also: group labels in tissue types, if necessary
if mapping_file is None:
atlas_vol = np.load(atlas_file)['vol_data'][np.newaxis, ...]
nb_labels = atlas_vol.shape[-1]
else:
atlas_full = np.load(atlas_file)['vol_data'][np.newaxis, ...]
mapping = np.load(mapping_file)['mapping'].astype('int').flatten()
assert len(mapping) == atlas_full.shape[-1], \
'mapping shape %d is inconsistent with atlas shape %d' % (len(mapping), atlas_full.shape[-1])
nb_labels = 1 + np.max(mapping)
atlas_vol = np.zeros(
[1, *atlas_full.shape[1:-1],
nb_labels.astype('int')])
for j in range(np.max(mapping.shape)):
atlas_vol[0, ...,
mapping[j]] = atlas_vol[0, ...,
mapping[j]] + atlas_full[0, ...,
j]
vol_size = atlas_vol.shape[1:-1]
# load guesses for means and variances
init_mu = np.load(init_stats)['init_mu']
init_sigma = np.load(init_stats)['init_std']
# prepare data files
train_vol_names = glob.glob(os.path.join(data_dir, '*.npz'))
random.shuffle(train_vol_names)
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]]
# prepare model and log folders
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
log_dir = os.path.join(model_dir, 'logs')
if not os.path.isdir(log_dir):
os.mkdir(log_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):
# prepare the model
model = networks.cvpr2018_net_probatlas(vol_size,
nf_enc,
nf_dec,
nb_labels,
diffeomorphic=True,
full_size=False,
stat_post_warp=stat_post_warp,
warp_method=warp_method,
init_mu=init_mu,
init_sigma=init_sigma)
# 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))
# 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)
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
# tensorBoard callback
tensorboard = TensorBoard(log_dir=log_dir,
histogram_freq=0,
write_graph=True,
write_images=False)
# compile loss and parameters
def data_loss(_, yp):
m = tf.cast(model.inputs[0] > 0, tf.float32)
return -K.sum(yp * m) / K.sum(m)
if warp_method != 'WARP':
reg_param = 0
# 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, tensorboard],
steps_per_epoch=steps_per_epoch,
verbose=1)
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)
def train(
fixed_dir,
moving_dir,
fixedval_dir,
movingval_dir,
model_dir,
gpu_id,
lr,
nb_epochs,
img_size,
lambda_param,
steps_per_epoch,
batch_size,
):
fixed_training_file = open(fixed_dir, 'r')
moving_traning_file = open(moving_dir, 'r')
fixed_val_file = open(fixedval_dir, 'r')
moving_val_file = open(movingval_dir, 'r')
def str_strip(file):
list = []
for f in file.readlines():
list.append(f.strip())
return list
fixed_image_names = str_strip(fixed_training_file)
moving_image_names = str_strip(moving_traning_file)
fixedval_image_names = str_strip(fixed_val_file)
movingval_image_names = str_strip(moving_val_file)
assert len(fixed_image_names) > 0, "fixed路径中找不到训练数据"
assert len(moving_image_names) > 0, "moving路径中找不到训练数据"
assert len(fixedval_image_names) > 0, "fixed路径中找不到训练数据"
assert len(movingval_image_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('网络输入维度', len(unet.inputs))
moving_input_tensor = unet.inputs[0]
fixed_input_tensor = unet.inputs[1]
#---------网络输出----------------
print('网络输出维度:', unet.output)
#-------------转换成流场维度---------------
disp_tensor = keras.layers.Conv2D(ndims,
kernel_size=3,
padding='same',
name='disp')(unet.output)
# （512，512，2）
# check
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 = [losses.NCC().loss, losses.Grad('l2').loss]
loss = ['mse', losses.Grad('l2').loss]
# usually, we have to balance the two losses by a hyper-parameter.
loss_weights = [1, lambda_param]
vxm_model.compile(optimizer=Adam(lr=lr),
loss=loss,
loss_weights=loss_weights)
vxm_model.summary()
#----------------------准备训练数据-------------------------------
x_train_path = fixed_image_names
y_train_path = moving_image_names
x_val_path = fixedval_image_names
y_val_path = movingval_image_names
train_generator = datagenerators.my_data_generator(x_train_path,
y_train_path,
batch_size=batch_size,
img_size=img_size)
valid_generator = datagenerators.my_data_generator(x_val_path,
y_train_path,
batch_size=1,
img_size=img_size)
input_sample, output_sample = next(train_generator)
# ---------------可视化---------------------
slices_2d = [f[0, ..., 0] for f in input_sample + output_sample]
titles = [
'input_moving', 'input_fixed', 'output_moved_ground_truth', 'zero'
]
neuron.plot.slices(slices_2d,
titles=titles,
cmaps=['gray'],
do_colorbars=True)
checkpoint = ModelCheckpoint(filepath='models/my_model{epoch:03d}.h5',
monitor='loss',
verbose=1,
save_best_only='True',
mode='min',
period=1)
tensorboard = TensorBoard(log_dir='logs/',
histogram_freq=0,
write_graph=True,
write_images=True)
#-------------------------开始训练---------------------------
hist = vxm_model.fit_generator(train_generator,
epochs=nb_epochs,
steps_per_epoch=steps_per_epoch,
verbose=2,
callbacks=[checkpoint, tensorboard],
validation_data=valid_generator,
validation_steps=2990)
vxm_model.save_weights(os.path.join(model_dir, 'my_model.h5'))
with open('data/SAR_model_hist.pickle', 'wb') as file_pi:
pickle.dump(hist.history, file_pi)
def plot_history(hist, loss_name='loss', path='data/'):
"""
Quick function to plot the history
"""
plt.figure()
plt.plot(hist.epoch, hist.history[loss_name], '.-')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.savefig(path)
plt.show()
plot_history(hist, 'loss', 'data/loss.jpg')
plot_history(hist, 'val_loss', 'data/val_loss.jpg')
#----------------直接测试一张图配准--------------------------
val_generator = datagenerators.my_data_generator(x_train_path,
y_train_path,
batch_size=1,
img_size=img_size)
val_input, _ = next(val_generator)
val_pred = vxm_model.predict(val_input)
# ---------------可视化---------------------
slices_2d = [f[0, ..., 0] for f in val_input + val_pred]
titles = [
'input_moving', 'input_fixed', 'predicted_moved', 'deformation_x'
]
neuron.plot.slices(slices_2d,
titles=titles,
cmaps=['gray'],
do_colorbars=True)
neuron.plot.flow([val_pred[1].squeeze()], width=5)
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 配准后")
