-
Notifications
You must be signed in to change notification settings - Fork 1.5k
/
model_targets.py
109 lines (82 loc) · 3.22 KB
/
model_targets.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
import numpy as np
import torch
import torchvision
class ClassifierOutputTarget:
def __init__(self, category):
self.category = category
def __call__(self, model_output):
if len(model_output.shape) == 1:
return model_output[self.category]
return model_output[:, self.category]
class ClassifierOutputSoftmaxTarget:
def __init__(self, category):
self.category = category
def __call__(self, model_output):
if len(model_output.shape) == 1:
return torch.softmax(model_output, dim=-1)[self.category]
return torch.softmax(model_output, dim=-1)[:, self.category]
class BinaryClassifierOutputTarget:
def __init__(self, category):
self.category = category
def __call__(self, model_output):
if self.category == 1:
sign = 1
else:
sign = -1
return model_output * sign
class SoftmaxOutputTarget:
def __init__(self):
pass
def __call__(self, model_output):
return torch.softmax(model_output, dim=-1)
class RawScoresOutputTarget:
def __init__(self):
pass
def __call__(self, model_output):
return model_output
class SemanticSegmentationTarget:
""" Gets a binary spatial mask and a category,
And return the sum of the category scores,
of the pixels in the mask. """
def __init__(self, category, mask):
self.category = category
self.mask = torch.from_numpy(mask)
if torch.cuda.is_available():
self.mask = self.mask.cuda()
if torch.backends.mps.is_available():
self.mask = self.mask.to("mps")
def __call__(self, model_output):
return (model_output[self.category, :, :] * self.mask).sum()
class FasterRCNNBoxScoreTarget:
""" For every original detected bounding box specified in "bounding boxes",
assign a score on how the current bounding boxes match it,
1. In IOU
2. In the classification score.
If there is not a large enough overlap, or the category changed,
assign a score of 0.
The total score is the sum of all the box scores.
"""
def __init__(self, labels, bounding_boxes, iou_threshold=0.5):
self.labels = labels
self.bounding_boxes = bounding_boxes
self.iou_threshold = iou_threshold
def __call__(self, model_outputs):
output = torch.Tensor([0])
if torch.cuda.is_available():
output = output.cuda()
elif torch.backends.mps.is_available():
output = output.to("mps")
if len(model_outputs["boxes"]) == 0:
return output
for box, label in zip(self.bounding_boxes, self.labels):
box = torch.Tensor(box[None, :])
if torch.cuda.is_available():
box = box.cuda()
elif torch.backends.mps.is_available():
box = box.to("mps")
ious = torchvision.ops.box_iou(box, model_outputs["boxes"])
index = ious.argmax()
if ious[0, index] > self.iou_threshold and model_outputs["labels"][index] == label:
score = ious[0, index] + model_outputs["scores"][index]
output = output + score
return output