def predict_img(net,
full_img,
device,
scale_factor=1,
out_threshold=0.5):
net.eval()
img = torch.from_numpy(BasicDataset.preprocess(full_img, scale_factor))
img = img.unsqueeze(0)
img = img.to(device=device, dtype=torch.float32)
with torch.no_grad():
output = net(img)
if net.n_classes > 1:
probs = F.softmax(output, dim=1)
else:
probs = torch.sigmoid(output)
probs = probs.squeeze(0)
tf = transforms.Compose(
[
transforms.ToPILImage(),
transforms.Resize(full_img.size[1]),
transforms.ToTensor()
]
)
probs = tf(probs.cpu())
full_mask = probs.squeeze().cpu().numpy()
return full_mask > out_threshold
def segment_img(img_path, net, device, scale_factor=1, out_threshold=0.5):
"""
Segement out the mask (target vein region) of the input image (specified by image path)
"""
print(f"\nPerforming segmentation on ---> {img_path} ...")
img = torch.from_numpy(BasicDataset.preprocess(Image.open(img_path).convert('L'), scale_factor))
img = img.unsqueeze(0) # add dimension
img = img.to(device=device, dtype=torch.float32)
begin_time = time.time()
with torch.no_grad():
output = net(img)
end_time = time.time()
inference_time = end_time - begin_time
print(f'inference_time: {inference_time}s')
# transform to image numpy array
if net.n_classes > 1:
probs = F.softmax(output, dim=1)
else:
probs = torch.sigmoid(output)
probs = probs.squeeze(0) #cuda tensor
tf = transforms.Compose(
[
transforms.ToPILImage(),
# transforms.Resize(full_img.size[1]),
transforms.ToTensor()
]
)
probs = tf(probs.cpu())
full_mask = probs.squeeze().cpu().numpy()
if np.count_nonzero(full_mask) == 0:
print("No veins segmented out on this image!")
return full_mask > out_threshold, inference_time
image_folder = args.image_folder + '/' + args.dataset_name
for i, fn in enumerate(os.listdir(image_folder)):
if not fn.endswith('.jpg'):
continue
count += 1
# input image
img_path = os.path.join(image_folder, fn)
target_path = img_path.replace('imgs', 'masks')
# single image prediction
mask, inference_time = segment_img(img_path, net, device, args.scale, args.mask_threshold)
# target mask
# FloatTensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
FloatTensor = torch.cuda.FloatTensor if args.device == 'cuda' else torch.FloatTensor
target = BasicDataset.preprocess(Image.open(target_path), args.scale)
# validation score computation
val_score = dice_coeff(FloatTensor(mask), FloatTensor(target).squeeze(0)).item()
print(f'validation_score (dice_coeff): {val_score}')
val_score_total += val_score
# prediction speed
inference_time_total += inference_time
if args.save_results:
output_path = os.path.join(output_dir, fn)
mask_to_image(mask).save(output_path)
print(f"saved predicted mask to ---> {output_path}")
if args.viz: