def eval_model(model, dataloaders, log_dir="./log/", logger=None, opt=None):
since = time.time()
if False: #opt.do_seg:
# eval lung segmentation
logger.info("-" * 8 + "eval lung segmentation" + "-" * 8)
model.eval()
all_dices = []
all_dices_au = []
for batch_idx, (inputs,
labels) in enumerate(dataloaders["tgt_lung_seg_val"],
0):
annotation = dataloaders["tgt_lung_seg_val"].dataset.annotations[
batch_idx]
img_dir = annotation.strip().split(',')[0]
img_name = Path(img_dir).name
inputs = inputs.to(device)
# adjust labels
labels[labels == opt.xray_mask_value_dict["lung"]] = 1
labels = labels[:, -1].to(device)
labels = torch.stack([labels == c for c in range(2)], dim=1)
with torch.set_grad_enabled(False):
if opt.use_aux:
_, _, seg_logits, _, seg_logits_au = model(inputs)
else:
_, _, seg_logits, _, _ = model(inputs)
seg_probs = torch.softmax(seg_logits, dim=1)
predicted_mask = probs2one_hot(seg_probs.detach())
# change the infection to Lung
predicted_mask_lung = predicted_mask[:, :-1]
predicted_mask_lung[:, -1] += predicted_mask[:, -1]
dices = dice_coef(
predicted_mask_lung,
labels.detach().type_as(predicted_mask)).cpu().numpy()
all_dices.append(dices) # [(B,C)]
predicted_mask_lung = predicted_mask_lung.squeeze().cpu(
).numpy() # 3xwxh
mask_inone = (np.zeros_like(predicted_mask_lung[0]) +
predicted_mask_lung[1] * 255).astype(np.uint8)
# save dir:
save_dir = os.path.join(opt.logs, "tgt_lung_seg_val", "eval")
#
if not os.path.exists(save_dir):
os.makedirs(save_dir)
cv2.imwrite(os.path.join(save_dir, img_name), mask_inone)
###################################################au
if opt.use_aux:
seg_probs_au = torch.softmax(seg_logits_au, dim=1)
predicted_mask_au = probs2one_hot(seg_probs_au.detach())
# change the infection to Lung
predicted_mask_lung_au = predicted_mask_au[:, :-1]
predicted_mask_lung_au[:, -1] += predicted_mask_au[:, -1]
dices_au = dice_coef(
predicted_mask_lung_au,
labels.detach().type_as(
predicted_mask_au)).cpu().numpy()
all_dices_au.append(dices_au) # [(B,C)]
predicted_mask_lung_au = predicted_mask_lung_au.squeeze(
).cpu().numpy() # 3xwxh
mask_inone_au = (np.zeros_like(predicted_mask_lung_au[0]) +
predicted_mask_lung_au[1] * 255).astype(
np.uint8)
# save dir:
save_dir_au = os.path.join(opt.logs, "tgt_lung_seg_val_au",
"eval")
#
if not os.path.exists(save_dir_au):
os.makedirs(save_dir_au)
cv2.imwrite(os.path.join(save_dir_au, img_name),
mask_inone_au)
avg_dice = np.mean(np.concatenate(all_dices, 0), 0) #
logger.info(
"tgt_lung_seg_val:[%d/%d],dice0:%.03f,dice1:%.03f,dice:%.03f" %
(batch_idx, len(dataloaders['tgt_lung_seg_val'].dataset) //
inputs.shape[0], avg_dice[0], avg_dice[1],
np.mean(np.concatenate(all_dices, 0))))
if opt.use_aux:
avg_dice_au = np.mean(np.concatenate(all_dices_au, 0), 0) #
logger.info(
"tgt_lung_seg_val_au:[%d/%d],dice0:%.03f,dice1:%.03f,dice:%.03f"
%
(batch_idx, len(dataloaders['tgt_lung_seg_val'].dataset) //
inputs.shape[0], avg_dice_au[0], avg_dice_au[1],
np.mean(np.concatenate(all_dices_au, 0))))
if True:
# eval infection segmentation and cls
logger.info("-" * 8 + "eval infection cls" + "-" * 8)
model.eval()
val_gt = []
val_cls_pred = []
val_cls_probs = [] # for VOC
val_seg_pred = []
val_seg_probs = [] # for VOC
val_seg_probs_au = []
val_seg_pred_au = [] # for VOC
for batch_idx, (inputs,
labels) in enumerate(dataloaders["tgt_cls_val"], 0):
inputs = inputs.to(device)
# adjust label
val_gt.append(labels.cpu().data.numpy())
with torch.set_grad_enabled(False):
annotation = dataloaders["tgt_cls_val"].dataset.annotations[
batch_idx]
img_dir = annotation.strip().split(',')[0]
img_name = Path(img_dir).name
if opt.use_aux:
cls_logits, _, seg_logits, _, seg_logits_au = model(inputs)
else:
cls_logits, _, seg_logits, _, _ = model(inputs)
if opt.do_seg:
seg_probs = torch.softmax(seg_logits, dim=1)
val_seg_probs.append(seg_probs[:, -1:].detach().cpu().view(
seg_probs.shape[0], 1, -1).max(-1)[0])
predicted_mask_onehot = probs2one_hot(seg_probs.detach())
# for save
predicted_mask = predicted_mask_onehot.squeeze().cpu(
).numpy() # 3xwxh
mask_inone = (np.zeros_like(predicted_mask[0]) +
predicted_mask[1] * 128 +
predicted_mask[2] * 255).astype(np.uint8)
# save dir:
save_dir = os.path.join(opt.logs, "tgt_cls_val", "eval")
#
if not os.path.exists(save_dir):
os.makedirs(save_dir)
cv2.imwrite(os.path.join(save_dir, img_name), mask_inone)
# seg2cls
preds_cls_seg = (predicted_mask_onehot[:,
-1:].sum(-1).sum(-1)
> 0).cpu().numpy().astype(np.uint8)
val_seg_pred.append(preds_cls_seg)
if opt.do_seg and opt.use_aux:
seg_probs_au = torch.softmax(seg_logits_au, dim=1)
val_seg_probs_au.append(
seg_probs_au[:, -1:].detach().cpu().view(
seg_probs_au.shape[0], 1, -1).max(-1)[0])
predicted_mask_onehot_au = probs2one_hot(
seg_probs_au.detach())
# for save
predicted_mask_au = predicted_mask_onehot_au.squeeze().cpu(
).numpy() # 3xwxh
mask_inone_au = (np.zeros_like(predicted_mask_au[0]) +
predicted_mask_au[1] * 128 +
predicted_mask_au[2] * 255).astype(
np.uint8)
# save dir:
save_dir_au = os.path.join(opt.logs, "tgt_cls_val_au",
"eval")
#
if not os.path.exists(save_dir_au):
os.makedirs(save_dir_au)
cv2.imwrite(os.path.join(save_dir_au, img_name),
mask_inone_au)
# seg2cls
preds_cls_seg_au = (
predicted_mask_onehot_au[:, -1:].sum(-1).sum(-1) >
0).cpu().numpy().astype(np.uint8)
val_seg_pred_au.append(preds_cls_seg_au)
# cls
#print(cls_logits)
if opt.do_cls:
probs_cls = torch.softmax(cls_logits, dim=1)
val_cls_probs.append(probs_cls[...,
1:].detach().cpu().numpy())
preds_cls = (probs_cls[..., 1:] > 0.5).type(torch.long)
val_cls_pred.append(preds_cls.cpu().data.numpy())
if not os.path.exists(os.path.join(opt.logs, "cf")):
os.makedirs(os.path.join(opt.logs, "cf"))
val_gt = np.concatenate(val_gt, axis=0)
if opt.do_cls:
val_cls_pred = np.concatenate(val_cls_pred, axis=0)
val_cls_probs = np.concatenate(val_cls_probs, axis=0)
save_cf_png_dir = os.path.join(opt.logs, "cf", "eval_cls_cf.png")
save_metric_dir = os.path.join(opt.logs, "eval_metric_cls.txt")
result_str = get_results(val_gt, val_cls_pred, val_cls_probs,
save_cf_png_dir, save_metric_dir)
logger.info("tgt_cls_val:[cls]: %s" % (result_str))
if opt.do_seg:
val_seg_pred = np.concatenate(val_seg_pred, axis=0)
val_seg_probs = np.concatenate(val_seg_probs, axis=0)
# seg2cls
save_cf_png_dir = os.path.join(opt.logs, "cf", "eval_seg_cf.png")
save_metric_dir = os.path.join(opt.logs, "eval_metric_seg.txt")
result_str = get_results(val_gt, val_seg_pred, val_seg_probs,
save_cf_png_dir, save_metric_dir)
logger.info("tgt_seg_val:[seg2cls]: %s" % (result_str))
if opt.do_seg and opt.use_aux:
val_seg_pred_au = np.concatenate(val_seg_pred_au, axis=0)
val_seg_probs_au = np.concatenate(val_seg_probs_au, axis=0)
# seg2cls
save_cf_png_dir_au = os.path.join(opt.logs, "cf",
"eval_seg_au_cf.png")
save_metric_dir_au = os.path.join(opt.logs,
"eval_metric_seg_au.txt")
result_str_au = get_results(val_gt, val_seg_pred_au,
val_seg_probs_au, save_cf_png_dir_au,
save_metric_dir_au)
logger.info("tgt_seg_au_val:[seg2cls]: %s" % (result_str_au))
time_elapsed = time.time() - since
logger.info("Eval complete in {:.0f}m {:.0f}s".format(
time_elapsed // 60, time_elapsed % 60))
print(file)
image = np.load(os.path.join(root, 'in_npy', file))
gt = np.load(os.path.join(root, 'gt_npy', file))
if len(np.unique(gt)) > 0:
#print('infering {} of shape {} and classes {}, max {} and min {} '.format( file, image.shape, np.unique(gt), image.max(), image.min()))
image = image.reshape(-1, 1, 256, 256) / 255.00
image = torch.tensor(image, dtype=torch.float)
image = Variable(image, requires_grad=True)
pred = net(image)
pred = F.softmax(pred, dim=1).to('cpu')
predicted_output = probs2one_hot(pred.detach())
#print(predicted_output.to('cpu')[:,:2:].shape,class2one_hot(torch.tensor(gt).to('cpu'), n_classes).shape )
#np.save(os.path.join(path, 'predictions', '{}'.format(file)), pred.to('cpu').detach().numpy())
#dice = dice_coef(predicted_output.to('cpu'), class2one_hot(torch.tensor(gt).to('cpu'), n_classes))[:,n,]
dice = dice_coef(
predicted_output.to('cpu'),
class2one_hot(torch.tensor(gt).to('cpu'), n_classes))[:, n, ]
hauss = haussdorf(predicted_output,
class2one_hot(torch.tensor(gt), n_classes))[:,
n, ]
'''
fig, ax = plt.subplots()
ax.imshow(np.argmax(predicted_output.detach().numpy(), axis=1)[0], cmap=plt.cm.gray)
#r, contours= Get_contour_characteristics(np.argmax(predicted_output.detach().numpy(), axis=1)[0])
g, contours = Get_contour_characteristics(np.array(gt).round())
total_summ = 0
for n, contour in enumerate(contours):
ax.plot(contour[:, 1].astype(int), contour[:, 0].astype(int),color='red', linewidth=-1)
def do_epoch(mode: str, net: Any, device: Any, loader: DataLoader, epc: int,
loss_fns: List[Callable], loss_weights: List[float], C: int,
savedir: str = "", optimizer: Any = None,
metric_axis: List[int] = [1], compute_haussdorf: bool = False) \
-> Tuple[Tensor, Tensor, Tensor, Tensor]:
assert mode in ["train", "val"]
L: int = len(loss_fns)
if mode == "train":
net.train()
desc = f">> Training ({epc})"
elif mode == "val":
net.eval()
desc = f">> Validation ({epc})"
total_iteration, total_images = len(loader), len(loader.dataset)
all_dices: Tensor = torch.zeros((total_images, C),
dtype=torch.float32,
device=device)
batch_dices: Tensor = torch.zeros((total_iteration, C),
dtype=torch.float32,
device=device)
loss_log: Tensor = torch.zeros((total_iteration),
dtype=torch.float32,
device=device)
haussdorf_log: Tensor = torch.zeros((total_images, C),
dtype=torch.float32,
device=device)
tq_iter = tqdm_(enumerate(loader), total=total_iteration, desc=desc)
done: int = 0
for j, data in tq_iter:
data[1:] = [e.to(device)
for e in data[1:]] # Move all tensors to device
filenames, image, target = data[:3]
labels = data[3:3 + L]
bounds = data[3 + L:]
assert len(labels) == len(bounds)
B = len(image)
# Reset gradients
if optimizer:
optimizer.zero_grad()
# Forward
pred_logits: Tensor = net(image)
pred_probs: Tensor = F.softmax(pred_logits, dim=1)
predicted_mask: Tensor = probs2one_hot(
pred_probs.detach()) # Used only for dice computation
assert len(bounds) == len(loss_fns) == len(loss_weights)
ziped = zip(loss_fns, labels, loss_weights, bounds)
losses = [
w * loss_fn(pred_probs, label, bound)
for loss_fn, label, w, bound in ziped
]
loss = reduce(add, losses)
assert loss.shape == (), loss.shape
# Backward
if optimizer:
loss.backward()
optimizer.step()
# Compute and log metrics
loss_log[j] = loss.detach()
sm_slice = slice(done, done + B) # Values only for current batch
dices: Tensor = dice_coef(predicted_mask, target.detach())
assert dices.shape == (B, C), (dices.shape, B, C)
all_dices[sm_slice, ...] = dices
if B > 1 and mode == "val":
batch_dice: Tensor = dice_batch(predicted_mask, target.detach())
assert batch_dice.shape == (C, ), (batch_dice.shape, B, C)
batch_dices[j] = batch_dice
if compute_haussdorf:
haussdorf_res: Tensor = haussdorf(predicted_mask.detach(),
target.detach())
assert haussdorf_res.shape == (B, C)
haussdorf_log[sm_slice] = haussdorf_res
# Save images
if savedir:
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UserWarning)
predicted_class: Tensor = probs2class(pred_probs)
save_images(predicted_class, filenames, savedir, mode, epc)
# Logging
big_slice = slice(0,
done + B) # Value for current and previous batches
dsc_dict = {
f"DSC{n}": all_dices[big_slice, n].mean()
for n in metric_axis
}
hauss_dict = {
f"HD{n}": haussdorf_log[big_slice, n].mean()
for n in metric_axis
} if compute_haussdorf else {}
batch_dict = {
f"bDSC{n}": batch_dices[:j, n].mean()
for n in metric_axis
} if B > 1 and mode == "val" else {}
mean_dict = {
"DSC": all_dices[big_slice, metric_axis].mean(),
"HD": haussdorf_log[big_slice, metric_axis].mean()
} if len(metric_axis) > 1 else {}
stat_dict = {
**dsc_dict,
**hauss_dict,
**mean_dict,
**batch_dict, "loss": loss_log[:j].mean()
}
nice_dict = {k: f"{v:.3f}" for (k, v) in stat_dict.items()}
tq_iter.set_postfix(nice_dict)
done += B
print(f"{desc} " + ', '.join(f"{k}={v}" for (k, v) in nice_dict.items()))
return loss_log, all_dices, batch_dices, haussdorf_log
def main(argv):
# turn off log message
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.FATAL)
# check folder
if not os.path.exists(FLAGS.dir):
raise Exception("model dirctory is not existed!")
if not os.path.exists(os.path.join(FLAGS.dir, 'dice')):
os.makedirs(os.path.join(FLAGS.dir, 'dice'))
# get ground truth list
ground_truth_list = io.load_list(FLAGS.ground_truth)
# load ground truth
ground_truth = io.load_data_from_path(ground_truth_list, dtype='int32')
# get tfrecord list
test_data_list = glob.glob(FLAGS.indir + '/*')
# load test data
test_set = tf.data.TFRecordDataset(test_data_list)
test_set = test_set.map(lambda x: utils._parse_function_val_test(
x, image_size=FLAGS.image_size),
num_parallel_calls=os.cpu_count())
test_set = test_set.repeat()
test_set = test_set.batch(FLAGS.batch_size)
test_iter = test_set.make_one_shot_iterator()
test_data = test_iter.get_next()
# initializer
init_op = tf.group(tf.initializers.global_variables(),
tf.initializers.local_variables())
with tf.Session(config=utils.config(index=FLAGS.gpu_index)) as sess:
# set network
kwargs = {
'sess': sess,
'latent_dim': FLAGS.latent_dim,
'scale_lambda': FLAGS.scale_lambda,
'scale_kappa': FLAGS.scale_kappa,
'scale_psi': FLAGS.scale_psi,
'image_size': FLAGS.image_size,
'points_num': FLAGS.points_num,
'k_size': FLAGS.k_size,
'encoder_layer': encoder_layer,
'points_encoder_layer': points_encoder_layer,
'generator_layer': generator_layer,
'discriminator_layer': discriminator_layer,
'code_discriminator_layer': code_discriminator_layer,
'is_training': False
}
Model = conditional_alphaGAN(**kwargs)
sess.run(init_op)
# print parameters
utils.cal_parameter()
# test
dice_list = []
Model.restore_model(FLAGS.dir +
'/model/model_{}'.format(FLAGS.model_index))
for i in range(FLAGS.num_of_test):
_, test_points_batch, _ = sess.run(test_data)
np.random.seed(4)
tbar = tqdm(range(FLAGS.num_of_generate // FLAGS.batch_size),
ascii=True)
for j in tbar:
z = np.random.normal(0.,
1.,
size=[FLAGS.batch_size, FLAGS.latent_dim])
# z = utils.truncated_noise_sample(FLAGS.batch_size, FLAGS.latent_dim, truncation=2.0)
generate_batch = Model.generate_sample(z, test_points_batch)
# save logodds
generate_batch_ = np.asarray(generate_batch)
generate_batch_ = generate_batch_[0, :, :, :]
for image_index in range(generate_batch_.shape[0]):
gen = generate_batch_[image_index][:, :, :, 0]
io.write_mhd_and_raw(
gen,
'{}.mhd'.format(
os.path.join(
FLAGS.dir, 'dice', '{}'.format(i),
'{}'.format(j * FLAGS.batch_size +
image_index))),
spacing=[1, 1, 1],
origin=[0, 0, 0],
compress=True)
if j is 0:
data = np.asarray(generate_batch)[0]
label = np.where(data > 0.5, 0, 1)
label = label.astype(np.int8)
pa = np.sum(label, axis=0)
else:
data = np.asarray(generate_batch)[0]
label_ = np.where(data > 0.5, 0, 1)
label_ = label_.astype(np.int8)
pa = pa + np.sum(label_, axis=0)
pa = pa / float(FLAGS.num_of_generate)
pa = pa.astype(np.float32)
# output image
io.write_mhd_and_raw(pa,
'{}_{}.mhd'.format(
os.path.join(FLAGS.dir, 'dice', 'PA'), i),
spacing=[1, 1, 1],
origin=[0, 0, 0],
compress=True)
# dice
gt = ground_truth[i]
gt = gt.astype(np.float32)
dice = utils.dice_coef(gt, pa)
dice_list.append([round(dice, 6)])
print(dice)
print('dice = %f' % np.mean(dice_list))
# write csv
io.write_csv(
dice_list,
os.path.join(FLAGS.dir, 'dice',
'dice_{}.csv'.format(FLAGS.model_index)), 'dice')
def do_epoch(args,
mode: str,
net: Any,
device: Any,
loader: DataLoader,
epc: int,
loss_fns: List[Callable],
loss_weights: List[float],
loss_fns_source: List[Callable],
loss_weights_source: List[float],
new_w: int,
num_steps: int,
C: int,
metric_axis: List[int],
savedir: str = "",
optimizer: Any = None,
target_loader: Any = None):
assert mode in ["train", "val"]
L: int = len(loss_fns)
indices = torch.tensor(metric_axis, device=device)
if mode == "train":
net.train()
desc = f">> Training ({epc})"
elif mode == "val":
net.eval()
# net.train()
desc = f">> Validation ({epc})"
total_it_s, total_images = len(loader), len(loader.dataset)
total_it_t, total_images_t = len(target_loader), len(target_loader.dataset)
total_iteration = max(total_it_s, total_it_t)
# Lazy add lines below because we will be cycling until the biggest length is reached
total_images = max(total_images, total_images_t)
total_images_t = total_images
pho = 1
dtype = eval(args.dtype)
all_dices: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
all_inter_card: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
all_card_gt: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
all_card_pred: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
loss_log: Tensor = torch.zeros((total_images), dtype=dtype, device=device)
loss_inf: Tensor = torch.zeros((total_images), dtype=dtype, device=device)
loss_cons: Tensor = torch.zeros((total_images), dtype=dtype, device=device)
loss_fs: Tensor = torch.zeros((total_images), dtype=dtype, device=device)
posim_log: Tensor = torch.zeros((total_images), dtype=dtype, device=device)
haussdorf_log: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
all_grp: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
dice_3d_log: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
dice_3d_sd_log: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
if args.source_metrics == True:
all_dices_s: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
all_inter_card_s: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
all_card_gt_s: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
all_card_pred_s: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
all_grp_s: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
dice_3d_s_log: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
dice_3d_s_sd_log: Tensor = torch.zeros((total_images, C),
dtype=dtype,
device=device)
# if len(loader)>len(target_loader):
# tq_iter = tqdm_(enumerate(zip(loader, cycle(target_loader))), total=total_iteration, desc=desc)
# elif len(loader)<len(target_loader):
# tq_iter = tqdm_(enumerate(zip(cycle(loader), target_loader)), total=total_iteration, desc=desc)
# else:
# tq_iter = tqdm_(enumerate(zip(loader, target_loader)), total=total_iteration, desc=desc)
tq_iter = tqdm_(enumerate(zip(loader, target_loader)),
total=total_iteration,
desc=desc)
#tq_iter = tqdm_(enumerate(target_loader), total=total_iteration, desc=desc)
done: int = 0
#ratio_losses = 0
n_warmup = 0
mult_lw = [pho**(epc - n_warmup + 1)] * len(loss_weights)
#if epc > 100:
# mult_lw = [pho ** 100] * len(loss_weights)
mult_lw[0] = 1
loss_weights = [a * b for a, b in zip(loss_weights, mult_lw)]
losses_vec, source_vec, target_vec, baseline_target_vec = [], [], [], []
pen_count = 0
with warnings.catch_warnings():
warnings.simplefilter("ignore")
for j, (source_data, target_data) in tq_iter:
#for j, target_data in tq_iter:
source_data[1:] = [e.to(device) for e in source_data[1:]
] # Move all tensors to device
filenames_source, source_image, source_gt = source_data[:3]
target_data[1:] = [e.to(device) for e in target_data[1:]
] # Move all tensors to device
filenames_target, target_image, target_gt = target_data[:3]
labels = target_data[3:3 + L]
labels_source = source_data[3:3 + L]
bounds = target_data[3 + L:]
bounds_source = source_data[3 + L:]
assert len(labels) == len(bounds), len(bounds)
if args.mix == False:
assert filenames_source == filenames_target
#print(filenames_source,filenames_target)
B = len(target_image)
# Reset gradients
if optimizer:
#adjust_learning_rate(optimizer, 1, args.l_rate, num_steps, args.power)
optimizer.zero_grad()
# Forward
with torch.set_grad_enabled(mode == "train"):
pred_logits: Tensor = net(target_image)
pred_logits_source: Tensor = net(source_image)
pred_probs: Tensor = F.softmax(pred_logits, dim=1)
pred_probs_source: Tensor = F.softmax(pred_logits_source,
dim=1)
if new_w > 0:
pred_probs = resize(pred_probs, new_w)
labels = [resize(label, new_w) for label in labels]
target = resize(target, new_w)
predicted_mask: Tensor = probs2one_hot(
pred_probs) # Used only for dice computation
predicted_mask_source: Tensor = probs2one_hot(
pred_probs_source) # Used only for dice computation
#print(torch.sum(predicted_mask, dim=[2,3]).cpu().numpy())
#print(list(map(lambda n: [int(f) for f in n], np.around(torch.sum(pred_probs, dim=[2,3]).detach().cpu().numpy()))))
assert len(bounds) == len(loss_fns) == len(loss_weights)
if epc < n_warmup:
loss_weights = [0] * len(loss_weights)
loss: Tensor = torch.zeros(1, requires_grad=True).to(device)
loss_vec = []
loss_k = []
for loss_fn, label, w, bound in zip(loss_fns, labels, loss_weights,
bounds):
if w > 0:
if args.lin_aug_w:
if epc < 70:
w = w * (epc + 1) / 70
loss_b = loss_fn(pred_probs, label, bound)
loss = loss + w * loss_b
#pen_count += count_b.detach()
#print(count_b.detach())
loss_k.append(w * loss_b.detach())
#for loss_fn, label, w, bound in zip(loss_fns_source, [source_gt], loss_weights_source, torch.randn(1)):
#for loss_fn, label, w, bound in zip(loss_fns_source, labels_source, loss_weights_source, torch.randn(1)):
for loss_fn, label, w, bound in zip(loss_fns_source, labels_source,
loss_weights_source,
bounds_source):
if w > 0:
loss_b = loss_fn(pred_probs_source, label, bound)
loss = loss + w * loss_b
loss_k.append(w * loss_b.detach())
#print(loss_k)
# Backward
if optimizer:
loss.backward()
optimizer.step()
# Compute and log metrics
#dices: Tensor = dice_coef(predicted_mask.detach(), target.detach())
# baseline_dices: Tensor = dice_coef(labels[0].detach(), target.detach())
#batch_dice: Tensor = dice_batch(predicted_mask.detach(), target.detach())
# assert batch_dice.shape == (C,) and dices.shape == (B, C), (batch_dice.shape, dices.shape, B, C)
dices, inter_card, card_gt, card_pred = dice_coef(
predicted_mask.detach(), target_gt.detach())
assert dices.shape == (B, C), (dices.shape, B, C)
sm_slice = slice(done, done + B) # Values only for current batch
all_dices[sm_slice, ...] = dices
# # for 3D dice
all_grp[sm_slice, ...] = int(
re.split('_', filenames_target[0])[1]) * torch.ones([1, C])
all_inter_card[sm_slice, ...] = inter_card
all_card_gt[sm_slice, ...] = card_gt
all_card_pred[sm_slice, ...] = card_pred
# 3D dice on source
if args.source_metrics == True:
dices_s, inter_card_s, card_gt_s, card_pred_s = dice_coef(
predicted_mask_source.detach(), source_gt.detach())
all_grp_s[sm_slice, ...] = int(
re.split('_', filenames_source[0])[1]) * torch.ones([1, C])
all_inter_card_s[sm_slice, ...] = inter_card_s
all_card_gt_s[sm_slice, ...] = card_gt_s
all_card_pred_s[sm_slice, ...] = card_pred_s
#loss_log[sm_slice] = loss.detach()
loss_inf[sm_slice] = loss_k[0]
if len(loss_k) > 1:
loss_cons[sm_slice] = loss_k[1]
else:
loss_cons[sm_slice] = 0
if len(loss_k) > 2:
loss_fs[sm_slice] = loss_k[2]
else:
loss_fs[sm_slice] = 0
#posim_log[sm_slice] = torch.einsum("bcwh->b", [target_gt[:, 1:, :, :]]).detach() > 0
#haussdorf_res: Tensor = haussdorf(predicted_mask.detach(), target_gt.detach(), dtype)
#assert haussdorf_res.shape == (B, C)
#haussdorf_log[sm_slice] = haussdorf_res
# # Save images
if savedir:
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UserWarning)
warnings.simplefilter("ignore")
predicted_class: Tensor = probs2class(pred_probs)
#save_images_p(predicted_class, filenames_target, args.dataset, mode, epc, False)
save_images(predicted_class, filenames_target, savedir,
mode, epc, True)
# Logging
big_slice = slice(0, done +
B) # Value for current and previous batches
stat_dict = {
"dice": torch.index_select(all_dices, 1, indices).mean(),
"loss": loss_log[big_slice].mean()
}
nice_dict = {k: f"{v:.4f}" for (k, v) in stat_dict.items()}
done += B
tq_iter.set_postfix(nice_dict)
print(f"{desc} " + ', '.join(f"{k}={v}"
for (k, v) in nice_dict.items()))
#dice_posim = torch.masked_select(all_dices[:, -1], posim_log.type(dtype=torch.uint8)).mean()
# dice3D gives back the 3d dice mai on images
# if not args.debug:
# dice_3d_log_o, dice_3d_sd_log_o = dice3d(args.workdir, f"iter{epc:03d}", mode, "Subj_\\d+_",args.dataset + mode + '/CT_GT', C)
dice_3d_log, dice_3d_sd_log = dice3dn(all_grp, all_inter_card, all_card_gt,
all_card_pred, metric_axis, True)
if args.source_metrics == True:
dice_3d_s_log, dice_3d_s_sd_log = dice3dn(all_grp_s, all_inter_card_s,
all_card_gt_s,
all_card_pred_s, metric_axis,
True)
print("mean 3d_dice over all patients:", dice_3d_log)
#source_vec = [ dice_3d_s, dice_3d_sd_s, haussdorf_log_s]
dice_2d = torch.index_select(all_dices, 1, indices).mean().cpu().numpy()
target_vec = [dice_3d_log, dice_3d_sd_log, dice_2d]
if args.source_metrics == True:
source_vec = [dice_3d_s_log, dice_3d_s_sd_log]
else:
source_vec = [0, 0]
#losses_vec = [loss_log.mean().item()]
losses_vec = [
loss_inf.mean().item(),
loss_cons.mean().item(),
loss_fs.mean().item()
]
return losses_vec, target_vec, source_vec
def do_epoch(mode: str, net: Any, device: Any, loaders: List[DataLoader], epc: int,
list_loss_fns: List[List[Callable]], list_loss_weights: List[List[float]], C: int,
savedir: str = "", optimizer: Any = None,
metric_axis: List[int] = [1], compute_haussdorf: bool = False, compute_miou: bool = False,
temperature: float = 1) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tuple[None, Tensor]]:
assert mode in ["train", "val"]
if mode == "train":
net.train()
desc = f">> Training ({epc})"
elif mode == "val":
net.eval()
desc = f">> Validation ({epc})"
total_iteration: int = sum(len(loader) for loader in loaders) # U
total_images: int = sum(len(loader.dataset) for loader in loaders) # D
n_loss: int = max(map(len, list_loss_fns))
all_dices: Tensor = torch.zeros((total_images, C), dtype=torch.float32, device=device)
batch_dices: Tensor = torch.zeros((total_iteration, C), dtype=torch.float32, device=device)
loss_log: Tensor = torch.zeros((total_iteration, n_loss), dtype=torch.float32, device=device)
haussdorf_log: Tensor = torch.zeros((total_images, C), dtype=torch.float32, device=device)
iiou_log: Tensor = torch.zeros((total_images, C), dtype=torch.float32, device=device)
intersections: Tensor = torch.zeros((total_images, C), dtype=torch.float32, device=device)
unions: Tensor = torch.zeros((total_images, C), dtype=torch.float32, device=device)
few_axis: bool = len(metric_axis) <= 3
done_img: int = 0
done_batch: int = 0
tq_iter = tqdm_(total=total_iteration, desc=desc)
for i, (loader, loss_fns, loss_weights) in enumerate(zip(loaders, list_loss_fns, list_loss_weights)):
L: int = len(loss_fns)
for data in loader:
data[1:] = [e.to(device) for e in data[1:]] # Move all tensors to device
filenames, image, target = data[:3]
assert not target.requires_grad
labels = data[3:3 + L]
bounds = data[3 + L:]
assert len(labels) == len(bounds)
B = len(image)
# Reset gradients
if optimizer:
optimizer.zero_grad()
# Forward
pred_logits: Tensor = net(image)
pred_probs: Tensor = F.softmax(temperature * pred_logits, dim=1)
predicted_mask: Tensor = probs2one_hot(pred_probs.detach()) # Used only for dice computation
assert not predicted_mask.requires_grad
assert len(bounds) == len(loss_fns) == len(loss_weights) == len(labels)
ziped = zip(loss_fns, labels, loss_weights, bounds)
losses = [w * loss_fn(pred_probs, label, bound) for loss_fn, label, w, bound in ziped]
loss = reduce(add, losses)
assert loss.shape == (), loss.shape
# if epc >= 1 and False:
# import matplotlib.pyplot as plt
# _, axes = plt.subplots(nrows=1, ncols=3)
# axes[0].imshow(image[0, 0].cpu().numpy(), cmap='gray')
# axes[0].contour(target[0, 1].cpu().numpy(), cmap='rainbow')
# pred_np = pred_probs[0, 1].detach().cpu().numpy()
# axes[1].imshow(pred_np)
# bins = np.linspace(0, 1, 50)
# axes[2].hist(pred_np.flatten(), bins)
# print(bounds)
# print(bounds[2].cpu().numpy())
# print(bounds[2][0, 1].cpu().numpy())
# print(pred_np.sum())
# plt.show()
# Backward
if optimizer:
loss.backward()
optimizer.step()
# Compute and log metrics
# loss_log[done_batch] = loss.detach()
for j in range(len(loss_fns)):
loss_log[done_batch, j] = losses[j].detach()
sm_slice = slice(done_img, done_img + B) # Values only for current batch
dices: Tensor = dice_coef(predicted_mask, target)
assert dices.shape == (B, C), (dices.shape, B, C)
all_dices[sm_slice, ...] = dices
if B > 1 and mode == "val":
batch_dice: Tensor = dice_batch(predicted_mask, target)
assert batch_dice.shape == (C,), (batch_dice.shape, B, C)
batch_dices[done_batch] = batch_dice
if compute_haussdorf:
haussdorf_res: Tensor = haussdorf(predicted_mask, target)
assert haussdorf_res.shape == (B, C)
haussdorf_log[sm_slice] = haussdorf_res
if compute_miou:
IoUs: Tensor = iIoU(predicted_mask, target)
assert IoUs.shape == (B, C), IoUs.shape
iiou_log[sm_slice] = IoUs
intersections[sm_slice] = inter_sum(predicted_mask, target)
unions[sm_slice] = union_sum(predicted_mask, target)
# Save images
if savedir:
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UserWarning)
predicted_class: Tensor = probs2class(pred_probs)
save_images(predicted_class, filenames, savedir, mode, epc)
# Logging
big_slice = slice(0, done_img + B) # Value for current and previous batches
dsc_dict = {f"DSC{n}": all_dices[big_slice, n].mean() for n in metric_axis} if few_axis else {}
hauss_dict = {f"HD{n}": haussdorf_log[big_slice, n].mean() for n in metric_axis} \
if compute_haussdorf and few_axis else {}
batch_dict = {f"bDSC{n}": batch_dices[:done_batch, n].mean() for n in metric_axis} \
if B > 1 and mode == "val" and few_axis else {}
miou_dict = {f"iIoU": iiou_log[big_slice, metric_axis].mean(),
f"mIoU": (intersections.sum(dim=0) / (unions.sum(dim=0) + 1e-10)).mean()} \
if compute_miou else {}
if len(metric_axis) > 1:
mean_dict = {"DSC": all_dices[big_slice, metric_axis].mean()}
if compute_haussdorf:
mean_dict["HD"] = haussdorf_log[big_slice, metric_axis].mean()
else:
mean_dict = {}
stat_dict = {**miou_dict, **dsc_dict, **hauss_dict, **mean_dict, **batch_dict,
"loss": loss_log[:done_batch].mean()}
nice_dict = {k: f"{v:.3f}" for (k, v) in stat_dict.items()}
done_img += B
done_batch += 1
tq_iter.set_postfix({**nice_dict, "loader": str(i)})
tq_iter.update(1)
tq_iter.close()
print(f"{desc} " + ', '.join(f"{k}={v}" for (k, v) in nice_dict.items()))
if compute_miou:
mIoUs: Tensor = (intersections.sum(dim=0) / (unions.sum(dim=0) + 1e-10))
assert mIoUs.shape == (C,), mIoUs.shape
else:
mIoUs = None
if not few_axis and False:
print(f"DSC: {[f'{all_dices[:, n].mean():.3f}' for n in metric_axis]}")
print(f"iIoU: {[f'{iiou_log[:, n].mean():.3f}' for n in metric_axis]}")
if mIoUs:
print(f"mIoU: {[f'{mIoUs[n]:.3f}' for n in metric_axis]}")
return loss_log, all_dices, batch_dices, haussdorf_log, mIoUs
def runInference(args: argparse.Namespace, pred_folder: str):
# print('>>> Loading the data')
device = torch.device("cuda") if torch.cuda.is_available(
) and not args.cpu else torch.device("cpu")
C: int = args.num_classes
# Let's just reuse some code
png_transform = transforms.Compose([
lambda img: np.array(img)[np.newaxis, ...],
lambda nd: nd / 255, # max <= 1
lambda nd: torch.tensor(nd, dtype=torch.float32)
])
gt_transform = transforms.Compose([
lambda img: np.array(img)[np.newaxis, ...],
lambda nd: torch.tensor(nd, dtype=torch.int64),
partial(class2one_hot, C=C),
itemgetter(0)
])
bounds_gen = [(lambda *a: torch.zeros(C, 1, 2)) for _ in range(2)]
folders: List[Path] = [
Path(pred_folder),
Path(pred_folder),
Path(args.gt_folder)
] # First one is dummy
names: List[str] = map_(lambda p: str(p.name), folders[0].glob("*.png"))
are_hots = [False, True, True]
dt_set = SliceDataset(
names,
folders,
transforms=[png_transform, gt_transform, gt_transform],
debug=False,
C=C,
are_hots=are_hots,
in_memory=False,
bounds_generators=bounds_gen)
sampler = PatientSampler(dt_set, args.grp_regex)
loader = DataLoader(dt_set, batch_sampler=sampler, num_workers=11)
# print('>>> Computing the metrics')
total_iteration, total_images = len(loader), len(loader.dataset)
metrics = {
"all_dices":
torch.zeros((total_images, C), dtype=torch.float64, device=device),
"batch_dices":
torch.zeros((total_iteration, C), dtype=torch.float64, device=device),
"sizes":
torch.zeros((total_images, 1), dtype=torch.float64, device=device)
}
desc = f">> Computing"
tq_iter = tqdm_(enumerate(loader), total=total_iteration, desc=desc)
done: int = 0
for j, (filenames, _, pred, gt, _) in tq_iter:
B = len(pred)
pred = pred.to(device)
gt = gt.to(device)
assert simplex(pred) and sset(pred, [0, 1])
assert simplex(gt) and sset(gt, [0, 1])
dices: Tensor = dice_coef(pred, gt)
b_dices: Tensor = dice_batch(pred, gt)
assert dices.shape == (B, C)
assert b_dices.shape == (C, ), b_dices.shape
sm_slice = slice(done, done + B) # Values only for current batch
metrics["all_dices"][sm_slice, ...] = dices
metrics["sizes"][sm_slice, :] = torch.einsum("bwh->b",
gt[:, 1, ...])[..., None]
metrics["batch_dices"][j] = b_dices
done += B
print(f">>> {pred_folder}")
for key, v in metrics.items():
print(key, map_("{:.4f}".format, v.mean(dim=0)))
from tensorflow.python.ops import math_ops
from tensorflow.keras.models import load_model
from data_processing import get_data_orig
from utils import customized_loss_fn, dice_coef
# image size
width = 1024
height = 1024
# data preprocessing (histgram-equalization and stretching)
preprocessing = True
images, soma_labels, vessel_labels = get_data_orig('../data/Test',
size=(width, height),
preprocessing=preprocessing)
images = np.expand_dims(images, axis=-1)
soma_labels = np.expand_dims(soma_labels, axis=-1)
vessel_labels = np.expand_dims(vessel_labels, axis=-1)
# network prediction
model = load_model('./weights/model-weights-epoch-10.hdf5')
out = model.predict([images, soma_labels, vessel_labels])
soma_labels_pred, vessel_labels_pred = out[0][0], out[0][1]
# network performance evaluation
soma_dice = dice_coef(np.squeeze(soma_labels), np.squeeze(soma_labels_pred))
vessel_dice = dice_coef(np.squeeze(vessel_labels),
np.squeeze(vessel_labels_pred))
print("Boosting network performance for SOMA: ", np.mean(soma_dice))
print("Boosting network for VESSEL: ", np.mean(vessel_dice))
# ================================================================
unet_seg = model.Res_Unet(images)
prior_mvn = model.Prior_net(images, latent_dim, 'prior_dist')
z_prior = prior_mvn.sample()
seg_prior = model.Fcomb(unet_seg, z_prior, 'prior')
posterior_mvn = model.Posterior_net(images, y_true, latent_dim, 'post_dist')
z_posterior = posterior_mvn.sample()
seg = model.Fcomb(unet_seg, z_posterior, 'posterior')
# ================================================================
# LOSS
# ================================================================
dice_coef = utils.dice_coef(y_true, unet_seg)
loss, kl, ce = utils.elbo(y_true, seg, prior_mvn, posterior_mvn, 2)
lambda_ = 1e-6
l2_norms = [tf.nn.l2_loss(v) for v in tf.trainable_variables()]
l2_norm = tf.reduce_sum(l2_norms)
reg_loss = loss + lambda_ * l2_norm
optimizer = utils.optimize(reg_loss)
# ================================================================
# SAVER
# ================================================================
tf.add_to_collection('saved_variables', value=images)
def do_epoch(mode: str, args, net, device, use_cuda, loader, optimizer,
num_classes, epoch):
totalImages = len(loader)
if mode == "train":
net.train()
desc = f">> Training ({epoch})"
elif mode == "val":
net.eval()
desc = f">> Validation ({epoch})"
total_iteration, total_images = len(loader), len(loader.dataset)
all_dices: Tensor = torch.zeros((total_images, num_classes),
dtype=eval(args.dtype),
device=device)
batch_dices: Tensor = torch.zeros((total_iteration, num_classes),
dtype=eval(args.dtype),
device=device)
loss_log: Tensor = torch.zeros((total_images),
dtype=eval(args.dtype),
device=device)
entropy_log: Tensor = torch.zeros((total_images),
dtype=eval(args.dtype),
device=device)
KL_log: Tensor = torch.zeros((total_images),
dtype=eval(args.dtype),
device=device)
tq_iter = tqdm_(enumerate(loader), total=total_iteration, desc=desc)
done: int = 0
for j, data in tq_iter:
image_f, image_i, image_d, image_o, image_w, image_c, labels, img_names = data
#image_f=image_f.type(torch.FloatTensor)/65535.
#image_f = image_f.type(torch.FloatTensor)/65535.
#image_i = image_i.type(torch.FloatTensor)/65535.
#image_d = image_d.type(torch.FloatTensor)/65535.
#image_o = image_o.type(torch.FloatTensor)/65535.
#image_w = image_w.type(torch.FloatTensor)/65535.
#image_c = image_c.type(torch.FloatTensor)/65535.
MRI: Tensor = torch.zeros((1, 6, image_f.size()[2], image_f.size()[3]),
dtype=eval(args.dtype))
MRI = torch.cat((image_f, image_i, image_d, image_o, image_w, image_c),
dim=1)
MRI = MRI.type(
torch.FloatTensor
) / 65535.0 #.type(eval(args.dtype)) #.type(torch.FloatTensor)
targets = torch.cat((1 - labels, labels),
dim=1) #.type(torch.LongTensor)
B = len(image_f)
#print(type(labels))
#MRI = torch.cat((image_f,image_i,image_d,image_w),dim=1)
if use_cuda:
MRI, targets = MRI.to(device), targets.to(device)
# forward
outputs = net(MRI)
pred_probs = F.softmax(outputs, dim=1)
predicted_mask = probs2one_hot(pred_probs)
entropy = crossEntropy_f(pred_probs, targets)
pred_probs_aver: Tensor = torch.sum(pred_probs, dim=(2, 3))
pred_probs_aver = pred_probs_aver / torch.sum(targets).float()
target_aver: Tensor = torch.sum(targets, dim=(2, 3)).float()
target_aver = target_aver / torch.sum(targets).float()
KL_loss = args.lam * kl(target_aver, pred_probs_aver)
loss = entropy + KL_loss
if mode == "train":
# zero the parameter gradients8544
optimizer.zero_grad()
# backward + optimize
loss.backward()
optimizer.step()
# Compute and log metrics
dices: Tensor = dice_coef(predicted_mask.detach(),
targets.type(torch.cuda.IntTensor).detach())
batch_dice: Tensor = dice_batch(
predicted_mask.detach(),
targets.type(torch.cuda.IntTensor).detach())
assert batch_dice.shape == (num_classes, ) and dices.shape == (
B, num_classes), (batch_dice.shape, dices.shape, B, num_classes)
sm_slice = slice(done, done + B) # Values only for current batch
all_dices[sm_slice, ...] = dices
entropy_log[sm_slice] = entropy.detach()
loss_log[sm_slice] = loss.detach()
KL_log[sm_slice] = KL_loss.detach()
batch_dices[j] = batch_dice
# Logging
big_slice = slice(0,
done + B) # Value for current and previous batches
stat_dict = {
"dice": all_dices[big_slice, -1].mean(),
"total loss": loss_log[big_slice].mean(),
"entropy loss": entropy_log[big_slice].mean(),
"KL loss": KL_log[big_slice].mean(),
"b dice": batch_dices[:j + 1, -1].mean()
}
nice_dict = {k: f"{v:.4f}" for (k, v) in stat_dict.items()}
done += B
tq_iter.set_postfix(nice_dict)
return loss_log, entropy_log, KL_log, all_dices, batch_dices
def do_epoch(
mode: str,
net: Any,
device: Any,
loaders: list[DataLoader],
epc: int,
list_loss_fns: list[list[Callable]],
list_loss_weights: list[list[float]],
K: int,
savedir: str = "",
optimizer: Any = None,
metric_axis: list[int] = [1],
compute_3d_dice: bool = False,
temperature: float = 1) -> Tuple[Tensor, Tensor, Optional[Tensor]]:
assert mode in ["train", "val", "dual"]
if mode == "train":
net.train()
desc = f">> Training ({epc})"
elif mode == "val":
net.eval()
desc = f">> Validation ({epc})"
total_iteration: int = sum(len(loader) for loader in loaders) # U
total_images: int = sum(len(loader.dataset) for loader in loaders) # D
n_loss: int = max(map(len, list_loss_fns))
all_dices: Tensor = torch.zeros((total_images, K),
dtype=torch.float32,
device=device)
loss_log: Tensor = torch.zeros((total_iteration, n_loss),
dtype=torch.float32,
device=device)
three_d_dices: Optional[Tensor]
if compute_3d_dice:
three_d_dices = torch.zeros((total_iteration, K),
dtype=torch.float32,
device=device)
else:
three_d_dices = None
done_img: int = 0
done_batch: int = 0
tq_iter = tqdm_(total=total_iteration, desc=desc)
for i, (loader, loss_fns, loss_weights) in enumerate(
zip(loaders, list_loss_fns, list_loss_weights)):
for data in loader:
# t0 = time()
image: Tensor = data["images"].to(device)
target: Tensor = data["gt"].to(device)
filenames: list[str] = data["filenames"]
assert not target.requires_grad
labels: list[Tensor] = [e.to(device) for e in data["labels"]]
B, C, *_ = image.shape
# Reset gradients
if optimizer:
optimizer.zero_grad()
# Forward
pred_logits: Tensor = net(image)
pred_probs: Tensor = F.softmax(temperature * pred_logits, dim=1)
predicted_mask: Tensor = probs2one_hot(
pred_probs.detach()) # Used only for dice computation
assert not predicted_mask.requires_grad
assert len(loss_fns) == len(loss_weights) == len(labels)
ziped = zip(loss_fns, labels, loss_weights)
losses = [
w * loss_fn(pred_probs, label) for loss_fn, label, w in ziped
]
loss = reduce(add, losses)
assert loss.shape == (), loss.shape
# Backward
if optimizer:
loss.backward()
optimizer.step()
# Compute and log metrics
for j in range(len(loss_fns)):
loss_log[done_batch, j] = losses[j].detach()
sm_slice = slice(done_img,
done_img + B) # Values only for current batch
dices: Tensor = dice_coef(predicted_mask, target)
assert dices.shape == (B, K), (dices.shape, B, K)
all_dices[sm_slice, ...] = dices
if compute_3d_dice:
three_d_DSC: Tensor = dice_batch(predicted_mask, target)
assert three_d_DSC.shape == (K, )
three_d_dices[done_batch] = three_d_DSC # type: ignore
# Save images
if savedir:
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UserWarning)
predicted_class: Tensor = probs2class(pred_probs)
save_images(predicted_class, filenames, savedir, mode, epc)
# Logging
big_slice = slice(0, done_img +
B) # Value for current and previous batches
dsc_dict: dict = {f"DSC{n}": all_dices[big_slice, n].mean() for n in metric_axis} | \
({f"3d_DSC{n}": three_d_dices[:done_batch, n].mean() for n in metric_axis}
if three_d_dices is not None else {})
loss_dict = {
f"loss_{i}": loss_log[:done_batch].mean(dim=0)[i]
for i in range(n_loss)
}
stat_dict = dsc_dict | loss_dict
nice_dict = {k: f"{v:.3f}" for (k, v) in stat_dict.items()}
done_img += B
done_batch += 1
tq_iter.set_postfix({**nice_dict, "loader": str(i)})
tq_iter.update(1)
tq_iter.close()
print(f"{desc} " + ', '.join(f"{k}={v}" for (k, v) in nice_dict.items()))
return (loss_log.detach().cpu(), all_dices.detach().cpu(),
three_d_dices.detach().cpu()
if three_d_dices is not None else None)
def runInference(args: argparse.Namespace):
# print('>>> Loading the data')
# device = torch.device("cuda") if torch.cuda.is_available() and not args.cpu else torch.device("cpu")
device = torch.device("cpu")
C: int = args.num_classes
# Let's just reuse some code
png_transform = transforms.Compose([
lambda img: np.array(img)[np.newaxis, ...],
lambda nd: nd / 255, # max <= 1
lambda nd: torch.tensor(nd, dtype=torch.float32)
])
gt_transform = transforms.Compose([
lambda img: np.array(img)[np.newaxis, ...],
lambda nd: torch.tensor(nd, dtype=torch.int64),
partial(class2one_hot, C=C),
itemgetter(0)
])
bounds_gen = [(lambda *a: torch.zeros(C, 1, 2)) for _ in range(2)]
metrics = None
pred_folders = sorted(list(Path(args.pred_root).glob('iter*')))
assert len(pred_folders) == args.epochs, (len(pred_folders), args.epochs)
for epoch, pred_folder in enumerate(pred_folders):
if args.do_only and epoch not in args.do_only:
continue
# First one is dummy:
folders: List[Path] = [Path(pred_folder, 'val'), Path(pred_folder, 'val'), Path(args.gt_folder)]
names: List[str] = map_(lambda p: str(p.name), folders[0].glob("*.png"))
are_hots = [False, True, True]
# spacing_dict = pickle.load(open(Path(args.gt_folder, "..", "spacing.pkl"), 'rb'))
spacing_dict = None
dt_set = SliceDataset(names,
folders,
transforms=[png_transform, gt_transform, gt_transform],
debug=False,
C=C,
are_hots=are_hots,
in_memory=False,
spacing_dict=spacing_dict,
bounds_generators=bounds_gen,
quiet=True)
loader = DataLoader(dt_set,
num_workers=2)
# print('>>> Computing the metrics')
total_iteration, total_images = len(loader), len(loader.dataset)
if not metrics:
metrics = {"all_dices": torch.zeros((args.epochs, total_images, C), dtype=torch.float64, device=device),
"hausdorff": torch.zeros((args.epochs, total_images, C), dtype=torch.float64, device=device)}
desc = f">> Computing"
tq_iter = tqdm_(enumerate(loader), total=total_iteration, desc=desc)
done: int = 0
for j, (filenames, _, pred, gt, _) in tq_iter:
B = len(pred)
pred = pred.to(device)
gt = gt.to(device)
assert simplex(pred) and sset(pred, [0, 1])
assert simplex(gt) and sset(gt, [0, 1])
dices: Tensor = dice_coef(pred, gt)
assert dices.shape == (B, C)
haussdorf_res: Tensor = haussdorf(pred, gt)
assert haussdorf_res.shape == (B, C)
sm_slice = slice(done, done + B) # Values only for current batch
metrics["all_dices"][epoch, sm_slice, ...] = dices
metrics["hausdorff"][epoch, sm_slice, ...] = haussdorf_res
done += B
for key, v in metrics.items():
print(epoch, key, map_("{:.4f}".format, v[epoch].mean(dim=0)))
if metrics:
savedir: Path = Path(args.save_folder)
for k, e in metrics.items():
np.save(Path(savedir, f"{k}.npy"), e.cpu().numpy())
with tf.name_scope('cross_entropy'):
cross_entropy_loss = tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(labels=label, logits=pred))
# define optimzer
global_step = tf.Variable(0, name='global_step', trainable=False)
with tf.name_scope('learning_rate'):
learning_rate = tf.train.exponential_decay(0.0001,
global_step,
iter_epoch * 3,
0.5,
staircase=True)
train_step = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(
cross_entropy_loss, global_step=global_step)
# compute dice score for simple evaluation during training
with tf.name_scope('dice_eval'):
dice_evaluator = tf.reduce_mean(dice_coef(label, pred))
# visualization conv1
conv1_kernel = tf.global_variables()[0] # get the first convolutional kernel
# normalize to the rgb space
x_min = tf.reduce_min(conv1_kernel)
x_max = tf.reduce_max(conv1_kernel)
kernel_0_to_1 = (conv1_kernel - x_min) / (x_max - x_min)
kernel_transposed = tf.transpose(kernel_0_to_1, [3, 0, 1, 2]) # transpose
# deine summary for tensorboard
tf.summary.scalar('cross_entropy_loss', cross_entropy_loss)
tf.summary.image('prediction', pred)
tf.summary.scalar('learning_rate', learning_rate)
tf.summary.scalar('dice_score', dice_evaluator)
tf.summary.image('conv1/kernel', kernel_transposed)
def do_epoch(args, mode: str, net: Any, device: Any, epc: int,
loss_fns: List[Callable], loss_weights: List[float],
new_w:int, C: int, metric_axis:List[int], savedir: str = "",
optimizer: Any = None, target_loader: Any = None, best_dice3d_val:Any=None):
assert mode in ["train", "val"]
L: int = len(loss_fns)
indices = torch.tensor(metric_axis,device=device)
if mode == "train":
net.train()
desc = f">> Training ({epc})"
elif mode == "val":
net.eval()
# net.train()
desc = f">> Validation ({epc})"
total_it_t, total_images_t = len(target_loader), len(target_loader.dataset)
total_iteration = total_it_t
total_images = total_images_t
if args.debug:
total_iteration = 10
pho=1
dtype = eval(args.dtype)
all_dices: Tensor = torch.zeros((total_images, C), dtype=dtype, device=device)
all_sizes: Tensor = torch.zeros((total_images, C), dtype=dtype, device=device)
all_sizes: Tensor = torch.zeros((total_images, C), dtype=dtype, device=device)
all_gt_sizes: Tensor = torch.zeros((total_images, C), dtype=dtype, device=device)
all_sizes2: Tensor = torch.zeros((total_images, C), dtype=dtype, device=device)
all_inter_card: Tensor = torch.zeros((total_images, C), dtype=dtype, device=device)
all_card_gt: Tensor = torch.zeros((total_images, C), dtype=dtype, device=device)
all_card_pred: Tensor = torch.zeros((total_images, C), dtype=dtype, device=device)
all_gt = []
all_pred = []
if args.do_hd:
all_gt: Tensor = torch.zeros((total_images, args.wh, args.wh), dtype=dtype)
all_pred: Tensor = torch.zeros((total_images, args.wh, args.wh), dtype=dtype)
loss_log: Tensor = torch.zeros((total_images), dtype=dtype, device=device)
loss_cons: Tensor = torch.zeros((total_images), dtype=dtype, device=device)
loss_se: Tensor = torch.zeros((total_images), dtype=dtype, device=device)
loss_tot: Tensor = torch.zeros((total_images), dtype=dtype, device=device)
posim_log: Tensor = torch.zeros((total_images), dtype=dtype, device=device)
haussdorf_log: Tensor = torch.zeros((total_images, C), dtype=dtype, device=device)
all_grp: Tensor = torch.zeros((total_images, C), dtype=dtype, device=device)
all_pnames = np.zeros([total_images]).astype('U256')
#dice_3d_log: Tensor = torch.zeros((1, C), dtype=dtype, device=device)
dice_3d_log, dice_3d_sd_log = 0, 0
#dice_3d_sd_log: Tensor = torch.zeros((1, C), dtype=dtype, device=device)
hd_3d_log, asd_3d_log, hd_3d_sd_log, asd_3d_sd_log= 0, 0, 0, 0
tq_iter = tqdm_(enumerate(target_loader), total=total_iteration, desc=desc)
done: int = 0
n_warmup = args.n_warmup
mult_lw = [pho ** (epc - n_warmup + 1)] * len(loss_weights)
mult_lw[0] = 1
loss_weights = [a * b for a, b in zip(loss_weights, mult_lw)]
losses_vec, source_vec, target_vec, baseline_target_vec = [], [], [], []
pen_count = 0
with warnings.catch_warnings():
warnings.simplefilter("ignore")
count_losses = 0
for j, target_data in tq_iter:
target_data[1:] = [e.to(device) for e in target_data[1:]] # Move all tensors to device
filenames_target, target_image, target_gt = target_data[:3]
#print("target", filenames_target)
labels = target_data[3:3+L]
bounds = target_data[3+L:]
filenames_target = [f.split('.nii')[0] for f in filenames_target]
assert len(labels) == len(bounds), len(bounds)
B = len(target_image)
# Reset gradients
if optimizer:
optimizer.zero_grad()
# Forward
with torch.set_grad_enabled(mode == "train"):
pred_logits: Tensor = net(target_image)
pred_probs: Tensor = F.softmax(pred_logits, dim=1)
if new_w > 0:
pred_probs = resize(pred_probs, new_w)
labels = [resize(label, new_w) for label in labels]
target = resize(target, new_w)
predicted_mask: Tensor = probs2one_hot(pred_probs) # Used only for dice computation
assert len(bounds) == len(loss_fns) == len(loss_weights)
if epc < n_warmup:
loss_weights = [0]*len(loss_weights)
loss: Tensor = torch.zeros(1, requires_grad=True).to(device)
loss_vec = []
loss_kw = []
for loss_fn,label, w, bound in zip(loss_fns,labels, loss_weights, bounds):
if w > 0:
if eval(args.target_losses)[0][0]=="EntKLProp":
loss_1, loss_cons_prior,est_prop = loss_fn(pred_probs, label, bound)
loss = loss_1 + loss_cons_prior
else:
loss = loss_fn(pred_probs, label, bound)
loss = w*loss
loss_1 = loss
loss_kw.append(loss_1.detach())
# Backward
if optimizer:
loss.backward()
optimizer.step()
dices, inter_card, card_gt, card_pred = dice_coef(predicted_mask.detach(), target_gt.detach())
assert dices.shape == (B, C), (dices.shape, B, C)
sm_slice = slice(done, done + B) # Values only for current batch
all_dices[sm_slice, ...] = dices
if eval(args.target_losses)[0][0] in ["EntKLProp"]:
all_sizes[sm_slice, ...] = torch.round(est_prop.detach()*target_image.shape[2]*target_image.shape[3])
all_sizes2[sm_slice, ...] = torch.sum(predicted_mask,dim=(2,3))
all_gt_sizes[sm_slice, ...] = torch.sum(target_gt,dim=(2,3))
all_grp[sm_slice, ...] = torch.FloatTensor(get_subj_nb(filenames_target)).unsqueeze(1).repeat(1,C)
all_pnames[sm_slice] = filenames_target
all_inter_card[sm_slice, ...] = inter_card
all_card_gt[sm_slice, ...] = card_gt
all_card_pred[sm_slice, ...] = card_pred
if args.do_hd:
all_pred[sm_slice, ...] = probs2class(predicted_mask[:,:,:,:]).cpu().detach()
all_gt[sm_slice, ...] = probs2class(target_gt).detach()
loss_se[sm_slice] = loss_kw[0]
if len(loss_kw)>1:
loss_cons[sm_slice] = loss_kw[1]
loss_tot[sm_slice] = loss_kw[1]+loss_kw[0]
else:
loss_cons[sm_slice] = 0
loss_tot[sm_slice] = loss_kw[0]
# # Save images
if savedir and args.saveim and mode =="val":
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UserWarning)
warnings.simplefilter("ignore")
predicted_class: Tensor = probs2class(pred_probs)
save_images(predicted_class, filenames_target, savedir, mode, epc, False)
if args.entmap:
ent_map = torch.einsum("bcwh,bcwh->bwh", [-pred_probs, (pred_probs+1e-10).log()])
save_images_ent(ent_map, filenames_target, savedir,'ent_map', epc)
# Logging
big_slice = slice(0, done + B) # Value for current and previous batches
stat_dict = {**{f"DSC{n}": all_dices[big_slice, n].mean() for n in metric_axis},
**{f"SZ{n}": all_sizes[big_slice, n].mean() for n in metric_axis},
**({f"DSC_source{n}": all_dices_s[big_slice, n].mean() for n in metric_axis}
if args.source_metrics else {})}
size_dict = {**{f"SZ{n}": all_sizes[big_slice, n].mean() for n in metric_axis}}
nice_dict = {k: f"{v:.4f}" for (k, v) in stat_dict.items()}
done += B
tq_iter.set_postfix(nice_dict)
if args.dice_3d and (mode == 'val'):
dice_3d_log, dice_3d_sd_log,asd_3d_log, asd_3d_sd_log,hd_3d_log, hd_3d_sd_log = dice3d(all_grp, all_inter_card, all_card_gt, all_card_pred,all_pred,all_gt,all_pnames,metric_axis,args.pprint,args.do_hd,args.do_asd,best_dice3d_val)
dice_2d = torch.index_select(all_dices, 1, indices).mean().cpu().numpy().item()
target_vec = [dice_3d_log, dice_3d_sd_log,asd_3d_log, asd_3d_sd_log,hd_3d_log,hd_3d_sd_log,dice_2d]
size_mean = torch.index_select(all_sizes2, 1, indices).mean(dim=0).cpu().numpy()
size_gt_mean = torch.index_select(all_gt_sizes, 1, indices).mean(dim=0).cpu().numpy()
mask_pos = torch.index_select(all_sizes2, 1, indices)!=0
gt_pos = torch.index_select(all_gt_sizes, 1, indices)!=0
size_mean_pos = torch.index_select(all_sizes2, 1, indices).sum(dim=0).cpu().numpy()/mask_pos.sum(dim=0).cpu().numpy()
gt_size_mean_pos = torch.index_select(all_gt_sizes, 1, indices).sum(dim=0).cpu().numpy()/gt_pos.sum(dim=0).cpu().numpy()
size_mean2 = torch.index_select(all_sizes2, 1, indices).mean(dim=0).cpu().numpy()
losses_vec = [loss_se.mean().item(),loss_cons.mean().item(),loss_tot.mean().item(),size_mean.mean(),size_mean_pos.mean(),size_gt_mean.mean(),gt_size_mean_pos.mean()]
if not epc%50:
df_t = pd.DataFrame({
"val_ids":all_pnames,
"proposal_size":all_sizes2.cpu()})
df_t.to_csv(Path(savedir,mode+str(epc)+"sizes.csv"), float_format="%.4f", index_label="epoch")
return losses_vec, target_vec,source_vec
def do_epoch(
mode: str,
net: Any,
device: Any,
loaders: List[DataLoader],
epc: int,
list_loss_fns: List[List[Callable]],
list_loss_weights: List[List[float]],
K: int,
savedir: str = "",
optimizer: Any = None,
metric_axis: List[int] = [1],
compute_hausdorff: bool = False,
compute_miou: bool = False,
compute_3d_dice: bool = False,
temperature: float = 1
) -> Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor],
Optional[Tensor]]:
assert mode in ["train", "val"]
if mode == "train":
net.train()
desc = f">> Training ({epc})"
elif mode == "val":
net.eval()
desc = f">> Validation ({epc})"
total_iteration: int = sum(len(loader) for loader in loaders) # U
total_images: int = sum(len(loader.dataset) for loader in loaders) # D
n_loss: int = max(map(len, list_loss_fns))
all_dices: Tensor = torch.zeros((total_images, K),
dtype=torch.float32,
device=device)
loss_log: Tensor = torch.zeros((total_iteration, n_loss),
dtype=torch.float32,
device=device)
iiou_log: Optional[Tensor]
intersections: Optional[Tensor]
unions: Optional[Tensor]
if compute_miou:
iiou_log = torch.zeros((total_images, K),
dtype=torch.float32,
device=device)
intersections = torch.zeros((total_images, K),
dtype=torch.float32,
device=device)
unions = torch.zeros((total_images, K),
dtype=torch.float32,
device=device)
else:
iiou_log = None
intersections = None
unions = None
three_d_dices: Optional[Tensor]
if compute_3d_dice:
three_d_dices = torch.zeros((total_iteration, K),
dtype=torch.float32,
device=device)
else:
three_d_dices = None
hausdorff_log: Optional[Tensor]
if compute_hausdorff:
hausdorff_log = torch.zeros((total_images, K),
dtype=torch.float32,
device=device)
else:
hausdorff_log = None
few_axis: bool = len(metric_axis) <= 3
done_img: int = 0
done_batch: int = 0
tq_iter = tqdm_(total=total_iteration, desc=desc)
for i, (loader, loss_fns, loss_weights) in enumerate(
zip(loaders, list_loss_fns, list_loss_weights)):
for data in loader:
image: Tensor = data["images"].to(device)
target: Tensor = data["gt"].to(device)
spacings: Tensor = data["spacings"] # Keep that one on CPU
assert not target.requires_grad
labels: List[Tensor] = [e.to(device) for e in data["labels"]]
bounds: List[Tensor] = [e.to(device) for e in data["bounds"]]
box_priors: List[List[Tuple[
Tensor, Tensor]]] # one more level for the batch
box_priors = [[(m.to(device), b.to(device)) for (m, b) in B]
for B in data["box_priors"]]
assert len(labels) == len(bounds)
B, C, *_ = image.shape
samplings: List[List[Tuple[slice]]] = data["samplings"]
assert len(samplings) == B
assert len(samplings[0][0]) == len(
image[0, 0].shape), (samplings[0][0], image[0, 0].shape)
probs_receptacle: Tensor = -torch.ones_like(
target, dtype=torch.float32) # -1 for unfilled
mask_receptacle: Tensor = -torch.ones_like(
target, dtype=torch.int32) # -1 for unfilled
# Use the sampling coordinates of the first batch item
assert not (len(samplings[0]) > 1 and
B > 1), samplings # No subsampling if batch size > 1
loss_sub_log: Tensor = torch.zeros(
(len(samplings[0]), len(loss_fns)),
dtype=torch.float32,
device=device)
for k, sampling in enumerate(samplings[0]):
img_sampling = [slice(0, B), slice(0, C)] + list(sampling)
label_sampling = [slice(0, B), slice(0, K)] + list(sampling)
assert len(img_sampling) == len(image.shape), (img_sampling,
image.shape)
sub_img = image[img_sampling]
# Reset gradients
if optimizer:
optimizer.zero_grad()
# Forward
pred_logits: Tensor = net(sub_img)
pred_probs: Tensor = F.softmax(temperature * pred_logits,
dim=1)
predicted_mask: Tensor = probs2one_hot(
pred_probs.detach()) # Used only for dice computation
assert not predicted_mask.requires_grad
probs_receptacle[label_sampling] = pred_probs[...]
mask_receptacle[label_sampling] = predicted_mask[...]
assert len(bounds) == len(loss_fns) == len(
loss_weights) == len(labels)
ziped = zip(loss_fns, labels, loss_weights, bounds)
losses = [
w * loss_fn(pred_probs, label[label_sampling], bound,
box_priors)
for loss_fn, label, w, bound in ziped
]
loss = reduce(add, losses)
assert loss.shape == (), loss.shape
# Backward
if optimizer:
loss.backward()
optimizer.step()
# Compute and log metrics
for j in range(len(loss_fns)):
loss_sub_log[k, j] = losses[j].detach()
reduced_loss_sublog: Tensor = loss_sub_log.sum(dim=0)
assert reduced_loss_sublog.shape == (len(loss_fns), ), (
reduced_loss_sublog.shape, len(loss_fns))
loss_log[done_batch, ...] = reduced_loss_sublog[...]
del loss_sub_log
sm_slice = slice(done_img,
done_img + B) # Values only for current batch
dices: Tensor = dice_coef(mask_receptacle, target)
assert dices.shape == (B, K), (dices.shape, B, K)
all_dices[sm_slice, ...] = dices
if compute_3d_dice:
three_d_DSC: Tensor = dice_batch(mask_receptacle, target)
assert three_d_DSC.shape == (K, )
three_d_dices[done_batch] = three_d_DSC # type: ignore
if compute_hausdorff:
hausdorff_res: Tensor
try:
hausdorff_res = hausdorff(mask_receptacle, target,
spacings)
except RuntimeError:
hausdorff_res = torch.zeros((B, K), device=device)
assert hausdorff_res.shape == (B, K)
hausdorff_log[sm_slice] = hausdorff_res # type: ignore
if compute_miou:
IoUs: Tensor = iIoU(mask_receptacle, target)
assert IoUs.shape == (B, K), IoUs.shape
iiou_log[sm_slice] = IoUs # type: ignore
intersections[sm_slice] = inter_sum(mask_receptacle,
target) # type: ignore
unions[sm_slice] = union_sum(mask_receptacle,
target) # type: ignore
# if False and target[0, 1].sum() > 0: # Useful template for quick and dirty inspection
# import matplotlib.pyplot as plt
# from pprint import pprint
# from mpl_toolkits.axes_grid1 import ImageGrid
# from utils import soft_length
# print(data["filenames"])
# pprint(data["bounds"])
# pprint(soft_length(mask_receptacle))
# fig = plt.figure()
# fig.clear()
# grid = ImageGrid(fig, 211, nrows_ncols=(1, 2))
# grid[0].imshow(data["images"][0, 0], cmap="gray")
# grid[0].contour(data["gt"][0, 1], cmap='jet', alpha=.75, linewidths=2)
# grid[1].imshow(data["images"][0, 0], cmap="gray")
# grid[1].contour(mask_receptacle[0, 1], cmap='jet', alpha=.75, linewidths=2)
# plt.show()
# Save images
if savedir:
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UserWarning)
predicted_class: Tensor = probs2class(pred_probs)
save_images(predicted_class, data["filenames"], savedir,
mode, epc)
# Logging
big_slice = slice(0, done_img +
B) # Value for current and previous batches
dsc_dict: Dict
if few_axis:
dsc_dict = {
**{
f"DSC{n}": all_dices[big_slice, n].mean()
for n in metric_axis
},
**({
f"3d_DSC{n}": three_d_dices[:done_batch, n].mean()
for n in metric_axis
} if three_d_dices is not None else {})
}
else:
dsc_dict = {}
# dsc_dict = {f"DSC{n}": all_dices[big_slice, n].mean() for n in metric_axis} if few_axis else {}
hauss_dict = {f"HD{n}": hausdorff_log[big_slice, n].mean() for n in metric_axis} \
if hausdorff_log is not None and few_axis else {}
miou_dict = {f"iIoU": iiou_log[big_slice, metric_axis].mean(),
f"mIoU": (intersections.sum(dim=0) / (unions.sum(dim=0) + 1e-10)).mean()} \
if iiou_log is not None and intersections is not None and unions is not None else {}
if len(metric_axis) > 1:
mean_dict = {"DSC": all_dices[big_slice, metric_axis].mean()}
if hausdorff_log:
mean_dict["HD"] = hausdorff_log[big_slice,
metric_axis].mean()
else:
mean_dict = {}
stat_dict = {
**miou_dict,
**dsc_dict,
**hauss_dict,
**mean_dict, "loss": loss_log[:done_batch].mean()
}
nice_dict = {k: f"{v:.3f}" for (k, v) in stat_dict.items()}
done_img += B
done_batch += 1
tq_iter.set_postfix({**nice_dict, "loader": str(i)})
tq_iter.update(1)
tq_iter.close()
print(f"{desc} " + ', '.join(f"{k}={v}" for (k, v) in nice_dict.items()))
mIoUs: Optional[Tensor]
if intersections and unions:
mIoUs = (intersections.sum(dim=0) / (unions.sum(dim=0) + 1e-10))
assert mIoUs.shape == (K, ), mIoUs.shape
else:
mIoUs = None
if not few_axis and False:
print(f"DSC: {[f'{all_dices[:, n].mean():.3f}' for n in metric_axis]}")
print(f"iIoU: {[f'{iiou_log[:, n].mean():.3f}' for n in metric_axis]}")
if mIoUs:
print(f"mIoU: {[f'{mIoUs[n]:.3f}' for n in metric_axis]}")
return (
loss_log.detach().cpu(), all_dices.detach().cpu(),
hausdorff_log.detach().cpu() if hausdorff_log is not None else None,
mIoUs.detach().cpu() if mIoUs is not None else None,
three_d_dices.detach().cpu() if three_d_dices is not None else None)
def do_epoch(mode: str, net: Any, device: Any, loaders: list[DataLoader], epc: int,
list_loss_fns: list[list[Callable]], list_loss_weights: list[list[float]], K: int,
savedir: Path = None, optimizer: Any = None,
metric_axis: list[int] = [1], requested_metrics: list[str] = None,
temperature: float = 1) -> dict[str, Tensor]:
assert mode in ["train", "val", "dual"]
if requested_metrics is None:
requested_metrics = []
if mode == "train":
net.train()
desc = f">> Training ({epc})"
elif mode == "val":
net.eval()
desc = f">> Validation ({epc})"
elif mode == "dual":
net.eval()
desc = f">> Dual ({epc})"
total_iteration: int = sum(len(loader) for loader in loaders) # U
total_images: int = sum(len(loader.dataset) for loader in loaders) # D
n_loss: int = max(map(len, list_loss_fns))
epoch_metrics: dict[str, Tensor]
epoch_metrics = {"dice": torch.zeros((total_images, K), dtype=torch.float32, device=device),
"loss": torch.zeros((total_iteration, n_loss), dtype=torch.float32, device=device)}
if "3d_dsc" in requested_metrics:
epoch_metrics["3d_dsc"] = torch.zeros((total_iteration, K), dtype=torch.float32, device=device)
few_axis: bool = len(metric_axis) <= 4
# time_log: np.ndarray = np.ndarray(total_iteration, dtype=np.float32)
done_img: int = 0
done_batch: int = 0
tq_iter = tqdm_(total=total_iteration, desc=desc)
for i, (loader, loss_fns, loss_weights) in enumerate(zip(loaders, list_loss_fns, list_loss_weights)):
for data in loader:
# t0 = time()
image: Tensor = data["images"].to(device)
target: Tensor = data["gt"].to(device)
filenames: list[str] = data["filenames"]
assert not target.requires_grad
labels: list[Tensor] = [e.to(device) for e in data["labels"]]
bounds: list[Tensor] = [e.to(device) for e in data["bounds"]]
assert len(labels) == len(bounds)
B, C, *_ = image.shape
samplings: list[list[Tuple[slice]]] = data["samplings"]
assert len(samplings) == B
assert len(samplings[0][0]) == len(image[0, 0].shape), (samplings[0][0], image[0, 0].shape)
probs_receptacle: Tensor = - torch.ones_like(target, dtype=torch.float32) # -1 for unfilled
mask_receptacle: Tensor = - torch.ones_like(target, dtype=torch.int32) # -1 for unfilled
# Use the sampling coordinates of the first batch item
assert not (len(samplings[0]) > 1 and B > 1), samplings # No subsampling if batch size > 1
loss_sub_log: Tensor = torch.zeros((len(samplings[0]), len(loss_fns)),
dtype=torch.float32, device=device)
for k, sampling in enumerate(samplings[0]):
img_sampling = [slice(0, B), slice(0, C)] + list(sampling)
label_sampling = [slice(0, B), slice(0, K)] + list(sampling)
assert len(img_sampling) == len(image.shape), (img_sampling, image.shape)
sub_img = image[img_sampling]
# Reset gradients
if optimizer:
optimizer.zero_grad()
# Forward
pred_logits: Tensor = net(sub_img)
pred_probs: Tensor = F.softmax(temperature * pred_logits, dim=1)
# Used only for dice computation:
predicted_mask: Tensor = probs2one_hot(pred_probs.detach())
assert not predicted_mask.requires_grad
probs_receptacle[label_sampling] = pred_probs[...]
mask_receptacle[label_sampling] = predicted_mask[...]
assert len(bounds) == len(loss_fns) == len(loss_weights) == len(labels)
ziped = zip(loss_fns, labels, loss_weights, bounds)
losses = [w * loss_fn(pred_probs, label[label_sampling], bound, filenames)
for loss_fn, label, w, bound in ziped]
loss = reduce(add, losses)
assert loss.shape == (), loss.shape
# Backward
if optimizer:
loss.backward()
optimizer.step()
# Compute and log metrics
for j in range(len(loss_fns)):
loss_sub_log[k, j] = losses[j].detach()
reduced_loss_sublog: Tensor = loss_sub_log.sum(dim=0)
assert reduced_loss_sublog.shape == (len(loss_fns),), (reduced_loss_sublog.shape, len(loss_fns))
epoch_metrics["loss"][done_batch, ...] = reduced_loss_sublog[...]
del loss_sub_log
sm_slice = slice(done_img, done_img + B) # Values only for current batch
dices: Tensor = dice_coef(mask_receptacle, target)
assert dices.shape == (B, K), (dices.shape, B, K)
epoch_metrics["dice"][sm_slice, ...] = dices
if "3d_dsc" in requested_metrics:
three_d_DSC: Tensor = dice_batch(mask_receptacle, target)
assert three_d_DSC.shape == (K,)
epoch_metrics["3d_dsc"][done_batch] = three_d_DSC # type: ignore
# Save images
if savedir:
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=UserWarning)
predicted_class: Tensor = probs2class(pred_probs)
save_images(predicted_class,
data["filenames"],
savedir / f"iter{epc:03d}" / mode)
# Logging
big_slice = slice(0, done_img + B) # Value for current and previous batches
stat_dict: dict[str, Any] = {}
# The order matters for the final display -- it is easy to change
if few_axis:
stat_dict |= {f"DSC{n}": epoch_metrics["dice"][big_slice, n].mean()
for n in metric_axis}
if "3d_dsc" in requested_metrics:
stat_dict |= {f"3d_DSC{n}": epoch_metrics["3d_dsc"][:done_batch, n].mean()
for n in metric_axis}
if len(metric_axis) > 1:
stat_dict |= {"DSC": epoch_metrics["dice"][big_slice, metric_axis].mean()}
stat_dict |= {f"loss_{i}": epoch_metrics["loss"][:done_batch].mean(dim=0)[i]
for i in range(n_loss)}
nice_dict = {k: f"{v:.3f}" for (k, v) in stat_dict.items()}
# t1 = time()
# time_log[done_batch] = (t1 - t0)
done_img += B
done_batch += 1
tq_iter.set_postfix({**nice_dict, "loader": str(i)})
tq_iter.update(1)
tq_iter.close()
print(f"{desc} " + ', '.join(f"{k}={v}" for (k, v) in nice_dict.items()))
return {k: v.detach().cpu() for (k, v) in epoch_metrics.items()}