def visualize(args, path):
model = ProbabilisticUnet(input_channels=1, num_classes=1, num_filters=[32,64,128,192], latent_dim=2, no_convs_fcomb=4, beta=10.0)
model.to(device)
model.load_state_dict(torch.load(path))
task_dir = args.task
testset = MedicalDataset(task_dir=task_dir, mode='test')
testloader = data.DataLoader(testset, batch_size=1, shuffle=False)
model.eval()
with torch.no_grad():
while testset.iteration < args.test_iteration:
x, y = testset.next()
x, y = torch.from_numpy(x).unsqueeze(0).cuda(), torch.from_numpy(y).unsqueeze(0).cuda()
#output = torch.nn.Sigmoid()(model(x))
#output = torch.round(output)
output = model.forward(x,y,training=True)
output = torch.round(output)
#elbo = model.elbo(y)
# reg_loss = l2_regularisation(model.posterior) + l2_regularisation(model.prior) + l2_regularisation(model.fcomb.layers)
# valid_loss = -elbo + 1e-5 * reg_loss
print (x.size(), y.size(), output.size())
grid = torch.cat((x,y,output), dim=0)
torchvision.utils.save_image(grid, './save/'+testset.task_dir+'prediction'+str(testset.iteration)+'.png', nrow=8, padding=2, pad_value=1)
# eval_indices = indices
eval_sampler = SubsetRandomSampler(eval_indices)
eval_loader = DataLoader(dataset, batch_size=1, sampler=eval_sampler)
print("Number of test patches:", (len(eval_indices)))
# model
net = ProbabilisticUnet(input_channels=1,
num_classes=1,
num_filters=[32, 64, 128, 192],
latent_dim=2,
no_convs_fcomb=4,
beta=10.0)
if LOAD_MODEL_FROM is not None:
import os
net.load_state_dict(
torch.load(os.path.join("./saved_checkpoints/", LOAD_MODEL_FROM)))
net.to(device)
net.eval()
def energy_distance(seg_samples, gt_seg_modes, num_samples=2):
num_modes = 4 # fixed for LIDC
# if num_samples != len(seg_samples) or num_samples != len(gt_seg_modes):
# raise ValueError
d_matrix_YS = np.zeros(shape=(num_modes, num_samples), dtype=np.float32)
d_matrix_YY = np.zeros(shape=(num_modes, num_modes), dtype=np.float32)
d_matrix_SS = np.zeros(shape=(num_samples, num_samples), dtype=np.float32)
batch_size=batch_size_val,
sampler=test_sampler)
print("Number of test patches:", len(test_indices))
# network
net = ProbabilisticUnet(input_channels=1,
num_classes=1,
num_filters=[32, 64, 128, 192],
latent_dim=2,
no_convs_fcomb=4,
beta=10.0)
net.cuda()
# load pretrained model
cpk_name = os.path.join(cpk_directory, 'model_dict.pth')
net.load_state_dict(torch.load(cpk_name))
net.eval()
with torch.no_grad():
for step, (patch, mask, _) in enumerate(test_loader):
if step >= save_batches_n:
break
patch = patch.cuda()
mask = mask.cuda()
mask = torch.unsqueeze(mask, 1)
output_samples = []
for i in range(samples_per_example):
net.forward(patch, mask, training=True)
output_samples.append(
torch.sigmoid(net.sample()).detach().cpu().numpy())