mirror of
https://github.com/Laurent2916/REVA-QCAV.git
synced 2024-09-19 19:45:28 +00:00
36 lines
1.3 KiB
Python
36 lines
1.3 KiB
Python
|
import torch
|
||
|
|
import torch.nn.functional as F
|
||
|
|
from tqdm import tqdm
|
||
|
|
|
||
|
|
from dice_score import multiclass_dice_coeff
|
||
|
|
|
||
|
|
|
||
|
|
def evaluate(net, dataloader, device):
|
||
|
|
net.eval()
|
||
|
|
num_val_batches = len(dataloader)
|
||
|
|
dice_score = 0
|
||
|
|
|
||
|
|
# iterate over the validation set
|
||
|
|
for batch in tqdm(dataloader, total=num_val_batches, desc='Validation round', unit='batch', leave=False):
|
||
|
|
image, mask_true = batch['image'], batch['mask']
|
||
|
|
# move images and labels to correct device and type
|
||
|
|
image = image.to(device=device, dtype=torch.float32)
|
||
|
|
mask_true = mask_true.to(device=device, dtype=torch.long)
|
||
|
|
mask_true = F.one_hot(mask_true, net.n_classes).permute(0, 3, 1, 2).float()
|
||
|
|
|
||
|
|
with torch.no_grad():
|
||
|
|
# predict the mask
|
||
|
|
mask_pred = net(image)
|
||
|
|
|
||
|
|
# convert to one-hot format
|
||
|
|
if net.n_classes == 1:
|
||
|
|
mask_pred = (F.sigmoid(mask_pred) > 0).float()
|
||
|
|
else:
|
||
|
|
mask_pred = F.one_hot(mask_pred.argmax(dim=1), net.n_classes).permute(0, 3, 1, 2).float()
|
||
|
|
|
||
|
|
# compute the Dice score, ignoring background
|
||
|
|
dice_score += multiclass_dice_coeff(mask_pred[:, :1, ...], mask_true[:, :1, ...], reduce_batch_first=False)
|
||
|
|
|
||
|
|
net.train()
|
||
|
|
return dice_score / num_val_batches
|