def getnet(): model_file = 'C:/Users/cisguest/Downloads/voxelmorph-master/voxelmorph-master/my_models/1500.h5' nf_enc = [16, 32, 32, 32] nf_dec = [32, 32, 32, 32, 16, 3] # atlas_vol = np.load(fixed)['vol'][np.newaxis, ..., np.newaxis] # vol_size = atlas_vol.shape[1:-1] vol_size = (480, 640) bidir = 0 net = networks.miccai2018_net(vol_size, nf_enc, nf_dec, bidir=bidir) net.load_weights(model_file) return net #def main(): # net = getnet() # mapImages('C:/Users/cisguest/Downloads/iris/dark iris/301.png','C:/Users/cisguest/Downloads/iris/dark iris/301.png',net) # #if __name__ == "__main__": # main() #%% #[moved, warp] = net.predict([movIm, fixIm]) #plt.figure(figsize=(10,5),dpi=250) #plt.subplot(221) #plt.imshow(warp[0,:,:,0]) #plt.colorbar() #plt.subplot(222) #plt.imshow(warp[0,:,:,1]) #plt.colorbar() #plt.subplot(223) #plt.imshow(warp[0,:,:,2]) #plt.colorbar() #plt.subplot(224) #plt.imshow(warp[0,:,:,3]) #plt.colorbar() # ##%% # #plt.figure(figsize=(10,5),dpi=250) #plt.subplot(211) #plt.imshow(np.sqrt(((warp[0,:,:,0])**2+(warp[0,:,:,1])**2))) #plt.colorbar() #plt.title('Mag') #plt.subplot(212) #plt.imshow(np.arctan2(warp[0,:,:,0],warp[0,:,:,1])) #plt.colorbar() #plt.title('Phase')
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, 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 """ # load atlas from provided files. The atlas we used is 160x192x224. atlas_vol = np.load(atlas_file)['vol'][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" # Diffeomorphic network architecture used in MICCAI 2018 paper nf_enc = [16, 32, 32, 32] nf_dec = [32, 32, 32, 32, 16, 3] # 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): # 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.miccai2018_net(vol_size, nf_enc, nf_dec, bidir=bidir) # load initial weights if load_model_file is not None and load_model_file != "": model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # compile # note: best to supply vol_shape here than to let tf figure it out. 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] # 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) miccai2018_gen = datagenerators.miccai2018_gen(train_example_gen, atlas_vol_bs, batch_size=batch_size, bidir=bidir) # 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 mg_model.compile(optimizer=Adam(lr=lr), loss=model_losses, loss_weights=loss_weights) mg_model.fit_generator(miccai2018_gen, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback], steps_per_epoch=steps_per_epoch, verbose=1)
def test( gpu_id, model_dir, iter_num, compute_type='GPU', # GPU or CPU vol_size=(160, 192, 224), nf_enc=[16, 32, 32, 32], nf_dec=[32, 32, 32, 32, 16, 3], save_file=None): """ test via segmetnation propagation works by iterating over some iamge files, registering them to atlas, propagating the warps, then computing Dice with atlas segmentations """ # 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): # if testing miccai run, should be xy indexing. net = networks.miccai2018_net(vol_size, nf_enc, nf_dec, use_miccai_int=False, indexing='ij') net.load_weights(os.path.join(model_dir, str(iter_num) + '.h5')) # compose diffeomorphic flow output model diff_net = keras.models.Model(net.inputs, net.get_layer('diffflow').output) # 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) # prepare a matrix of dice values dice_vals = np.zeros((len(good_labels), n_batches)) for k in range(n_batches): # get data vol_name, seg_name = test_brain_strings[k].split(",") X_vol, X_seg = datagenerators.load_example_by_name(vol_name, seg_name) # predict transform with tf.device(gpu): pred = diff_net.predict([X_vol, atlas_vol]) # Warp segments with flow if compute_type == 'CPU': flow = pred[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])[0, ..., 0] # compute Volume Overlap (Dice) dice_vals[:, k] = dice(warp_seg, atlas_seg, labels=good_labels) print('%3d %5.3f %5.3f' % (k, np.mean( dice_vals[:, k]), np.mean(np.mean(dice_vals[:, :k + 1])))) if save_file is not None: sio.savemat(save_file, { 'dice_vals': dice_vals, 'labels': good_labels })
def test( gpu_id, model_dir, iter_num, compute_type='GPU', # GPU or CPU vol_size=(160, 192, 224), nf_enc=[16, 32, 32, 32], nf_dec=[32, 32, 32, 32, 16, 3], save_file=None): """ test via segmetnation propagation works by iterating over some iamge files, registering them to atlas, propagating the warps, then computing Dice with atlas segmentations """ # 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): # if testing miccai run, should be xy indexing. net = networks.miccai2018_net(vol_size, nf_enc, nf_dec, use_miccai_int=False, indexing='ij') net.load_weights(os.path.join(model_dir, str(iter_num) + '.h5')) # compose diffeomorphic flow output model diff_net = keras.models.Model(net.inputs, net.get_layer('diffflow').output) # 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) # get data X_vol = nib.load(r'D:\users\zzx\data\2018nor\pre\01.nii').get_data() X_vol = X_vol[np.newaxis, ..., np.newaxis] atlas_vol = nib.load(r'D:\users\zzx\data\2018nor\pre\01_a.nii').get_data() atlas_vol = atlas_vol[np.newaxis, ..., np.newaxis] X_mask = nib.load( r'D:\users\zzx\data\2018mask\pre\pig01_pre_final_r.nii').get_data() X_mask = X_mask[np.newaxis, ..., np.newaxis] X_mask[X_mask == np.max(X_mask)] = 1 X_mask[X_mask != 1] = 0 atlas_mask = nib.load( r'D:\users\zzx\data\2018mask\aft\pig01_02_final_r.nii').get_data() atlas_mask[atlas_mask == np.max(atlas_mask)] = 1 atlas_mask[atlas_mask != 1] = 0 ## feature point # X_feapt = np.zeros((160,192,224)) # X_feapt[128,51,165] = 1 # X_feapt = X_feapt[np.newaxis,...,np.newaxis] # predict transform with tf.device(gpu): pred = diff_net.predict([X_vol, atlas_vol]) # Warp segments with flow if compute_type == 'CPU': flow = pred[0, :, :, :, :] warp_seg = util.warp_seg(X_mask, flow, grid=grid, xx=xx, yy=yy, zz=zz) else: # GPU warp_mask = nn_trf_model.predict([X_mask, pred])[0, ..., 0] warp_vol = nn_trf_model.predict([X_vol, pred])[0, ..., 0] # pred_point1 = nn_trf_model.predict([X_feapt, pred])[0,...,0] print(X_vol.shape) # warp_vol = nib.Nifti1Image(warp_vol,np.eye(4)) warp_vol = nib.Nifti1Image(warp_vol, np.eye(4)) nib.save(warp_vol, r'D:\users\zzx\data\2018warp\1w.nii') # compute Volume Overlap (Dice) # X_mask = X_mask[0,...,0] # print(X_mask.shape, atlas_mask.shape,pred_point1.shape,np.where(pred_point1 != 0)) dice_vals = dice(warp_mask, atlas_mask) # print('%3d %5.3f %5.3f' % (k, np.mean(dice_vals[:, k]), np.mean(np.mean(dice_vals[:, :k+1])))) print(dice_vals) if save_file is not None: sio.savemat(save_file, {'dice_vals': dice_vals})
def train(data_dir, atlas_file, model_dir, gpu_id, lr, nb_epochs, prior_lambda, image_sigma, steps_per_epoch, batch_size, load_model_file, bidir, bool_cc, initial_epoch=0): """ model training function :param data_dir: folder with npz files for each subject. :param atlas_file: atlas filename. So far we support npz file with a 'vol' variable :param model_dir: model folder to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param nb_epochs: number of training iterations :param prior_lambda: the prior_lambda, the scalar in front of the smoothing laplacian, in MICCAI paper :param image_sigma: the image sigma in MICCAI paper :param steps_per_epoch: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size :param load_model_file: optional h5 model file to initialize with :param bidir: logical whether to use bidirectional cost function :param bool_cc: Train CC or MICCAI version """ # load atlas from provided files. The atlas we used is 160x192x224. #atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis] vm_dir = '/home/jdram/voxelmorph/' base = np.load( os.path.join(vm_dir, "data", "ts12_dan_a88_fin_o_trim_adpc_002661_256.npy")) monitor = np.load( os.path.join(vm_dir, "data", "ts12_dan_a05_fin_o_trim_adpc_002682_256.npy")) #base = np.load(os.path.join(vm_dir, "data","ts12_dan_a88_fin_o_trim_adpc_002661_abs.npy")) #monitor = np.load(os.path.join(vm_dir, "data","ts12_dan_a05_fin_o_trim_adpc_002682_abs.npy")) #vol_size = (64, 64, 64) vol_size = (64, 64, 256 - 64) #vol_size = (128, 128, 256) # prepare data files # for the CVPR and MICCAI papers, we have data arranged in train/validate/test folders # inside each folder is a /vols/ and a /asegs/ folder with the volumes # and segmentations. All of our papers use npz formated data. #train_vol_names = glob.glob(os.path.join(data_dir, '*.npy')) #random.shuffle(train_vol_names) # shuffle volume list #assert len(train_vol_names) > 0, "Could not find any training data" # Diffeomorphic network architecture used in MICCAI 2018 paper nf_enc = [32, 64, 64, 64] nf_dec = [64, 64, 64, 64, 32, 3] # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) tf.reset_default_graph() if bool_cc: pre_net = "cc_" else: if bidir: pre_net = "miccai_bidir_" else: pre_net = "miccai_" # gpu handling gpu = '/device:GPU:%d' % int(gpu_id) # gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # prepare the model with tf.device(gpu): # prepare the model # in the CVPR layout, the model takes in [image_1, image_2] and outputs [warped_image_1, flow] # in the experiments, we use image_2 as atlas if bool_cc: model = networks.cvpr2018_net(vol_size, nf_enc, nf_dec) else: model = networks.miccai2018_net(vol_size, nf_enc, nf_dec, bidir=bidir, vel_resize=.5) # load initial weights if load_model_file is not None and load_model_file != "": print('loading', load_model_file) model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, f'{pre_net}{initial_epoch:02d}.h5')) model.summary() if bool_cc: model_losses = [losses.NCC().loss, losses.Grad('l2').loss] loss_weights = [1.0, 0.01] # recommend 1.0 for ncc, 0.01 for mse else: flow_vol_shape = model.outputs[-1].shape[1:-1] loss_class = losses.Miccai2018(image_sigma, prior_lambda, flow_vol_shape=flow_vol_shape) if bidir: model_losses = [ loss_class.recon_loss, loss_class.recon_loss, loss_class.kl_loss ] loss_weights = [0.5, 0.5, 1] else: model_losses = [loss_class.recon_loss, loss_class.kl_loss] loss_weights = [1, 1] segy_gen = datagenerators.segy_gen(base, monitor, batch_size=batch_size) # prepare callbacks save_file_name = os.path.join(model_dir, pre_net + '{epoch:02d}.h5') with tf.device(gpu): # fit generator save_callback = ModelCheckpoint(save_file_name, period=5) csv_cb = CSVLogger(f'{pre_net}log.csv') nan_cb = TerminateOnNaN() rlr_cb = ReduceLROnPlateau(monitor='loss', verbose=1) els_cb = EarlyStopping(monitor='loss', patience=15, verbose=1, restore_best_weights=True) cbs = [save_callback, csv_cb, nan_cb, rlr_cb, els_cb] mg_model = model # compile mg_model.compile(optimizer=Adam(lr=lr), loss=model_losses, loss_weights=loss_weights) mg_model.fit( [base, monitor], [monitor, np.zeros_like(base)], initial_epoch=initial_epoch, batch_size=8, epochs=nb_epochs, callbacks=cbs, #steps_per_epoch=steps_per_epoch, verbose=1)
def test( gpu_id, model_dir, iter_num, compute_type='GPU', # GPU or CPU vol_size=(160, 192, 224), nf_enc=[16, 32, 32, 32], nf_dec=[32, 32, 32, 32, 16, 3], save_file=None): """ test via segmetnation propagation works by iterating over some iamge files, registering them to atlas, propagating the warps, then computing Dice with atlas segmentations """ # 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): # if testing miccai run, should be xy indexing. net = networks.miccai2018_net(vol_size, nf_enc, nf_dec, use_miccai_int=True, indexing='xy') net.load_weights(os.path.join(model_dir, str(iter_num) + '.h5')) # compose diffeomorphic flow output model diff_net = keras.models.Model(net.inputs, net.get_layer('diffflow').output) # NN transfer model nn_trf_model = networks.nn_trf(vol_size) # if CPU, prepare grid if compute_type == 'CPU': grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4) # prepare a matrix of dice values dice_vals = np.zeros((len(good_labels), n_batches)) for k in range(n_batches): # get data vol_name, seg_name = test_brain_strings[k].split(",") X_vol, X_seg = datagenerators.load_example_by_name(vol_name, seg_name) # predict transform with tf.device(gpu): pred = diff_net.predict([X_vol, atlas_vol]) # Warp segments with flow if compute_type == 'CPU': flow = pred[0, :, :, :, :] warp_seg = util.warp_seg(X_seg, flow, grid=grid, xx=xx, yy=yy, zz=zz) else: # GPU # Rigid registration only by GPU flow = pred[0, :, :, :, :] # Compute A(all about coordinate computation) x = np.linspace(0, 160 - 16, sample_num) x = x.astype(np.int32) y = np.linspace(0, 190 - 19, sample_num) y = y.astype(np.int32) z = np.linspace(0, 220 - 22, sample_num) z = z.astype(np.int32) index = np.rollaxis(np.array(np.meshgrid(x, y, z)), 0, 4) x = index[:, :, :, 0] y = index[:, :, :, 1] z = index[:, :, :, 2] # Y in formula x_flow = np.arange(vol_size[0]) y_flow = np.arange(vol_size[1]) z_flow = np.arange(vol_size[2]) grid = np.rollaxis(np.array((np.meshgrid(y_flow, x_flow, z_flow))), 0, 4) # original coordinate grid_x = grid_sample(x, y, z, grid[:, :, :, 0], sample_num) grid_y = grid_sample(x, y, z, grid[:, :, :, 1], sample_num) grid_z = grid_sample(x, y, z, grid[:, :, :, 2], sample_num) # X (10,10,10) sample = flow + grid sample_x = grid_sample(x, y, z, sample[:, :, :, 0], sample_num) sample_y = grid_sample(x, y, z, sample[:, :, :, 1], sample_num) sample_z = grid_sample(x, y, z, sample[:, :, :, 2], sample_num) # Y (10,10,10) sum_x = np.sum(flow[:, :, :, 0]) sum_y = np.sum(flow[:, :, :, 1]) sum_z = np.sum(flow[:, :, :, 2]) ave_x = sum_x / (vol_size[0] * vol_size[1] * vol_size[2]) ave_y = sum_y / (vol_size[0] * vol_size[1] * vol_size[2]) ave_z = sum_z / (vol_size[0] * vol_size[1] * vol_size[2]) # formula Y = np.zeros((10, 10, 10, grid_dimension)) X = np.zeros((10, 10, 10, grid_dimension)) T = np.array([ave_x, ave_y, ave_z, 1]) # (4,1) # R = np.zeros((10, 10, 10, grid_dimension, grid_dimension)) for i in np.arange(10): for j in np.arange(10): for z in np.arange(10): Y[i, j, z, :] = np.array([ sample_x[i, j, z], sample_y[i, j, z], sample_z[i, j, z], 1 ]) for i in np.arange(10): for j in np.arange(10): for z in np.arange(10): X[i, j, z, :] = np.array([ grid_x[i, j, z], grid_y[i, j, z], grid_z[i, j, z], 1 ]) X = X.reshape((1000, grid_dimension)) Y = Y.reshape((1000, grid_dimension)) R = np.dot( np.dot(np.linalg.pinv(np.dot(np.transpose(X), X)), np.transpose(X)), Y) # R # build new grid(Use R to do the spatial transform) shifted_x = np.arange(vol_size[0]) shifted_y = np.arange(vol_size[1]) shifted_z = np.arange(vol_size[2]) print(shifted_x.shape) print(shifted_y.shape) print(shifted_z.shape) shifted_grid = np.rollaxis( np.array((np.meshgrid(shifted_y, shifted_x, shifted_z))), 0, 4) print(shifted_grid.shape) for i in np.arange(vol_size[0]): for j in np.arange(vol_size[1]): for z in np.arange(vol_size[2]): coordinates = np.dot( R, np.array([i, j, z, 1]).reshape(4, 1)) + T.reshape( 4, 1) print("voxel." + '(' + str(i) + ',' + str(j) + ',' + str(z) + ')') shifted_grid[i, j, z, 0] = coordinates[0] shifted_grid[i, j, z, 1] = coordinates[1] shifted_grid[i, j, z, 2] = coordinates[2] # interpolation xx = np.arange(vol_size[1]) yy = np.arange(vol_size[0]) zz = np.arange(vol_size[2]) warp_seg = interpn((yy, xx, zz), X_seg[0, :, :, :, 0], shifted_grid, method='nearest', bounds_error=False, fill_value=0) # CVPR grid = np.rollaxis(np.array(np.meshgrid(xx, yy, zz)), 0, 4) sample = flow + grid sample = np.stack( (sample[:, :, :, 1], sample[:, :, :, 0], sample[:, :, :, 2]), 3) warp_seg2 = interpn((yy, xx, zz), X_seg[0, :, :, :, 0], sample, method='nearest', bounds_error=False, fill_value=0) # compute dice vals, _ = dice(warp_seg, atlas_seg, labels=labels, nargout=2) vals2, _ = dice(X_seg[0, :, :, :, 0], atlas_seg, labels=labels, nargout=2) vals3, _ = dice(warp_seg2, atlas_seg, labels=labels, nargout=2) print("dice before:") print(np.mean(vals2), np.std(vals2)) print("dice after deformable registration:") print(np.mean(vals3), np.std(vals3)) print("dice after rigid registration:") print(np.mean(vals), np.std(vals)) warp_seg = nn_trf_model.predict([X_seg, pred])[0, ..., 0] # compute Volume Overlap (Dice) dice_vals[:, k] = dice(warp_seg, atlas_seg, labels=good_labels) print('%3d %5.3f %5.3f' % (k, np.mean( dice_vals[:, k]), np.mean(np.mean(dice_vals[:, :k + 1])))) if save_file is not None: sio.savemat(save_file, { 'dice_vals': dice_vals, 'labels': good_labels })
def train(model_dir, gpu_id, lr, n_iterations, alpha, image_sigma, model_save_iter, batch_size=1): """ model training function :param model_dir: model folder to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param n_iterations: number of training iterations :param alpha: the alpha, the scalar in front of the smoothing laplacian, in MICCAI paper :param image_sigma: the image sigma in MICCAI 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 """ # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) print(model_dir) # 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)) # Diffeomorphic network architecture used in MICCAI 2018 paper nf_enc = [16, 32, 32, 32] nf_dec = [32, 32, 32, 32, 16, 3] # prepare the model # in the CVPR layout, the model takes in [image_1, image_2] and outputs [warped_image_1, velocity_stats] # in the experiments, we use image_2 as atlas with tf.device(gpu): # miccai 2018 used xy indexing. model = networks.miccai2018_net(vol_size, nf_enc, nf_dec, use_miccai_int=True, indexing='xy') # compile model_losses = [losses.kl_l2loss(image_sigma), losses.kl_loss(alpha)] model.compile(optimizer=Adam(lr=lr), loss=model_losses) # save first iteration model.save(os.path.join(model_dir, str(0) + '.h5')) train_example_gen = datagenerators.example_gen(train_vol_names) zeros = np.zeros((1, *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(1, n_iterations): # get_data X = next(train_example_gen)[0] # train with tf.device(gpu): train_loss = model.train_on_batch([X, atlas_vol], [atlas_vol, zeros]) if not isinstance(train_loss, list): train_loss = [train_loss] # print print_loss(step, 0, train_loss) # save model with tf.device(gpu): if (step % model_save_iter == 0) or step < 10: model.save(os.path.join(model_dir, str(step) + '.h5'))
def train(data_dir, model, model_name, gpu_id, lr, nb_epochs, reg_param, steps_per_epoch, batch_size, load_model_file, atlas_file, max_clip, distance, patch_size, use_ssc, use_gaussian_kernel, use_fixed_var, use_miccai, 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 = nib.load(atlas_file).get_data()[np.newaxis, ..., np.newaxis] atlas_vol = atlas_vol / np.max(atlas_vol) * max_clip # atlas_vol = nib.load('../data/t1_atlas.nii').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]] 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) print(gpu) 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 use_miccai: print('miccai: therefore diffeomorphic') model = networks.miccai2018_net(vol_size, nf_enc, nf_dec) else: 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) if use_miccai: data_gen = datagenerators.cvpr2018_gen(train_example_gen, atlas_vol_bs, batch_size=batch_size) else: data_gen = datagenerators.miccai2018_gen(train_example_gen, atlas_vol_bs, batch_size=batch_size) # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') loss_function = losses.mind(distance, patch_size, use_ssc=use_ssc, use_gaussian_kernel=use_gaussian_kernel, use_fixed_var=use_fixed_var) # 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, verbose=1) mg_model = model # compile mg_model.compile(optimizer=Adam(lr=lr), loss=[loss_function, losses.Grad('l2').loss], loss_weights=[1.0, reg_param]) # fit mg_model.fit_generator(data_gen, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback], steps_per_epoch=steps_per_epoch, verbose=1)
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( gpu_id, iter_num, compute_type='GPU', # GPU or CPU vol_size=(160, 192, 224), nf_enc=[16, 32, 32, 32], nf_dec=[32, 32, 32, 32, 16, 3], save_file=None): """ test by segmentation, compute dice between atlas_seg and warp_seg :param gpu_id: gpu id :param iter_num: specify the model to read :param compute_type: CPU/GPU :param vol_size: volume size :param nf_enc: number of encoder :param nf_dec: number of decoder :param save_file: None :return: None """ # 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): # if testing miccai run, should be xy indexing. net = networks.miccai2018_net(vol_size, nf_enc, nf_dec, use_miccai_int=True, indexing='xy') model_dir = "/home/ys895/rigid_diff_model/" net.load_weights(os.path.join(model_dir, str(iter_num) + '.h5')) # compose diffeomorphic flow output model diff_net = keras.models.Model(net.inputs, net.get_layer('diffflow').output) # NN transfer model nn_trf_model = networks.nn_trf(vol_size) # if CPU, prepare grid if compute_type == 'CPU': grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4) # prepare a matrix of dice values dice_vals = np.zeros((len(good_labels), n_batches)) for k in range(n_batches): # get data vol_name, seg_name = test_brain_strings[k].split(",") X_vol, X_seg = datagenerators.load_example_by_name(vol_name, seg_name) orig_vol = X_vol orig_seg = X_seg theta = 0 beta = 5 omega = 0 X_seg = rotate_img(X_seg[0, :, :, :, 0], theta=theta, beta=beta, omega=omega) X_vol = rotate_img(X_vol[0, :, :, :, 0], theta=theta, beta=beta, omega=omega) X_seg = X_seg.reshape((1, ) + X_seg.shape + (1, )) X_vol = X_vol.reshape((1, ) + X_vol.shape + (1, )) sample_num = 30 grid_dimension = 4 # predict transform with tf.device(gpu): pred = diff_net.predict([X_vol, atlas_vol]) # Warp segments with flow if compute_type == 'CPU': flow = pred[0, :, :, :, :] warp_seg = util.warp_seg(X_seg, flow, grid=grid, xx=xx, yy=yy, zz=zz) else: # GPU flow = pred[0, :, :, :, :] # sample coordinate(sample_num * sample_num * sample_num) x = np.linspace(0, (vol_size[0] / sample_num) * (sample_num - 1), sample_num) x = x.astype(np.int32) y = np.linspace(0, (vol_size[1] / sample_num) * (sample_num - 1), sample_num) y = y.astype(np.int32) z = np.linspace(0, (vol_size[2] / sample_num) * (sample_num - 1), sample_num) z = z.astype(np.int32) index = np.rollaxis(np.array(np.meshgrid(y, x, z)), 0, 4) x = index[:, :, :, 1] y = index[:, :, :, 0] z = index[:, :, :, 2] # Y in formula x_flow = np.arange(vol_size[0]) y_flow = np.arange(vol_size[1]) z_flow = np.arange(vol_size[2]) grid = np.rollaxis(np.array((np.meshgrid(y_flow, x_flow, z_flow))), 0, 4) # original coordinate grid_x = grid_sample(x, y, z, grid[:, :, :, 1], sample_num) grid_y = grid_sample(x, y, z, grid[:, :, :, 0], sample_num) grid_z = grid_sample(x, y, z, grid[:, :, :, 2], sample_num) # X (10,10,10) sample = flow + grid sample_x = grid_sample(x, y, z, sample[:, :, :, 1], sample_num) sample_y = grid_sample(x, y, z, sample[:, :, :, 0], sample_num) sample_z = grid_sample(x, y, z, sample[:, :, :, 2], sample_num) # Y (10,10,10) sum_x = np.sum(flow[:, :, :, 1]) sum_y = np.sum(flow[:, :, :, 0]) sum_z = np.sum(flow[:, :, :, 2]) ave_x = sum_x / (vol_size[0] * vol_size[1] * vol_size[2]) ave_y = sum_y / (vol_size[0] * vol_size[1] * vol_size[2]) ave_z = sum_z / (vol_size[0] * vol_size[1] * vol_size[2]) # formula Y = np.zeros((sample_num, sample_num, sample_num, grid_dimension)) X = np.zeros((sample_num, sample_num, sample_num, grid_dimension)) T = np.array([ave_x, ave_y, ave_z, 1]) # (4,1) print(T) for i in np.arange(sample_num): for j in np.arange(sample_num): for z in np.arange(sample_num): Y[i, j, z, :] = np.array([ sample_x[i, j, z], sample_y[i, j, z], sample_z[i, j, z], 1 ]) #Y[i, j, z, :] = Y[i, j, z, :] - np.array([ave_x, ave_y, ave_z, 0]) # amend: Y` = Y - T for i in np.arange(sample_num): for j in np.arange(sample_num): for z in np.arange(sample_num): X[i, j, z, :] = np.array([ grid_x[i, j, z], grid_y[i, j, z], grid_z[i, j, z], 1 ]) X = X.reshape( (sample_num * sample_num * sample_num, grid_dimension)) Y = Y.reshape( (sample_num * sample_num * sample_num, grid_dimension)) R = np.dot( np.dot(np.linalg.pinv(np.dot(np.transpose(X), X)), np.transpose(X)), Y) # R(4, 4) print(R) beta = -(beta / 180) * math.pi R = np.array([[math.cos(beta), 0, -math.sin(beta), 0], [0, 1, 0, 0], [math.sin(beta), 0, math.cos(beta), 0], [0, 0, 0, 1]]) #R = R.transpose() # build new grid(Use R to do the spatial transform) shifted_x = np.arange(vol_size[0]) shifted_y = np.arange(vol_size[1]) shifted_z = np.arange(vol_size[2]) shifted_grid = np.rollaxis( np.array((np.meshgrid(shifted_y, shifted_x, shifted_z))), 0, 4) # some required matrixs T1 = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [ -int(vol_size[0] / 2), -int(vol_size[1] / 2), -int(vol_size[2] / 2), 1 ]]) T2 = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [ int(vol_size[0] / 2), int(vol_size[1] / 2), int(vol_size[2] / 2), 1 ]]) for i in np.arange(vol_size[0]): for j in np.arange(vol_size[1]): for z in np.arange(vol_size[2]): #coordinates = np.dot(R, np.array([i, j, z, 1]).reshape(4, 1)) + T.reshape(4, 1) coordinates = np.dot( np.dot( np.dot( np.array([i, j, z, 1]).reshape(1, 4), T1), R), T2) # new implementation # print("voxel." + '(' + str(i) + ',' + str(j) + ',' + str(z) + ')') shifted_grid[i, j, z, 1] = coordinates[0, 0] shifted_grid[i, j, z, 0] = coordinates[0, 1] shifted_grid[i, j, z, 2] = coordinates[0, 2] # interpolation xx = np.arange(vol_size[1]) yy = np.arange(vol_size[0]) zz = np.arange(vol_size[2]) shifted_grid = np.stack( (shifted_grid[:, :, :, 1], shifted_grid[:, :, :, 0], shifted_grid[:, :, :, 2]), 3 ) # notice: the shifted_grid is reverse in x and y, so this step is used for making it back. warp_seg = interpn((yy, xx, zz), X_seg[0, :, :, :, 0], shifted_grid, method='nearest', bounds_error=False, fill_value=0) # rigid registration warp_vol = interpn((yy, xx, zz), X_vol[0, :, :, :, 0], shifted_grid, method='nearest', bounds_error=False, fill_value=0) # rigid registration # compute Volume Overlap (Dice) dice_vals[:, k] = dice(warp_seg, orig_seg[0, :, :, :, 0], labels=good_labels) print('%3d %5.3f %5.3f' % (k, np.mean( dice_vals[:, k]), np.mean(np.mean(dice_vals[:, :k + 1])))) if save_file is not None: sio.savemat(save_file, { 'dice_vals': dice_vals, 'labels': good_labels }) # specify slice num_slice = 90 plt.figure() plt.subplot(1, 3, 1) plt.imshow(orig_vol[0, :, num_slice, :, 0]) plt.subplot(1, 3, 2) plt.imshow(X_vol[0, :, num_slice, :, 0]) plt.subplot(1, 3, 3) plt.imshow(warp_vol[:, num_slice, :]) plt.savefig("slice" + str(num_slice) + '_' + str(k) + ".png") plt.figure() plt.subplot(1, 3, 1) plt.imshow(flow[:, num_slice, :, 1]) plt.subplot(1, 3, 2) plt.imshow(flow[:, num_slice, :, 0]) plt.subplot(1, 3, 3) plt.imshow(flow[:, num_slice, :, 2]) plt.savefig("flow.png")