SACuP/eval.py at main · AIRLABkhu/SACuP · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
import torch
import numpy as np
from sklearn.metrics import confusion_matrix
from visualization import *

def compute_IoU(cm):
    '''
    Adapted from:
        https://github.com/davidtvs/PyTorch-ENet/blob/master/metric/iou.py
        https://github.com/tensorflow/tensorflow/blob/v2.3.0/tensorflow/python/keras/metrics.py#L2716-L2844
    '''

    sum_over_row = cm.sum(axis=0)
    sum_over_col = cm.sum(axis=1)
    true_positives = np.diag(cm)

    # sum_over_row + sum_over_col = 2 * true_positives + false_positives + false_negatives.
    denominator = sum_over_row + sum_over_col - true_positives

    iou = true_positives / denominator

    return iou, np.nanmean(iou)


def eval_net_loader(net, val_loader, n_classes, device='cpu', epoch=0):
    # for validation vis_epoch folder
    if not os.path.exists(os.path.join('./vis/val/',str(epoch))):
        os.makedirs(os.path.join('./vis/val/',str(epoch)))

    net.eval()
    labels = np.arange(n_classes)
    cf = np.zeros((n_classes, n_classes))

    for i, sample_batch in enumerate(val_loader):
        imgs = sample_batch['image']
        true_masks = sample_batch['mask']

        imgs = imgs.to(device)
        true_masks = true_masks.to(device)

        outputs = net(imgs)
        probs = torch.softmax(outputs, dim=1)
        preds = torch.argmax(probs, dim=1)

        for j in range(len(true_masks)):
            true = true_masks[j].cpu().detach().numpy().flatten()
            pred = preds[j].cpu().detach().numpy().flatten()
            cf += confusion_matrix(true, pred, labels=labels)
            # visualization
            pred_colorization(imgs.cpu().detach().numpy(), true, pred, epoch, i)

    class_iou, mean_iou = compute_IoU(cf)

    return class_iou, mean_iou, cf


def IoU(mask_true, mask_pred, n_classes=2):
    labels = np.arange(n_classes)
    cm = confusion_matrix(mask_true.flatten(), mask_pred.flatten(), labels=labels)

    return compute_IoU(cm)