def test(model_name,
gpu_id,
nf_enc=[16, 32, 32, 32],
nf_dec=[32, 32, 32, 32, 32, 16, 16]):
"""
test
nf_enc and nf_dec
#nf_dec = [32,32,32,32,32,16,16,3]
# This needs to be changed. Ideally, we could just call load_model, and we wont have to
# specify the # of channels here, but the load_model is not working with the custom loss...
"""
# load subject test
print("load_data start")
X_train, y_train = load_data(data_dir='../data',
mode='test',
fixed='joyoungje')
vol_size = y_train.shape[1:-1]
# gpu handling
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# load weights of model
with tf.device(gpu):
net = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
print("model load weights")
net.load_weights(model_name)
# NN transfer model
nn_trf_model = networks.nn_trf(vol_size, indexing='ij')
# # if CPU, prepare grid
# if compute_type == 'CPU':
# grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4)
with tf.device(gpu):
print("model predict")
pred = net.predict([X_train, y_train])
print("nn_tft_model.predict")
X_warp = nn_trf_model.predict([X_train, pred[1]])[0, ..., 0]
reshape_y_train = y_train.reshape(y_train.shape[1:-1])
vals = dice(pred[0].reshape(pred[0].shape[1:-1]), reshape_y_train)
dice_mean = np.mean(vals)
dice_std = np.std(vals)
print('Dice mean over structures: {:.2f} ({:.2f})'.format(
dice_mean, dice_std))
def test(gpu_id, compute_type='GPU', vol_size=(256, 256, 256), flow=None):
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
flow = sitk.GetArrayFromImage(sitk.ReadImage(flow))
flow = np.reshape(flow, (1, ) + np.shape(flow))
flow_tensor = tf.convert_to_tensor(flow)
print(flow_tensor.shape)
# load weights of model
with tf.device(gpu):
nn_trf_model = networks.nn_trf(vol_size, indexing='ij')
warped_seg = nn_trf_model.predict([X_seg, flow])[0, ..., 0]
print(warped_seg.shape)
result_seg = sitk.GetImageFromArray(warped_seg)
result_seg.CopyInformation(mov)
sitk.WriteImage(result_seg, '../data/results/warped_seg.nii.gz')
def test(
gpu_id,
model_dir,
iter_num,
compute_type='GPU', # GPU or CPU
vol_size=(160, 192, 224),
nf_enc=[16, 32, 32, 32],
nf_dec=[32, 32, 32, 32, 16, 3],
save_file=None):
"""
test via segmetnation propagation
works by iterating over some iamge files, registering them to atlas,
propagating the warps, then computing Dice with atlas segmentations
"""
# GPU handling
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# load weights of model
with tf.device(gpu):
# if testing miccai run, should be xy indexing.
net = networks.miccai2018_net(vol_size,
nf_enc,
nf_dec,
use_miccai_int=False,
indexing='ij')
net.load_weights(os.path.join(model_dir, str(iter_num) + '.h5'))
# compose diffeomorphic flow output model
diff_net = keras.models.Model(net.inputs,
net.get_layer('diffflow').output)
# NN transfer model
nn_trf_model = networks.nn_trf(vol_size, indexing='ij')
# if CPU, prepare grid
if compute_type == 'CPU':
grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4)
# prepare a matrix of dice values
dice_vals = np.zeros((len(good_labels), n_batches))
for k in range(n_batches):
# get data
vol_name, seg_name = test_brain_strings[k].split(",")
X_vol, X_seg = datagenerators.load_example_by_name(vol_name, seg_name)
# predict transform
with tf.device(gpu):
pred = diff_net.predict([X_vol, atlas_vol])
# Warp segments with flow
if compute_type == 'CPU':
flow = pred[0, :, :, :, :]
warp_seg = util.warp_seg(X_seg,
flow,
grid=grid,
xx=xx,
yy=yy,
zz=zz)
else: # GPU
warp_seg = nn_trf_model.predict([X_seg, pred])[0, ..., 0]
# compute Volume Overlap (Dice)
dice_vals[:, k] = dice(warp_seg, atlas_seg, labels=good_labels)
print('%3d %5.3f %5.3f' % (k, np.mean(
dice_vals[:, k]), np.mean(np.mean(dice_vals[:, :k + 1]))))
if save_file is not None:
sio.savemat(save_file, {
'dice_vals': dice_vals,
'labels': good_labels
})
def test(
gpu_id,
model_dir,
iter_num,
compute_type='GPU', # GPU or CPU
vol_size=(160, 192, 224),
nf_enc=[16, 32, 32, 32],
nf_dec=[32, 32, 32, 32, 16, 3],
save_file=None):
"""
test via segmetnation propagation
works by iterating over some iamge files, registering them to atlas,
propagating the warps, then computing Dice with atlas segmentations
"""
# GPU handling
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# load weights of model
with tf.device(gpu):
# if testing miccai run, should be xy indexing.
net = networks.miccai2018_net(vol_size,
nf_enc,
nf_dec,
use_miccai_int=False,
indexing='ij')
net.load_weights(os.path.join(model_dir, str(iter_num) + '.h5'))
# compose diffeomorphic flow output model
diff_net = keras.models.Model(net.inputs,
net.get_layer('diffflow').output)
# NN transfer model
nn_trf_model = networks.nn_trf(vol_size, indexing='ij')
# if CPU, prepare grid
if compute_type == 'CPU':
grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4)
# get data
X_vol = nib.load(r'D:\users\zzx\data\2018nor\pre\01.nii').get_data()
X_vol = X_vol[np.newaxis, ..., np.newaxis]
atlas_vol = nib.load(r'D:\users\zzx\data\2018nor\pre\01_a.nii').get_data()
atlas_vol = atlas_vol[np.newaxis, ..., np.newaxis]
X_mask = nib.load(
r'D:\users\zzx\data\2018mask\pre\pig01_pre_final_r.nii').get_data()
X_mask = X_mask[np.newaxis, ..., np.newaxis]
X_mask[X_mask == np.max(X_mask)] = 1
X_mask[X_mask != 1] = 0
atlas_mask = nib.load(
r'D:\users\zzx\data\2018mask\aft\pig01_02_final_r.nii').get_data()
atlas_mask[atlas_mask == np.max(atlas_mask)] = 1
atlas_mask[atlas_mask != 1] = 0
## feature point
# X_feapt = np.zeros((160,192,224))
# X_feapt[128,51,165] = 1
# X_feapt = X_feapt[np.newaxis,...,np.newaxis]
# predict transform
with tf.device(gpu):
pred = diff_net.predict([X_vol, atlas_vol])
# Warp segments with flow
if compute_type == 'CPU':
flow = pred[0, :, :, :, :]
warp_seg = util.warp_seg(X_mask, flow, grid=grid, xx=xx, yy=yy, zz=zz)
else: # GPU
warp_mask = nn_trf_model.predict([X_mask, pred])[0, ..., 0]
warp_vol = nn_trf_model.predict([X_vol, pred])[0, ..., 0]
# pred_point1 = nn_trf_model.predict([X_feapt, pred])[0,...,0]
print(X_vol.shape)
# warp_vol = nib.Nifti1Image(warp_vol,np.eye(4))
warp_vol = nib.Nifti1Image(warp_vol, np.eye(4))
nib.save(warp_vol, r'D:\users\zzx\data\2018warp\1w.nii')
# compute Volume Overlap (Dice)
# X_mask = X_mask[0,...,0]
# print(X_mask.shape, atlas_mask.shape,pred_point1.shape,np.where(pred_point1 != 0))
dice_vals = dice(warp_mask, atlas_mask)
# print('%3d %5.3f %5.3f' % (k, np.mean(dice_vals[:, k]), np.mean(np.mean(dice_vals[:, :k+1]))))
print(dice_vals)
if save_file is not None:
sio.savemat(save_file, {'dice_vals': dice_vals})
def test(
model_name,
gpu_id,
compute_type='GPU', # GPU or CPU
nf_enc=[16, 32, 32, 32],
nf_dec=[32, 32, 32, 32, 32, 16, 16]):
"""
test
nf_enc and nf_dec
#nf_dec = [32,32,32,32,32,16,16,3]
# This needs to be changed. Ideally, we could just call load_model, and we wont have to
# specify the # of channels here, but the load_model is not working with the custom loss...
"""
# Anatomical labels we want to evaluate
labels = sio.loadmat('../data/labels.mat')['labels'][0]
atlas = np.load('../data/atlas_norm.npz')
atlas_vol = atlas['vol'][np.newaxis, ..., np.newaxis]
atlas_seg = atlas['seg']
vol_size = atlas_vol.shape[1:-1]
# gpu handling
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# load weights of model
with tf.device(gpu):
net = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
net.load_weights(model_name)
# NN transfer model
nn_trf_model = networks.nn_trf(vol_size, indexing='ij')
# if CPU, prepare grid
if compute_type == 'CPU':
grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4)
# load subject test
X_vol, X_seg = datagenerators.load_example_by_name('../data/test_vol.npz',
'../data/test_seg.npz')
with tf.device(gpu):
pred = net.predict([X_vol, atlas_vol])
# Warp segments with flow
if compute_type == 'CPU':
flow = pred[1][0, :, :, :, :]
warp_seg = util.warp_seg(X_seg,
flow,
grid=grid,
xx=xx,
yy=yy,
zz=zz)
else: # GPU
warp_seg = nn_trf_model.predict([X_seg, pred[1]])[0, ..., 0]
vals, _ = dice(warp_seg, atlas_seg, labels=labels, nargout=2)
dice_mean = np.mean(vals)
dice_std = np.std(vals)
print('Dice mean over structures: {:.2f} ({:.2f})'.format(
dice_mean, dice_std))
def test(
gpu_id,
model_dir,
iter_num,
compute_type='GPU', # GPU or CPU
vol_size=(160, 192, 224),
nf_enc=[16, 32, 32, 32],
nf_dec=[32, 32, 32, 32, 16, 3],
save_file=None):
"""
test via segmetnation propagation
works by iterating over some iamge files, registering them to atlas,
propagating the warps, then computing Dice with atlas segmentations
"""
# GPU handling
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# load weights of model
with tf.device(gpu):
# if testing miccai run, should be xy indexing.
net = networks.miccai2018_net(vol_size,
nf_enc,
nf_dec,
use_miccai_int=True,
indexing='xy')
net.load_weights(os.path.join(model_dir, str(iter_num) + '.h5'))
# compose diffeomorphic flow output model
diff_net = keras.models.Model(net.inputs,
net.get_layer('diffflow').output)
# NN transfer model
nn_trf_model = networks.nn_trf(vol_size)
# if CPU, prepare grid
if compute_type == 'CPU':
grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4)
# prepare a matrix of dice values
dice_vals = np.zeros((len(good_labels), n_batches))
for k in range(n_batches):
# get data
vol_name, seg_name = test_brain_strings[k].split(",")
X_vol, X_seg = datagenerators.load_example_by_name(vol_name, seg_name)
# predict transform
with tf.device(gpu):
pred = diff_net.predict([X_vol, atlas_vol])
# Warp segments with flow
if compute_type == 'CPU':
flow = pred[0, :, :, :, :]
warp_seg = util.warp_seg(X_seg,
flow,
grid=grid,
xx=xx,
yy=yy,
zz=zz)
else: # GPU
# Rigid registration only by GPU
flow = pred[0, :, :, :, :]
# Compute A(all about coordinate computation)
x = np.linspace(0, 160 - 16, sample_num)
x = x.astype(np.int32)
y = np.linspace(0, 190 - 19, sample_num)
y = y.astype(np.int32)
z = np.linspace(0, 220 - 22, sample_num)
z = z.astype(np.int32)
index = np.rollaxis(np.array(np.meshgrid(x, y, z)), 0, 4)
x = index[:, :, :, 0]
y = index[:, :, :, 1]
z = index[:, :, :, 2]
# Y in formula
x_flow = np.arange(vol_size[0])
y_flow = np.arange(vol_size[1])
z_flow = np.arange(vol_size[2])
grid = np.rollaxis(np.array((np.meshgrid(y_flow, x_flow, z_flow))),
0, 4) # original coordinate
grid_x = grid_sample(x, y, z, grid[:, :, :, 0], sample_num)
grid_y = grid_sample(x, y, z, grid[:, :, :, 1], sample_num)
grid_z = grid_sample(x, y, z, grid[:, :, :, 2],
sample_num) # X (10,10,10)
sample = flow + grid
sample_x = grid_sample(x, y, z, sample[:, :, :, 0], sample_num)
sample_y = grid_sample(x, y, z, sample[:, :, :, 1], sample_num)
sample_z = grid_sample(x, y, z, sample[:, :, :, 2],
sample_num) # Y (10,10,10)
sum_x = np.sum(flow[:, :, :, 0])
sum_y = np.sum(flow[:, :, :, 1])
sum_z = np.sum(flow[:, :, :, 2])
ave_x = sum_x / (vol_size[0] * vol_size[1] * vol_size[2])
ave_y = sum_y / (vol_size[0] * vol_size[1] * vol_size[2])
ave_z = sum_z / (vol_size[0] * vol_size[1] * vol_size[2])
# formula
Y = np.zeros((10, 10, 10, grid_dimension))
X = np.zeros((10, 10, 10, grid_dimension))
T = np.array([ave_x, ave_y, ave_z, 1]) # (4,1)
# R = np.zeros((10, 10, 10, grid_dimension, grid_dimension))
for i in np.arange(10):
for j in np.arange(10):
for z in np.arange(10):
Y[i, j, z, :] = np.array([
sample_x[i, j, z], sample_y[i, j, z],
sample_z[i, j, z], 1
])
for i in np.arange(10):
for j in np.arange(10):
for z in np.arange(10):
X[i, j, z, :] = np.array([
grid_x[i, j, z], grid_y[i, j, z], grid_z[i, j, z],
1
])
X = X.reshape((1000, grid_dimension))
Y = Y.reshape((1000, grid_dimension))
R = np.dot(
np.dot(np.linalg.pinv(np.dot(np.transpose(X), X)),
np.transpose(X)), Y) # R
# build new grid(Use R to do the spatial transform)
shifted_x = np.arange(vol_size[0])
shifted_y = np.arange(vol_size[1])
shifted_z = np.arange(vol_size[2])
print(shifted_x.shape)
print(shifted_y.shape)
print(shifted_z.shape)
shifted_grid = np.rollaxis(
np.array((np.meshgrid(shifted_y, shifted_x, shifted_z))), 0, 4)
print(shifted_grid.shape)
for i in np.arange(vol_size[0]):
for j in np.arange(vol_size[1]):
for z in np.arange(vol_size[2]):
coordinates = np.dot(
R,
np.array([i, j, z, 1]).reshape(4, 1)) + T.reshape(
4, 1)
print("voxel." + '(' + str(i) + ',' + str(j) + ',' +
str(z) + ')')
shifted_grid[i, j, z, 0] = coordinates[0]
shifted_grid[i, j, z, 1] = coordinates[1]
shifted_grid[i, j, z, 2] = coordinates[2]
# interpolation
xx = np.arange(vol_size[1])
yy = np.arange(vol_size[0])
zz = np.arange(vol_size[2])
warp_seg = interpn((yy, xx, zz),
X_seg[0, :, :, :, 0],
shifted_grid,
method='nearest',
bounds_error=False,
fill_value=0)
# CVPR
grid = np.rollaxis(np.array(np.meshgrid(xx, yy, zz)), 0, 4)
sample = flow + grid
sample = np.stack(
(sample[:, :, :, 1], sample[:, :, :, 0], sample[:, :, :, 2]),
3)
warp_seg2 = interpn((yy, xx, zz),
X_seg[0, :, :, :, 0],
sample,
method='nearest',
bounds_error=False,
fill_value=0)
# compute dice
vals, _ = dice(warp_seg, atlas_seg, labels=labels, nargout=2)
vals2, _ = dice(X_seg[0, :, :, :, 0],
atlas_seg,
labels=labels,
nargout=2)
vals3, _ = dice(warp_seg2, atlas_seg, labels=labels, nargout=2)
print("dice before:")
print(np.mean(vals2), np.std(vals2))
print("dice after deformable registration:")
print(np.mean(vals3), np.std(vals3))
print("dice after rigid registration:")
print(np.mean(vals), np.std(vals))
warp_seg = nn_trf_model.predict([X_seg, pred])[0, ..., 0]
# compute Volume Overlap (Dice)
dice_vals[:, k] = dice(warp_seg, atlas_seg, labels=good_labels)
print('%3d %5.3f %5.3f' % (k, np.mean(
dice_vals[:, k]), np.mean(np.mean(dice_vals[:, :k + 1]))))
if save_file is not None:
sio.savemat(save_file, {
'dice_vals': dice_vals,
'labels': good_labels
})
def test(data_dir, fixed_image, label, device, load_model_file, DLR_model):
assert DLR_model in [
'VM', 'FAIM'
], 'DLR_model should be one of VM or FAIM, found %s' % LBR_model
# prepare data files
# inside the folder are npz files with the 'vol' and 'label'.
test_vol_names = glob.glob(os.path.join(data_dir, '*.npz'))
assert len(test_vol_names) > 0, "Could not find any testing data"
fixed_vol = np.load(fixed_image)['vol'][np.newaxis, ..., np.newaxis]
fixed_seg = np.load(fixed_image)['label']
vol_size = fixed_vol.shape[1:-1]
label = np.load(label)
# 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'
# load weights of model
with tf.device(device):
net = networks.AAN(vol_size, DLR_model)
net.load_weights(load_model_file)
# NN transfer model
nn_trf_model_nearest = networks.nn_trf(vol_size,
interp_method='nearest',
indexing='ij')
nn_trf_model_linear = networks.nn_trf(vol_size,
interp_method='linear',
indexing='ij')
dice_result = []
for test_image in test_vol_names:
X_vol, X_seg, x_boundary = datagenerators.load_example_by_name(
test_image, return_boundary=True)
with tf.device(device):
pred = net.predict([X_vol, fixing_vol, x_boundary])
warp_vol = nn_trf_model_linear.predict([X_vol, pred[1]])[0, ..., 0]
warp_seg = nn_trf_model_nearest.predict([X_seg, pred[1]])[0, ...,
0]
vals, _ = dice(warp_seg, fixing_seg, label, nargout=2)
dice_result.append(vals)
print('Dice mean: {:.3f} ({:.3f})'.format(np.mean(vals), np.std(vals)))
dice_result = np.array(dice_result)
print('Average dice mean: {:.3f} ({:.3f})'.format(np.mean(dice_result),
np.std(dice_result)))
def test(expr_name,
epoch,
gpu_id,
sample_num,
dataset,
nf_enc=[16, 32, 32, 32],
nf_dec=[32, 32, 32, 32, 32, 16, 16]):
"""
test
nf_enc and nf_dec
#nf_dec = [32,32,32,32,32,16,16,3]
# This needs to be changed. Ideally, we could just call load_model, and we wont have to
# specify the # of channels here, but the load_model is not working with the custom loss...
"""
# load subject test
print("load_data start")
X_train, y_train = load_data(data_dir='../data',
sample_num=sample_num,
mode=dataset,
fixed='joyoungje')
vol_size = y_train.shape[1:-1]
# gpu handling
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# load weights of model
with tf.device(gpu):
net = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
print("model load weights")
net.load_weights(
os.path.join("../models", expr_name, "{}.h5".format(epoch)))
# NN transfer model
nn_trf_model = networks.nn_trf(vol_size, indexing='ij')
# # if CPU, prepare grid
# if compute_type == 'CPU':
# grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4)
with tf.device(gpu):
print("model predict")
pred = net.predict([X_train, y_train])
#print("nn_tft_model.predict")
#X_warp = nn_trf_model.predict([X_train, pred[1]])[0,...,0]
warp_img = nib.Nifti1Image(pred[0].reshape(pred[0].shape[1:-1]),
affine=np.eye(4))
nib.save(
warp_img,
os.path.join(
'../result', '{}_epoch{}_{}{}_registration.nii.gz'.format(
expr_name, epoch, dataset, sample_num)))
moving_img = nib.Nifti1Image(X_train.reshape(X_train.shape[1:-1]),
affine=np.eye(4))
nib.save(
moving_img,
os.path.join(
'../result',
'{}_epoch{}_{}{}_original.nii.gz'.format(expr_name, epoch, dataset,
sample_num)))
fixed_img = nib.Nifti1Image(y_train.reshape(y_train.shape[1:-1]),
affine=np.eye(4))
nib.save(
fixed_img,
os.path.join(
'../result',
'{}_epoch{}_{}{}_reference.nii.gz'.format(expr_name, epoch,
dataset, sample_num)))
def test(expr_name, epoch, gpu_id,
nf_enc=(16, 32, 32, 32), nf_dec=(32, 32, 32, 32, 32, 16, 16)):
"""
test
nf_enc and nf_dec
#nf_dec = [32,32,32,32,32,16,16,3]
# This needs to be changed. Ideally, we could just call load_model, and we wont have to
# specify the # of channels here, but the load_model is not working with the custom loss...
"""
vol_size = (256, 256, 144)
# gpu handling
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# load weights of model
with tf.device(gpu):
net = networks.cvpr2018_net(vol_size, nf_enc, nf_dec)
print("model load weights")
net.load_weights(os.path.join("../models", expr_name, "{:04d}.h5".format(epoch)))
# NN transfer model
nn_trf_model = networks.nn_trf(vol_size, indexing='ij')
# load subject test
data_dir = '../../../dataset/urinary'
except_list = [6, 97, 125, 198, 199, 228, 271]
# vol_names = [filename for filename in os.listdir(data_dir) if (int(filename.split("_")[-1].split('.')[0]) <= 10) and
# (int(filename.split("_")[-1].split('.')[0]) not in except_list)]
vol_names = [filename for filename in os.listdir(data_dir) if int(filename.split("_")[-1].split(".")[0]) in normal]
vol_names.sort()
print("data length:", len(vol_names))
vol_names = vol_names[:10]
generator = eval_gen(data_dir=data_dir, vol_names=vol_names)
# # if CPU, prepare grid
# if compute_type == 'CPU':
# grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4)
if not os.path.isdir(os.path.join('../result', '{}_epoch{}'.format(expr_name, epoch))):
os.mkdir(os.path.join('../result', '{}_epoch{}'.format(expr_name, epoch)))
for pre, delay, vol_name in generator:
with tf.device(gpu):
pred = net.predict([pre, delay])
X_warp = nn_trf_model.predict([pre, pred[1]])[0, ..., 0]
pre = np.flip(pre, 3)[0, ..., 0]
delay = np.flip(delay, 3)[0, ..., 0]
X_warp = np.flip(X_warp, 2)
pre = np.rot90(pre, 3)
delay = np.rot90(delay, 3)
X_warp = np.rot90(X_warp, 3)
pre = pre * 4096 - 1024
pre = pre.astype(np.int16)
delay = delay * 4096 - 1024
delay = delay.astype(np.int16)
X_warp = X_warp * 4096 - 1024
X_warp = X_warp.astype(np.int16)
warp_img = nib.Nifti1Image(X_warp, affine=np.eye(4))
nib.save(warp_img, os.path.join('../result', '{}_epoch{}'.format(expr_name, epoch),
"{}_registration.nii.gz".format(vol_name)))
moving_img = nib.Nifti1Image(pre, affine=np.eye(4))
nib.save(moving_img, os.path.join('../result', '{}_epoch{}'.format(expr_name, epoch),
"{}_pre.nii.gz".format(vol_name)))
fixed_img = nib.Nifti1Image(delay, affine=np.eye(4))
nib.save(fixed_img, os.path.join('../result', '{}_epoch{}'.format(expr_name, epoch),
"{}_delay.nii.gz".format(vol_name)))
print("successfully done!")
def test(
gpu_id,
iter_num,
compute_type='GPU', # GPU or CPU
vol_size=(160, 192, 224),
nf_enc=[16, 32, 32, 32],
nf_dec=[32, 32, 32, 32, 16, 3],
save_file=None):
"""
test by segmentation, compute dice between atlas_seg and warp_seg
:param gpu_id: gpu id
:param iter_num: specify the model to read
:param compute_type: CPU/GPU
:param vol_size: volume size
:param nf_enc: number of encoder
:param nf_dec: number of decoder
:param save_file: None
:return: None
"""
# GPU handling
gpu = '/gpu:' + str(gpu_id)
os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
config.allow_soft_placement = True
set_session(tf.Session(config=config))
# load weights of model
with tf.device(gpu):
# if testing miccai run, should be xy indexing.
net = networks.miccai2018_net(vol_size,
nf_enc,
nf_dec,
use_miccai_int=True,
indexing='xy')
model_dir = "/home/ys895/rigid_diff_model/"
net.load_weights(os.path.join(model_dir, str(iter_num) + '.h5'))
# compose diffeomorphic flow output model
diff_net = keras.models.Model(net.inputs,
net.get_layer('diffflow').output)
# NN transfer model
nn_trf_model = networks.nn_trf(vol_size)
# if CPU, prepare grid
if compute_type == 'CPU':
grid, xx, yy, zz = util.volshape2grid_3d(vol_size, nargout=4)
# prepare a matrix of dice values
dice_vals = np.zeros((len(good_labels), n_batches))
for k in range(n_batches):
# get data
vol_name, seg_name = test_brain_strings[k].split(",")
X_vol, X_seg = datagenerators.load_example_by_name(vol_name, seg_name)
orig_vol = X_vol
orig_seg = X_seg
theta = 0
beta = 5
omega = 0
X_seg = rotate_img(X_seg[0, :, :, :, 0],
theta=theta,
beta=beta,
omega=omega)
X_vol = rotate_img(X_vol[0, :, :, :, 0],
theta=theta,
beta=beta,
omega=omega)
X_seg = X_seg.reshape((1, ) + X_seg.shape + (1, ))
X_vol = X_vol.reshape((1, ) + X_vol.shape + (1, ))
sample_num = 30
grid_dimension = 4
# predict transform
with tf.device(gpu):
pred = diff_net.predict([X_vol, atlas_vol])
# Warp segments with flow
if compute_type == 'CPU':
flow = pred[0, :, :, :, :]
warp_seg = util.warp_seg(X_seg,
flow,
grid=grid,
xx=xx,
yy=yy,
zz=zz)
else: # GPU
flow = pred[0, :, :, :, :]
# sample coordinate(sample_num * sample_num * sample_num)
x = np.linspace(0, (vol_size[0] / sample_num) * (sample_num - 1),
sample_num)
x = x.astype(np.int32)
y = np.linspace(0, (vol_size[1] / sample_num) * (sample_num - 1),
sample_num)
y = y.astype(np.int32)
z = np.linspace(0, (vol_size[2] / sample_num) * (sample_num - 1),
sample_num)
z = z.astype(np.int32)
index = np.rollaxis(np.array(np.meshgrid(y, x, z)), 0, 4)
x = index[:, :, :, 1]
y = index[:, :, :, 0]
z = index[:, :, :, 2]
# Y in formula
x_flow = np.arange(vol_size[0])
y_flow = np.arange(vol_size[1])
z_flow = np.arange(vol_size[2])
grid = np.rollaxis(np.array((np.meshgrid(y_flow, x_flow, z_flow))),
0, 4) # original coordinate
grid_x = grid_sample(x, y, z, grid[:, :, :, 1], sample_num)
grid_y = grid_sample(x, y, z, grid[:, :, :, 0], sample_num)
grid_z = grid_sample(x, y, z, grid[:, :, :, 2],
sample_num) # X (10,10,10)
sample = flow + grid
sample_x = grid_sample(x, y, z, sample[:, :, :, 1], sample_num)
sample_y = grid_sample(x, y, z, sample[:, :, :, 0], sample_num)
sample_z = grid_sample(x, y, z, sample[:, :, :, 2],
sample_num) # Y (10,10,10)
sum_x = np.sum(flow[:, :, :, 1])
sum_y = np.sum(flow[:, :, :, 0])
sum_z = np.sum(flow[:, :, :, 2])
ave_x = sum_x / (vol_size[0] * vol_size[1] * vol_size[2])
ave_y = sum_y / (vol_size[0] * vol_size[1] * vol_size[2])
ave_z = sum_z / (vol_size[0] * vol_size[1] * vol_size[2])
# formula
Y = np.zeros((sample_num, sample_num, sample_num, grid_dimension))
X = np.zeros((sample_num, sample_num, sample_num, grid_dimension))
T = np.array([ave_x, ave_y, ave_z, 1]) # (4,1)
print(T)
for i in np.arange(sample_num):
for j in np.arange(sample_num):
for z in np.arange(sample_num):
Y[i, j, z, :] = np.array([
sample_x[i, j, z], sample_y[i, j, z],
sample_z[i, j, z], 1
])
#Y[i, j, z, :] = Y[i, j, z, :] - np.array([ave_x, ave_y, ave_z, 0]) # amend: Y` = Y - T
for i in np.arange(sample_num):
for j in np.arange(sample_num):
for z in np.arange(sample_num):
X[i, j, z, :] = np.array([
grid_x[i, j, z], grid_y[i, j, z], grid_z[i, j, z],
1
])
X = X.reshape(
(sample_num * sample_num * sample_num, grid_dimension))
Y = Y.reshape(
(sample_num * sample_num * sample_num, grid_dimension))
R = np.dot(
np.dot(np.linalg.pinv(np.dot(np.transpose(X), X)),
np.transpose(X)), Y) # R(4, 4)
print(R)
beta = -(beta / 180) * math.pi
R = np.array([[math.cos(beta), 0, -math.sin(beta), 0],
[0, 1, 0, 0], [math.sin(beta), 0,
math.cos(beta), 0], [0, 0, 0, 1]])
#R = R.transpose()
# build new grid(Use R to do the spatial transform)
shifted_x = np.arange(vol_size[0])
shifted_y = np.arange(vol_size[1])
shifted_z = np.arange(vol_size[2])
shifted_grid = np.rollaxis(
np.array((np.meshgrid(shifted_y, shifted_x, shifted_z))), 0, 4)
# some required matrixs
T1 = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0],
[
-int(vol_size[0] / 2), -int(vol_size[1] / 2),
-int(vol_size[2] / 2), 1
]])
T2 = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0],
[
int(vol_size[0] / 2),
int(vol_size[1] / 2),
int(vol_size[2] / 2), 1
]])
for i in np.arange(vol_size[0]):
for j in np.arange(vol_size[1]):
for z in np.arange(vol_size[2]):
#coordinates = np.dot(R, np.array([i, j, z, 1]).reshape(4, 1)) + T.reshape(4, 1)
coordinates = np.dot(
np.dot(
np.dot(
np.array([i, j, z, 1]).reshape(1, 4), T1),
R), T2) # new implementation
# print("voxel." + '(' + str(i) + ',' + str(j) + ',' + str(z) + ')')
shifted_grid[i, j, z, 1] = coordinates[0, 0]
shifted_grid[i, j, z, 0] = coordinates[0, 1]
shifted_grid[i, j, z, 2] = coordinates[0, 2]
# interpolation
xx = np.arange(vol_size[1])
yy = np.arange(vol_size[0])
zz = np.arange(vol_size[2])
shifted_grid = np.stack(
(shifted_grid[:, :, :, 1], shifted_grid[:, :, :, 0],
shifted_grid[:, :, :, 2]), 3
) # notice: the shifted_grid is reverse in x and y, so this step is used for making it back.
warp_seg = interpn((yy, xx, zz),
X_seg[0, :, :, :, 0],
shifted_grid,
method='nearest',
bounds_error=False,
fill_value=0) # rigid registration
warp_vol = interpn((yy, xx, zz),
X_vol[0, :, :, :, 0],
shifted_grid,
method='nearest',
bounds_error=False,
fill_value=0) # rigid registration
# compute Volume Overlap (Dice)
dice_vals[:, k] = dice(warp_seg,
orig_seg[0, :, :, :, 0],
labels=good_labels)
print('%3d %5.3f %5.3f' % (k, np.mean(
dice_vals[:, k]), np.mean(np.mean(dice_vals[:, :k + 1]))))
if save_file is not None:
sio.savemat(save_file, {
'dice_vals': dice_vals,
'labels': good_labels
})
# specify slice
num_slice = 90
plt.figure()
plt.subplot(1, 3, 1)
plt.imshow(orig_vol[0, :, num_slice, :, 0])
plt.subplot(1, 3, 2)
plt.imshow(X_vol[0, :, num_slice, :, 0])
plt.subplot(1, 3, 3)
plt.imshow(warp_vol[:, num_slice, :])
plt.savefig("slice" + str(num_slice) + '_' + str(k) + ".png")
plt.figure()
plt.subplot(1, 3, 1)
plt.imshow(flow[:, num_slice, :, 1])
plt.subplot(1, 3, 2)
plt.imshow(flow[:, num_slice, :, 0])
plt.subplot(1, 3, 3)
plt.imshow(flow[:, num_slice, :, 2])
plt.savefig("flow.png")
