def test_networks_unet(): device = 'cuda:0' if torch.cuda.is_available() else 'cpu' model = unet(3, 1) model.to(device) assert torchsummary.summary(model, image_size)
def cremi_unet(name='unet', sample_to_isotropy=False): in_shape = (84, 268, 268) n_channels = 12 # These values reproduce jans network initial_fmaps = 12 fmap_increase = 5 downsample_factors = [[1, 3, 3], [1, 3, 3], [3, 3, 3]] if sample_to_isotropy else \ [[1, 3, 3], [1, 3, 3], [1, 3, 3]] raw = tf.placeholder(tf.float32, shape=in_shape) raw_batched = tf.reshape(raw, ( 1, 1, ) + in_shape) unet = networks.unet(raw_batched, initial_fmaps, fmap_increase, downsample_factors) affs_batched = networks.conv_pass(unet, kernel_size=1, num_fmaps=n_channels, num_repetitions=1, activation='sigmoid') output_shape_batched = affs_batched.get_shape().as_list() output_shape = output_shape_batched[1:] # strip the batch dimension affs = tf.reshape(affs_batched, output_shape) gt_affs = tf.placeholder(tf.float32, shape=output_shape) loss_weights = tf.placeholder(tf.float32, shape=output_shape) loss = tf.losses.mean_squared_error(gt_affs, affs, loss_weights) tf.summary.scalar('loss_total', loss) opt = tf.train.AdamOptimizer(learning_rate=0.5e-4, beta1=0.95, beta2=0.999, epsilon=1e-8) optimizer = opt.minimize(loss) #for trainable in tf.trainable_variables(): # networks.tf_var_summary(trainable) merged = tf.summary.merge_all() tf.train.export_meta_graph(filename='%s.meta' % name) names = { 'raw': raw.name, 'affs': affs.name, 'gt_affs': gt_affs.name, 'loss_weights': loss_weights.name, 'loss': loss.name, 'optimizer': optimizer.name, 'summary': merged.name } with open('net_io_names.json', 'w') as f: json.dump(names, f)
def train(model_name, gpu_id): model_dir = '../models/' + model_name if not os.path.isdir(model_dir): os.mkdir(model_dir) 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)) with tf.device(gpu): model = networks.unet(vol_size, nf_enc, nf_dec) model.compile(optimizer=Adam(lr=lr), loss=[ losses.cc3D(), losses.gradientLoss('l2')], loss_weights=[1.0, reg_param]) # model.load_weights('../models/udrnet2/udrnet1_1/120000.h5') train_example_gen = datagenerators.example_gen(train_vol_names) zero_flow = np.zeros((1, vol_size[0], vol_size[1], vol_size[2], 3)) for step in xrange(0, n_iterations): X = train_example_gen.next()[0] train_loss = model.train_on_batch( [X, atlas_vol], [atlas_vol, zero_flow]) if not isinstance(train_loss, list): train_loss = [train_loss] printLoss(step, 1, train_loss) if(step % model_save_iter == 0): model.save(model_dir + '/' + str(step) + '.h5')
def train(model, gpu_id, lr, n_iterations, reg_param, model_save_iter, load_iter): model_dir = '/home/ys895/MAS3_Models' if not os.path.isdir(model_dir): os.mkdir(model_dir) 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)) # UNET filters nf_enc = [16, 32, 32, 32] if (model == 'vm1'): nf_dec = [32, 32, 32, 32, 8, 8, 3] else: nf_dec = [32, 32, 32, 32, 32, 16, 16, 3] with tf.device(gpu): model = networks.unet(vol_size, nf_enc, nf_dec) if (load_iter != 0): model.load_weights('/home/ys895/MAS3_Models/' + str(load_iter) + '.h5') model.compile(optimizer=Adam(lr=lr), loss=[losses.cc3D(), losses.gradientLoss('l2')], loss_weights=[1.0, reg_param]) # model.load_weights('../models/udrnet2/udrnet1_1/120000.h5') # return the data, add one more dimension into the data train_example_gen = datagenerators.example_gen(train_vol_names) zero_flow = np.zeros((1, vol_size[0], vol_size[1], vol_size[2], 3)) # In this part, the code inputs the data into the model # Before this part, the model was set for step in range(1, n_iterations + 1): # choose randomly one of the atlas from the atlas_list rand_num = random.randint(0, list_num - 1) atlas_vol = atlas_list[rand_num] #Parameters for training : X(train_vol) ,atlas_vol(atlas) ,zero_flow X = train_example_gen.__next__()[0] train_loss = model.train_on_batch([atlas_vol, X], [X, zero_flow]) if not isinstance(train_loss, list): train_loss = [train_loss] printLoss(step, 1, train_loss) if (step % model_save_iter == 0): model.save(model_dir + '/' + str(load_iter + step) + '.h5')
def model_factory(model_name, x, dropout, is_training, weight_decay, crop, num_input_bands, num_classes, crop_size): if model_name == 'dilated_grsl': logits = dilated_grsl(x, is_training, weight_decay, crop, num_input_bands, num_classes) elif model_name == 'dilated_icpr_rate6': logits = dilated_icpr_rate6(x, is_training, weight_decay, crop, num_input_bands, num_classes) elif model_name == 'dilated_icpr_rate6_densely': logits = dilated_icpr_rate6_densely(x, is_training, weight_decay, crop, num_input_bands, num_classes) elif model_name == 'dilated_grsl_rate8': logits = dilated_grsl_rate8(x, is_training, weight_decay, crop, num_input_bands, num_classes) elif model_name == 'fcn_25_1_4x': logits = fcn_25_1_4x(x, dropout, is_training, crop, weight_decay, num_input_bands, num_classes) elif model_name == 'fcn_25_2_2x': logits = fcn_25_2_2x(x, dropout, is_training, crop, weight_decay, num_input_bands, num_classes) elif model_name == 'fcn_25_3_2x_icpr': logits = fcn_25_3_2x_icpr(x, dropout, is_training, crop, weight_decay, num_input_bands, num_classes) elif model_name == 'fcn_50_1_8x': logits = fcn_50_1_8x(x, dropout, is_training, crop, weight_decay, num_input_bands, num_classes) elif model_name == 'fcn_50_2_4x': logits = fcn_50_2_4x(x, dropout, is_training, crop, weight_decay, num_input_bands, num_classes) elif model_name == 'fcn_50_3_2x': logits = fcn_50_3_2x(x, dropout, is_training, crop, weight_decay, num_input_bands, num_classes) elif model_name == 'pixelwise': logits = pixelwise(x, dropout, is_training, weight_decay, crop, num_input_bands, num_classes) elif model_name == 'segnet': logits = segnet(x, dropout, is_training, weight_decay, crop, num_input_bands, num_classes, crop_size) elif model_name == 'segnet_4': logits = segnet_4(x, dropout, is_training, weight_decay, crop, num_input_bands, num_classes, crop_size) elif model_name == 'unet': logits = unet(x, dropout, is_training, weight_decay, crop, num_input_bands, num_classes, crop_size) elif model_name == 'deeplabv3+': logits = deeplab(x, dropout, is_training, weight_decay, crop, num_input_bands, num_classes, crop_size) else: raise NotImplementedError('Network not identified: ' + model_name) return logits
def _network_graph_def(self, verbose=False): """ defines the network graph :param verbose: if True prints messages as it defines layers :return: """ # placeholders for input and ground truth with tf.name_scope('input'): self.x = tf.placeholder(name='x', shape=(None, None, None, 1), dtype=tf.float32) self.y = tf.placeholder(name='y', shape=(None, None, None, 1), dtype=tf.float32) # by default we use batch norm in train_mode (i.e with current if self.name == 'UNET': self.sigmoided_logits, self.logits, _, _ = unet( self.x, 1, self.num_layers, self.feature_maps_root, self.norm, True, verbose=verbose) elif self.name == 'iUNET': self.sigmoided_logits, self.logits, _, _ = iunet( self.x, 1, self.num_layers, self.feature_maps_root, self.norm, True, self.n_iterations, verbose=verbose) elif self.name == 'SHN': self.sigmoided_logits, self.logits, _, _ = shn( self.x, 1, self.num_layers, self.feature_maps_root, self.norm, True, self.n_modules, verbose=verbose) self.output = self.sigmoided_logits if self.name == 'UNET' else self.sigmoided_logits[ -1] self.variables = [ v for v in tf.global_variables() if self.name in v.name ]
def test(model_name, iter_num, gpu_id, vol_size=(160,192,224), nf_enc=[16,32,32,32], nf_dec=[32,32,32,32,32,16,16,3]): """ test nf_enc and nf_dec #nf_dec = [32,32,32,32,32,16,16,3] # This needs to be changed. Ideally, we could just call load_model, and we wont have to # specify the # of channels here, but the load_model is not working with the custom loss... """ gpu = '/gpu:' + str(gpu_id) # Anatomical labels we want to evaluate labels = sio.loadmat('../data/labels.mat')['labels'][0] atlas = np.load('../data/atlas_norm.npz') atlas_vol = atlas['vol'] atlas_seg = atlas['seg'] atlas_vol = np.reshape(atlas_vol, (1,)+atlas_vol.shape+(1,)) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # load weights of model with tf.device(gpu): net = networks.unet(vol_size, nf_enc, nf_dec) net.load_weights('../models/' + model_name + '/' + str(iter_num) + '.h5') xx = np.arange(vol_size[1])#192 yy = np.arange(vol_size[0])#160 zz = np.arange(vol_size[2])#224 grid = np.rollaxis(np.array(np.meshgrid(xx, yy, zz)), 0, 4)#(160,192,224,3) it stores the co-ordinate of the original position of the point in the X_vol, X_seg = datagenerators.load_example_by_name('../data/test_vol.npz', '../data/test_seg.npz') with tf.device(gpu): pred = net.predict([X_vol, atlas_vol]) # Warp segments with flow flow = pred[1][0, :, :, :, :]#(160,192,224,3) sample = flow+grid#add the original position with the shift flow the dimension is: (160,192,224,3) sample = np.stack((sample[:, :, :, 1], sample[:, :, :, 0], sample[:, :, :, 2]), 3) warp_seg = interpn((yy, xx, zz), X_seg[0, :, :, :, 0], sample, method='nearest', bounds_error=False, fill_value=0) vals, _ = dice(warp_seg, atlas_seg, labels=labels, nargout=2) print(np.mean(vals), np.std(vals))
def test(model_name, iter_num, gpu_id, vol_size=(160,192,224), nf_enc=[16,32,32,32], nf_dec=[32,32,32,32,32,16,16,3]): """ test nf_enc and nf_dec #nf_dec = [32,32,32,32,32,16,16,3] # This needs to be changed. Ideally, we could just call load_model, and we wont have to # specify the # of channels here, but the load_model is not working with the custom loss... """ gpu = '/gpu:' + str(gpu_id) # Anatomical labels we want to evaluate labels = sio.loadmat('../data/labels.mat')['labels'][0] atlas = np.load('../data/atlas_norm.npz') atlas_vol = atlas['vol'] atlas_seg = atlas['seg'] atlas_vol = np.reshape(atlas_vol, (1,)+atlas_vol.shape+(1,)) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # load weights of model with tf.device(gpu): net = networks.unet(vol_size, nf_enc, nf_dec) net.load_weights('../models/' + model_name + '/' + str(iter_num) + '.h5') xx = np.arange(vol_size[1]) yy = np.arange(vol_size[0]) zz = np.arange(vol_size[2]) grid = np.rollaxis(np.array(np.meshgrid(xx, yy, zz)), 0, 4) X_vol, X_seg = datagenerators.load_example_by_name('../data/test_vol.npz', '../data/test_seg.npz') with tf.device(gpu): pred = net.predict([X_vol, atlas_vol]) # Warp segments with flow flow = pred[1][0, :, :, :, :] sample = flow+grid sample = np.stack((sample[:, :, :, 1], sample[:, :, :, 0], sample[:, :, :, 2]), 3) warp_seg = interpn((yy, xx, zz), X_seg[0, :, :, :, 0], sample, method='nearest', bounds_error=False, fill_value=0) vals, _ = dice(warp_seg, atlas_seg, labels=labels, nargout=2) print(np.mean(vals), np.std(vals))
def train(model,save_name, gpu_id, lr, n_iterations, reg_param, model_save_iter): model_dir = '../models/' + save_name if not os.path.isdir(model_dir): os.mkdir(model_dir) 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)) # UNET filters nf_enc = [16,32,32,32] if(model == 'vm1'): nf_dec = [32,32,32,32,8,8,3] else: nf_dec = [32,32,32,32,32,16,16,3] with tf.device(gpu): model = networks.unet(vol_size, nf_enc, nf_dec) model.compile(optimizer=Adam(lr=lr), loss=[losses.cc3D(), losses.gradientLoss('l2')], loss_weights=[1.0, reg_param]) # model.load_weights('../models/udrnet2/udrnet1_1/120000.h5') zeroflow = np.zeros((1, vol_size[0], vol_size[1], vol_size[2], 3)) for step in range(0, n_iterations): sub = np.load(train_pairs[step % (noftrain ** 2)][0]) sub = np.reshape(sub, (1,) + sub.shape + (1,)) tmp = np.load(train_pairs[step % (noftrain ** 2)][1]) tmp = np.reshape(tmp, (1,) + tmp.shape + (1,)) train_loss = model.train_on_batch([sub, tmp], [tmp, zeroflow]) printLoss(step, train_loss, keras.get_value(model.optimizer.lr)) if(step % model_save_iter == 0): model.save(model_dir + '/' + str(step) + '.h5') if(step % (2*(noftrain ** 2)) == 0 and step > 0): keras.set_value(model.optimizer.lr, keras.get_value(model.optimizer.lr) / 2)
def train(model,save_name, gpu_id, lr, n_iterations, reg_param, model_save_iter): model_dir = '../models/' + save_name if not os.path.isdir(model_dir): os.mkdir(model_dir) 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)) # UNET filters nf_enc = [16,32,32,32] if(model == 'vm1'): nf_dec = [32,32,32,32,8,8,3] else: nf_dec = [32,32,32,32,32,16,16,3] with tf.device(gpu): model = networks.unet(vol_size, nf_enc, nf_dec) model.compile(optimizer=Adam(lr=lr), loss=[ losses.cc3D(), losses.gradientLoss('l2')], loss_weights=[1.0, reg_param]) # model.load_weights('../models/udrnet2/udrnet1_1/120000.h5') train_example_gen = datagenerators.example_gen(train_vol_names) zero_flow = np.zeros((1, vol_size[0], vol_size[1], vol_size[2], 3)) for step in range(0, n_iterations): X = train_example_gen.__next__()[0] train_loss = model.train_on_batch( [X, atlas_vol], [atlas_vol, zero_flow]) if not isinstance(train_loss, list): train_loss = [train_loss] printLoss(step, 1, train_loss) if(step % model_save_iter == 0): model.save(model_dir + '/' + str(step) + '.h5')
def train(model, model_dir, gpu_id, lr, n_iterations, reg_param, model_save_iter, batch_size=1): """ model training function :param model: either vm1 or vm2 (based on CVPR 2018 paper) :param model_dir: the model directory to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param n_iterations: number of training iterations :param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper) :param model_save_iter: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size """ # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) # GPU handling os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # UNET filters for voxelmorph-1 and voxelmorph-2, # these are architectures presented in CVPR 2018 nf_enc = [16, 32, 32, 32] if model == 'vm1': nf_dec = [32, 32, 32, 32, 8, 8] else: nf_dec = [32, 32, 32, 32, 32, 16, 16] # prepare the model # in the CVPR layout, the model takes in [image_1, image_2] and outputs [warped_image_1, flow] # in the experiments, we use image_2 as atlas model = networks.unet(vol_size, nf_enc, nf_dec) model.compile(optimizer=Adam(lr=lr), loss=[losses.cc3D(), losses.gradientLoss('l2')], loss_weights=[1.0, reg_param]) # if you'd like to initialize the data, you can do it here: # model.load_weights(os.path.join(model_dir, '120000.h5')) # prepare data for training train_example_gen = datagenerators.example_gen(train_vol_names) zero_flow = np.zeros([batch_size, *vol_size, 3]) # train. Note: we use train_on_batch and design out own print function as this has enabled # faster development and debugging, but one could also use fit_generator and Keras callbacks. for step in range(0, n_iterations): # get data X = next(train_example_gen)[0] # train train_loss = model.train_on_batch([X, atlas_vol], [atlas_vol, zero_flow]) if not isinstance(train_loss, list): train_loss = [train_loss] # print the loss. print_loss(step, 1, train_loss) # save model if step % model_save_iter == 0: model.save(os.path.join(model_dir, str(step) + '.h5'))
if __name__ == "__main__": print(networks.__file__) raw = tf.placeholder(tf.float32, shape=(43, 430, 430)) raw_batched = tf.reshape(raw, ( 1, 1, ) + (43, 430, 430)) #unet = networks.unet(raw_batched, 12, 6, [[1, 3, 3], [1, 3, 3], [3, 3, 3]], anisotropy=10) unet, fov, anisotropy = networks.unet( raw_batched, 12, 6, [[1, 3, 3], [1, 3, 3], [3, 3, 3]], [[(1, 3, 3), (1, 3, 3)], [(1, 3, 3), (1, 3, 3)], [(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)]], [[(1, 3, 3), (1, 3, 3)], [(1, 3, 3), (1, 3, 3)], [(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)]], anisotropy=[10, 1, 1], fov=[10, 1, 1]) # raw = tf.placeholder(tf.float32, shape=(132,)*3) # raw_batched = tf.reshape(raw, (1, 1,) + (132,)*3) # # unet = networks.unet(raw_batched, 24, 3, [[2, 2, 2], [2, 2, 2], [2, 2, 2]]) dist_batched, fov = networks.conv_pass(unet, kernel_size=[[1, 1, 1]], num_fmaps=1, activation=None,
def test(load_iters, gpu_id, vol_size=(160,192,224), nf_enc=[16,32,32,32], nf_dec=[32,32,32,32,32,16,16,3], sample_num = 10, grid_dimension = 4): """ Test of the rigid registration by calculating the dice score between the atlas's segmentation and warped image's segmentation :param iter_num: iteration number :param gpu_id: gpu id :param vol_size: volume's size :param nf_enc: number of encode :param nf_dec: number of decoder :param model_name: load model's name :param sample_num: sample grid's dimension, this can be changed to improve the performance :param grid_dimension: R(in the formula)'s dimension :return: None """ gpu = '/gpu:' + str(gpu_id) # Anatomical labels we want to evaluate labels = sio.loadmat('../data/labels.mat')['labels'][0] atlas = np.load('../data/atlas_norm.npz') atlas_vol = atlas['vol'] atlas_seg = atlas['seg'] atlas_vol = np.reshape(atlas_vol, (1,)+atlas_vol.shape+(1,)) os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # load weights of model with tf.device(gpu): net = networks.unet(vol_size, nf_enc, nf_dec) net.load_weights('../rigid_model/' + load_iters + '.h5', by_name=True) X_vol, X_seg = datagenerators.load_example_by_name('../data/test_vol.npz', '../data/test_seg.npz') orig_vol = X_vol theta = 0 beta = 4 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,)) with tf.device(gpu): pred = net.predict([X_vol, atlas_vol]) # get flow flow = pred[1][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) #R = np.zeros((10, 10, 10, grid_dimension, grid_dimension)) 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, :] - T# 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 print(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]) shifted_grid = np.rollaxis(np.array((np.meshgrid(shifted_y, shifted_x, shifted_z))), 0, 4) 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, 1] = coordinates[0] shifted_grid[i, j, z, 0] = 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]) 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 # 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)# deformable registration # 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)) # plot #fig1, axs1 = nplt.slices(warp_seg[100, :, :], do_colorbars=True) #fig1.savefig('warp_seg100.png') #fig2, axs2 = nplt.slices(warp_seg[130, :, :], do_colorbars=True) #fig2.savefig('warp_seg130.png') #fig3, axs3 = nplt.slices(atlas_seg[100, :, :], do_colorbars=True) #fig3.savefig('atlas_seg100.png') #fig4, axs4 = nplt.slices(atlas_seg[130, :, :], do_colorbars=True) #fig4.savefig('atlas_seg130.png') # 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")
def test(model_name, iter_num, gpu_id, n_test, filename, vol_size=(160, 192, 224), nf_enc=[16, 32, 32, 32], nf_dec=[32, 32, 32, 32, 32, 16, 16]): """ test nf_enc and nf_dec #nf_dec = [32,32,32,32,32,16,16,3] # This needs to be changed. Ideally, we could just call load_model, and we wont have to # specify the # of channels here, but the load_model is not working with the custom loss... """ start_time = time.time() gpu = '/gpu:' + str(gpu_id) print(gpu) os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # Anatomical labels we want to evaluate labels = sio.loadmat('../data/labels.mat')['labels'][0] atlas = np.load('../data/atlas_norm.npz') atlas_vol = atlas['vol'] atlas_seg = atlas['seg'] atlas_vol = np.reshape(atlas_vol, (1, ) + atlas_vol.shape + (1, )) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # load weights of model with tf.device(gpu): net = networks.unet(vol_size, nf_enc, nf_dec) net.load_weights('../models/' + model_name + '/' + str(iter_num) + '.h5') seg_path = '../models/seg_pretrained/0.h5' feature_model, num_features = networks.segmenter_feature_model(seg_path) with open('seg_feature_stats.txt', 'rb') as file: feature_stats = pickle.loads( file.read()) # use `pickle.loads` to do the reverse xx = np.arange(vol_size[1]) yy = np.arange(vol_size[0]) zz = np.arange(vol_size[2]) grid = np.rollaxis(np.array(np.meshgrid(xx, yy, zz)), 0, 4) percentile = 99 dice_means = [] results = {} for step in range(0, n_test): res = {} vol_name, seg_name = test_brain_strings[step].split(",") X_vol, X_seg = datagenerators.load_example_by_name(vol_name, seg_name) with tf.device(gpu): pred = net.predict([X_vol, atlas_vol]) warped_image = np.transpose(pred[0][0, :, :, :, :], (2, 0, 1, 3)) pred_ac_features = feature_model.predict([warped_image]) orig_ac_features = feature_model.predict( [np.transpose(X_vol[0, :, :, :, :], (2, 0, 1, 3))]) # Warp segments with flow flow = pred[1][0, :, :, :, :] sample = flow + grid sample = np.stack( (sample[:, :, :, 1], sample[:, :, :, 0], sample[:, :, :, 2]), 3) warp_seg = interpn((yy, xx, zz), X_seg[0, :, :, :, 0], sample, method='nearest', bounds_error=False, fill_value=0) vals, _ = dice(warp_seg, atlas_seg, labels=labels, nargout=2) # print(np.mean(vals), np.std(vals)) mean = np.mean(vals) std = np.std(vals) res['dice_mean'] = mean res['dice_std'] = std for i in range(len(pred_ac_features)): normalized_pred = normalize_percentile(pred_ac_features[i], percentile, feature_stats, i, twod=True) normalized_orig = normalize_percentile(orig_ac_features[i], percentile, feature_stats, i, twod=True) for j in range(normalized_pred.shape[-1]): pred_feature = normalized_pred[:, :, :, j] orig_feature = normalized_orig[:, :, :, j] append_to_dict(res, 'l1_diff', np.mean(np.abs(pred_feature - orig_feature))) append_to_dict(res, 'l2_diff', np.mean(np.square(pred_feature - orig_feature))) append_to_dict(res, 'pred_mean', np.mean(pred_feature)) append_to_dict(res, 'pred_std', np.std(pred_feature)) append_to_dict(res, 'pred_99pc', np.percentile(pred_feature, 99)) append_to_dict(res, 'pred_1pc', np.percentile(pred_feature, 1)) append_to_dict(res, 'orig_mean', np.mean(orig_feature)) append_to_dict(res, 'orig_std', np.std(orig_feature)) append_to_dict(res, 'orig_99pc', np.percentile(orig_feature, 99)) append_to_dict(res, 'orig_1pc', np.percentile(orig_feature, 1)) dice_means.append(mean) results[vol_name] = res print(step, mean, std) print('time:', time.time() - start_time) print('average dice:', np.mean(dice_means)) print('time taken:', time.time() - start_time) # for key, value in results.items(): # print(key) # print(value) with open(filename, 'wb') as file: file.write( pickle.dumps(results)) # use `pickle.loads` to do the reverse
if __name__ == "__main__": print(networks.__file__) raw = tf.placeholder(tf.float32, shape=(43, 430, 430)) raw_batched = tf.reshape(raw, ( 1, 1, ) + (43, 430, 430)) #unet = networks.unet(raw_batched, 12, 6, [[1, 3, 3], [1, 3, 3], [3, 3, 3]], anisotropy=10) unet, fov, voxel_size = networks.unet( raw_batched, 12, 6, [[1, 3, 3], [1, 3, 3], [3, 3, 3]], [[(1, 3, 3), (1, 3, 3)], [(1, 3, 3), (1, 3, 3)], [(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)]], [[(1, 3, 3), (1, 3, 3)], [(1, 3, 3), (1, 3, 3)], [(3, 3, 3), (3, 3, 3)], [(3, 3, 3), (3, 3, 3)]], voxel_size=(10, 1, 1), fov=(10, 1, 1)) # raw = tf.placeholder(tf.float32, shape=(132,)*3) # raw_batched = tf.reshape(raw, (1, 1,) + (132,)*3) # # unet = networks.unet(raw_batched, 24, 3, [[2, 2, 2], [2, 2, 2], [2, 2, 2]]) dist_batched, fov = networks.conv_pass(unet, kernel_size=[[1, 1, 1]], num_fmaps=1, activation=None,
def train(model,save_name, gpu_id, lr, n_iterations, reg_param, model_save_iter, alpha): model_dir = '../models/' + save_name if not os.path.isdir(model_dir): os.mkdir(model_dir) 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)) # UNET filters nf_enc = [16,32,32,32] if(model == 'vm1'): nf_dec = [32,32,32,32,8,8,3] else: nf_dec = [32,32,32,32,32,16,16,3] nf = [0, 8, 16, 32, 64] with tf.device(gpu): deformer = networks.unet(vol_size, nf_enc, nf_dec) #deformer.compile(optimizer=Adam(lr=lr), loss=[losses.cc3D(), losses.gradientLoss('l2')], loss_weights=[1.0, reg_param]) # model.load_weights('../models/udrnet2/udrnet1_1/120000.h5') discriminator = networks.similarity_net(vol_size, nf) discriminator.compile(optimizer=Adam(lr), loss='binary_crossentropy') discriminator.trainable = False # Build GAN model src = Input(shape=vol_size + (1,)) tgt = Input(shape=vol_size + (1,)) [warp, df] = deformer([src,tgt]) #pdb.set_trace() sim_p = discriminator([warp, tgt]) GAN = Model([src, tgt], [sim_p, df]) GAN.compile(optimizer=Adam(lr), loss=['binary_crossentropy',losses.gradientLoss('l2')], loss_weights=[1.0, reg_param]) GAN.summary() sz = sim_p.shape.as_list() sz[0] = 1 p_one = np.ones(sz) p_zero = np.zeros(sz) zeroflow = np.zeros((1, vol_size[0], vol_size[1], vol_size[2], 3)) for step in range(0, n_iterations): print 'Epoch ' + str(step // (n_pairs)) + ', Total iterations ' + str(step) + ', lr= ' + str(keras.get_value(discriminator.optimizer.lr)) + ':' sub = np.load(train_pairs[step % (n_pairs)][0]) sub = np.reshape(sub, (1,) + sub.shape + (1,)) tmp = np.load(train_pairs[step % (n_pairs)][1]) tmp = np.reshape(tmp, (1,) + tmp.shape + (1,)) ref_sub = np.load(ref_pairs[step % (n_ref)][0]) ref_sub = np.reshape(ref_sub, (1,) + ref_sub.shape + (1,)) ref_tmp = np.load(ref_pairs[step % (n_ref)][1]) ref_tmp = np.reshape(ref_tmp, (1,) + ref_tmp.shape + (1,)) ## ---------------- Train deformer --------------------------------## keras.set_value(GAN.optimizer.lr, lr) g_loss = GAN.train_on_batch([sub, tmp],[p_one, zeroflow]) print ' Train deformer: ' + str(g_loss[1]) print ' Regularization: ' + str(g_loss[2]) ## ---------------- Train discriminator for reference --------------## fused = alpha * sub + (1-alpha) * tmp d_loss1 = discriminator.test_on_batch([fused, tmp], p_one) ## --------------------Tricks of blancing Registor and Discriminator-------------# if d_loss1 > 0.6: keras.set_value(discriminator.optimizer.lr, lr) elif d_loss1 > 0.4: keras.set_value(discriminator.optimizer.lr, lr * 0.1) elif d_loss1 > 0.2: keras.set_value(discriminator.optimizer.lr, lr * 0.01) else: keras.set_value(discriminator.optimizer.lr, lr * 0) d_loss1 = discriminator.train_on_batch([fused, tmp], p_one) print ' Test discriminator for positive sample: ' + str(d_loss1) ## ---------------- Train discriminator for deformer --------------## [warped, deform] = deformer.predict([sub, tmp]) d_loss0 = discriminator.test_on_batch([warped, tmp], p_zero) if d_loss0 > 0.6: keras.set_value(discriminator.optimizer.lr, lr) elif d_loss0 > 0.4: keras.set_value(discriminator.optimizer.lr, lr * 0.1) elif d_loss0 > 0.2: keras.set_value(discriminator.optimizer.lr, lr * 0.01) else: keras.set_value(discriminator.optimizer.lr, lr * 0) d_loss0 = discriminator.train_on_batch([warped, tmp], p_zero) print ' Test discriminator for negative sample: ' + str(d_loss0) if(step % model_save_iter == 0): deformer.save(model_dir + '/' + str(step) + '.h5') #if(step % (20 * n_pairs) == 0 and step > 0): #lr = lr / 2 #alpha = np.abs(alpha - 0.05) sys.stdout.flush()
import networks import tensorflow as tf import json if __name__ == "__main__": raw = tf.placeholder(tf.float32, shape=(196, ) * 3) raw_batched = tf.reshape(raw, ( 1, 1, ) + (196, ) * 3) unet = networks.unet(raw_batched, 12, 6, [[2, 2, 2], [2, 2, 2], [3, 3, 3]]) affs_batched = networks.conv_pass(unet, kernel_size=1, num_fmaps=3, num_repetitions=1, activation='sigmoid') output_shape_batched = affs_batched.get_shape().as_list() output_shape = output_shape_batched[1:] # strip the batch dimension affs = tf.reshape(affs_batched, output_shape) gt_affs = tf.placeholder(tf.float32, shape=output_shape) loss_weights = tf.placeholder(tf.float32, shape=output_shape) loss = tf.losses.mean_squared_error(gt_affs, affs, loss_weights) tf.summary.scalar('loss_total', loss)
def test(iter_num, gpu_id, vol_size=(160,192,224), nf_enc=[16,32,32,32], nf_dec=[32,32,32,32,32,16,16,3]): gpu = '/gpu:' + str(gpu_id) # Anatomical labels we want to evaluate labels = sio.loadmat('../data/labels.mat')['labels'][0] # read atlas atlas_vol1, atlas_seg1 = datagenerators.load_example_by_name('/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/990114_vc722.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/990114_vc722.npz')# [1,160,192,224,1] atlas_seg1 = atlas_seg1[0,:,:,:,0]# reduce the dimension to [160,192,224] atlas_vol2, atlas_seg2 = datagenerators.load_example_by_name('/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/990210_vc792.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/990210_vc792.npz') atlas_seg2 = atlas_seg2[0, :, :, :, 0] #gpu = '/gpu:' + str(gpu_id) os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # load weights of model with tf.device(gpu): net = networks.unet(vol_size, nf_enc, nf_dec) net.load_weights('/home/ys895/MAS2_Models/'+str(iter_num)+'.h5') #net.load_weights('../models/' + model_name + '/' + str(iter_num) + '.h5') xx = np.arange(vol_size[1]) yy = np.arange(vol_size[0]) zz = np.arange(vol_size[2]) grid = np.rollaxis(np.array(np.meshgrid(xx, yy, zz)), 0, 4) # (160, 192, 224, 3) #X_vol, X_seg = datagenerators.load_example_by_name('../data/test_vol.npz', '../data/test_seg.npz') X_vol1, X_seg1 = datagenerators.load_example_by_name('/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/981216_vc681.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/981216_vc681.npz') X_vol2, X_seg2 = datagenerators.load_example_by_name('/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/990205_vc783.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/990205_vc783.npz') X_vol3, X_seg3 = datagenerators.load_example_by_name('/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/990525_vc1024.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/990525_vc1024.npz') X_vol4, X_seg4 = datagenerators.load_example_by_name('/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/991025_vc1379.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/991025_vc1379.npz') X_vol5, X_seg5 = datagenerators.load_example_by_name('/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/991122_vc1463.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/991122_vc1463.npz') # change the direction of the atlas data and volume data # pred[0].shape (1, 160, 192, 224, 1) # pred[1].shape (1, 160, 192, 224, 3) # X1 with tf.device(gpu): pred1 = net.predict([atlas_vol1, X_vol1]) # Warp segments with flow flow1 = pred1[1][0, :, :, :, :]# (1, 160, 192, 224, 3) sample1 = flow1+grid sample1 = np.stack((sample1[:, :, :, 1], sample1[:, :, :, 0], sample1[:, :, :, 2]), 3) warp_seg1 = interpn((yy, xx, zz), atlas_seg1[ :, :, : ], sample1, method='nearest', bounds_error=False, fill_value=0) # (160, 192, 224) # label fusion: get the final warp_seg warp_seg = np.empty((160, 192, 224)) for x in range(0,160): for y in range(0,192): for z in range(0,224): warp_seg = np.array(warp_seg1[x, y, z]) vals, _ = dice(warp_seg, X_seg1[0, :, :, :, 0], labels=labels, nargout=2) mean1 = np.mean(vals) # X2 with tf.device(gpu): pred1 = net.predict([atlas_vol1, X_vol2]) # Warp segments with flow flow1 = pred1[1][0, :, :, :, :]# (1, 160, 192, 224, 3) sample1 = flow1+grid sample1 = np.stack((sample1[:, :, :, 1], sample1[:, :, :, 0], sample1[:, :, :, 2]), 3) warp_seg1 = interpn((yy, xx, zz), atlas_seg1[ :, :, : ], sample1, method='nearest', bounds_error=False, fill_value=0) # (160, 192, 224) # label fusion: get the final warp_seg warp_seg = np.empty((160, 192, 224)) for x in range(0,160): for y in range(0,192): for z in range(0,224): warp_seg = np.array(warp_seg1[x, y, z]) vals, _ = dice(warp_seg, X_seg2[0,:,:,:,0], labels=labels, nargout=2) mean2 = np.mean(vals) #print(np.mean(vals), np.std(vals)) # X3 with tf.device(gpu): pred1 = net.predict([atlas_vol1, X_vol3]) # Warp segments with flow flow1 = pred1[1][0, :, :, :, :]# (1, 160, 192, 224, 3) sample1 = flow1+grid sample1 = np.stack((sample1[:, :, :, 1], sample1[:, :, :, 0], sample1[:, :, :, 2]), 3) warp_seg1 = interpn((yy, xx, zz), atlas_seg1[ :, :, : ], sample1, method='nearest', bounds_error=False, fill_value=0) # (160, 192, 224) # label fusion: get the final warp_seg warp_seg = np.empty((160, 192, 224)) for x in range(0,160): for y in range(0,192): for z in range(0,224): warp_seg = np.array(warp_seg1[x, y, z]) vals, _ = dice(warp_seg, X_seg3[0, :, :, :, 0], labels=labels, nargout=2) mean3 = np.mean(vals) # X4 with tf.device(gpu): pred1 = net.predict([atlas_vol1, X_vol4]) # Warp segments with flow flow1 = pred1[1][0, :, :, :, :]# (1, 160, 192, 224, 3) sample1 = flow1+grid sample1 = np.stack((sample1[:, :, :, 1], sample1[:, :, :, 0], sample1[:, :, :, 2]), 3) warp_seg1 = interpn((yy, xx, zz), atlas_seg1[ :, :, : ], sample1, method='nearest', bounds_error=False, fill_value=0) # (160, 192, 224) # label fusion: get the final warp_seg warp_seg = np.empty((160, 192, 224)) for x in range(0,160): for y in range(0,192): for z in range(0,224): warp_seg = np.array(warp_seg1[x, y, z]) vals, _ = dice(warp_seg, X_seg4[0, :, :, :, 0], labels=labels, nargout=2) mean4 = np.mean(vals) # X5 with tf.device(gpu): pred1 = net.predict([atlas_vol1, X_vol5]) # Warp segments with flow flow1 = pred1[1][0, :, :, :, :]# (1, 160, 192, 224, 3) sample1 = flow1+grid sample1 = np.stack((sample1[:, :, :, 1], sample1[:, :, :, 0], sample1[:, :, :, 2]), 3) warp_seg1 = interpn((yy, xx, zz), atlas_seg1[ :, :, : ], sample1, method='nearest', bounds_error=False, fill_value=0) # (160, 192, 224) # label fusion: get the final warp_seg warp_seg = np.empty((160, 192, 224)) for x in range(0,160): for y in range(0,192): for z in range(0,224): warp_seg = np.array(warp_seg1[x,y,z]) #print(warp_arr) #warp_seg[x,y,z] = stats.mode(warp_arr)[0] vals, _ = dice(warp_seg, X_seg5[0, :, :, :, 0], labels=labels, nargout=2) mean5 = np.mean(vals) # compute mean of dice score sum = mean1 + mean2 + mean3 + mean4 + mean5 mean_dice = sum/5 print(mean_dice)
def test(iter_num, gpu_id, vol_size=(160, 192, 224), nf_enc=[16, 32, 32, 32], nf_dec=[32, 32, 32, 32, 32, 16, 16, 3]): gpu = '/gpu:' + str(gpu_id) # Anatomical labels we want to evaluate labels = sio.loadmat('../data/labels.mat')['labels'][0] # read atlas atlas_vol1, atlas_seg1 = datagenerators.load_example_by_name( '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/990114_vc722.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/990114_vc722.npz' ) # [1,160,192,224,1] atlas_seg1 = atlas_seg1[0, :, :, :, 0] # reduce the dimension to [160,192,224] atlas_seg1 = keras.utils.to_categorical(atlas_seg1) atlas_vol2, atlas_seg2 = datagenerators.load_example_by_name( '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/990210_vc792.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/990210_vc792.npz' ) atlas_seg2 = atlas_seg2[0, :, :, :, 0] atlas_seg2 = keras.utils.to_categorical(atlas_seg2) atlas_vol3, atlas_seg3 = datagenerators.load_example_by_name( '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/990405_vc922.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/990405_vc922.npz' ) atlas_seg3 = atlas_seg3[0, :, :, :, 0] atlas_seg3 = keras.utils.to_categorical(atlas_seg3) atlas_vol4, atlas_seg4 = datagenerators.load_example_by_name( '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/991006_vc1337.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/991006_vc1337.npz' ) atlas_seg4 = atlas_seg4[0, :, :, :, 0] atlas_seg4 = keras.utils.to_categorical(atlas_seg4) atlas_vol5, atlas_seg5 = datagenerators.load_example_by_name( '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/991120_vc1456.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/991120_vc1456.npz' ) atlas_seg5 = atlas_seg5[0, :, :, :, 0] atlas_seg5 = keras.utils.to_categorical(atlas_seg5) #atlas = np.load('../data/atlas_norm.npz') #atlas_vol = atlas['vol'] #print('the size of atlas:') #print(atlas_vol.shape) #atlas_seg = atlas['seg'] #atlas_vol = np.reshape(atlas_vol, (1,)+atlas_vol.shape+(1,)) #gpu = '/gpu:' + str(gpu_id) os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # load weights of model with tf.device(gpu): net = networks.unet(vol_size, nf_enc, nf_dec) net.load_weights('/home/ys895/MAS_Models/' + str(iter_num) + '.h5') #net.load_weights('../models/' + model_name + '/' + str(iter_num) + '.h5') xx = np.arange(vol_size[1]) yy = np.arange(vol_size[0]) zz = np.arange(vol_size[2]) grid = np.rollaxis(np.array(np.meshgrid(xx, yy, zz)), 0, 4) # (160, 192, 224, 3) #X_vol, X_seg = datagenerators.load_example_by_name('../data/test_vol.npz', '../data/test_seg.npz') X_vol1, X_seg1 = datagenerators.load_example_by_name( '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/981216_vc681.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/981216_vc681.npz' ) X_vol2, X_seg2 = datagenerators.load_example_by_name( '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/990205_vc783.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/990205_vc783.npz' ) X_vol3, X_seg3 = datagenerators.load_example_by_name( '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/990525_vc1024.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/990525_vc1024.npz' ) X_vol4, X_seg4 = datagenerators.load_example_by_name( '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/991025_vc1379.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/991025_vc1379.npz' ) X_vol5, X_seg5 = datagenerators.load_example_by_name( '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/vols/991122_vc1463.npz', '/home/ys895/resize256/resize256-crop_x32/FromEugenio_prep/labels/991122_vc1463.npz' ) # change the direction of the atlas data and volume data # pred[0].shape (1, 160, 192, 224, 1) # pred[1].shape (1, 160, 192, 224, 3) # X1 with tf.device(gpu): pred1 = net.predict([atlas_vol1, X_vol1]) pred2 = net.predict([atlas_vol2, X_vol1]) pred3 = net.predict([atlas_vol3, X_vol1]) pred4 = net.predict([atlas_vol4, X_vol1]) pred5 = net.predict([atlas_vol5, X_vol1]) # Warp segments with flow flow1 = pred1[1][0, :, :, :, :] # (1, 160, 192, 224, 3) flow2 = pred2[1][0, :, :, :, :] flow3 = pred3[1][0, :, :, :, :] flow4 = pred4[1][0, :, :, :, :] flow5 = pred5[1][0, :, :, :, :] sample1 = flow1 + grid sample1 = np.stack( (sample1[:, :, :, 1], sample1[:, :, :, 0], sample1[:, :, :, 2]), 3) sample2 = flow2 + grid sample2 = np.stack( (sample2[:, :, :, 1], sample2[:, :, :, 0], sample2[:, :, :, 2]), 3) sample3 = flow3 + grid sample3 = np.stack( (sample3[:, :, :, 1], sample3[:, :, :, 0], sample3[:, :, :, 2]), 3) sample4 = flow4 + grid sample4 = np.stack( (sample4[:, :, :, 1], sample4[:, :, :, 0], sample4[:, :, :, 2]), 3) sample5 = flow5 + grid sample5 = np.stack( (sample5[:, :, :, 1], sample5[:, :, :, 0], sample5[:, :, :, 2]), 3) warp_seg1 = interpn((yy, xx, zz), atlas_seg1[:, :, :, :], sample1, method='linear', bounds_error=False, fill_value=0) # (160, 192, 224) warp_seg2 = interpn((yy, xx, zz), atlas_seg2[:, :, :, :], sample2, method='linear', bounds_error=False, fill_value=0) warp_seg3 = interpn((yy, xx, zz), atlas_seg3[:, :, :, :], sample3, method='linear', bounds_error=False, fill_value=0) warp_seg4 = interpn((yy, xx, zz), atlas_seg4[:, :, :, :], sample4, method='linear', bounds_error=False, fill_value=0) warp_seg5 = interpn((yy, xx, zz), atlas_seg5[:, :, :, :], sample5, method='linear', bounds_error=False, fill_value=0) # label fusion: get the final warp_seg warp_seg = np.empty((160, 192, 224, atlas_seg1.shape[3])) warp_seg = (warp_seg1 + warp_seg2 + warp_seg3 + warp_seg4 + warp_seg5) / 5 warp_seg = np.argmax(warp_seg, axis=3) vals, _ = dice(warp_seg, X_seg1[0, :, :, :, 0], labels=labels, nargout=2) mean1 = np.mean(vals) var1 = np.std(vals) # X2 with tf.device(gpu): pred1 = net.predict([atlas_vol1, X_vol2]) pred2 = net.predict([atlas_vol2, X_vol2]) pred3 = net.predict([atlas_vol3, X_vol2]) pred4 = net.predict([atlas_vol4, X_vol2]) pred5 = net.predict([atlas_vol5, X_vol2]) # Warp segments with flow flow1 = pred1[1][0, :, :, :, :] # (1, 160, 192, 224, 3) flow2 = pred2[1][0, :, :, :, :] flow3 = pred3[1][0, :, :, :, :] flow4 = pred4[1][0, :, :, :, :] flow5 = pred5[1][0, :, :, :, :] sample1 = flow1 + grid sample1 = np.stack( (sample1[:, :, :, 1], sample1[:, :, :, 0], sample1[:, :, :, 2]), 3) sample2 = flow2 + grid sample2 = np.stack( (sample2[:, :, :, 1], sample2[:, :, :, 0], sample2[:, :, :, 2]), 3) sample3 = flow3 + grid sample3 = np.stack( (sample3[:, :, :, 1], sample3[:, :, :, 0], sample3[:, :, :, 2]), 3) sample4 = flow4 + grid sample4 = np.stack( (sample4[:, :, :, 1], sample4[:, :, :, 0], sample4[:, :, :, 2]), 3) sample5 = flow5 + grid sample5 = np.stack( (sample5[:, :, :, 1], sample5[:, :, :, 0], sample5[:, :, :, 2]), 3) warp_seg1 = interpn((yy, xx, zz), atlas_seg1[:, :, :, :], sample1, method='linear', bounds_error=False, fill_value=0) # (160, 192, 224) warp_seg2 = interpn((yy, xx, zz), atlas_seg2[:, :, :, :], sample2, method='linear', bounds_error=False, fill_value=0) warp_seg3 = interpn((yy, xx, zz), atlas_seg3[:, :, :, :], sample3, method='linear', bounds_error=False, fill_value=0) warp_seg4 = interpn((yy, xx, zz), atlas_seg4[:, :, :, :], sample4, method='linear', bounds_error=False, fill_value=0) warp_seg5 = interpn((yy, xx, zz), atlas_seg5[:, :, :, :], sample5, method='linear', bounds_error=False, fill_value=0) # label fusion: get the final warp_seg warp_seg = np.empty((160, 192, 224, atlas_seg1.shape[3])) warp_seg = (warp_seg1 + warp_seg2 + warp_seg3 + warp_seg4 + warp_seg5) / 5 warp_seg = np.argmax(warp_seg, axis=3) vals, _ = dice(warp_seg, X_seg2[0, :, :, :, 0], labels=labels, nargout=2) mean2 = np.mean(vals) var2 = np.std(vals) #print(np.mean(vals), np.std(vals)) # X3 with tf.device(gpu): pred1 = net.predict([atlas_vol1, X_vol3]) pred2 = net.predict([atlas_vol2, X_vol3]) pred3 = net.predict([atlas_vol3, X_vol3]) pred4 = net.predict([atlas_vol4, X_vol3]) pred5 = net.predict([atlas_vol5, X_vol3]) # Warp segments with flow flow1 = pred1[1][0, :, :, :, :] # (1, 160, 192, 224, 3) flow2 = pred2[1][0, :, :, :, :] flow3 = pred3[1][0, :, :, :, :] flow4 = pred4[1][0, :, :, :, :] flow5 = pred5[1][0, :, :, :, :] sample1 = flow1 + grid sample1 = np.stack( (sample1[:, :, :, 1], sample1[:, :, :, 0], sample1[:, :, :, 2]), 3) sample2 = flow2 + grid sample2 = np.stack( (sample2[:, :, :, 1], sample2[:, :, :, 0], sample2[:, :, :, 2]), 3) sample3 = flow3 + grid sample3 = np.stack( (sample3[:, :, :, 1], sample3[:, :, :, 0], sample3[:, :, :, 2]), 3) sample4 = flow4 + grid sample4 = np.stack( (sample4[:, :, :, 1], sample4[:, :, :, 0], sample4[:, :, :, 2]), 3) sample5 = flow5 + grid sample5 = np.stack( (sample5[:, :, :, 1], sample5[:, :, :, 0], sample5[:, :, :, 2]), 3) warp_seg1 = interpn((yy, xx, zz), atlas_seg1[:, :, :, :], sample1, method='linear', bounds_error=False, fill_value=0) warp_seg2 = interpn((yy, xx, zz), atlas_seg2[:, :, :, :], sample2, method='linear', bounds_error=False, fill_value=0) warp_seg3 = interpn((yy, xx, zz), atlas_seg3[:, :, :, :], sample3, method='linear', bounds_error=False, fill_value=0) warp_seg4 = interpn((yy, xx, zz), atlas_seg4[:, :, :, :], sample4, method='linear', bounds_error=False, fill_value=0) warp_seg5 = interpn((yy, xx, zz), atlas_seg5[:, :, :, :], sample5, method='linear', bounds_error=False, fill_value=0) # label fusion: get the final warp_seg warp_seg = np.empty((160, 192, 224, atlas_seg1.shape[3])) warp_seg = (warp_seg1 + warp_seg2 + warp_seg3 + warp_seg4 + warp_seg5) / 5 warp_seg = np.argmax(warp_seg, axis=3) vals, _ = dice(warp_seg, X_seg3[0, :, :, :, 0], labels=labels, nargout=2) mean3 = np.mean(vals) var3 = np.std(vals) # X4 with tf.device(gpu): pred1 = net.predict([atlas_vol1, X_vol4]) pred2 = net.predict([atlas_vol2, X_vol4]) pred3 = net.predict([atlas_vol3, X_vol4]) pred4 = net.predict([atlas_vol4, X_vol4]) pred5 = net.predict([atlas_vol5, X_vol4]) # Warp segments with flow flow1 = pred1[1][0, :, :, :, :] # (1, 160, 192, 224, 3) flow2 = pred2[1][0, :, :, :, :] flow3 = pred3[1][0, :, :, :, :] flow4 = pred4[1][0, :, :, :, :] flow5 = pred5[1][0, :, :, :, :] sample1 = flow1 + grid sample1 = np.stack( (sample1[:, :, :, 1], sample1[:, :, :, 0], sample1[:, :, :, 2]), 3) sample2 = flow2 + grid sample2 = np.stack( (sample2[:, :, :, 1], sample2[:, :, :, 0], sample2[:, :, :, 2]), 3) sample3 = flow3 + grid sample3 = np.stack( (sample3[:, :, :, 1], sample3[:, :, :, 0], sample3[:, :, :, 2]), 3) sample4 = flow4 + grid sample4 = np.stack( (sample4[:, :, :, 1], sample4[:, :, :, 0], sample4[:, :, :, 2]), 3) sample5 = flow5 + grid sample5 = np.stack( (sample5[:, :, :, 1], sample5[:, :, :, 0], sample5[:, :, :, 2]), 3) warp_seg1 = interpn((yy, xx, zz), atlas_seg1[:, :, :, :], sample1, method='linear', bounds_error=False, fill_value=0) warp_seg2 = interpn((yy, xx, zz), atlas_seg2[:, :, :, :], sample2, method='linear', bounds_error=False, fill_value=0) warp_seg3 = interpn((yy, xx, zz), atlas_seg3[:, :, :, :], sample3, method='linear', bounds_error=False, fill_value=0) warp_seg4 = interpn((yy, xx, zz), atlas_seg4[:, :, :, :], sample4, method='linear', bounds_error=False, fill_value=0) warp_seg5 = interpn((yy, xx, zz), atlas_seg5[:, :, :, :], sample5, method='linear', bounds_error=False, fill_value=0) # label fusion: get the final warp_seg warp_seg = np.empty((160, 192, 224, atlas_seg1.shape[3])) warp_seg = (warp_seg1 + warp_seg2 + warp_seg3 + warp_seg4 + warp_seg5) / 5 warp_seg = np.argmax(warp_seg, axis=3) vals, _ = dice(warp_seg, X_seg4[0, :, :, :, 0], labels=labels, nargout=2) mean4 = np.mean(vals) var4 = np.std(vals) # X5 with tf.device(gpu): pred1 = net.predict([atlas_vol1, X_vol5]) pred2 = net.predict([atlas_vol2, X_vol5]) pred3 = net.predict([atlas_vol3, X_vol5]) pred4 = net.predict([atlas_vol4, X_vol5]) pred5 = net.predict([atlas_vol5, X_vol5]) # Warp segments with flow flow1 = pred1[1][0, :, :, :, :] # (1, 160, 192, 224, 3) flow2 = pred2[1][0, :, :, :, :] flow3 = pred3[1][0, :, :, :, :] flow4 = pred4[1][0, :, :, :, :] flow5 = pred5[1][0, :, :, :, :] sample1 = flow1 + grid sample1 = np.stack( (sample1[:, :, :, 1], sample1[:, :, :, 0], sample1[:, :, :, 2]), 3) sample2 = flow2 + grid sample2 = np.stack( (sample2[:, :, :, 1], sample2[:, :, :, 0], sample2[:, :, :, 2]), 3) sample3 = flow3 + grid sample3 = np.stack( (sample3[:, :, :, 1], sample3[:, :, :, 0], sample3[:, :, :, 2]), 3) sample4 = flow4 + grid sample4 = np.stack( (sample4[:, :, :, 1], sample4[:, :, :, 0], sample4[:, :, :, 2]), 3) sample5 = flow5 + grid sample5 = np.stack( (sample5[:, :, :, 1], sample5[:, :, :, 0], sample5[:, :, :, 2]), 3) warp_seg1 = interpn((yy, xx, zz), atlas_seg1[:, :, :, :], sample1, method='linear', bounds_error=False, fill_value=0) warp_seg2 = interpn((yy, xx, zz), atlas_seg2[:, :, :, :], sample2, method='linear', bounds_error=False, fill_value=0) warp_seg3 = interpn((yy, xx, zz), atlas_seg3[:, :, :, :], sample3, method='linear', bounds_error=False, fill_value=0) warp_seg4 = interpn((yy, xx, zz), atlas_seg4[:, :, :, :], sample4, method='linear', bounds_error=False, fill_value=0) warp_seg5 = interpn((yy, xx, zz), atlas_seg5[:, :, :, :], sample5, method='linear', bounds_error=False, fill_value=0) # label fusion: get the final warp_seg warp_seg = np.empty((160, 192, 224, atlas_seg1.shape[3])) warp_seg = (warp_seg1 + warp_seg2 + warp_seg3 + warp_seg4 + warp_seg5) / 5 warp_seg = np.argmax(warp_seg, axis=3) vals, _ = dice(warp_seg, X_seg5[0, :, :, :, 0], labels=labels, nargout=2) mean5 = np.mean(vals) var5 = np.std(vals) # compute mean of dice score sum = mean1 + mean2 + mean3 + mean4 + mean5 mean_dice = sum / 5 var = (var1 + var2 + var3 + var4 + var5) / 5 print(str(mean_dice) + ',' + str(var))
def train(model, pretrained_path, model_name, gpu_id, lr, n_iterations, autoencoder_iters, autoencoder_model, autoencoder_num_downsample, ac_coef, use_normalize, norm_percentile, use_mse, reg_param, gamma, model_save_iter, batch_size=1): """ model training function :param model: either vm1 or vm2 (based on CVPR 2018 paper) :param model_dir: the model directory to save to :param gpu_id: integer specifying the gpu to use :param lr: learning rate :param n_iterations: number of training iterations :param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper) :param model_save_iter: frequency with which to save models :param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size """ os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' restrict_GPU_tf(str(gpu_id)) restrict_GPU_keras(str(gpu_id)) train_labels = sio.loadmat('../data/labels.mat')['labels'][0] n_labels = train_labels.shape[0] atlas_vol = atlas['vol'] atlas_seg = atlas['seg'] atlas_vol = np.reshape(atlas_vol, (1, ) + atlas_vol.shape + (1, )) atlas_seg = np.reshape(atlas_seg, (1, ) + atlas_seg.shape + (1, )) atlas_seg = datagenerators.split_seg_into_channels(atlas_seg, train_labels) atlas_seg = datagenerators.downsample(atlas_seg) model_dir = "../models/" + model_name # prepare model folder if not os.path.isdir(model_dir): os.mkdir(model_dir) # GPU handling os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # UNET filters for voxelmorph-1 and voxelmorph-2, # these are architectures presented in CVPR 2018 nf_enc = [16, 32, 32, 32] if model == 'vm1': nf_dec = [32, 32, 32, 32, 8, 8] else: nf_dec = [32, 32, 32, 32, 32, 16, 16] # prepare the model # in the CVPR layout, the model takes in [image_1, image_2] and outputs [warped_image_1, flow] # in the experiments, we use image_2 as atlas autoencoder_path = '../models/%s/%s.h5' % (autoencoder_model, autoencoder_iters) mse_coef = 1.0 if use_mse else 0 print('mse coef', mse_coef) model = networks.unet(vol_size, nf_enc, nf_dec, use_seg=True, n_seg=len(train_labels)) model.compile(optimizer=Adam(lr=lr), loss=['mse', losses.gradientLoss('l2'), losses.diceLoss], loss_weights=[mse_coef, reg_param, gamma]) # if you'd like to initialize the data, you can do it here: if pretrained_path != None: model.load_weights(pretrained_path) # prepare data for training train_example_gen = datagenerators.example_gen(train_vol_names, return_segs=True, seg_dir=train_seg_dir) zero_flow = np.zeros([batch_size, *vol_size, 3]) # train. Note: we use train_on_batch and design out own print function as this has enabled # faster development and debugging, but one could also use fit_generator and Keras callbacks. for step in range(0, n_iterations): # get data X = next(train_example_gen) X_seg = X[1] X_seg = datagenerators.split_seg_into_channels(X_seg, train_labels) X_seg = datagenerators.downsample(X_seg) # train train_loss = model.train_on_batch([X[0], atlas_vol, X_seg], [atlas_vol, zero_flow, atlas_seg]) if not isinstance(train_loss, list): train_loss = [train_loss] # print the loss. print_loss(step, 1, train_loss) # save model if step % model_save_iter == 0: model.save(os.path.join(model_dir, str(step) + '.h5'))
validation_dataset = ImageDataset(input_dir=validation_directory, transform=True) #Validation Dataset num_epochs = n_iters / (len(train_dataset) / batch_size) num_epochs = int(num_epochs) train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True) validation_loader = torch.utils.data.DataLoader( dataset=validation_dataset, batch_size=validation_batch_size, shuffle=False) model = unet() iteri = 0 iter_new = 0 #checking if checkpoints exist to resume training and create it if not if os.path.exists(checkpoints_directory_unet) and len( os.listdir(checkpoints_directory_unet)): checkpoints = os.listdir(checkpoints_directory_unet) checkpoints.sort(key=lambda x: int((x.split('_')[2]).split('.')[0])) model = torch.load(checkpoints_directory_unet + '/' + checkpoints[-1]) #changed to checkpoints iteri = int(re.findall(r'\d+', checkpoints[-1])[0]) # changed to checkpoints iter_new = iteri print("Resuming from iteration " + str(iteri)) elif not os.path.exists(checkpoints_directory_unet):
if __name__ == "__main__": raw = tf.placeholder(tf.float32, shape=(132, ) * 3) pred = tf.placeholder(tf.float32, shape=(132, ) * 3) raw_batched = tf.reshape(raw, ( 1, 1, ) + (132, ) * 3) pred_reshaped = tf.reshape(raw, ( 1, 1, ) + (132, ) * 3) unet = networks.unet(raw_batched, 24, 3, [[2, 2, 2], [2, 2, 2], [2, 2, 2]]) affs_batched = networks.conv_pass(unet, kernel_size=1, num_fmaps=3, num_repetitions=1, activation='sigmoid') output_shape_batched = affs_batched.get_shape().as_list() output_shape = output_shape_batched[1:] # strip the batch dimension affs = tf.reshape(affs_batched, output_shape) gt_affs = tf.placeholder(tf.float32, shape=output_shape) loss_weights = tf.placeholder(tf.float32, shape=output_shape)
def test(model_name, iter_num, gpu_id, vol_size=(160, 224, 192), nf_enc=[16, 32, 32, 32], nf_dec=[32, 32, 32, 32, 8, 8, 3]): """ test nf_enc and nf_dec #nf_dec = [32,32,32,32,32,16,16,3] # This needs to be changed. Ideally, we could just call load_model, and we wont have to # specify the # of channels here, but the load_model is not working with the custom loss... """ gpu = '/gpu:' + str(gpu_id) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # load weights of model with tf.device(gpu): net = networks.unet(vol_size, nf_enc, nf_dec) net.load_weights('../models/' + model_name + '/' + str(iter_num) + '.h5') xx = np.arange(vol_size[1]) yy = np.arange(vol_size[0]) zz = np.arange(vol_size[2]) grid = np.rollaxis(np.array(np.meshgrid(xx, yy, zz)), 0, 4) sum_dice = 0 for i in range(0, noftest): X_vol = np.load(subjcet_filepaths[i]) X_vol = np.reshape(X_vol, (1, ) + X_vol.shape + (1, )) X_seg = np.load(seg_filepaths[i]) #X_seg = np.reshape(X_seg, (1,) + X_seg.shape + (1,)) for j in range(0, noftest): Y_vol = np.load(subjcet_filepaths[j]) Y_vol = np.reshape(Y_vol, (1, ) + Y_vol.shape + (1, )) Y_seg = np.load(seg_filepaths[j]) #Y_seg = np.reshape(Y_seg, (1,) + Y_seg.shape + (1,)) with tf.device(gpu): pred = net.predict([X_vol, Y_vol]) # Warp segments with flow flow = pred[1][0, :, :, :, :] sample = flow + grid sample = np.stack( (sample[:, :, :, 1], sample[:, :, :, 0], sample[:, :, :, 2]), 3) warp_seg = interpn((yy, xx, zz), X_seg, sample, method='nearest', bounds_error=False, fill_value=0) #print 'X Image size:', X_seg.shape #print 'warped Image size:', warp_seg.shape #print 'Y Image size:', Y_seg.shape vals, lnames = dice(warp_seg, Y_seg, nargout=2) print '%dth image to %dth image --- ' % ( i + startoftest, j + startoftest), 'mean: ', np.mean(vals), 'std: ', np.std(vals) sum_dice = sum_dice + np.mean(vals) mat_Deform = pred[1][0, :, :, :, :] img_Deform = sitk.GetImageFromArray(mat_Deform) print 'img_Deform shape, ', mat_Deform.shape outputfilename = output_dir + 'deformations/' + 'Deform%02dto%02d.mha' % ( i + startoftest, j + startoftest) sitk.WriteImage(img_Deform, outputfilename) mat_warpimage = pred[0][0, :, :, :, 0] img_Warp = sitk.GetImageFromArray(mat_warpimage) caster = sitk.CastImageFilter() caster.SetOutputPixelType(sitk.sitkInt16) img_Warp = caster.Execute(img_Warp) outputfilename = output_dir + 'warp_images/' + 'na%02dto%02d.mha' % ( i + startoftest, j + startoftest) sitk.WriteImage(img_Warp, outputfilename) img_Warpseg = sitk.GetImageFromArray(warp_seg) img_Warpseg = caster.Execute(img_Warpseg) outputfilename = output_dir + 'warp_images/' + 'seg%02dto%02d.mha' % ( i + startoftest, j + startoftest) sitk.WriteImage(img_Warpseg, outputfilename) print 'Average Dice ratio: ', sum_dice / (noftest**2)
with tf.variable_scope('autocontext') as scope: # phase 1 raw_0 = tf.placeholder(tf.float32, shape=shape_0) raw_0_batched = tf.reshape(raw_0, (1, 1) + shape_0) input_0 = tf.concat([raw_0_batched, affs_0_batched], 1) if ignore: keep_raw = tf.ones_like(raw_0_batched) ignore_aff = tf.zeros_like(affs_0_batched) ignore_mask = tf.concat([keep_raw, ignore_aff], 1) input_0 = networks.ignore(input_0, ignore_mask) unet = networks.unet(input_0, 24, 3, [[2, 2, 2], [2, 2, 2], [2, 2, 2]]) affs_1_batched = networks.conv_pass(unet, kernel_size=1, num_fmaps=3, num_repetitions=1, activation='sigmoid') affs_1 = tf.reshape(affs_1_batched, (3, ) + shape_1) gt_affs_1 = tf.placeholder(tf.float32, shape=(3, ) + shape_1) loss_weights_1 = tf.placeholder(tf.float32, shape=(3, ) + shape_1) loss_1 = tf.losses.mean_squared_error(gt_affs_1, affs_1, loss_weights_1) # phase 2
# Anatomical labels we want to evaluate labels = sio.loadmat('../data/labels.mat')['labels'][0] atlas = np.load('../data/atlas_norm.npz') atlas_vol = atlas['vol'] atlas_seg = atlas['seg'] atlas_vol = np.reshape(atlas_vol, (1,)+atlas_vol.shape+(1,)) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # load weights of model with tf.device(gpu): net = networks.unet(vol_size, nf_enc, nf_dec) # net.load_weights('../models/' + model_name + '/' + str(iter_num) + '.h5') net.load_weights(model_name) xx = np.arange(vol_size[1]) yy = np.arange(vol_size[0]) zz = np.arange(vol_size[2]) grid = np.rollaxis(np.array(np.meshgrid(xx, yy, zz)), 0, 4) X_vol, X_seg = datagenerators.load_example_by_name('../data/test_vol.npz', '../data/test_seg.npz') with tf.device(gpu): pred = net.predict([X_vol, atlas_vol]) # Warp segments with flow flow = pred[1][0, :, :, :, :]
def test(model_name, gpu_id, iter_num, vol_size=(160, 192, 224), nf_enc=[16, 32, 32, 32], nf_dec=[32, 32, 32, 32, 8, 8, 3]): """ test nf_enc and nf_dec #nf_dec = [32,32,32,32,32,16,16,3] # This needs to be changed. Ideally, we could just call load_model, and we wont have to # specify the # of channels here, but the load_model is not working with the custom loss... """ gpu = '/gpu:' + str(gpu_id) # Test file and anatomical labels we want to evaluate test_brain_file = open('../data/test_examples.txt') test_brain_strings = test_brain_file.readlines() test_brain_strings = [x.strip() for x in test_brain_strings] good_labels = sio.loadmat('../data/test_labels.mat')['labels'][0] atlas = np.load('../data/atlas_norm.npz') atlas_vol = atlas['vol'] atlas_seg = atlas['seg'] atlas_vol = np.reshape(atlas_vol, (1, ) + atlas_vol.shape + (1, )) config = tf.ConfigProto() config.gpu_options.allow_growth = True config.allow_soft_placement = True set_session(tf.Session(config=config)) # load weights of model with tf.device(gpu): net = networks.unet(vol_size, nf_enc, nf_dec) net.load_weights('../models/' + model_name + '/' + str(iter_num) + '.h5') n_batches = len(test_brain_strings) xx = np.arange(vol_size[1]) yy = np.arange(vol_size[0]) zz = np.arange(vol_size[2]) grid = np.rollaxis(np.array(np.meshgrid(xx, yy, zz)), 0, 4) dice_vals = np.zeros((len(good_labels), n_batches)) np.random.seed(17) for k in range(0, n_batches): vol_name, seg_name = test_brain_strings[k].split(",") X_vol, X_seg = datagenerators.load_example_by_name(vol_name, seg_name) with tf.device(gpu): pred = net.predict([X_vol, atlas_vol]) # Warp segments with flow flow = pred[1][0, :, :, :, :] sample = flow + grid sample = np.stack( (sample[:, :, :, 1], sample[:, :, :, 0], sample[:, :, :, 2]), 3) warp_seg = interpn((yy, xx, zz), X_seg[0, :, :, :, 0], sample, method='nearest', bounds_error=False, fill_value=0) vals, labels = dice(warp_seg, atlas_seg, labels=good_labels, nargout=2) dice_vals[:, k] = vals print np.mean(dice_vals[:, k])
import networks import tensorflow as tf import json if __name__ == "__main__": raw = tf.placeholder(tf.float32, shape=(84, 268,268)) raw_batched = tf.reshape(raw, (1, 1,) + (84, 268, 268)) unet = networks.unet(raw_batched, 12, 6, [[1, 3, 3], [1, 3, 3], [1, 3, 3]]) # raw = tf.placeholder(tf.float32, shape=(132,)*3) # raw_batched = tf.reshape(raw, (1, 1,) + (132,)*3) # # unet = networks.unet(raw_batched, 24, 3, [[2, 2, 2], [2, 2, 2], [2, 2, 2]]) dist_batched = networks.conv_pass( unet, kernel_size=1, num_fmaps=2, num_repetitions=1, activation=None) syn_dist, bdy_dist = tf.unstack(dist_batched, 2, axis=1) #bdy_dist_batched = networks.conv_pass(unet, # kernel_size=1, # num_fmaps=1, # num_repetitions=1, # activation=None) output_shape = syn_dist.get_shape().as_list()