def register(gpu_id, moving, fixed, model_file, out_img, out_warp): """ register moving and fixed. """ assert model_file, "A model file is necessary" assert out_img or out_warp, "output image or warp file needs to be specified" # GPU handling if gpu_id is not None: 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)) else: gpu = '/cpu:0' # load data mov_nii = nib.load(moving) mov = mov_nii.get_data()[np.newaxis, ..., np.newaxis] fix_nii = nib.load(fixed) fix = fix_nii.get_data()[np.newaxis, ..., np.newaxis] with tf.device(gpu): # load model custom_objects = { 'SpatialTransformer': nrn_layers.SpatialTransformer, 'VecInt': nrn_layers.VecInt, 'Sample': networks.Sample, 'Rescale': networks.RescaleDouble, 'Resize': networks.ResizeDouble, 'Negate': networks.Negate, 'recon_loss': losses.Miccai2018(0.02, 10).recon_loss, # values shouldn't matter 'kl_loss': losses.Miccai2018(0.02, 10).kl_loss # values shouldn't matter } net = keras.models.load_model(model_file, custom_objects=custom_objects) # register [moved, warp] = net.predict([mov, fix]) # output image if out_img is not None: img = nib.Nifti1Image(moved[0, ..., 0], mov_nii.affine) nib.save(img, out_img) # output warp if out_warp is not None: img = nib.Nifti1Image(warp[0, ...], mov_nii.affine) nib.save(img, out_warp)
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 train(data_dir, atlas_file, model_dir, model, gpu_id, lr, nb_epochs, prior_lambda, image_sigma, mean_lambda, steps_per_epoch, batch_size, load_model_file, bidir, atlas_wt, bias_mult, smooth_pen_layer, data_loss, reg_param, ncc_win, initial_epoch=0): """ model training function :param data_dir: folder with npz files for each subject. :param atlas_file: atlas filename. So far we support npz file with a 'vol' variable :param model_dir: model folder to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param nb_epochs: number of training iterations :param prior_lambda: the prior_lambda, the scalar in front of the smoothing laplacian, in MICCAI paper :param image_sigma: the image sigma in MICCAI paper :param steps_per_epoch: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size :param load_model_file: optional h5 model file to initialize with :param bidir: logical whether to use bidirectional cost function """ # prepare data files # we have data arranged in train/validate/test folders # inside each folder is a /vols/ and a /asegs/ folder with the volumes # and segmentations. All of our papers use npz formated data. train_vol_names = glob.glob(data_dir) train_vol_names = [f for f in train_vol_names if 'ADNI' not in f] random.shuffle(train_vol_names) # shuffle volume list assert len(train_vol_names) > 0, "Could not find any training data" # data generator train_example_gen = datagenerators.example_gen(train_vol_names, batch_size=batch_size) # prepare the initial weights for the atlas "layer" if atlas_file is None or atlas_file == "": nb_atl_creation = 100 print('creating "atlas" by averaging %d subjects' % nb_atl_creation) x_avg = 0 for _ in range(nb_atl_creation): x_avg += next(train_example_gen)[0][0,...,0] x_avg /= nb_atl_creation x_avg = x_avg[np.newaxis,...,np.newaxis] atlas_vol = x_avg else: atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis] vol_size = atlas_vol.shape[1:-1] # Diffeomorphic network architecture used in MICCAI 2018 paper nf_enc = [16,32,32,32] nf_dec = [32,32,32,32,16,3] if model == 'm1': pass elif model == 'm1double': nf_enc = [f*2 for f in nf_enc] nf_dec = [f*2 for f in nf_dec] # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) assert data_loss in ['mse', 'cc', 'ncc'], 'Loss should be one of mse or cc, found %s' % data_loss if data_loss in ['ncc', 'cc']: data_loss = losses.NCC(win=[ncc_win]*3).loss else: data_loss = lambda y_t, y_p: K.mean(K.square(y_t-y_p)) # gpu handling gpu = '/gpu:' + str(gpu_id) os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # prepare the model with tf.device(gpu): # the MICCAI201 model takes in [image_1, image_2] and outputs [warped_image_1, velocity_stats] # in these experiments, we use image_2 as atlas model = networks.img_atlas_diff_model(vol_size, nf_enc, nf_dec, atl_mult=1, bidir=bidir, smooth_pen_layer=smooth_pen_layer) # compile mean_layer_loss = lambda _, y_pred: mean_lambda * K.mean(K.square(y_pred)) flow_vol_shape = model.outputs[-2].shape[1:-1] loss_class = losses.Miccai2018(image_sigma, prior_lambda, flow_vol_shape=flow_vol_shape) if bidir: model_losses = [data_loss, lambda _,y_p: data_loss(model.get_layer('atlas').output, y_p), mean_layer_loss, losses.Grad('l2').loss] loss_weights = [atlas_wt, 1-atlas_wt, 1, reg_param] else: model_losses = [loss_class.recon_loss, loss_class.kl_loss, mean_layer_loss] loss_weights = [1, 1, 1] model.compile(optimizer=Adam(lr=lr), loss=model_losses, loss_weights=loss_weights) # set initial weights in model model.get_layer('atlas').set_weights([atlas_vol[0,...]]) # load initial weights. # note this overloads the img_param weights if load_model_file is not None and len(load_model_file) > 0: model.load_weights(load_model_file, by_name=True) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # atlas_generator specific to this model. Once we're convinced of this, move to datagenerators def atl_gen(gen): zero_flow = np.zeros([batch_size, *vol_size, len(vol_size)]) zero_flow_half = np.zeros([batch_size] + [f//2 for f in vol_size] + [len(vol_size)]) while 1: x2 = next(train_example_gen)[0] # TODO: note this is the opposite of train_miccai and it might be confusing. yield ([atlas_vol, x2], [x2, atlas_vol, zero_flow, zero_flow]) atlas_gen = atl_gen(train_example_gen) # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') save_callback = ModelCheckpoint(save_file_name) # fit generator with tf.device(gpu): model.fit_generator(atlas_gen, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback], steps_per_epoch=steps_per_epoch, verbose=1)
def train(data_dir, atlas_file, model_dir, gpu_id, lr, nb_epochs, prior_lambda, image_sigma, steps_per_epoch, batch_size, load_model_file, bidir, bool_cc, initial_epoch=0): """ model training function :param data_dir: folder with npz files for each subject. :param atlas_file: atlas filename. So far we support npz file with a 'vol' variable :param model_dir: model folder to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param nb_epochs: number of training iterations :param prior_lambda: the prior_lambda, the scalar in front of the smoothing laplacian, in MICCAI paper :param image_sigma: the image sigma in MICCAI paper :param steps_per_epoch: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size :param load_model_file: optional h5 model file to initialize with :param bidir: logical whether to use bidirectional cost function :param bool_cc: Train CC or MICCAI version """ # load atlas from provided files. The atlas we used is 160x192x224. #atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis] vm_dir = '/home/jdram/voxelmorph/' base = np.load( os.path.join(vm_dir, "data", "ts12_dan_a88_fin_o_trim_adpc_002661_256.npy")) monitor = np.load( os.path.join(vm_dir, "data", "ts12_dan_a05_fin_o_trim_adpc_002682_256.npy")) #base = np.load(os.path.join(vm_dir, "data","ts12_dan_a88_fin_o_trim_adpc_002661_abs.npy")) #monitor = np.load(os.path.join(vm_dir, "data","ts12_dan_a05_fin_o_trim_adpc_002682_abs.npy")) #vol_size = (64, 64, 64) vol_size = (64, 64, 256 - 64) #vol_size = (128, 128, 256) # prepare data files # for the CVPR and MICCAI papers, we have data arranged in train/validate/test folders # inside each folder is a /vols/ and a /asegs/ folder with the volumes # and segmentations. All of our papers use npz formated data. #train_vol_names = glob.glob(os.path.join(data_dir, '*.npy')) #random.shuffle(train_vol_names) # shuffle volume list #assert len(train_vol_names) > 0, "Could not find any training data" # Diffeomorphic network architecture used in MICCAI 2018 paper nf_enc = [32, 64, 64, 64] nf_dec = [64, 64, 64, 64, 32, 3] # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) tf.reset_default_graph() if bool_cc: pre_net = "cc_" else: if bidir: pre_net = "miccai_bidir_" else: pre_net = "miccai_" # gpu handling gpu = '/device:GPU:%d' % int(gpu_id) # gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # prepare the model with tf.device(gpu): # prepare the model # in the CVPR layout, the model takes in [image_1, image_2] and outputs [warped_image_1, flow] # in the experiments, we use image_2 as atlas if bool_cc: model = networks.cvpr2018_net(vol_size, nf_enc, nf_dec) else: model = networks.miccai2018_net(vol_size, nf_enc, nf_dec, bidir=bidir, vel_resize=.5) # load initial weights if load_model_file is not None and load_model_file != "": print('loading', load_model_file) model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, f'{pre_net}{initial_epoch:02d}.h5')) model.summary() if bool_cc: model_losses = [losses.NCC().loss, losses.Grad('l2').loss] loss_weights = [1.0, 0.01] # recommend 1.0 for ncc, 0.01 for mse else: flow_vol_shape = model.outputs[-1].shape[1:-1] loss_class = losses.Miccai2018(image_sigma, prior_lambda, flow_vol_shape=flow_vol_shape) if bidir: model_losses = [ loss_class.recon_loss, loss_class.recon_loss, loss_class.kl_loss ] loss_weights = [0.5, 0.5, 1] else: model_losses = [loss_class.recon_loss, loss_class.kl_loss] loss_weights = [1, 1] segy_gen = datagenerators.segy_gen(base, monitor, batch_size=batch_size) # prepare callbacks save_file_name = os.path.join(model_dir, pre_net + '{epoch:02d}.h5') with tf.device(gpu): # fit generator save_callback = ModelCheckpoint(save_file_name, period=5) csv_cb = CSVLogger(f'{pre_net}log.csv') nan_cb = TerminateOnNaN() rlr_cb = ReduceLROnPlateau(monitor='loss', verbose=1) els_cb = EarlyStopping(monitor='loss', patience=15, verbose=1, restore_best_weights=True) cbs = [save_callback, csv_cb, nan_cb, rlr_cb, els_cb] mg_model = model # compile mg_model.compile(optimizer=Adam(lr=lr), loss=model_losses, loss_weights=loss_weights) mg_model.fit( [base, monitor], [monitor, np.zeros_like(base)], initial_epoch=initial_epoch, batch_size=8, epochs=nb_epochs, callbacks=cbs, #steps_per_epoch=steps_per_epoch, verbose=1)
def train(data_dir, 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 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 model :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 """ vol_size = (160, 192, 224) # Fisrt 30 images are used as training set, and the rest 10 images are testing set. constrain = list(range(31, 41, 1)) moving_path = [] list_dirs = os.walk(data_dir) for root, dirs, files in list_dirs: for f in files: if f.endswith(".hdr") and f.startswith("l"): for c in constrain: if "l{}_".format(str(c)) in str(f) or "_l{}.".format( str(c)) in str(f): break else: moving_path.append(os.path.join(data_dir, f)) train_vol_names = moving_path random.shuffle(train_vol_names) # shuffle volume list assert len(train_vol_names) > 0, "Could not find any training data" 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): model = networks_lpba40.miccai2018_net(vol_size, nf_enc, nf_dec, bidir=bidir) # load initial weights if load_model_file is not None and load_model_file != "": print("load file from {}".format(load_model_file)) model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # compile 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_class.kl_loss_1, loss_class.kl_loss_2, loss_class.kl_loss_3, loss_class.mse_loss, loss_class.mse_loss, loss_class.mse_loss ] loss_weights = [1, 1, 0.5, 0.5, 0.5, 1, 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_lpba40.example_gen_lpba40( train_vol_names, batch_size=batch_size, image_path=data_dir) miccai2018_gen_lpba40 = datagenerators_lpba40.miccai2018_gen_lpba40( train_example_gen, vol_size, 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_lpba40, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback], steps_per_epoch=steps_per_epoch, verbose=1)
def train(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 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 """ # Diffeomorphic network architecture used in VoxelMorph 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): model = networks.Reg_Net(vol_size, nf_enc, nf_dec, bidir=bidir) flow = model.get_layer('flow-int').output [src, tgt] = model.inputs [y, flow_parameters] = model.outputs model = Model(inputs=[src, tgt], outputs=[y, flow_parameters, flow]) # 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[-2].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, losses_user.Sparse_Loss( tissue=src, loss_weights=1, ).loss ] loss_weights = [1, 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) data_gen = datagenerators.data_generator_vertices(train_vol_names, vol_size) # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') log_dir = "./logs/" + model_dir[9:] + datetime.datetime.now().strftime( "%Y%m%d-%H%M%S") tensorboard_callback = keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=0) # 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(data_gen, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback, tensorboard_callback], steps_per_epoch=steps_per_epoch, verbose=1)
def register(gpu_id, moving, fixed, model_file, out_img, out_warp): """ register moving and fixed. """ assert model_file, "A model file is necessary" assert out_img or out_warp, "output image or warp file needs to be specified" # GPU handling if gpu_id is not None: 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)) else: gpu = '/cpu:0' # load data if_normalize_moving = True mov_nii = nib.load(moving) mov = mov_nii.get_data()[np.newaxis, ..., np.newaxis] if if_normalize_moving: # Typically the CTs that we used for other stuff are not clipped+normalized, that is required # hard-coded range and norm, the same as in "convert_niftis_to_npz.py" clip_range = [0, 100] # HU mov_orig = mov mov[mov < clip_range[0]] = clip_range[0] mov[mov > clip_range[1]] = clip_range[1] mov = mov / np.max(mov) fix_nii = nib.load(fixed) fix = fix_nii.get_data()[np.newaxis, ..., np.newaxis] print('\n===\nLoading "moving" data (as in your input CT) from {}'.format( moving)) print(' mov.shape:', mov.shape) print(' mov (input):', np.amin(mov_orig), ',', np.amax(mov_orig), ']') print(' mov (intensity range, after clipping+normalization):', np.amin(mov), ',', np.amax(mov), ']') print('Loading "fixed" data (as in your CT atlas/template) from {}'.format( fixed)) print(' fix.shape:', fix.shape) print(' fix (intensity range):', np.amin(fix), ',', np.amax(fix), ']') image_sigma = 0.02 prior_lambda = 10 flow_vol_shape = None # Petteri: this was the default # PETTERI: Without this, not exactly sure why the shape needs to be 1/2 of the input shape # TODO! check from paper, or inspect code a bit better flow_vol_shape = (fix.shape[1] / 2, fix.shape[2] / 2, fix.shape[3] / 2) with tf.device(gpu): # load model custom_objects = { 'SpatialTransformer': nrn_layers.SpatialTransformer, 'VecInt': nrn_layers.VecInt, 'Sample': networks.Sample, 'Rescale': networks.RescaleDouble, 'Resize': networks.ResizeDouble, 'Negate': networks.Negate, 'recon_loss': losses.Miccai2018(image_sigma, prior_lambda, flow_vol_shape=flow_vol_shape ).recon_loss, # values shouldn't matter 'kl_loss': losses.Miccai2018(image_sigma, prior_lambda, flow_vol_shape=flow_vol_shape). kl_loss # values shouldn't matter } print('Loading pre-trained model for registration from {}'.format( model_file)) print('\n===\n') net = keras.models.load_model(model_file, custom_objects=custom_objects) # register [moved, warp] = net.predict([mov, fix]) # output image if out_img is not None: img = nib.Nifti1Image(moved[0, ..., 0], mov_nii.affine) nib.save(img, out_img) # output warp if out_warp is not None: img = nib.Nifti1Image(warp[0, ...], mov_nii.affine) nib.save(img, out_warp)