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)