def test(model_name,
gpu_id,
nf_enc=[16, 32, 32, 32],
nf_dec=[32, 32, 32, 32, 32, 16, 16]):
"""
test
nf_enc and nf_dec
#nf_dec = [32,32,32,32,32,16,16,3]
# This needs to be changed. Ideally, we could just call load_model, and we wont have to
# specify the # of channels here, but the load_model is not working with the custom loss...
"""
# load subject test
print("load_data start")
X_train, y_train = load_data(data_dir='../data',
mode='test',
fixed='joyoungje')
vol_size = y_train.shape[1:-1]
# 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))
# load weights of model
with tf.device(gpu):
net = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
print("model load weights")
net.load_weights(model_name)
# NN transfer model
nn_trf_model = networks.nn_trf(vol_size, indexing='ij')
# # if CPU, prepare grid
# if compute_type == 'CPU':
# grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4)
with tf.device(gpu):
print("model predict")
pred = net.predict([X_train, y_train])
print("nn_tft_model.predict")
X_warp = nn_trf_model.predict([X_train, pred[1]])[0, ..., 0]
reshape_y_train = y_train.reshape(y_train.shape[1:-1])
vals = dice(pred[0].reshape(pred[0].shape[1:-1]), reshape_y_train)
dice_mean = np.mean(vals)
dice_std = np.std(vals)
print('Dice mean over structures: {:.2f} ({:.2f})'.format(
dice_mean, dice_std))
def train(model,
pretrained_path,
model_name,
gpu_id,
lr,
n_iterations,
use_mi,
gamma,
num_bins,
patch_size,
max_clip,
reg_param,
model_save_iter,
local_mi,
sigma_ratio,
batch_size=1):
"""
model training function
: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 model_save_iter: frequency with which to save models
:param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size
"""
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
restrict_GPU_tf(str(gpu_id))
restrict_GPU_keras(str(gpu_id))
train_labels = sio.loadmat('../data/labels.mat')['labels'][0]
n_labels = train_labels.shape[0]
normalized_atlas_vol = atlas_vol / np.max(atlas_vol) * max_clip
atlas_seg = datagenerators.split_seg_into_channels(seg, train_labels)
atlas_seg = datagenerators.downsample(atlas_seg)
model_dir = "../models/" + model_name
# prepare model folder
if not os.path.isdir(model_dir):
os.mkdir(model_dir)
# GPU handling
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))
# 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]
else:
nf_dec = [32, 32, 32, 32, 32, 16, 16]
# 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
bin_centers = np.linspace(0, max_clip, num_bins * 2 + 1)[1::2]
loss_function = losses.mutualInformation(bin_centers,
max_clip=max_clip,
local_mi=local_mi,
patch_size=patch_size,
sigma_ratio=sigma_ratio)
model = networks.cvpr2018_net(vol_size,
nf_enc,
nf_dec,
use_seg=True,
n_seg=len(train_labels))
model.compile(optimizer=Adam(lr=lr),
loss=[
loss_function,
losses.gradientLoss('l2'),
sparse_categorical_crossentropy
],
loss_weights=[1 if use_mi else 0, reg_param, gamma])
# if you'd like to initialize the data, you can do it here:
if pretrained_path != None and pretrained_path != '':
model.load_weights(pretrained_path)
# prepare data for training
train_example_gen = datagenerators.example_gen(train_vol_names,
return_segs=True,
seg_dir=train_seg_dir)
zero_flow = np.zeros([batch_size, *vol_size, 3])
# train. Note: we use train_on_batch and design out own print function as this has enabled
# faster development and debugging, but one could also use fit_generator and Keras callbacks.
for step in range(0, n_iterations):
# get data
X = next(train_example_gen)
X_seg = X[1]
X_seg = datagenerators.split_seg_into_channels(X_seg, train_labels)
X_seg = datagenerators.downsample(X_seg)
# train
train_loss = model.train_on_batch(
[X[0], normalized_atlas_vol, X_seg],
[normalized_atlas_vol, zero_flow, atlas_seg])
if not isinstance(train_loss, list):
train_loss = [train_loss]
# print the loss.
print_loss(step, 1, train_loss)
# save model
if step % model_save_iter == 0:
model.save(os.path.join(model_dir, str(step) + '.h5'))
def test(
model_name,
gpu_id,
compute_type='GPU', # GPU or CPU
nf_enc=[16, 32, 32, 32],
nf_dec=[32, 32, 32, 32, 32, 16, 16]):
"""
test
nf_enc and nf_dec
#nf_dec = [32,32,32,32,32,16,16,3]
# This needs to be changed. Ideally, we could just call load_model, and we wont have to
# specify the # of channels here, but the load_model is not working with the custom loss...
"""
# Anatomical labels we want to evaluate
labels = sio.loadmat('../data/labels.mat')['labels'][0]
atlas = np.load('../data/atlas_norm.npz')
atlas_vol = atlas['vol'][np.newaxis, ..., np.newaxis]
atlas_seg = atlas['seg']
vol_size = atlas_vol.shape[1:-1]
# 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))
# load weights of model
with tf.device(gpu):
net = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
net.load_weights(model_name)
# NN transfer model
nn_trf_model = networks.nn_trf(vol_size, indexing='ij')
# if CPU, prepare grid
if compute_type == 'CPU':
grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4)
# load subject test
X_vol, X_seg = datagenerators.load_example_by_name('../data/test_vol.npz',
'../data/test_seg.npz')
with tf.device(gpu):
pred = net.predict([X_vol, atlas_vol])
# Warp segments with flow
if compute_type == 'CPU':
flow = pred[1][0, :, :, :, :]
warp_seg = util.warp_seg(X_seg,
flow,
grid=grid,
xx=xx,
yy=yy,
zz=zz)
else: # GPU
warp_seg = nn_trf_model.predict([X_seg, pred[1]])[0, ..., 0]
vals, _ = dice(warp_seg, atlas_seg, labels=labels, nargout=2)
dice_mean = np.mean(vals)
dice_std = np.std(vals)
print('Dice mean over structures: {:.2f} ({:.2f})'.format(
dice_mean, dice_std))
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)
labels = sio.loadmat('../data/labels.mat')['labels'][0]
# -> [ 2 3 4 7 8 10 11 12 13 14 15 16 17 18 24 28 31 41 42 43 46 47 49 50 51 52 53 54 60 63]
atlas = np.load('../data/atlas_norm.npz')
atlas_vol = atlas['vol']
atlas_seg = atlas['seg']
atlas_vol = np.reshape(atlas_vol, (1,)+atlas_vol.shape+(1,)) # 1x160x192x224x1
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# load weights of model
with tf.device(gpu):
net = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
# net.load_weights('../models/' + model_name + '/' + str(iter_num) + '.h5')
net.load_weights(model_name)
xx = np.arange(vol_size[1])
yy = np.arange(vol_size[0])
zz = np.arange(vol_size[2])
grid = np.rollaxis(np.array(np.meshgrid(xx, yy, zz)), 0, 4)
X_vol, X_seg = datagenerators.load_example_by_name('../data/test_vol.npz', '../data/test_seg.npz')
# 1x160x192x224x1, 1x160x192x224x1
with tf.device(gpu):
pred = net.predict([X_vol, atlas_vol])
# Warp segments with flow
def test(expr_name,
epoch,
gpu_id,
sample_num,
dataset,
nf_enc=[16, 32, 32, 32],
nf_dec=[32, 32, 32, 32, 32, 16, 16]):
"""
test
nf_enc and nf_dec
#nf_dec = [32,32,32,32,32,16,16,3]
# This needs to be changed. Ideally, we could just call load_model, and we wont have to
# specify the # of channels here, but the load_model is not working with the custom loss...
"""
# load subject test
print("load_data start")
X_train, y_train = load_data(data_dir='../data',
sample_num=sample_num,
mode=dataset,
fixed='joyoungje')
vol_size = y_train.shape[1:-1]
# 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))
# load weights of model
with tf.device(gpu):
net = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
print("model load weights")
net.load_weights(
os.path.join("../models", expr_name, "{}.h5".format(epoch)))
# NN transfer model
nn_trf_model = networks.nn_trf(vol_size, indexing='ij')
# # if CPU, prepare grid
# if compute_type == 'CPU':
# grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4)
with tf.device(gpu):
print("model predict")
pred = net.predict([X_train, y_train])
#print("nn_tft_model.predict")
#X_warp = nn_trf_model.predict([X_train, pred[1]])[0,...,0]
warp_img = nib.Nifti1Image(pred[0].reshape(pred[0].shape[1:-1]),
affine=np.eye(4))
nib.save(
warp_img,
os.path.join(
'../result', '{}_epoch{}_{}{}_registration.nii.gz'.format(
expr_name, epoch, dataset, sample_num)))
moving_img = nib.Nifti1Image(X_train.reshape(X_train.shape[1:-1]),
affine=np.eye(4))
nib.save(
moving_img,
os.path.join(
'../result',
'{}_epoch{}_{}{}_original.nii.gz'.format(expr_name, epoch, dataset,
sample_num)))
fixed_img = nib.Nifti1Image(y_train.reshape(y_train.shape[1:-1]),
affine=np.eye(4))
nib.save(
fixed_img,
os.path.join(
'../result',
'{}_epoch{}_{}{}_reference.nii.gz'.format(expr_name, epoch,
dataset, sample_num)))
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(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 test(model_name,
epoch,
gpu_id,
n_test,
invert_images,
max_clip,
indexing,
use_miccai,
atlas_file,
atlas_seg_file,
normalize_atlas,
vol_size=(160, 192, 224),
nf_enc=[16, 32, 32, 32],
nf_dec=[32, 32, 32, 32, 32, 16, 16]):
start_time = time.time()
good_labels = sio.loadmat('../data/labels.mat')['labels'][0]
# setup
gpu = '/gpu:' + str(gpu_id)
# print(gpu)
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
restrict_GPU_tf(str(gpu_id))
restrict_GPU_keras(str(gpu_id))
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
atlas_vol = nib.load(atlas_file).get_data()[np.newaxis, ..., np.newaxis]
atlas_seg = nib.load(atlas_seg_file).get_data()
if normalize_atlas:
atlas_vol = atlas_vol / np.max(atlas_vol) * max_clip
sz = atlas_seg.shape
z_inp1 = tf.placeholder(tf.float32, sz)
z_inp2 = tf.placeholder(tf.float32, sz)
z_out = losses.kdice(z_inp1, z_inp2, good_labels)
kdice_fn = K.function([z_inp1, z_inp2], [z_out])
# load weights of model
with tf.device(gpu):
if use_miccai:
net = networks.miccai2018_net(vol_size, nf_enc, nf_dec)
net.load_weights('../models/' + model_name + '/' + str(epoch) +
'.h5')
trf_model = networks.trf_core(
(vol_size[0] // 2, vol_size[1] // 2, vol_size[2] // 2),
nb_feats=len(good_labels) + 1,
indexing=indexing)
else:
net = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
net.load_weights('../models/' + model_name + '/' + str(epoch) +
'.h5')
trf_model = networks.trf_core(vol_size,
nb_feats=len(good_labels) + 1,
indexing=indexing)
dice_means = []
dice_stds = []
for step in range(0, n_test):
# get data
if n_test == 1:
X_vol = nib.load('../t1_atlas.nii').get_data()[np.newaxis, ...,
np.newaxis]
X_seg = nib.load('../t1_atlas_seg.nii').get_data()[np.newaxis, ...,
np.newaxis]
else:
vol_name, seg_name = test_brain_strings[step].split(",")
X_vol, X_seg = datagenerators.load_example_by_name(
vol_name, seg_name)
if invert_images:
X_vol = max_clip - X_vol
with tf.device(gpu):
pred = net.predict([X_vol, atlas_vol])
all_labels = np.unique(X_seg)
for l in all_labels:
if l not in good_labels:
X_seg[X_seg == l] = 0
for i in range(len(good_labels)):
X_seg[X_seg == good_labels[i]] = i + 1
seg_onehot = tf.keras.utils.to_categorical(
X_seg[0, :, :, :, 0], num_classes=len(good_labels) + 1)
warp_seg_onehot = trf_model.predict(
[seg_onehot[tf.newaxis, :, :, :, :], pred[1]])
warp_seg = np.argmax(warp_seg_onehot[0, :, :, :, :], axis=3)
warp_seg_correct = np.zeros(warp_seg.shape)
for i in range(len(good_labels)):
warp_seg_correct[warp_seg == i + 1] = good_labels[i]
dice = kdice_fn([warp_seg_correct, atlas_seg])
mean = np.mean(dice)
std = np.std(dice)
dice_means.append(mean)
dice_stds.append(std)
print(step, mean, std)
print('average dice:', np.mean(dice_means))
print('std over patients:', np.std(dice_means))
print('average std over regions:', np.mean(dice_stds))
print('time taken:', time.time() - start_time)
def test(expr_name, epoch, gpu_id,
nf_enc=(16, 32, 32, 32), nf_dec=(32, 32, 32, 32, 32, 16, 16)):
"""
test
nf_enc and nf_dec
#nf_dec = [32,32,32,32,32,16,16,3]
# This needs to be changed. Ideally, we could just call load_model, and we wont have to
# specify the # of channels here, but the load_model is not working with the custom loss...
"""
vol_size = (256, 256, 144)
# 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))
# load weights of model
with tf.device(gpu):
net = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
print("model load weights")
net.load_weights(os.path.join("../models", expr_name, "{:04d}.h5".format(epoch)))
# NN transfer model
nn_trf_model = networks.nn_trf(vol_size, indexing='ij')
# load subject test
data_dir = '../../../dataset/urinary'
except_list = [6, 97, 125, 198, 199, 228, 271]
# vol_names = [filename for filename in os.listdir(data_dir) if (int(filename.split("_")[-1].split('.')[0]) <= 10) 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.sort()
print("data length:", len(vol_names))
vol_names = vol_names[:10]
generator = eval_gen(data_dir=data_dir, vol_names=vol_names)
# # if CPU, prepare grid
# if compute_type == 'CPU':
# grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4)
if not os.path.isdir(os.path.join('../result', '{}_epoch{}'.format(expr_name, epoch))):
os.mkdir(os.path.join('../result', '{}_epoch{}'.format(expr_name, epoch)))
for pre, delay, vol_name in generator:
with tf.device(gpu):
pred = net.predict([pre, delay])
X_warp = nn_trf_model.predict([pre, pred[1]])[0, ..., 0]
pre = np.flip(pre, 3)[0, ..., 0]
delay = np.flip(delay, 3)[0, ..., 0]
X_warp = np.flip(X_warp, 2)
pre = np.rot90(pre, 3)
delay = np.rot90(delay, 3)
X_warp = np.rot90(X_warp, 3)
pre = pre * 4096 - 1024
pre = pre.astype(np.int16)
delay = delay * 4096 - 1024
delay = delay.astype(np.int16)
X_warp = X_warp * 4096 - 1024
X_warp = X_warp.astype(np.int16)
warp_img = nib.Nifti1Image(X_warp, affine=np.eye(4))
nib.save(warp_img, os.path.join('../result', '{}_epoch{}'.format(expr_name, epoch),
"{}_registration.nii.gz".format(vol_name)))
moving_img = nib.Nifti1Image(pre, affine=np.eye(4))
nib.save(moving_img, os.path.join('../result', '{}_epoch{}'.format(expr_name, epoch),
"{}_pre.nii.gz".format(vol_name)))
fixed_img = nib.Nifti1Image(delay, affine=np.eye(4))
nib.save(fixed_img, os.path.join('../result', '{}_epoch{}'.format(expr_name, epoch),
"{}_delay.nii.gz".format(vol_name)))
print("successfully done!")
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)
