806e98c744
Former-commit-id: 0181b044ec0497d95c80e9e7c5112537806e0afc [formerly 1f62452a8fa1bbf6cfa5c82347c58cc1886ae796] Former-commit-id: 31dba554cbd0126009890ac54dfd6b6d38f2a47c
41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
import torch
|
|
from tqdm import tqdm
|
|
|
|
import wandb
|
|
from src.utils.dice import dice_coeff
|
|
|
|
|
|
def evaluate(net, dataloader, device):
|
|
net.eval()
|
|
num_val_batches = len(dataloader)
|
|
dice_score = 0
|
|
|
|
# iterate over the validation set
|
|
with tqdm(dataloader, total=len(dataloader.dataset), desc="val", unit="img", leave=False) as pbar:
|
|
for images, masks_true in dataloader:
|
|
# move images and labels to correct device
|
|
images = images.to(device=device)
|
|
masks_true = masks_true.unsqueeze(1).float().to(device=device)
|
|
|
|
# forward, predict the mask
|
|
with torch.inference_mode():
|
|
masks_pred = net(images)
|
|
masks_pred = (torch.sigmoid(masks_pred) > 0.5).float()
|
|
|
|
# compute the Dice score
|
|
dice_score += dice_coeff(masks_pred, masks_true, reduce_batch_first=False)
|
|
|
|
# update progress bar
|
|
pbar.update(images.shape[0])
|
|
|
|
# save some images to wandb
|
|
table = wandb.Table(columns=["image", "mask", "prediction"])
|
|
for img, mask, pred in zip(images.to("cpu"), masks_true.to("cpu"), masks_pred.to("cpu")):
|
|
table.add_data(wandb.Image(img), wandb.Image(mask), wandb.Image(pred))
|
|
wandb.log({"predictions_table": table}, commit=False)
|
|
|
|
net.train()
|
|
|
|
# Fixes a potential division by zero error
|
|
return dice_score / num_val_batches if num_val_batches else dice_score
|