FasterAUPRO/aupro_efficient.py at main · AlirezaSalehy/FasterAUPRO · 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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import numpy as np
from skimage import measure
import torch

# primary ref: https://github.com/gudovskiy/cflow-ad/blob/master/train.py
def cal_pro_score_opp(masks, amaps, max_step=200, expect_fpr=0.3):
    binary_amaps = np.zeros_like(amaps, dtype=bool)
    min_th, max_th = amaps.min(), amaps.max()
    delta = (max_th - min_th) / max_step
    pros, fprs, ths = [], [], []

    regionprops_list = [measure.regionprops(measure.label(mask)) for mask in masks]
    coords_list = [[(region.coords[:, 0], region.coords[:, 1], len(region.coords)) for region in regionprops] for regionprops in regionprops_list]
    inverse_masks = 1 - masks
    tn_pixel = inverse_masks.sum() # Pixels that truly has the label of 0
    for th in np.arange(min_th, max_th, delta):
        binary_amaps[amaps <= th], binary_amaps[amaps > th] = 0, 1
        pro = []

        for binary_amap, regions_coords in zip(binary_amaps, coords_list):
            for coords in regions_coords:
                tp_pixels = binary_amap[coords[0], coords[1]].sum()
                pro.append(tp_pixels / coords[2])

        fp_pixels = np.logical_and(inverse_masks, binary_amaps).sum()
        fpr = fp_pixels / tn_pixel
        pros.append(np.mean(pro))
        fprs.append(fpr)
        ths.append(th)

    pros, fprs, ths = np.array(pros), np.array(fprs), np.array(ths)
    idxes = fprs < expect_fpr
    fprs = fprs[idxes]
    pros = pros[idxes]
    if np.ptp(fprs) == 0:
        pro_auc = np.mean(pros)
    else:
        fprs = (fprs - fprs.min()) / (fprs.max() - fprs.min())
        pro_auc = auc(fprs, pros)
    return pro_auc

# GPU implementation using PyTorch
def cal_pro_score_gpu(masks, amaps, max_step=200, expect_fpr=0.3, device="cuda"):
    if not torch.is_tensor(amaps) or not torch.is_tensor(masks):
        amaps = torch.tensor(amaps)
        masks = torch.tensor(masks)
    amaps = amaps.to(device)
    masks = masks.to(device)

    binary_amaps = torch.zeros_like(amaps, dtype=torch.bool, device=device)
    min_th, max_th = amaps.min().item(), amaps.max().item()
    delta = (max_th - min_th) / max_step
    pros, fprs, ths = [], [], []

    regionprops_list = [measure.regionprops(measure.label(mask.cpu().numpy())) for mask in masks]
    coords_list = [[(region.coords[:, 0], region.coords[:, 1], len(region.coords)) for region in regionprops] for regionprops in regionprops_list]
    inverse_masks = 1 - masks
    tn_pixel = inverse_masks.sum().item() # Pixels that truly has the label of 0
    for th in np.arange(min_th, max_th, delta):
        binary_amaps[amaps <= th], binary_amaps[amaps > th] = 0, 1
        pro = []

        for binary_amap, regions_coords in zip(binary_amaps, coords_list):
            for coords in regions_coords:
                tp_pixels = binary_amap[coords[0], coords[1]].sum().item()
                pro.append(tp_pixels / coords[2])

        fp_pixels = torch.logical_and(inverse_masks, binary_amaps).sum().item()
        fpr = fp_pixels / tn_pixel
        pros.append(np.mean(pro))
        fprs.append(fpr)
        ths.append(th.item())

    pros, fprs, ths = np.array(pros), np.array(fprs), np.array(ths)
    idxes = fprs < expect_fpr
    fprs = fprs[idxes]
    pros = pros[idxes]
    if (fprs.max() - fprs.min()) == 0:
        pro_auc = np.mean(pros)
    else:
        fprs = (fprs - fprs.min()) / (fprs.max() - fprs.min())
        pro_auc = auc(fprs, pros)
    return pro_auc