def eval_net(net, model_arch, loader, device, plot=False, dir_checkpoint=None):
"""Evaluation with the iou coefficient"""
net.eval()
mask_type = torch.float32 if net.n_classes == 1 else torch.long
# the number of batch
n_val = len(loader)
tot_ce = 0
tot_iou = 0
with tqdm(total=n_val, desc='Validation round', unit='batch',
leave=False) as pbar:
for batch in loader:
imgs, true_masks = batch['image'], batch['mask']
imgs = imgs.to(device=device, dtype=torch.float32)
true_masks = true_masks.to(device=device, dtype=torch.long)
with torch.no_grad():
if model_arch == 'icnet':
mask_pred, pred_sub4, pred_sub8, pred_sub16 = net(imgs)
else:
mask_pred = net(imgs)
tot_ce += F.cross_entropy(mask_pred, true_masks.squeeze(1)).item()
tot_iou += iou_coeff(mask_pred, true_masks.squeeze(1)).item()
pbar.update()
if plot:
for i, c in enumerate(zip(imgs, mask_pred)):
plot_img_and_mask(c[0], c[1], dir_checkpoint)
return tot_ce / n_val, tot_iou / n_val
def eval_net(net, dataset, gpu=False, epoch=None, ids=None):
"""Evaluation without the densecrf with the dice coefficient"""
net.eval()
tot = 0
filename = ''
for i, b in enumerate(dataset):
img = b[0]
true_mask = b[1]
if ids:
filename = ids[i][0]
filename = filename.replace('/', '_')
filename = filename.replace('\\', '_')
filename = filename.replace(':', '_')
img = torch.from_numpy(img).unsqueeze(0)
true_mask = torch.from_numpy(true_mask).unsqueeze(0)
if gpu:
img = img.cuda()
true_mask = true_mask.cuda()
mask_pred = net(img)[0]
if True or i < 500:
last_masks_pred = mask_pred[0].data.cpu().numpy()
last_masks_pred *= 255
last_img = img[0][0].data.cpu().numpy()
last_true_masks = true_mask[0].data.cpu().numpy()
last_true_masks *= 255
img_last_masks_pred = Image.fromarray(last_masks_pred)
min_ = last_img.min()
max_ = last_img.max()
img_last_img = Image.fromarray(255 * (last_img - min_) /
(max_ - min_))
img_last_true_masks = Image.fromarray(last_true_masks)
data_vis.plot_img_and_mask(
img_last_img, img_last_masks_pred, img_last_true_masks,
f'{filename}_ep{epoch}_i{i:03}_eval.jpg')
gc.collect()
if i % 300 == 0:
print(i)
mask_pred = (mask_pred > 0.5).float()
tot += dice_coeff(mask_pred, true_mask).item()
return tot / (i + 1)
def predict(input,
output=None,
model='MODEL.pth',
viz=False,
no_save=False,
mask_threshold=0.5,
scale=0.5):
in_files = input
out_files = get_output_filenames(input, output)
net = UNet(n_channels=3, n_classes=1)
logging.info("Loading model {}".format(model))
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
net.to(device=device)
net.load_state_dict(torch.load(model, map_location=device))
logging.info("Model loaded !")
for i, fn in enumerate(in_files):
logging.info("\nPredicting image {} ...".format(fn))
img = Image.open(fn)
mask = predict_img(net=net,
full_img=img,
scale_factor=scale,
out_threshold=mask_threshold,
device=device)
if not no_save:
out_fn = out_files[i]
result = mask_to_image(mask)
result.save(out_files[i])
logging.info("Mask saved to {}".format(out_files[i]))
if viz:
logging.info(
"Visualizing results for image {}, close to continue ...".
format(fn))
plot_img_and_mask(img, mask)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
net.to(device=device)
net.load_state_dict(torch.load(args.model, map_location=device))
logging.info("Model loaded !")
for i, fn in enumerate(in_files):
logging.info("\nPredicting image {} ...".format(fn))
img = Image.open(fn)
mask = predict_img(net=net,
full_img=img,
scale_factor=args.scale,
out_threshold=args.mask_threshold,
device=device)
if not args.no_save:
out_fn = out_files[i]
result = mask_to_image(mask)
result.save(out_files[i])
logging.info("Mask saved to {}".format(out_files[i]))
if args.viz:
logging.info(
"Visualizing results for image {}, close to continue ...".
format(fn))
plot_img_and_mask(img, mask)
logging.info("Model loaded !")
for i, fn in enumerate(in_files):
logging.info("\nPredicting image {} ...".format(fn))
img = Image.open(fn)
centerlines, points = predict_img(net=net,
full_img=img,
scale_factor=args.scale,
out_threshold=args.mask_threshold,
device=device)
if not args.no_save:
centerlines_fn = centerlines_files[i]
points_fn = points_files[i]
centerlines_result = mask_to_image(centerlines)
points_result = mask_to_image(points)
centerlines_result.save(centerlines_files[i])
points_result.save(points_files[i])
logging.info("Mask saved to {}".format(centerlines_files[i]))
logging.info("Mask saved to {}".format(points_files[i]))
if args.viz:
logging.info("Visualizing results for image {}, close to continue ...".format(centerlines_fn))
logging.info("Visualizing results for image {}, close to continue ...".format(points_fn))
plot_img_and_mask(img, centerlines)
plot_img_and_mask(img, points)
end_time = time.time()
total_time += end_time - start_time
print("Time taking for predicting:", str(end_time-start_time))
if not args.no_eval:
dc = dice_coeff(pred.unsqueeze(0).unsqueeze(0), true_mask.unsqueeze(0)).item()
tot += dc
print("Dice Coefficient: " + str(dc))
if not args.no_save:
output_file = args.output + 'output_' + img_files[i]
result = mask_to_image(mask)
result.save(output_file)
# replace background based on mask and provided bg
removed_bg = datasetGenerator.apply_mask(img_files[i], mask)
new_img = datasetGenerator.replace_background(img_files[i], mask)
cv2.imwrite(args.output + 'new_' + img_files[i], new_img)
logging.info("Mask saved to {}".format(output_file))
if not args.no_viz:
logging.info("Visualizing results for image {}, close to continue ...".format(img_files[i]))
plot_img_and_mask(plt, ax, org_img, mask, removed_bg, new_img)
plt.savefig(args.output + 'new_{}.jpg'.format(i + 1))
i += 1
print('Averaging Dice Coefficient on ' + str(len(img_files)) + ' testing frames: ' + str(tot/len(img_files)))
print("Total execution time:", str(total_time))
