def train(model, train_generator, callbacks, learning_rate, number_of_epochs, regularization_parameter, steps_per_epoch, loss_function='mse'): gpu_options = tf.GPUOptions(allow_growth=True) sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options, log_device_placement=False)) K.set_session(sess) model.compile(optimizer=Adam(lr=learning_rate), loss=[loss_function, losses.Grad('l2').loss], loss_weights=[1.0, regularization_parameter]) model.fit_generator(train_generator, initial_epoch=0, epochs=number_of_epochs, callbacks=callbacks, steps_per_epoch=steps_per_epoch, verbose=1)
def train(data_dir, atlas_file, model_dir, model, gpu_id, lr, nb_epochs, prior_lambda, image_sigma, mean_lambda, steps_per_epoch, batch_size, load_model_file, bidir, atlas_wt, bias_mult, smooth_pen_layer, data_loss, reg_param, ncc_win, initial_epoch=0): """ model training function :param data_dir: folder with npz files for each subject. :param atlas_file: atlas filename. So far we support npz file with a 'vol' variable :param model_dir: model folder to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param nb_epochs: number of training iterations :param prior_lambda: the prior_lambda, the scalar in front of the smoothing laplacian, in MICCAI paper :param image_sigma: the image sigma in MICCAI paper :param steps_per_epoch: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size :param load_model_file: optional h5 model file to initialize with :param bidir: logical whether to use bidirectional cost function """ # prepare data files # we have data arranged in train/validate/test folders # inside each folder is a /vols/ and a /asegs/ folder with the volumes # and segmentations. All of our papers use npz formated data. train_vol_names = glob.glob(data_dir) train_vol_names = [f for f in train_vol_names if 'ADNI' not in f] random.shuffle(train_vol_names) # shuffle volume list assert len(train_vol_names) > 0, "Could not find any training data" # prepare data generation specific to conditional templates train_vols_dir = os.path.join(data_dir, 'vols') train_vol_basenames = [ f for f in os.listdir(train_vols_dir) if ('ADNI' in f or 'ABIDE' in f) ] train_vol_names = [ os.path.join(train_vols_dir, f) for f in train_vol_basenames ] # csv pruning # our csv is a file of the form: # file,age,sex # ADNI_ADNI-1.5T-FS-5.3-Long_63196.long.094_S_1314_base_mri_talairach_norm.npz,81.0,2 csv_file_path = os.path.join(data_dir, 'combined_pheno.csv') train_atr_dct = load_pheno_csv(csv_file_path) train_vol_basenames = [ f for f in train_vol_basenames if f in list(train_atr_dct.keys()) ] train_vol_names = [ os.path.join(train_vols_dir, f) for f in train_vol_basenames ] # replace keys with full path for key in list(train_atr_dct.keys()): if key in train_vol_basenames: train_atr_dct[os.path.join(data_dir, 'vols', key)] = train_atr_dct[key] train_atr_dct.pop(key, None) # prepare the initial weights for the atlas "layer" if atlas_file is None or atlas_file == "": nb_atl_creation = 100 print('creating "atlas" by averaging %d subjects' % nb_atl_creation) x_avg = 0 for _ in range(nb_atl_creation): x_avg += next(train_example_gen)[0][0, ..., 0] x_avg /= nb_atl_creation x_avg = x_avg[np.newaxis, ..., np.newaxis] atlas_vol = x_avg else: atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis] vol_size = atlas_vol.shape[1:-1] genobj = Generator(train_vol_names, atlas_vol, y_k=train_atr_dct) # Diffeomorphic network architecture used in MICCAI 2018 paper nf_enc = [16, 32, 32, 32] nf_dec = [32, 32, 32, 32, 16, 3] if model == 'm1': pass elif model == 'm1double': nf_enc = [f * 2 for f in nf_enc] nf_dec = [f * 2 for f in nf_dec] # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) assert data_loss in [ 'mse', 'cc', 'ncc' ], 'Loss should be one of mse or cc, found %s' % data_loss if data_loss in ['ncc', 'cc']: data_loss = losses.NCC(win=[ncc_win] * 3).loss else: data_loss = lambda y_t, y_p: K.mean(K.square(y_t - y_p)) # gpu handling gpu = '/gpu:' + str(gpu_id) os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # prepare the model with tf.device(gpu): # parameters for atlas construction. nb_conv_features = 4 cond_im_input_shape = [160, 192, 224, 4] # note the 8 features cond_nb_levels = 0 cond_conv_size = [3, 3, 3] extra_conv_layers = 3 pheno_input_shape = [2] bidir = True smooth_pen_layer = 'diffflow' model, mn = networks.cond_img_atlas_diff_model( vol_size, nf_enc, nf_dec, atl_mult=1, bidir=bidir, smooth_pen_layer=smooth_pen_layer, nb_conv_features=nb_conv_features, cond_im_input_shape=cond_im_input_shape, cond_nb_levels=cond_nb_levels, cond_conv_size=cond_conv_size, pheno_input_shape=pheno_input_shape, extra_conv_layers=extra_conv_layers, ret_vm=True, ) outputs = [model.outputs[f] for f in [0, 2, 3, 3]] # latest model used in paper model = keras.models.Model(model.inputs, outputs) # compile mean_layer_loss = lambda _, y_pred: mean_lambda * K.mean( K.square(y_pred)) model_losses = [ data_loss, # could be mse or ncc or a combination, etc mean_layer_loss, lambda _, yp: losses.Grad('l2').loss(_, yp), lambda _, yp: K.mean(K.square(yp)) ] # parameters used in paper msmag_param = 0.01 reg_param = 1 centroid_reg = 1 loss_weights = [atlas_wt, centroid_reg, reg_param, msmag_param] model.compile(optimizer=keras.optimizers.Adam(lr=lr), loss=model_losses, loss_weights=loss_weights) # load initial weights. # note this overloads the img_param weights if load_model_file is not None and len(load_model_file) > 0: model.load_weights(load_model_file, by_name=True) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') save_callback = ModelCheckpoint(save_file_name) # fit generator with tf.device(gpu): model.fit_generator(genobj.cond_mean_flow_x2(batch_size=batch_size), initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback], steps_per_epoch=steps_per_epoch, verbose=1)
def FAIM_loss(y_true, y_pred): return losses.Grad('l2').loss( y_true, y_pred) + 1e-5 * losses.NJ_loss(y_true, y_pred)
def train(data_dir, fixed_image, model_dir, device, lr, nb_epochs, AAN_param, steps_per_epoch, batch_size, load_model_file, initial_epoch, DLR_model): # prepare data files # inside the folder are npz files with the 'vol' and 'label'. train_vol_names = glob.glob(os.path.join(data_dir, '*.npz')) random.shuffle(train_vol_names) # shuffle volume list_ assert len(train_vol_names) > 0, "Could not find any training data" vol_size = [144, 192, 160] # load atlas from provided files, if atlas-based registration if fixed_image != './': fixed_vol = np.load(fixed_image)['vol'][np.newaxis, ..., np.newaxis] def FAIM_loss(y_true, y_pred): return losses.Grad('l2').loss( y_true, y_pred) + 1e-5 * losses.NJ_loss(y_true, y_pred) assert DLR_model in [ 'VM', 'FAIM' ], 'DLR_model should be one of VM or FAIM, found %s' % LBR_model if DLR_model == 'FAIM': reg_loss = FAIM_loss else: reg_loss = losses.Grad('l2').loss # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) # device handling if 'gpu' in device: if '0' in device: device = '/gpu:0' if '1' in device: device = '/gpu:1' config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) else: device = '/cpu:0' # prepare the model with tf.device(device): model = networks.AAN(vol_size, DLR_model) # load initial weights if load_model_file != './': print('loading', load_model_file) model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # data generator train_example_gen = datagenerators.example_gen(train_vol_names, batch_size=batch_size, return_boundary=True) if fixed_image != './': fixed_vol_bs = np.repeat(fixed_vol, batch_size, axis=0) data_gen = datagenerators.gen_atlas(train_example_gen, fixing_vol_bs, batch_size=batch_size) else: data_gen = datagenerators.gen_s2s(train_example_gen, batch_size=batch_size) # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') # fit generator with tf.device(device): save_callback = ModelCheckpoint(save_file_name) # compile model.compile( optimizer=Adam(lr=lr), loss=['mse', reg_loss, losses.Grad('l1').loss_with_boundary], loss_weights=[1.0, 0.01, AAN_param]) # fit model.fit_generator(data_gen, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback], steps_per_epoch=steps_per_epoch, verbose=1)
def train(data_dir, atlas_file, model_dir, model, gpu_id, lr, nb_epochs, prior_lambda, image_sigma, mean_lambda, steps_per_epoch, batch_size, load_model_file, bidir, atlas_wt, bias_mult, smooth_pen_layer, data_loss, reg_param, ncc_win, initial_epoch=0): """ model training function :param data_dir: folder with npz files for each subject. :param atlas_file: atlas filename. So far we support npz file with a 'vol' variable :param model_dir: model folder to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param nb_epochs: number of training iterations :param prior_lambda: the prior_lambda, the scalar in front of the smoothing laplacian, in MICCAI paper :param image_sigma: the image sigma in MICCAI paper :param steps_per_epoch: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size :param load_model_file: optional h5 model file to initialize with :param bidir: logical whether to use bidirectional cost function """ # prepare data files # we have data arranged in train/validate/test folders # inside each folder is a /vols/ and a /asegs/ folder with the volumes # and segmentations. All of our papers use npz formated data. train_vol_names = glob.glob(data_dir) train_vol_names = [f for f in train_vol_names if 'ADNI' not in f] random.shuffle(train_vol_names) # shuffle volume list assert len(train_vol_names) > 0, "Could not find any training data" # data generator train_example_gen = datagenerators.example_gen(train_vol_names, batch_size=batch_size) # prepare the initial weights for the atlas "layer" if atlas_file is None or atlas_file == "": nb_atl_creation = 100 print('creating "atlas" by averaging %d subjects' % nb_atl_creation) x_avg = 0 for _ in range(nb_atl_creation): x_avg += next(train_example_gen)[0][0,...,0] x_avg /= nb_atl_creation x_avg = x_avg[np.newaxis,...,np.newaxis] atlas_vol = x_avg else: atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis] vol_size = atlas_vol.shape[1:-1] # Diffeomorphic network architecture used in MICCAI 2018 paper nf_enc = [16,32,32,32] nf_dec = [32,32,32,32,16,3] if model == 'm1': pass elif model == 'm1double': nf_enc = [f*2 for f in nf_enc] nf_dec = [f*2 for f in nf_dec] # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) assert data_loss in ['mse', 'cc', 'ncc'], 'Loss should be one of mse or cc, found %s' % data_loss if data_loss in ['ncc', 'cc']: data_loss = losses.NCC(win=[ncc_win]*3).loss else: data_loss = lambda y_t, y_p: K.mean(K.square(y_t-y_p)) # gpu handling gpu = '/gpu:' + str(gpu_id) os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # prepare the model with tf.device(gpu): # the MICCAI201 model takes in [image_1, image_2] and outputs [warped_image_1, velocity_stats] # in these experiments, we use image_2 as atlas model = networks.img_atlas_diff_model(vol_size, nf_enc, nf_dec, atl_mult=1, bidir=bidir, smooth_pen_layer=smooth_pen_layer) # compile mean_layer_loss = lambda _, y_pred: mean_lambda * K.mean(K.square(y_pred)) flow_vol_shape = model.outputs[-2].shape[1:-1] loss_class = losses.Miccai2018(image_sigma, prior_lambda, flow_vol_shape=flow_vol_shape) if bidir: model_losses = [data_loss, lambda _,y_p: data_loss(model.get_layer('atlas').output, y_p), mean_layer_loss, losses.Grad('l2').loss] loss_weights = [atlas_wt, 1-atlas_wt, 1, reg_param] else: model_losses = [loss_class.recon_loss, loss_class.kl_loss, mean_layer_loss] loss_weights = [1, 1, 1] model.compile(optimizer=Adam(lr=lr), loss=model_losses, loss_weights=loss_weights) # set initial weights in model model.get_layer('atlas').set_weights([atlas_vol[0,...]]) # load initial weights. # note this overloads the img_param weights if load_model_file is not None and len(load_model_file) > 0: model.load_weights(load_model_file, by_name=True) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # atlas_generator specific to this model. Once we're convinced of this, move to datagenerators def atl_gen(gen): zero_flow = np.zeros([batch_size, *vol_size, len(vol_size)]) zero_flow_half = np.zeros([batch_size] + [f//2 for f in vol_size] + [len(vol_size)]) while 1: x2 = next(train_example_gen)[0] # TODO: note this is the opposite of train_miccai and it might be confusing. yield ([atlas_vol, x2], [x2, atlas_vol, zero_flow, zero_flow]) atlas_gen = atl_gen(train_example_gen) # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') save_callback = ModelCheckpoint(save_file_name) # fit generator with tf.device(gpu): model.fit_generator(atlas_gen, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback], steps_per_epoch=steps_per_epoch, verbose=1)
def train(data_dir, atlas_file, model_dir, gpu_id, lr, nb_epochs, prior_lambda, image_sigma, steps_per_epoch, batch_size, load_model_file, bidir, bool_cc, initial_epoch=0): """ model training function :param data_dir: folder with npz files for each subject. :param atlas_file: atlas filename. So far we support npz file with a 'vol' variable :param model_dir: model folder to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param nb_epochs: number of training iterations :param prior_lambda: the prior_lambda, the scalar in front of the smoothing laplacian, in MICCAI paper :param image_sigma: the image sigma in MICCAI paper :param steps_per_epoch: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size :param load_model_file: optional h5 model file to initialize with :param bidir: logical whether to use bidirectional cost function :param bool_cc: Train CC or MICCAI version """ # load atlas from provided files. The atlas we used is 160x192x224. #atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis] vm_dir = '/home/jdram/voxelmorph/' base = np.load( os.path.join(vm_dir, "data", "ts12_dan_a88_fin_o_trim_adpc_002661_256.npy")) monitor = np.load( os.path.join(vm_dir, "data", "ts12_dan_a05_fin_o_trim_adpc_002682_256.npy")) #base = np.load(os.path.join(vm_dir, "data","ts12_dan_a88_fin_o_trim_adpc_002661_abs.npy")) #monitor = np.load(os.path.join(vm_dir, "data","ts12_dan_a05_fin_o_trim_adpc_002682_abs.npy")) #vol_size = (64, 64, 64) vol_size = (64, 64, 256 - 64) #vol_size = (128, 128, 256) # prepare data files # for the CVPR and MICCAI papers, we have data arranged in train/validate/test folders # inside each folder is a /vols/ and a /asegs/ folder with the volumes # and segmentations. All of our papers use npz formated data. #train_vol_names = glob.glob(os.path.join(data_dir, '*.npy')) #random.shuffle(train_vol_names) # shuffle volume list #assert len(train_vol_names) > 0, "Could not find any training data" # Diffeomorphic network architecture used in MICCAI 2018 paper nf_enc = [32, 64, 64, 64] nf_dec = [64, 64, 64, 64, 32, 3] # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) tf.reset_default_graph() if bool_cc: pre_net = "cc_" else: if bidir: pre_net = "miccai_bidir_" else: pre_net = "miccai_" # gpu handling gpu = '/device:GPU:%d' % int(gpu_id) # gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # prepare the model with tf.device(gpu): # prepare the model # in the CVPR layout, the model takes in [image_1, image_2] and outputs [warped_image_1, flow] # in the experiments, we use image_2 as atlas if bool_cc: model = networks.cvpr2018_net(vol_size, nf_enc, nf_dec) else: model = networks.miccai2018_net(vol_size, nf_enc, nf_dec, bidir=bidir, vel_resize=.5) # load initial weights if load_model_file is not None and load_model_file != "": print('loading', load_model_file) model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, f'{pre_net}{initial_epoch:02d}.h5')) model.summary() if bool_cc: model_losses = [losses.NCC().loss, losses.Grad('l2').loss] loss_weights = [1.0, 0.01] # recommend 1.0 for ncc, 0.01 for mse else: flow_vol_shape = model.outputs[-1].shape[1:-1] loss_class = losses.Miccai2018(image_sigma, prior_lambda, flow_vol_shape=flow_vol_shape) if bidir: model_losses = [ loss_class.recon_loss, loss_class.recon_loss, loss_class.kl_loss ] loss_weights = [0.5, 0.5, 1] else: model_losses = [loss_class.recon_loss, loss_class.kl_loss] loss_weights = [1, 1] segy_gen = datagenerators.segy_gen(base, monitor, batch_size=batch_size) # prepare callbacks save_file_name = os.path.join(model_dir, pre_net + '{epoch:02d}.h5') with tf.device(gpu): # fit generator save_callback = ModelCheckpoint(save_file_name, period=5) csv_cb = CSVLogger(f'{pre_net}log.csv') nan_cb = TerminateOnNaN() rlr_cb = ReduceLROnPlateau(monitor='loss', verbose=1) els_cb = EarlyStopping(monitor='loss', patience=15, verbose=1, restore_best_weights=True) cbs = [save_callback, csv_cb, nan_cb, rlr_cb, els_cb] mg_model = model # compile mg_model.compile(optimizer=Adam(lr=lr), loss=model_losses, loss_weights=loss_weights) mg_model.fit( [base, monitor], [monitor, np.zeros_like(base)], initial_epoch=initial_epoch, batch_size=8, epochs=nb_epochs, callbacks=cbs, #steps_per_epoch=steps_per_epoch, verbose=1)
def train(data_dir, atlas_file, model, model_name, gpu_id, lr, nb_epochs, reg_param, steps_per_epoch, batch_size, load_model_file, data_loss, initial_epoch=0): """ model training function :param data_dir: folder with npz files for each subject. :param atlas_file: atlas filename. So far we support npz file with a 'vol' variable :param model: either vm1 or vm2 (based on CVPR 2018 paper) :param model_dir: the model directory to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param n_iterations: number of training iterations :param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper) :param steps_per_epoch: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size :param load_model_file: optional h5 model file to initialize with :param data_loss: data_loss: 'mse' or 'ncc """ # load atlas from provided files. The atlas we used is 160x192x224. # atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis] atlas_vol = nib.load(atlas_file).get_data()[np.newaxis, ..., np.newaxis] vol_size = atlas_vol.shape[1:-1] # prepare data files # for the CVPR and MICCAI papers, we have data arranged in train/validate/test folders # inside each folder is a /vols/ and a /asegs/ folder with the volumes # and segmentations. All of our papers use npz formated data. train_vol_names = glob.glob(os.path.join(data_dir, '*.npz')) random.shuffle(train_vol_names) # shuffle volume list assert len(train_vol_names) > 0, "Could not find any training data" # UNET filters for voxelmorph-1 and voxelmorph-2, # these are architectures presented in CVPR 2018 nf_enc = [16, 32, 32, 32] if model == 'vm1': nf_dec = [32, 32, 32, 32, 8, 8] elif model == 'vm2': nf_dec = [32, 32, 32, 32, 32, 16, 16] else: # 'vm2double': nf_enc = [f * 2 for f in nf_enc] nf_dec = [f * 2 for f in [32, 32, 32, 32, 32, 16, 16]] assert data_loss in [ 'mse', 'cc', 'ncc' ], 'Loss should be one of mse or cc, found %s' % data_loss if data_loss in ['ncc', 'cc']: data_loss = losses.NCC().loss model_dir = "../models/" + model_name # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) # GPU handling gpu = '/gpu:%d' % gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # prepare the model with tf.device(gpu): # prepare the model # in the CVPR layout, the model takes in [image_1, image_2] and outputs [warped_image_1, flow] # in the experiments, we use image_2 as atlas model = networks.cvpr2018_net(vol_size, nf_enc, nf_dec) # load initial weights if load_model_file is not None and load_model_file != '': print('loading', load_model_file) model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # data generator # nb_gpus = len(gpu_id.split(',')) # assert np.mod(batch_size, nb_gpus) == 0, \ # 'batch_size should be a multiple of the nr. of gpus. ' + \ # 'Got batch_size %d, %d gpus' % (batch_size, nb_gpus) nb_gpus = 1 train_example_gen = datagenerators.example_gen(train_vol_names, batch_size=batch_size) atlas_vol_bs = np.repeat(atlas_vol, batch_size, axis=0) cvpr2018_gen = datagenerators.cvpr2018_gen(train_example_gen, atlas_vol_bs, batch_size=batch_size) # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') # fit generator with tf.device(gpu): # multi-gpu support if nb_gpus > 1: save_callback = nrn_gen.ModelCheckpointParallel(save_file_name) mg_model = multi_gpu_model(model, gpus=nb_gpus) # single-gpu else: save_callback = ModelCheckpoint(save_file_name) mg_model = model # compile mg_model.compile(optimizer=Adam(lr=lr), loss=[data_loss, losses.Grad('l2').loss], loss_weights=[1.0, reg_param]) # fit mg_model.fit_generator(cvpr2018_gen, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback], steps_per_epoch=steps_per_epoch, verbose=1)
def test(root_dir, fixed_dir, moving_dir, model_dir, gpu_id, img_size): # # 定义模型,加载权重 # 输入维度 #GPU handling gpu = '/gpu:%d' % 0 # gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) ndims = 2 vol_shape = (img_size, img_size) nb_enc_features = [32, 32, 32, 32] # 下采样卷积核个数 nb_dec_features = [32, 32, 32, 32, 32, 16] # 上采样卷积核个数 # 网络定义U-net unet = networks.unet_core(vol_shape, nb_enc_features, nb_dec_features) # 输入 print('numer of inputs', len(unet.inputs)) moving_input_tensor = unet.inputs[0] fixed_input_tensor = unet.inputs[1] # 输出 print('output:', unet.output) # 转换为流场维度 disp_tensor = keras.layers.Conv2D(ndims, kernel_size=3, padding='same', name='disp')(unet.output) # 显示流场维度 print('displacement tensor:', disp_tensor) spatial_transformer = neuron.layers.SpatialTransformer( name='spatial_transformer') # 扭转图像 moved_image_tensor = spatial_transformer( [moving_input_tensor, disp_tensor]) inputs = [moving_input_tensor, fixed_input_tensor] outputs = [moved_image_tensor, disp_tensor] vxm_model = keras.models.Model(inputs, outputs) # losses. Keras recognizes the string 'mse' as mean squared error, so we don't have to code it #loss = ['mse', losses.Grad('l2').loss] loss = [losses.NCC().loss, losses.Grad('l2').loss] # 损失函数 lambda_param = 0.01 loss_weights = [1, lambda_param] #---------------加载模型权重------------------------- vxm_model.compile(optimizer='Adam', loss=loss, loss_weights=loss_weights) vxm_model.load_weights(model_dir) #--------------前向推理------------------------------------ fixed_vol_names = glob.glob(os.path.join(fixed_dir, '*.png')) moving_vol_names = glob.glob(os.path.join(moving_dir, '*.png')) print(fixed_vol_names, moving_vol_names) # fixed_vol_names.sort() # moving_vol_names.sort() data = datagenerators.my_data_generator(fixed_vol_names, moving_vol_names, batch_size=1, img_size=img_size) sample, _ = next(data) print('输入维度:', sample[0].shape) sample_pred = vxm_model.predict(sample) #---------------保存流场维度数据------------------------------ print('流场输出维度:', sample_pred[1].shape) #a=sample_pred[0].squeeze() slices_in = [sample_pred[1].squeeze()] np.save(root_dir + 'flow.npy', slices_in) u, v = slices_in[0][..., 0], slices_in[0][..., 1] #u, v = flow_smooth.smooth(u,v,1) # np.savetxt(root_dir+'a.csv',a,delimiter=",") np.savetxt(root_dir + 'u.csv', u, delimiter=",") np.savetxt(root_dir + 'v.csv', v, delimiter=",") # 保存偏移值 #------------输出配准点的列表-------txt--------------------------- txt = open(root_dir + 'match.txt', 'w') moving_img = Image.open(moving_vol_names[0]) # moving_pixeles=str(moving_vol_names[0]).split('_') # fixed_pixeles=str(fixed_vol_names[0]).split('_') # print(moving_pixeles) # print(fixed_pixeles) # x_moving=float(moving_pixeles[2]) # y_moving=float(moving_pixeles[3].strip('.png')) # 读取文件名中左上角初始坐标 # x_fixed=float(fixed_pixeles[2]) # y_fixed=float(fixed_pixeles[3].strip('.png')) # width,height=moving_img.size # print('待配准图片长宽:',width,height) # w_scale=width/img_size # h_scale=height/img_size # for i in range(img_size): # for j in range(img_size): # x1=str(i*w_scale+x_moving) # 坐标转换 # y1=str(j*h_scale+y_moving) # x2=str(i+u[i][j]+x_fixed) # y2=str(j-v[i][j]+y_fixed) # txt.write(x1+' '+y1+' '+x2+' '+y2) # txt.write('\n') width, height = moving_img.size print('待配准图片长宽:', width, height) w_scale = width / img_size h_scale = height / img_size for i in range(img_size): for j in range(img_size): # if u[i][j]!=0: x1 = str(i * w_scale) # 坐标转换 y1 = str(j * h_scale) x2 = str(i + u[i][j]) y2 = str(j - v[i][j]) txt.write(x1 + ' ' + y1 + ' ' + x2 + ' ' + y2) txt.write('\n') #---------------可视化配准图像和流场----------------------------- slices_2d = [f[0, ..., 0] for f in sample + sample_pred] # fixed=sample[1].squeeze() # moving=sample[0].squeeze() warped = sample_pred[0].squeeze() #print(warped.shape) plt.imshow(warped, cmap='Greys_r') plt.subplots_adjust(top=1, bottom=0, left=0, right=1, hspace=0, wspace=0) plt.margins(0, 0) plt.axis('off') plt.savefig(root_dir + 'warped_picture.png', transparent=True, dpi=300, pad_inches=0.0) # plt.show() titles = [ 'input_moving', 'input_fixed', 'predicted_moved', 'deformation_x' ] neuron.plot.slices(slices_2d, titles=titles, cmaps=['gray'], do_colorbars=True, path=root_dir + 'predict.png') neuron.plot.flow([sample_pred[1].squeeze()], width=5, path=root_dir + 'flow.png')
def train( data_dir, val_data_dir, atlas_file, val_atlas_file, model, model_dir, gpu_id, lr, nb_epochs, reg_param, gama_param, steps_per_epoch, batch_size, load_model_file, data_loss, seg_dir=None, # one file val_seg_dir=None, Sf_file=None, # one file val_Sf_file=None, auxi_label=None, initial_epoch=0): """ model training function :param data_dir: folder with npz files for each subject. :param atlas_file: atlas filename. So far we support npz file with a 'vol' variable :param model: either vm1 or vm2 (based on CVPR 2018 paper) :param model_dir: the model directory to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param n_iterations: number of training iterations :param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper) :param steps_per_epoch: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size :param load_model_file: optional h5 model file to initialize with :param data_loss: 'mse' or 'ncc :param auxi_label: whether to use auxiliary informmation during the training """ # load atlas from provided files. The atlas we used is 160x192x224. # atlas_file = 'D:/voxel/data/t064.tif' atlas = Image.open(atlas_file) # is a TiffImageFile _size is (628, 690) atlas_vol = np.array(atlas)[ np.newaxis, ..., np.newaxis] # is a ndarray, shape is (1, 690, 628, 1) # new = Image.fromarray(X) new.size is (628, 690) vol_size = atlas_vol.shape[1:-1] # (690, 628) print(vol_size) val_atlas = Image.open( val_atlas_file) # is a TiffImageFile _size is (628, 690) val_atlas_vol = np.array(val_atlas)[ np.newaxis, ..., np.newaxis] # is a ndarray, shape is (1, 690, 628, 1) # new = Image.fromarray(X) new.size is (628, 690) val_vol_size = val_atlas_vol.shape[1:-1] # (690, 628) print(val_vol_size) Sm = Image.open(seg_dir) # is a TiffImageFile _size is (628, 690) Sm_ = np.array(Sm)[np.newaxis, ..., np.newaxis] val_Sm = Image.open(val_seg_dir) # is a TiffImageFile _size is (628, 690) val_Sm_ = np.array(val_Sm)[np.newaxis, ..., np.newaxis] # prepare data files # for the CVPR and MICCAI papers, we have data arranged in train/validate/test folders # inside each folder is a /vols/ and a /asegs/ folder with the volumes # and segmentations. All of our papers use npz formated data. # data_dir = D:/voxel/data/01 train_vol_names = data_dir # glob.glob(os.path.join(data_dir, '*.tif')) # is a list contain file path(name) # random.shuffle(train_vol_names) # shuffle volume list tif assert len(train_vol_names) > 0, "Could not find any training data" val_vol_names = val_data_dir # glob.glob(os.path.join(data_dir, '*.tif')) # is a list contain file path(name) # random.shuffle(train_vol_names) # shuffle volume list tif assert len(val_vol_names) > 0, "Could not find any training data" # UNET filters for voxelmorph-1 and voxelmorph-2, # these are architectures presented in CVPR 2018 nf_enc = [16, 32, 32, 32] if model == 'vm1': nf_dec = [32, 32, 32, 32, 8, 8] elif model == 'vm2': nf_dec = [32, 32, 32, 32, 32, 16, 16] else: # 'vm2double': nf_enc = [f * 2 for f in nf_enc] nf_dec = [f * 2 for f in [32, 32, 32, 32, 32, 16, 16]] assert data_loss in [ 'mse', 'cc', 'ncc' ], 'Loss should be one of mse or cc, found %s' % data_loss if data_loss in ['ncc', 'cc']: data_loss = losses.NCC().loss if Sf_file is not None: Sf = Image.open(Sf_file) Sf_ = np.array(Sf)[np.newaxis, ..., np.newaxis] if val_Sf_file is not None: val_Sf = Image.open(val_Sf_file) val_Sf_ = np.array(val_Sf)[np.newaxis, ..., np.newaxis] # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) # GPU handling gpu = '/gpu:%d' % 0 # gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) #gpu = gpu_id # data generator nb_gpus = len(gpu_id.split(',')) # 1 assert np.mod(batch_size, nb_gpus) == 0, \ 'batch_size should be a multiple of the nr. of gpus. ' + \ 'Got batch_size %d, %d gpus' % (batch_size, nb_gpus) train_example_gen = datagenerators.example_gen( train_vol_names, batch_size=batch_size) # it is a list contain a ndarray atlas_vol_bs = np.repeat( atlas_vol, batch_size, axis=0) # is a ndarray, if batch_size is 2, shape is (2, 690, 628, 1) cvpr2018_gen = datagenerators.cvpr2018_gen(train_example_gen, atlas_vol_bs, batch_size=batch_size) val_example_gen = datagenerators.example_gen( val_vol_names, batch_size=batch_size) # it is a list contain a ndarray val_atlas_vol_bs = np.repeat( val_atlas_vol, batch_size, axis=0) # is a ndarray, if batch_size is 2, shape is (2, 690, 628, 1) val_cvpr2018_gen = datagenerators.cvpr2018_gen(val_example_gen, val_atlas_vol_bs, batch_size=batch_size) # prepare the model with tf.device(gpu): sess = tf.Session(config=tf.ConfigProto(log_device_placement=True)) # prepare the model # in the CVPR layout, the model takes in [image_1, image_2] and outputs [warped_image_1, flow] # in the experiments, we use image_2 as atlas model = networks.cvpr2018_net(vol_size, nf_enc, nf_dec) # load initial weights if load_model_file is not None: print('loading', load_model_file) model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # if auxi_label is not None: # print('yes') # loss_model= [data_loss, losses.Grad('l2').loss, losses.Lseg()._lseg(Sf_) ] ########################## # loss_weight= [1.0, reg_param, gama_param] # else: loss_model = [ data_loss, losses.Grad(gama_param, Sf_, Sm_, penalty='l2').loss ] # real gama: reg_param*gama_param loss_weight = [1.0, reg_param] # reg_param_tensor = tf.constant(5, dtype=tf.float32) metrics_2 = losses.Grad(gama_param, val_Sf_, val_Sm_, penalty='l2', flag_vali=True).loss # reg_param # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') # fit generator with tf.device(gpu): # multi-gpu support if nb_gpus > 1: save_callback = nrn_gen.ModelCheckpointParallel(save_file_name) mg_model = multi_gpu_model(model, gpus=nb_gpus) # single-gpu else: save_callback = ModelCheckpoint(save_file_name) mg_model = model # compile mg_model.compile(optimizer=Adam(lr=lr), loss=loss_model, loss_weights=loss_weight, metrics={'flow': metrics_2}) # fit history = mg_model.fit_generator(cvpr2018_gen, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback], steps_per_epoch=steps_per_epoch, validation_data=val_cvpr2018_gen, validation_steps=1, verbose=2) # plot print('model', mg_model.metrics_names) print('keys()', history.history.keys()) # print(metrics.name) plt.plot(history.history['loss']) # plt.plot(history.history['val_spatial_transformer_1_loss']) plt.title('cvpr_auxi_loss') plt.ylabel('loss') plt.xlabel('Epoch') plt.legend(['Train', 'Validation']) plt.show()
def train_unsupervised_segmentation(data_dir, atlas_file, mapping_file, model, model_dir, gpu_id, lr, nb_epochs, init_stats, reg_param, steps_per_epoch, batch_size, stat_post_warp, warp_method, load_model_file, initial_epoch=0): """ model training function :param data_dir: folder with npz files (coregistered, intensity normalized) :param atlas_file: file with probabilistic atlas (coregistered to images) :param mapping_file: file with mapping from labels to tissue types :param model: registration (voxelmorph) model: vm1, vm2, or vm2double :param model_dir: the model directory to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param nb_epochs: number of epochs :param init_stats: file with guesses for means and log-variances (vectors init_mu, init_sigma) :param reg_param: smoothness/reconstruction tradeoff parameter (lambda in the paper) :param steps_per_epoch: frequency with which to save models :param batch_size: default of 1. can be larger, depends on GPU memory and volume size :param stat_post_warp: set to 1 to use warped atlas to estimate Gaussian parameters :param warp_method: set to 'WARP' if you want to warp the atlas :param load_model_file: optional h5 model file to initialize with :param initial_epoch: initial epoch """ # load reference soft edge and corresponding mask from provided files # (we used size 160x192x224). # Also: group labels in tissue types, if necessary if mapping_file is None: atlas_vol = np.load(atlas_file)['vol_data'][np.newaxis, ...] nb_labels = atlas_vol.shape[-1] else: atlas_full = np.load(atlas_file)['vol_data'][np.newaxis, ...] mapping = np.load(mapping_file)['mapping'].astype('int').flatten() assert len(mapping) == atlas_full.shape[-1], \ 'mapping shape %d is inconsistent with atlas shape %d' % (len(mapping), atlas_full.shape[-1]) nb_labels = 1 + np.max(mapping) atlas_vol = np.zeros( [1, *atlas_full.shape[1:-1], nb_labels.astype('int')]) for j in range(np.max(mapping.shape)): atlas_vol[0, ..., mapping[j]] = atlas_vol[0, ..., mapping[j]] + atlas_full[0, ..., j] vol_size = atlas_vol.shape[1:-1] # load guesses for means and variances init_mu = np.load(init_stats)['init_mu'] init_sigma = np.load(init_stats)['init_std'] # prepare data files train_vol_names = glob.glob(os.path.join(data_dir, '*.npz')) random.shuffle(train_vol_names) assert len(train_vol_names) > 0, "Could not find any training data" # UNET filters for voxelmorph-1 and voxelmorph-2, # these are architectures presented in CVPR 2018 nf_enc = [16, 32, 32, 32] if model == 'vm1': nf_dec = [32, 32, 32, 32, 8, 8] elif model == 'vm2': nf_dec = [32, 32, 32, 32, 32, 16, 16] else: # 'vm2double': nf_enc = [f * 2 for f in nf_enc] nf_dec = [f * 2 for f in [32, 32, 32, 32, 32, 16, 16]] # prepare model and log folders if not os.path.isdir(model_dir): os.mkdir(model_dir) log_dir = os.path.join(model_dir, 'logs') if not os.path.isdir(log_dir): os.mkdir(log_dir) # GPU handling gpu = '/gpu:%d' % 0 # gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # prepare the model with tf.device(gpu): # prepare the model model = networks.cvpr2018_net_probatlas(vol_size, nf_enc, nf_dec, nb_labels, diffeomorphic=True, full_size=False, stat_post_warp=stat_post_warp, warp_method=warp_method, init_mu=init_mu, init_sigma=init_sigma) # load initial weights if load_model_file is not None: print('loading', load_model_file) model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # data generator nb_gpus = len(gpu_id.split(',')) assert np.mod(batch_size, nb_gpus) == 0, \ 'batch_size should be a multiple of the nr. of gpus. ' + \ 'Got batch_size %d, %d gpus' % (batch_size, nb_gpus) train_example_gen = datagenerators.example_gen(train_vol_names, batch_size=batch_size) atlas_vol_bs = np.repeat(atlas_vol, batch_size, axis=0) cvpr2018_gen = datagenerators.cvpr2018_gen(train_example_gen, atlas_vol_bs, batch_size=batch_size) # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') # fit generator with tf.device(gpu): # multi-gpu support if nb_gpus > 1: save_callback = nrn_gen.ModelCheckpointParallel(save_file_name) mg_model = multi_gpu_model(model, gpus=nb_gpus) # single-gpu else: save_callback = ModelCheckpoint(save_file_name) mg_model = model # tensorBoard callback tensorboard = TensorBoard(log_dir=log_dir, histogram_freq=0, write_graph=True, write_images=False) # compile loss and parameters def data_loss(_, yp): m = tf.cast(model.inputs[0] > 0, tf.float32) return -K.sum(yp * m) / K.sum(m) if warp_method != 'WARP': reg_param = 0 # compile mg_model.compile(optimizer=Adam(lr=lr), loss=[data_loss, losses.Grad('l2').loss], loss_weights=[1.0, reg_param]) # fit mg_model.fit_generator(cvpr2018_gen, initial_epoch=initial_epoch, epochs=nb_epochs, callbacks=[save_callback, tensorboard], steps_per_epoch=steps_per_epoch, verbose=1)
def train(data_dir, depth_size, model, model_dir, gpu_id, lr, nb_epochs, reg_param, steps_per_epoch, train_mode, load_model_file, data_loss, batch_size, initial_epoch=0): """ model training function :param data_dir: folder with npz files for each subject. :param atlas_file: atlas filename. So far we support npz file with a 'vol' variable :param model: either vm1 or vm2 (based on CVPR 2018 paper) :param model_dir: the model directory to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param n_iterations: number of training iterations :param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper) :param steps_per_epoch: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size :param load_model_file: optional h5 model file to initialize with :param data_loss: data_loss: 'mse' or 'ncc' """ # load data train_x, train_y = load_data(data_dir=data_dir, depth_size=depth_size, mode='train', fixed='first') # test_x, test_y = load_data(data_dir=data_dir, depth_size=depth_size, mode='test', fixed='first') vol_size = train_x[0].shape[1:-1] # (width, height, depth) # set encoder, decoder feature number nf_enc = [16, 32, 32, 32] if model == 'vm1': nf_dec = [32, 32, 32, 32, 8, 8] elif model == 'vm2': nf_dec = [32, 32, 32, 32, 32, 16, 16] else: # 'vm2double': nf_enc = [f * 2 for f in nf_enc] nf_dec = [f * 2 for f in [32, 32, 32, 32, 32, 16, 16]] # set loss function # Mean Squared Error, Cross-Correlation, Negative Cross-Correlation assert data_loss in [ 'mse', 'cc', 'ncc' ], 'Loss should be one of mse or cc, found %s' % data_loss if data_loss in ['ncc', 'cc']: data_loss = losses.NCC().loss # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) # GPU handling gpu = '/gpu:%d' % 0 # gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # prepare the model with tf.device(gpu): # in the CVPR layout, the model takes in [moving image, fixed image] and outputs [warped image, flow] model = networks.cvpr2018_net(tuple(vol_size), nf_enc, nf_dec) model.load_weights(load_model_file) # save first iteration model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # prepare callbacks save_file_name = os.path.join(model_dir, '{epoch:02d}.h5') # fit with tf.device(gpu): save_callback = ModelCheckpoint(save_file_name) mg_model = model # compile mg_model.compile(optimizer=Adam(lr=lr), loss=[data_loss, losses.Grad('l2').loss], loss_weights=[1.0, reg_param]) # fit mg_model.fit(x=train_x, y=train_y, batch_size=None, epochs=nb_epochs, verbose=1, callbacks=[save_callback], steps_per_epoch=steps_per_epoch)
def train(data_dir, model, model_dir, gpu_id, lr, nb_epochs, reg_param, steps_per_epoch, load_model_file, data_loss, window_size, batch_size): """ model training function :param data_dir: folder with npz files for each subject. :param model: either vm1 or vm2 (based on CVPR 2018 paper) :param model_dir: the model directory to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper) :param steps_per_epoch: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size :param load_model_file: optional h5 model file to initialize with :param data_loss: data_loss: 'mse' or 'ncc' """ vol_size = [256, 256, 144] # (width, height, depth) # set encoder, decoder feature number nf_enc = [16, 32, 32, 32] if model == 'vm1': nf_dec = [32, 32, 32, 32, 8, 8] elif model == 'vm2': nf_dec = [32, 32, 32, 32, 32, 16, 16] else: # 'vm2double': nf_enc = [f * 2 for f in nf_enc] nf_dec = [f * 2 for f in [32, 32, 32, 32, 32, 16, 16]] # set loss function # Mean Squared Error, Cross-Correlation, Negative Cross-Correlation assert data_loss in [ 'mse', 'cc', 'ncc' ], 'Loss should be one of mse or cc, found %s' % data_loss if data_loss in ['ncc', 'cc']: NCC = losses.NCC(win=window_size) data_loss = NCC.loss # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) # GPU handling gpu = '/gpu:%d' % 0 # gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # prepare the model with tf.device(gpu): # in the CVPR layout, the model takes in [moving image, fixed image] and outputs [warped image, flow] model = networks.cvpr2018_net(vol_size, nf_enc, nf_dec) if load_model_file is not None: model.load_weights(load_model_file) # save first iteration # model.save(os.path.join(model_dir, '%02d.h5' % initial_epoch)) # load data # path = "../../dataset/urinary" # vol_names = [filename for filename in os.listdir(data_dir) if (int(filename.split("_")[-1].split('.')[0]) < 206) and # (int(filename.split("_")[-1].split('.')[0]) not in except_list)] vol_names = [ filename for filename in os.listdir(data_dir) if int(filename.split("_")[-1].split(".")[0]) in normal ] # vol_names = [filename for filename in os.listdir(data_dir) if int(filename.split("_")[-1].split(".")[0]) in (9, 130, 128)] vol_names.sort() uro_gen = uro_generator(vol_names, data_dir, fixed='joyoungje') # test_path = os.path.join(data_dir, 'test') # test_vol_names = [filename for filename in os.listdir(test_path) if '.npz'] # test_gen = uro_generator(test_vol_names, test_path) # fit with tf.device(gpu): mg_model = model # compile mg_model.compile(optimizer=Adam(lr=lr), loss=[data_loss, losses.Grad('l2').loss], loss_weights=[1.0, reg_param]) # fit save_file_name = os.path.join(model_dir, '{epoch:04d}.h5') save_callback = ModelCheckpoint(save_file_name) mg_model.fit_generator(uro_gen, epochs=nb_epochs, verbose=1, callbacks=[save_callback], steps_per_epoch=steps_per_epoch)
def train( fixed_dir, moving_dir, fixedval_dir, movingval_dir, model_dir, gpu_id, lr, nb_epochs, img_size, lambda_param, steps_per_epoch, batch_size, ): fixed_training_file = open(fixed_dir, 'r') moving_traning_file = open(moving_dir, 'r') fixed_val_file = open(fixedval_dir, 'r') moving_val_file = open(movingval_dir, 'r') def str_strip(file): list = [] for f in file.readlines(): list.append(f.strip()) return list fixed_image_names = str_strip(fixed_training_file) moving_image_names = str_strip(moving_traning_file) fixedval_image_names = str_strip(fixed_val_file) movingval_image_names = str_strip(moving_val_file) assert len(fixed_image_names) > 0, "fixed路径中找不到训练数据" assert len(moving_image_names) > 0, "moving路径中找不到训练数据" assert len(fixedval_image_names) > 0, "fixed路径中找不到训练数据" assert len(movingval_image_names) > 0, "moving路径中找不到训练数据" # GPU handling gpu = '/gpu:%d' % 0 # gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) #-------------------------------模型定义----------------------------------- # 输入维度 ndims = 2 vol_shape = (img_size, img_size) nb_enc_features = [32, 32, 32, 32] # 下采样卷积核个数 nb_dec_features = [32, 32, 32, 32, 32, 16] # 上采样卷积核个数 # 网络定义U-net unet = networks.unet_core(vol_shape, nb_enc_features, nb_dec_features) #------网络输入-------------------- print('网络输入维度', len(unet.inputs)) moving_input_tensor = unet.inputs[0] fixed_input_tensor = unet.inputs[1] #---------网络输出---------------- print('网络输出维度:', unet.output) #-------------转换成流场维度--------------- disp_tensor = keras.layers.Conv2D(ndims, kernel_size=3, padding='same', name='disp')(unet.output) # (512,512,2) # check print('displacement tensor:', disp_tensor) #--------------------流场作用于原图扭曲------------------- spatial_transformer = neuron.layers.SpatialTransformer( name='spatial_transformer') moved_image_tensor = spatial_transformer( [moving_input_tensor, disp_tensor]) inputs = [moving_input_tensor, fixed_input_tensor] outputs = [moved_image_tensor, disp_tensor] vxm_model = keras.models.Model(inputs, outputs) # losses. Keras recognizes the string 'mse' as mean squared error, so we don't have to code it #loss = [losses.NCC().loss, losses.Grad('l2').loss] loss = ['mse', losses.Grad('l2').loss] # usually, we have to balance the two losses by a hyper-parameter. loss_weights = [1, lambda_param] vxm_model.compile(optimizer=Adam(lr=lr), loss=loss, loss_weights=loss_weights) vxm_model.summary() #----------------------准备训练数据------------------------------- x_train_path = fixed_image_names y_train_path = moving_image_names x_val_path = fixedval_image_names y_val_path = movingval_image_names train_generator = datagenerators.my_data_generator(x_train_path, y_train_path, batch_size=batch_size, img_size=img_size) valid_generator = datagenerators.my_data_generator(x_val_path, y_train_path, batch_size=1, img_size=img_size) input_sample, output_sample = next(train_generator) # ---------------可视化--------------------- slices_2d = [f[0, ..., 0] for f in input_sample + output_sample] titles = [ 'input_moving', 'input_fixed', 'output_moved_ground_truth', 'zero' ] neuron.plot.slices(slices_2d, titles=titles, cmaps=['gray'], do_colorbars=True) checkpoint = ModelCheckpoint(filepath='models/my_model{epoch:03d}.h5', monitor='loss', verbose=1, save_best_only='True', mode='min', period=1) tensorboard = TensorBoard(log_dir='logs/', histogram_freq=0, write_graph=True, write_images=True) #-------------------------开始训练--------------------------- hist = vxm_model.fit_generator(train_generator, epochs=nb_epochs, steps_per_epoch=steps_per_epoch, verbose=2, callbacks=[checkpoint, tensorboard], validation_data=valid_generator, validation_steps=2990) vxm_model.save_weights(os.path.join(model_dir, 'my_model.h5')) with open('data/SAR_model_hist.pickle', 'wb') as file_pi: pickle.dump(hist.history, file_pi) def plot_history(hist, loss_name='loss', path='data/'): """ Quick function to plot the history """ plt.figure() plt.plot(hist.epoch, hist.history[loss_name], '.-') plt.ylabel('loss') plt.xlabel('epoch') plt.savefig(path) plt.show() plot_history(hist, 'loss', 'data/loss.jpg') plot_history(hist, 'val_loss', 'data/val_loss.jpg') #----------------直接测试一张图配准-------------------------- val_generator = datagenerators.my_data_generator(x_train_path, y_train_path, batch_size=1, img_size=img_size) val_input, _ = next(val_generator) val_pred = vxm_model.predict(val_input) # ---------------可视化--------------------- slices_2d = [f[0, ..., 0] for f in val_input + val_pred] titles = [ 'input_moving', 'input_fixed', 'predicted_moved', 'deformation_x' ] neuron.plot.slices(slices_2d, titles=titles, cmaps=['gray'], do_colorbars=True) neuron.plot.flow([val_pred[1].squeeze()], width=5)
def test(root_dir, fixed_dir, moving_dir, model_dir, gpu_id, img_size): fixed_file = open(fixed_dir, 'r') moving_file = open(moving_dir, 'r') def str_strip(file): list = [] for f in file.readlines(): list.append(f.strip()) return list fixed_vol_names = str_strip(fixed_file) moving_vol_names = str_strip(moving_file) assert len(fixed_vol_names) > 0, "fixed路径中找不到训练数据" assert len(moving_vol_names) > 0, "moving路径中找不到训练数据" #GPU handling gpu = '/gpu:%d' % 0 # gpu_id os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) ndims = 2 vol_shape = (img_size, img_size) nb_enc_features = [32, 32, 32, 32] # 下采样卷积核个数 nb_dec_features = [32, 32, 32, 32, 32, 16] # 上采样卷积核个数 # 网络定义U-net unet = networks.unet_core(vol_shape, nb_enc_features, nb_dec_features) # 输入 print('numer of inputs', len(unet.inputs)) moving_input_tensor = unet.inputs[0] fixed_input_tensor = unet.inputs[1] # 输出 print('output:', unet.output) # 转换为流场维度 disp_tensor = keras.layers.Conv2D(ndims, kernel_size=3, padding='same', name='disp')(unet.output) # 显示流场维度 print('displacement tensor:', disp_tensor) spatial_transformer = neuron.layers.SpatialTransformer( name='spatial_transformer') # 扭转图像 moved_image_tensor = spatial_transformer( [moving_input_tensor, disp_tensor]) inputs = [moving_input_tensor, fixed_input_tensor] outputs = [moved_image_tensor, disp_tensor] vxm_model = keras.models.Model(inputs, outputs) # losses. Keras recognizes the string 'mse' as mean squared error, so we don't have to code it loss = ['mse', losses.Grad('l2').loss] # 损失函数 lambda_param = 0.01 loss_weights = [1, lambda_param] #---------------加载模型权重------------------------- vxm_model.compile(optimizer='Adam', loss=loss, loss_weights=loss_weights) vxm_model.load_weights(model_dir) #------------定义DICE函数------------------------------- def dice_coef(y_true, y_pred): y_true_f = y_true.flatten() # 将 y_true 拉伸为一维. y_pred_f = y_pred.flatten() intersection = sum(y_true_f * y_pred_f) return (2. * intersection) / (sum(y_true_f * y_true_f) + sum(y_pred_f * y_pred_f)) def compare_images(imageA, imageB, title): # 分别计算输入图片的MSE和SSIM指标值的大小 s = ssim(imageA, imageB) #return s # 创建figure fig = plt.figure(title) plt.suptitle("SSIM: %.2f" % (s)) # 显示第一张图片 ax = fig.add_subplot(1, 2, 1) plt.imshow(imageA, cmap=plt.cm.gray) plt.axis("off") # 显示第二张图片 ax = fig.add_subplot(1, 2, 2) plt.imshow(imageB, cmap=plt.cm.gray) plt.axis("off") plt.tight_layout() plt.show() #--------------前向推理DICE计算------------------------------------ length = len(fixed_vol_names) dice_before = 0 dice_after = 0 for i in range(300, length): data = datagenerators.my_data_generator([fixed_vol_names[i]], [moving_vol_names[i]], batch_size=1, img_size=img_size) sample, _ = next(data) sample_pred = vxm_model.predict(sample) fixed = sample[1].squeeze() moving = sample[0].squeeze() warped = sample_pred[0].squeeze() # dice_before+=dice_coef(fixed,moving) # dice_after+=dice_coef(fixed,warped) #compare_images(fixed, fixed, "基准 vs 基准") # dice_before+=compare_images(fixed, moving, "基准 vs 待配准") # dice_after+=compare_images(fixed, warped, " 基准vs 配准后") #compare_images(fixed, fixed, "基准 vs 待配准") compare_images(fixed, moving, "基准 vs 待配准") compare_images(fixed, warped, " 基准vs 配准后")