-
Notifications
You must be signed in to change notification settings - Fork 387
/
dice.py
116 lines (95 loc) · 4 KB
/
dice.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
# Copyright The PyTorch Lightning team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Tuple
import torch
from torchmetrics.utilities.data import to_categorical
from torchmetrics.utilities.distributed import reduce
def _stat_scores(
preds: torch.Tensor,
target: torch.Tensor,
class_index: int,
argmax_dim: int = 1,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Calculates the number of true positive, false positive, true negative
and false negative for a specific class
Args:
pred: prediction tensor
target: target tensor
class_index: class to calculate over
argmax_dim: if pred is a tensor of probabilities, this indicates the
axis the argmax transformation will be applied over
Return:
True Positive, False Positive, True Negative, False Negative, Support
Example:
>>> x = torch.tensor([1, 2, 3])
>>> y = torch.tensor([0, 2, 3])
>>> tp, fp, tn, fn, sup = _stat_scores(x, y, class_index=1)
>>> tp, fp, tn, fn, sup
(tensor(0), tensor(1), tensor(2), tensor(0), tensor(0))
"""
if preds.ndim == target.ndim + 1:
preds = to_categorical(preds, argmax_dim=argmax_dim)
tp = ((preds == class_index) * (target == class_index)).to(torch.long).sum()
fp = ((preds == class_index) * (target != class_index)).to(torch.long).sum()
tn = ((preds != class_index) * (target != class_index)).to(torch.long).sum()
fn = ((preds != class_index) * (target == class_index)).to(torch.long).sum()
sup = (target == class_index).to(torch.long).sum()
return tp, fp, tn, fn, sup
def dice_score(
pred: torch.Tensor,
target: torch.Tensor,
bg: bool = False,
nan_score: float = 0.0,
no_fg_score: float = 0.0,
reduction: str = 'elementwise_mean',
) -> torch.Tensor:
"""
Compute dice score from prediction scores
Args:
pred: estimated probabilities
target: ground-truth labels
bg: whether to also compute dice for the background
nan_score: score to return, if a NaN occurs during computation
no_fg_score: score to return, if no foreground pixel was found in target
reduction: a method to reduce metric score over labels.
- ``'elementwise_mean'``: takes the mean (default)
- ``'sum'``: takes the sum
- ``'none'``: no reduction will be applied
Return:
Tensor containing dice score
Example:
>>> pred = torch.tensor([[0.85, 0.05, 0.05, 0.05],
... [0.05, 0.85, 0.05, 0.05],
... [0.05, 0.05, 0.85, 0.05],
... [0.05, 0.05, 0.05, 0.85]])
>>> target = torch.tensor([0, 1, 3, 2])
>>> dice_score(pred, target)
tensor(0.3333)
"""
num_classes = pred.shape[1]
bg = (1 - int(bool(bg)))
scores = torch.zeros(num_classes - bg, device=pred.device, dtype=torch.float32)
for i in range(bg, num_classes):
if not (target == i).any():
# no foreground class
scores[i - bg] += no_fg_score
continue
# TODO: rewrite to use general `stat_scores`
tp, fp, tn, fn, sup = _stat_scores(preds=pred, target=target, class_index=i)
denom = (2 * tp + fp + fn).to(torch.float)
# nan result
score_cls = (2 * tp).to(torch.float) / denom if torch.is_nonzero(denom) else nan_score
scores[i - bg] += score_cls
return reduce(scores, reduction=reduction)