-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmetrics.py
177 lines (145 loc) · 6.23 KB
/
metrics.py
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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
"""
This file contains the metrics of model accuracy measurement.
"""
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
class AverageMeter:
"""Computes average values"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val):
self.val = val
self.sum = self.sum + self.val
self.count += 1
self.avg = self.sum / self.count
class EMAMeter:
"""Computes and stores an exponential moving average and current value"""
def __init__(self, momentum=0.98):
self.mom = momentum
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val):
self.val = val
self.sum = (self.sum * self.mom) + (val * (1 - self.mom))
self.count += 1
self.avg = self.sum / (1 - self.mom ** (self.count))
class IoU:
"""
Give a meter object, calculates and stores average IoU results
from batches of train or evaluation data.
"""
def __init__(self, meter):
self.meter = meter
def calculate(self, output, target):
#make target the same shape as output by unsqueezing
#the channel dimension, if needed,
if target.ndim == output.ndim - 1:
target = target.unsqueeze(1)
# now the dimension is like [16, 7, H, W]
#get the number of classes from the output channels
n_classes = output.size(1)
#get reshape size based on number of dimensions
#can exclude first 2 dims, which are always batch and channel
empty_dims = (1,) * (target.ndim - 2)
if n_classes > 1:
#softmax the output and argmax,
# The softmax function converts these scores into probabilities,
# making the sum of probabilities for each class at each pixel equal to 1.
output = nn.Softmax(dim=1)(output) #(B, NC, H, W)
#find the index (class) with the maximum probability for each pixel
#dimension (NC) is reduced to 1, as argmax returns the index of the maximum class probability for each pixel
max_idx = torch.argmax(output, 1, keepdim=True) #(B, 1, H, W)
#one hot encode the target and output
k = torch.arange(0, n_classes).view(1, n_classes, *empty_dims).to(target.device)
target = (target == k)
output_onehot = torch.zeros_like(output)
# should be (B, NC, H, W)
output_onehot.scatter_(1, max_idx, 1)
else:
#just sigmoid the output
output = (nn.Sigmoid()(output) > 0.5).long()
x_detached = output.detach() # Detach the tensor from the computation graph
x_numpy = x_detached.cpu().numpy()
# Plotting
plt.imshow(x_numpy[0][0], cmap='gray')
plt.title("Sample Torch Image")
plt.savefig("/kasthuri_pp/transunetoutput.png")
plt.close()
#cast target to the correct type for operations
target = target.type(output.dtype)
#multiply the tensors, everything that is still as 1 is
#part of the intersection (N,)
# dims should be (0,2,3)
dims = (0,) + tuple(range(2, target.ndimension()))
# total intersection across all images in the batch and across all pixels, but separately for each class
intersect = torch.sum(output * target, dims)
#compute the union, (N,)
union = torch.sum(output + target, dims) - intersect
#avoid division errors by adding a small epsilon
#if intersect and union are zero, then iou is 1
iou = (intersect + 1e-5) / (union + 1e-5)
return iou
def update(self, value):
self.meter.update(value)
def reset(self):
self.meter.reset()
def average(self):
return self.meter.avg
class ComposeMetrics:
"""
A class for composing metrics together.
Arguments:
----------
metrics_dict: A dictionary in which each key is the name of a metric and
each value is a metric object.
class_names: Optional list of names for each class index (e.g. mitochondria, lysosomes, etc.)
reset_on_print: Boolean. Whether to reset the results for each metric after the print
function is called. Default True.
"""
def __init__(self, metrics_dict, class_names=None, reset_on_print=True):
self.metrics_dict = metrics_dict
self.class_names = class_names
self.reset_on_print = reset_on_print
def evaluate(self, output, target):
#calculate all the metrics in the dict
# calculate IoU, each metric is iou(EMAmeter()) instance
for metric in self.metrics_dict.values():
value = metric.calculate(output, target)
metric.update(value)
def print(self):
names = []
values = []
print("items:", self.metrics_dict.items())
for name, metric in self.metrics_dict.items(): # is IoU and IoU metric object inside
avg_values = metric.average()
if isinstance(avg_values, torch.Tensor):
avg_values = avg_values.cpu()
# Now, avg_values is either an int or a CPU tensor
#we expect metric to be a tensor of size (n_classes,)
#we want to print the corresponding class names if given
if self.class_names is None:
self.class_names = [f'class_{i}' for i in range(len(avg_values))]
#print("class_name", self.class_names)
#print("avg", avg_values)
for class_name, val in zip(self.class_names, avg_values):
#print("name&value:", class_name)
#print(val)
names.append(f'{class_name}_{name}') # name is IoU
values.append(val.item())
if self.reset_on_print:
metric.reset()
for name, value in zip(names, values):
print(f'{name}: {value:.3f}')
mean_iou = np.array(values).mean()
print(f'Mean IoU: {mean_iou}')