def test(gpu, ref_dir, mov_dir, model, init_model_file):
"""
model training function
:param gpu: integer specifying the gpu to use
:param atlas_file: atlas filename. So far we support npz file with a 'vol' variable
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param init_model_file: the model directory to load from
"""
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
device = "cuda"
# Prepare the vm1 or vm2 model and send to device
nf_enc = [16, 32, 32, 32]
if model == "vm1":
nf_dec = [32, 32, 32, 32, 8, 8]
elif model == "vm2":
nf_dec = [32, 32, 32, 32, 32, 16, 16]
# Set up model
vol_size = [2912, 2912]
model = cvpr2018_net(vol_size, nf_enc, nf_dec)
model.to(device)
model.load_state_dict(
torch.load(init_model_file, map_location=lambda storage, loc: storage))
# set up
ref_vol_names = glob.glob(os.path.join(ref_dir, '*.npy'))
mov_vol_names = glob.glob(os.path.join(mov_dir, '*npy'))
nums = len(ref_vol_names)
for k in range(0, nums):
refs, movs = datagenerators.example_gen(ref_vol_names,
mov_vol_names,
batch_size=1)
input_fixed = torch.from_numpy(refs).to(device).float()
input_fixed = input_fixed.permute(0, 3, 1, 2)
input_moving = torch.from_numpy(movs).to(device).float()
input_moving = input_moving.permute(0, 3, 1, 2)
# Use this to warp segments
# trf = SpatialTransformer(input_fixed.shape[2:], mode='nearest')
# trf.to(device)
warp, flow = model(input_moving, input_fixed)
flow_save = sitk.GetImageFromArray(flow.cpu().detach().numpy())
# sitk.WriteImage(flow_save,'D:\peizhunsd\data\\flow_img\\' + str(k) + '.nii')
# 位移向量场的可视化
# addimage(input_fixed,input_moving,warp,k) # 可视化结果
dice_score = metrics.dice_score(warp, input_fixed)
print('相似性度量dice:', dice_score)
def test_npy(gpu, npyfile, csv_path, model, init_model_file):
"""
for npy FIRE
model training function
:param gpu: integer specifying the gpu to use
:param atlas_file: atlas filename. So far we support npz file with a 'vol' variable
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param init_model_file: the model directory to load from
"""
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
device = "cuda"
# Prepare the vm1 or vm2 model and send to device
nf_enc = [16, 32, 32, 32]
if model == "vm1":
nf_dec = [32, 32, 32, 32, 8, 8]
elif model == "vm2":
nf_dec = [32, 32, 32, 32, 32, 16, 16]
# Set up model
vol_size = [2912, 2912]
model = cvpr2018_net(vol_size, nf_enc, nf_dec)
model.to(device)
model.load_state_dict(
torch.load(init_model_file, map_location=lambda storage, loc: storage))
# set up
with open(csv_path, 'r') as f:
reader = csv.reader(f)
for row in reader[1:]:
fixed_id_name = npyfile + row[0] + '.npy'
moving_id_name = npyfile + row[1] + '.npy'
refs = np.load(fixed_id_name)[np.newaxis, ..., np.newaxis]
movs = np.load(moving_id_name)[np.newaxis, ..., np.newaxis]
input_fixed = torch.from_numpy(refs).to(device).float()
input_fixed = input_fixed.permute(0, 3, 1, 2)
input_moving = torch.from_numpy(movs).to(device).float()
input_moving = input_moving.permute(0, 3, 1, 2)
# Use this to warp segments
# trf = SpatialTransformer(input_fixed.shape[2:], mode='nearest')
# trf.to(device)
warp, flow = model(input_moving, input_fixed)
# 位移向量场的可视化
# addimage(input_fixed,input_moving,warp,k) # 可视化结果
dice_score = metrics.dice_score(warp, input_fixed)
print('相似性度量dice:', row[0], 'and', row[1], dice_score)
def test_hippocampusMRI(gpu, niifile, csv_path, model, init_model_file):
"""
for hippocampus MRI
model training function
:param gpu: integer specifying the gpu to use
:param atlas_file: atlas filename. So far we support npz file with a 'vol' variable
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param init_model_file: the model directory to load from
"""
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
device = "cuda"
# Prepare the vm1 or vm2 model and send to device
nf_enc = [16, 32, 32, 32]
if model == "vm1":
nf_dec = [32, 32, 32, 32, 8, 8]
elif model == "vm2":
nf_dec = [32, 32, 32, 32, 32, 16, 16]
# Set up model
vol_size = [2912, 2912]
model = cvpr2018_net(vol_size, nf_enc, nf_dec)
model.to(device)
model.load_state_dict(
torch.load(init_model_file, map_location=lambda storage, loc: storage))
# set up
with open(csv_path, 'r') as f:
reader = csv.reader(f)
for row in reader[1:]:
fixed_id_name = niifile + 'hippocampus_' + row[0] + '.nii.gz'
moving_id_name = niifile + 'hippocampus_' + row[1] + '.nii.gz'
X = sitk.ReadImage(fixed_id_name)
X = sitk.GetArrayFromImage(X)[np.newaxis, np.newaxis, ...]
Y = sitk.ReadImage(moving_id_name)
Y = sitk.GetArrayFromImage(Y)[np.newaxis, np.newaxis, ...]
input_fixed = torch.from_numpy(X).to(device).float()
input_fixed = input_fixed.permute(0, 3, 1, 2)
input_moving = torch.from_numpy(Y).to(device).float()
input_moving = input_moving.permute(0, 3, 1, 2)
warp, flow = model(input_moving, input_fixed)
dice_score = metrics.dice_score(warp, input_fixed)
print('相似性度量dice:', row[0], 'and', row[1], dice_score)
def test(gpu, size, data_dir, atlas_dir, model, init_model_file, saveDir,
nr_val_data):
"""
model testing function
:param gpu: integer specifying the gpu to use
:param size: integer related to desired size of the input images
:param data_dir: String describing the location of the data
:param atlas_dir: String where atlases are located
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param init_model_file: the model directory to load from
:param saveDir: String specifiying the direction to store the outputs
:param nr_val_data: the number of validation samples must corresponed to the valiable for the train function
"""
#set gpu
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
device = "cuda"
#set size
vol_size = np.array([size, size, size])
# Prepare the vm1 or vm2 model and send to device
nf_enc = [16, 32, 32, 32]
if model == "vm1":
nf_dec = [32, 32, 32, 32, 8, 8]
elif model == "vm2":
nf_dec = [32, 32, 32, 32, 32, 16, 16]
#sim_loss_fn = losses.ncc_loss if data_loss == "ncc" else losses.mse_loss
sim_loss_fn = losses.mse_loss
# Set up model
model = cvpr2018_net(vol_size, nf_enc, nf_dec)
model.to(device)
model.load_state_dict(
torch.load(init_model_file, map_location=lambda storage, loc: storage))
# Use this to warp segments
trf = SpatialTransformer(vol_size, mode='nearest')
trf.to(device)
test_strings = glob.glob(os.path.join(data_dir, '*.nii'))
test_strings = test_strings[:nr_val_data]
test_strings = [i for i in test_strings if "L1-L3" in i]
#test_strings = ['.\\Data\\Train\\64\\3YQ_unhealthyL2_L1-L3.nii','.\\Data\\Train\\64\\x1b_unhealthyL2_L1-L3.nii' ]
mean_val_loss = 0
#iteration over the test volumes
for k in range(0, len(test_strings)):
#load the data create fixed and moving image
X_vol, atlas_vol = datagenerators.load_example_by_name(
test_strings[k], atlas_dir, size)
input_fixed = torch.from_numpy(atlas_vol).to(device).float()
input_fixed = input_fixed.permute(0, 4, 1, 2, 3)
input_moving = torch.from_numpy(X_vol).to(device).float()
input_moving = input_moving.permute(0, 4, 1, 2, 3)
#produce the warp field
warp, flow = model(input_moving, input_fixed)
flow_vectors = flow[0]
shape = flow_vectors.shape
#plot the middle slice of the vector field
#flow_vectors = flow_vectors.permute(1, 2, 3, 0)
#flow_vectors = flow_vectors.detach().cpu().numpy()
#flow_vectors_middle = flow_vectors[:,:,int(shape[2]/2),0:2]
#print('shape')
#print(flow_vectors_middle.shape)
#flow_vectors_middle = flow_vectors_middle.squeeze()
#fig, axes = neuron.plot.flow([flow_vectors_middle], width=5,show = False)
#print(type(fig))
#fig.savefig(os.path.join(get_outputs_path(), test_strings[k][len(data_dir):-4] +'.png'))
# generate the new sample with the vector field
warped = trf(input_moving, flow).detach().cpu().numpy()
mean_val_loss += sim_loss_fn(warp, input_fixed)
warped = np.squeeze(warped)
#print(warped.shape)
#plot_middle_slices(warped)
# store the generated volume
img = nib.Nifti1Image(warped, np.eye(4))
nib.save(img, os.path.join(saveDir, test_strings[k][len(data_dir):]))
mean_val_loss /= len(test_strings)
print("Mean validation loss: ")
print(mean_val_loss)
def train(gpu, data_dir, atlas_file, lr, n_iter, data_loss, model, reg_param,
batch_size, n_save_iter, model_dir):
"""
model training function
:param gpu: integer specifying the gpu to use
:param data_dir: folder with npz files for each subject.
:param atlas_file: atlas filename. So far we support npz file with a 'vol' variable
:param lr: learning rate
:param n_iter: number of training iterations
:param data_loss: data_loss: 'mse' or 'ncc
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper)
:param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size
:param n_save_iter: Optional, default of 500. Determines how many epochs before saving model version.
:param model_dir: the model directory to save to
"""
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
device = "cuda"
# Produce the loaded atlas with dims.:160x192x224.
atlas_vol = np.load(atlas_file)['vol'][np.newaxis, ..., np.newaxis]
vol_size = atlas_vol.shape[1:-1]
# Get all the names of the training data
train_vol_names = glob.glob(os.path.join(data_dir, '*.npz'))
random.shuffle(train_vol_names)
# Prepare the vm1 or vm2 model and send to device
nf_enc = [16, 32, 32, 32]
if model == "vm1":
nf_dec = [32, 32, 32, 32, 8, 8]
elif model == "vm2":
nf_dec = [32, 32, 32, 32, 32, 16, 16]
else:
raise ValueError("Not yet implemented!")
model = cvpr2018_net(vol_size, nf_enc, nf_dec)
model.to(device)
# Set optimizer and losses
opt = Adam(model.parameters(), lr=lr)
sim_loss_fn = losses.ncc_loss if data_loss == "ncc" else losses.mse_loss
grad_loss_fn = losses.gradient_loss
# data generator
train_example_gen = datagenerators.example_gen(train_vol_names, batch_size)
# set up atlas tensor
atlas_vol_bs = np.repeat(atlas_vol, batch_size, axis=0)
input_fixed = torch.from_numpy(atlas_vol_bs).to(device).float()
input_fixed = input_fixed.permute(0, 4, 1, 2, 3)
# Training loop.
for i in range(n_iter):
# Save model checkpoint
if i % n_save_iter == 0:
save_file_name = os.path.join(model_dir, '%d.ckpt' % i)
torch.save(model.state_dict(), save_file_name)
# Generate the moving images and convert them to tensors.
moving_image = next(train_example_gen)[0]
input_moving = torch.from_numpy(moving_image).to(device).float()
input_moving = input_moving.permute(0, 4, 1, 2, 3)
# Run the data through the model to produce warp and flow field
warp, flow = model(input_moving, input_fixed)
# Calculate loss
recon_loss = sim_loss_fn(warp, input_fixed)
grad_loss = grad_loss_fn(flow)
loss = recon_loss + reg_param * grad_loss
print("%d,%f,%f,%f" %
(i, loss.item(), recon_loss.item(), grad_loss.item()),
flush=True)
# Backwards and optimize
opt.zero_grad()
loss.backward()
opt.step()
def train(gpu, data_dir, size, atlas_dir, lr, n_iter, data_loss, model,
reg_param, batch_size, n_save_iter, model_dir, nr_val_data):
"""
model training function
:param gpu: integer specifying the gpu to use
:param data_dir: folder with npz files for each subject.
:param size: int desired size of the volumes: [size,size,size]
:param atlas_dir: direction to atlas folder
:param lr: learning rate
:param n_iter: number of training iterations
:param data_loss: data_loss: 'mse' or 'ncc
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param reg_param: the smoothness/reconstruction tradeoff parameter (lambda in CVPR paper)
:param batch_size: Optional, default of 1. can be larger, depends on GPU memory and volume size
:param n_save_iter: Optional, default of 500. Determines how many epochs before saving model version.
:param model_dir: the model directory to save to
:param nr_val_data: number of validation examples that should be separated from the training data
"""
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
device = "cuda"
vol_size = np.array([size, size, size])
# Get all the names of the training data
vol_names = glob.glob(os.path.join(data_dir, '*.nii'))
#random.shuffle(vol_names)
#test_vol_names = vol_names[-nr_val_data:]
test_vol_names = vol_names[:nr_val_data]
#test_vol_names = [i for i in test_vol_names if "L2-L4" in i]
print(
'these volumes are separated from the data and serve as validation data : '
)
print(test_vol_names)
#train_vol_names = vol_names[:-nr_val_data]
train_vol_names = vol_names[nr_val_data:]
#train_vol_names = [i for i in train_vol_names if "L2-L4" in i]
random.shuffle(train_vol_names)
writer = SummaryWriter(get_outputs_path())
# Prepare the vm1 or vm2 model and send to device
nf_enc = [16, 32, 32, 32]
if model == "vm1":
nf_dec = [32, 32, 32, 32, 8, 8]
elif model == "vm2":
nf_dec = [32, 32, 32, 32, 32, 16, 16]
else:
raise ValueError("Not yet implemented!")
model = cvpr2018_net(vol_size, nf_enc, nf_dec)
model.to(device)
# Set optimizer and losses
opt = Adam(model.parameters(), lr=lr)
sim_loss_fn = losses.ncc_loss if data_loss == "ncc" else losses.mse_loss
grad_loss_fn = losses.gradient_loss
# data generator
train_example_gen = datagenerators.example_gen(train_vol_names, atlas_dir,
size, batch_size)
# Training loop.
for i in range(n_iter):
# Save model checkpoint and plot validation score
if i % n_save_iter == 0:
save_file_name = os.path.join(model_dir, '%d.ckpt' % i)
torch.save(model.state_dict(), save_file_name)
# load validation data
val_example_gen = datagenerators.example_gen(
test_vol_names, atlas_dir, size, 4)
val_data = next(val_example_gen)
val_fixed = torch.from_numpy(val_data[1]).to(device).float()
val_fixed = val_fixed.permute(0, 4, 1, 2, 3)
val_moving = torch.from_numpy(val_data[0]).to(device).float()
val_moving = val_moving.permute(0, 4, 1, 2, 3)
#create validation data for the model
val_warp, val_flow = model(val_moving, val_fixed)
#calculte validation score
val_recon_loss = sim_loss_fn(val_warp, val_fixed)
val_grad_loss = grad_loss_fn(val_flow)
val_loss = val_recon_loss + reg_param * val_grad_loss
#tensorboard
writer.add_scalar('Loss/Test', val_loss, i)
#prints
print('validation')
print("%d,%f,%f,%f" % (i, val_loss.item(), val_recon_loss.item(),
val_grad_loss.item()),
flush=True)
# Generate the moving images and convert them to tensors.
data_for_network = next(train_example_gen)
input_fixed = torch.from_numpy(data_for_network[1]).to(device).float()
input_fixed = input_fixed.permute(0, 4, 1, 2, 3)
input_moving = torch.from_numpy(data_for_network[0]).to(device).float()
input_moving = input_moving.permute(0, 4, 1, 2, 3)
# Run the data through the model to produce warp and flow field
warp, flow = model(input_moving, input_fixed)
print("warp_and_flow_field")
print(warp.size())
print(flow.size())
# Calculate loss
recon_loss = sim_loss_fn(warp, input_fixed)
grad_loss = grad_loss_fn(flow)
loss = recon_loss + reg_param * grad_loss
#tensorboard
writer.add_scalar('Loss/Train', loss, i)
print("%d,%f,%f,%f" %
(i, loss.item(), recon_loss.item(), grad_loss.item()),
flush=True)
# Backwards and optimize
opt.zero_grad()
loss.backward()
opt.step()
def test(gpu, atlas_file, model, init_model_file):
"""
model training function
:param gpu: integer specifying the gpu to use
:param atlas_file: atlas filename. So far we support npz file with a 'vol' variable
:param model: either vm1 or vm2 (based on CVPR 2018 paper)
:param init_model_file: the model directory to load from
"""
os.environ["CUDA_VISIBLE_DEVICES"] = gpu
device = "cuda"
# Produce the loaded atlas with dims.:160x192x224.
atlas = np.load(atlas_file)
atlas_vol = atlas['vol'][np.newaxis, ..., np.newaxis]
atlas_seg = atlas['seg']
vol_size = atlas_vol.shape[1:-1]
# Test file and anatomical labels we want to evaluate
test_file = open('../voxelmorph/data/val_examples.txt')
test_strings = test_file.readlines()
test_strings = [x.strip() for x in test_strings]
good_labels = sio.loadmat(
'../voxelmorph/data/test_labels.mat')['labels'][0]
# Prepare the vm1 or vm2 model and send to device
nf_enc = [16, 32, 32, 32]
if model == "vm1":
nf_dec = [32, 32, 32, 32, 8, 8]
elif model == "vm2":
nf_dec = [32, 32, 32, 32, 32, 16, 16]
# Set up model
model = cvpr2018_net(vol_size, nf_enc, nf_dec)
model.to(device)
model.load_state_dict(
torch.load(init_model_file, map_location=lambda storage, loc: storage))
# set up atlas tensor
input_fixed = torch.from_numpy(atlas_vol).to(device).float()
input_fixed = input_fixed.permute(0, 4, 1, 2, 3)
# Use this to warp segments
trf = SpatialTransformer(atlas_vol.shape[1:-1], mode='nearest')
trf.to(device)
for k in range(0, len(test_strings)):
vol_name, seg_name = test_strings[k].split(",")
X_vol, X_seg = datagenerators.load_example_by_name(vol_name, seg_name)
input_moving = torch.from_numpy(X_vol).to(device).float()
input_moving = input_moving.permute(0, 4, 1, 2, 3)
warp, flow = model(input_moving, input_fixed)
# Warp segment using flow
moving_seg = torch.from_numpy(X_seg).to(device).float()
moving_seg = moving_seg.permute(0, 4, 1, 2, 3)
warp_seg = trf(moving_seg, flow).detach().cpu().numpy()
vals, labels = dice(warp_seg, atlas_seg, labels=good_labels, nargout=2)
#dice_vals[:, k] = vals
#print(np.mean(dice_vals[:, k]))
print(np.mean(vals))