Skip to content

Commit cbe16eb

Browse files
yashika-gitpre-commit-ci[bot]Nic-Ma
authored
Mean IoU Implementation (#4663)
* Create meaniou.py Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * Create test_compute_meaniou.py Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * Update __init__.py Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * Update metrics.rst Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * Update test_compute_meaniou.py Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * Update meaniou.py Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * Update metrics.rst Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * Update __init__.py Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * Update meaniou.py Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * Update test_compute_meaniou.py Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * Update metrics.rst Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * Update test_compute_meaniou.py Signed-off-by: Yashika Jain <yashikajain201@gmail.com> * Update test_compute_meaniou.py Signed-off-by: Yashika Jain <yashikajain201@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Nic Ma <nma@nvidia.com>
1 parent 9580603 commit cbe16eb

4 files changed

Lines changed: 377 additions & 0 deletions

File tree

‎docs/source/metrics.rst‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,13 @@ Metrics
3939
.. autoclass:: DiceMetric
4040
:members:
4141

42+
`Mean IoU`
43+
----------
44+
.. autofunction:: compute_meaniou
45+
46+
.. autoclass:: MeanIoU
47+
:members:
48+
4249
`Generalized Dice Score`
4350
------------------------
4451
.. autofunction:: compute_generalized_dice

‎monai/metrics/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
from .generalized_dice import GeneralizedDiceScore, compute_generalized_dice
1616
from .hausdorff_distance import HausdorffDistanceMetric, compute_hausdorff_distance, compute_percent_hausdorff_distance
1717
from .meandice import DiceMetric, compute_meandice
18+
from .meaniou import MeanIoU, compute_meaniou
1819
from .metric import Cumulative, CumulativeIterationMetric, IterationMetric, Metric
1920
from .regression import MAEMetric, MSEMetric, PSNRMetric, RMSEMetric, SSIMMetric
2021
from .rocauc import ROCAUCMetric, compute_roc_auc

‎monai/metrics/meaniou.py‎

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
# Copyright (c) MONAI Consortium
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
# http://www.apache.org/licenses/LICENSE-2.0
6+
# Unless required by applicable law or agreed to in writing, software
7+
# distributed under the License is distributed on an "AS IS" BASIS,
8+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
9+
# See the License for the specific language governing permissions and
10+
# limitations under the License.
11+
12+
from typing import Union
13+
14+
import torch
15+
16+
from monai.metrics.utils import do_metric_reduction, ignore_background, is_binary_tensor
17+
from monai.utils import MetricReduction
18+
19+
from .metric import CumulativeIterationMetric
20+
21+
22+
class MeanIoU(CumulativeIterationMetric):
23+
"""
24+
Compute average IoU score between two tensors. It can support both multi-classes and multi-labels tasks.
25+
Input `y_pred` is compared with ground truth `y`.
26+
`y_pred` is expected to have binarized predictions and `y` should be in one-hot format. You can use suitable transforms
27+
in ``monai.transforms.post`` first to achieve binarized values.
28+
The `include_background` parameter can be set to ``False`` to exclude
29+
the first category (channel index 0) which is by convention assumed to be background. If the non-background
30+
segmentations are small compared to the total image size they can get overwhelmed by the signal from the
31+
background.
32+
`y_pred` and `y` can be a list of channel-first Tensor (CHW[D]) or a batch-first Tensor (BCHW[D]).
33+
34+
Args:
35+
include_background: whether to skip IoU computation on the first channel of
36+
the predicted output. Defaults to ``True``.
37+
reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values,
38+
available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``,
39+
``"mean_channel"``, ``"sum_channel"``}, default to ``"mean"``. if "none", will not do reduction.
40+
get_not_nans: whether to return the `not_nans` count, if True, aggregate() returns (metric, not_nans).
41+
Here `not_nans` count the number of not nans for the metric, thus its shape equals to the shape of the metric.
42+
ignore_empty: whether to ignore empty ground truth cases during calculation.
43+
If `True`, NaN value will be set for empty ground truth cases.
44+
If `False`, 1 will be set if the predictions of empty ground truth cases are also empty.
45+
46+
"""
47+
48+
def __init__(
49+
self,
50+
include_background: bool = True,
51+
reduction: Union[MetricReduction, str] = MetricReduction.MEAN,
52+
get_not_nans: bool = False,
53+
ignore_empty: bool = True,
54+
) -> None:
55+
super().__init__()
56+
self.include_background = include_background
57+
self.reduction = reduction
58+
self.get_not_nans = get_not_nans
59+
self.ignore_empty = ignore_empty
60+
61+
def _compute_tensor(self, y_pred: torch.Tensor, y: torch.Tensor): # type: ignore
62+
"""
63+
Args:
64+
y_pred: input data to compute, typical segmentation model output.
65+
It must be one-hot format and first dim is batch, example shape: [16, 3, 32, 32]. The values
66+
should be binarized.
67+
y: ground truth to compute mean IoU metric. It must be one-hot format and first dim is batch.
68+
The values should be binarized.
69+
70+
Raises:
71+
ValueError: when `y` is not a binarized tensor.
72+
ValueError: when `y_pred` has less than three dimensions.
73+
"""
74+
is_binary_tensor(y_pred, "y_pred")
75+
is_binary_tensor(y, "y")
76+
77+
dims = y_pred.ndimension()
78+
if dims < 3:
79+
raise ValueError(f"y_pred should have at least 3 dimensions (batch, channel, spatial), got {dims}.")
80+
# compute IoU (BxC) for each channel for each batch
81+
return compute_meaniou(
82+
y_pred=y_pred, y=y, include_background=self.include_background, ignore_empty=self.ignore_empty
83+
)
84+
85+
def aggregate(self, reduction: Union[MetricReduction, str, None] = None): # type: ignore
86+
"""
87+
Execute reduction logic for the output of `compute_meaniou`.
88+
89+
Args:
90+
reduction: define mode of reduction to the metrics, will only apply reduction on `not-nan` values,
91+
available reduction modes: {``"none"``, ``"mean"``, ``"sum"``, ``"mean_batch"``, ``"sum_batch"``,
92+
``"mean_channel"``, ``"sum_channel"``}, default to `self.reduction`. if "none", will not do reduction.
93+
94+
"""
95+
data = self.get_buffer()
96+
if not isinstance(data, torch.Tensor):
97+
raise ValueError("the data to aggregate must be PyTorch Tensor.")
98+
99+
# do metric reduction
100+
f, not_nans = do_metric_reduction(data, reduction or self.reduction)
101+
return (f, not_nans) if self.get_not_nans else f
102+
103+
104+
def compute_meaniou(
105+
y_pred: torch.Tensor, y: torch.Tensor, include_background: bool = True, ignore_empty: bool = True
106+
) -> torch.Tensor:
107+
"""Computes IoU score metric from full size Tensor and collects average.
108+
109+
Args:
110+
y_pred: input data to compute, typical segmentation model output.
111+
It must be one-hot format and first dim is batch, example shape: [16, 3, 32, 32]. The values
112+
should be binarized.
113+
y: ground truth to compute mean IoU metric. It must be one-hot format and first dim is batch.
114+
The values should be binarized.
115+
include_background: whether to skip IoU computation on the first channel of
116+
the predicted output. Defaults to True.
117+
ignore_empty: whether to ignore empty ground truth cases during calculation.
118+
If `True`, NaN value will be set for empty ground truth cases.
119+
If `False`, 1 will be set if the predictions of empty ground truth cases are also empty.
120+
121+
Returns:
122+
IoU scores per batch and per class, (shape [batch_size, num_classes]).
123+
124+
Raises:
125+
ValueError: when `y_pred` and `y` have different shapes.
126+
127+
"""
128+
129+
if not include_background:
130+
y_pred, y = ignore_background(y_pred=y_pred, y=y)
131+
132+
y = y.float()
133+
y_pred = y_pred.float()
134+
135+
if y.shape != y_pred.shape:
136+
raise ValueError(f"y_pred and y should have same shapes, got {y_pred.shape} and {y.shape}.")
137+
138+
# reducing only spatial dimensions (not batch nor channels)
139+
n_len = len(y_pred.shape)
140+
reduce_axis = list(range(2, n_len))
141+
intersection = torch.sum(y * y_pred, dim=reduce_axis)
142+
143+
y_o = torch.sum(y, reduce_axis)
144+
y_pred_o = torch.sum(y_pred, dim=reduce_axis)
145+
union = y_o + y_pred_o - intersection
146+
147+
if ignore_empty is True:
148+
return torch.where(y_o > 0, (intersection) / union, torch.tensor(float("nan"), device=y_o.device))
149+
return torch.where(union > 0, (intersection) / union, torch.tensor(1.0, device=y_o.device))

‎tests/test_compute_meaniou.py‎

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
# Copyright (c) MONAI Consortium
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
# http://www.apache.org/licenses/LICENSE-2.0
6+
# Unless required by applicable law or agreed to in writing, software
7+
# distributed under the License is distributed on an "AS IS" BASIS,
8+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
9+
# See the License for the specific language governing permissions and
10+
# limitations under the License.
11+
12+
import unittest
13+
14+
import numpy as np
15+
import torch
16+
from parameterized import parameterized
17+
18+
from monai.metrics import MeanIoU, compute_meaniou
19+
20+
# keep background
21+
TEST_CASE_1 = [ # y (1, 1, 2, 2), y_pred (1, 1, 2, 2), expected out (1, 1)
22+
{
23+
"y_pred": torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]),
24+
"y": torch.tensor([[[[1.0, 0.0], [1.0, 1.0]]]]),
25+
"include_background": True,
26+
},
27+
[[0.6667]],
28+
]
29+
30+
# remove background and not One-Hot target
31+
TEST_CASE_2 = [ # y (2, 3, 2, 2), y_pred (2, 3, 2, 2), expected out (2, 2) (no background)
32+
{
33+
"y_pred": torch.tensor(
34+
[
35+
[[[0.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [1.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]]],
36+
[[[0.0, 0.0], [0.0, 1.0]], [[1.0, 0.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 0.0]]],
37+
]
38+
),
39+
"y": torch.tensor(
40+
[
41+
[[[0.0, 0.0], [0.0, 1.0]], [[1.0, 0.0], [1.0, 0.0]], [[0.0, 1.0], [0.0, 0.0]]],
42+
[[[0.0, 0.0], [0.0, 1.0]], [[1.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [1.0, 0.0]]],
43+
]
44+
),
45+
"include_background": False,
46+
},
47+
[[0.3333, 0.0000], [0.5000, 0.5000]],
48+
]
49+
50+
# should return Nan for all labels=0 case and skip for MeanIoU
51+
TEST_CASE_3 = [
52+
{
53+
"y_pred": torch.tensor(
54+
[
55+
[[[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]], [[1.0, 1.0], [1.0, 1.0]]],
56+
[[[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]], [[1.0, 1.0], [1.0, 1.0]]],
57+
]
58+
),
59+
"y": torch.tensor(
60+
[
61+
[[[1.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]],
62+
[[[0.0, 1.0], [1.0, 0.0]], [[1.0, 0.0], [0.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]]],
63+
]
64+
),
65+
"include_background": True,
66+
},
67+
[[False, True, True], [False, False, True]],
68+
]
69+
70+
TEST_CASE_4 = [
71+
{"include_background": True, "reduction": "mean_batch", "get_not_nans": True},
72+
{
73+
"y_pred": torch.tensor(
74+
[
75+
[[[1.0, 1.0], [1.0, 0.0]], [[0.0, 1.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 1.0]]],
76+
[[[1.0, 0.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 0.0]]],
77+
]
78+
),
79+
"y": torch.tensor(
80+
[
81+
[[[1.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]],
82+
[[[0.0, 0.0], [0.0, 1.0]], [[1.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [1.0, 0.0]]],
83+
]
84+
),
85+
},
86+
[0.5416, 0.2500, 0.5000],
87+
]
88+
89+
TEST_CASE_5 = [
90+
{"include_background": True, "reduction": "mean", "get_not_nans": True},
91+
{
92+
"y_pred": torch.tensor(
93+
[
94+
[[[1.0, 1.0], [1.0, 0.0]], [[0.0, 1.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 1.0]]],
95+
[[[1.0, 0.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 0.0]]],
96+
]
97+
),
98+
"y": torch.tensor(
99+
[
100+
[[[1.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]],
101+
[[[0.0, 0.0], [0.0, 1.0]], [[1.0, 1.0], [0.0, 0.0]], [[0.0, 0.0], [1.0, 0.0]]],
102+
]
103+
),
104+
},
105+
0.5555,
106+
]
107+
108+
TEST_CASE_6 = [
109+
{"include_background": True, "reduction": "sum_batch", "get_not_nans": True},
110+
{
111+
"y_pred": torch.tensor(
112+
[
113+
[[[1.0, 1.0], [1.0, 0.0]], [[0.0, 1.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 1.0]]],
114+
[[[1.0, 0.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 0.0]]],
115+
]
116+
),
117+
"y": torch.tensor(
118+
[
119+
[[[1.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]],
120+
[[[1.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]],
121+
]
122+
),
123+
},
124+
[1.5000, 0.0000, 0.0000],
125+
]
126+
127+
TEST_CASE_7 = [
128+
{"include_background": True, "reduction": "mean", "get_not_nans": True},
129+
{
130+
"y_pred": torch.tensor(
131+
[
132+
[[[1.0, 1.0], [1.0, 0.0]], [[0.0, 1.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 1.0]]],
133+
[[[1.0, 0.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 0.0]]],
134+
]
135+
),
136+
"y": torch.tensor(
137+
[
138+
[[[1.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]],
139+
[[[1.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]],
140+
]
141+
),
142+
},
143+
0.7500,
144+
]
145+
146+
TEST_CASE_8 = [
147+
{"include_background": False, "reduction": "sum_batch", "get_not_nans": True},
148+
{
149+
"y_pred": torch.tensor(
150+
[
151+
[[[1.0, 1.0], [1.0, 0.0]], [[0.0, 1.0], [0.0, 0.0]], [[0.0, 1.0], [1.0, 1.0]]],
152+
[[[1.0, 0.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 1.0]], [[0.0, 1.0], [1.0, 0.0]]],
153+
]
154+
),
155+
"y": torch.tensor(
156+
[
157+
[[[1.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]],
158+
[[[1.0, 1.0], [1.0, 1.0]], [[0.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]],
159+
]
160+
),
161+
},
162+
[0.0000, 0.0000],
163+
]
164+
165+
TEST_CASE_9 = [
166+
{"y": torch.ones((2, 2, 3, 3)), "y_pred": torch.ones((2, 2, 3, 3))},
167+
[[1.0000, 1.0000], [1.0000, 1.0000]],
168+
]
169+
170+
TEST_CASE_10 = [
171+
{"y": [torch.ones((2, 3, 3)), torch.ones((2, 3, 3))], "y_pred": [torch.ones((2, 3, 3)), torch.ones((2, 3, 3))]},
172+
[[1.0000, 1.0000], [1.0000, 1.0000]],
173+
]
174+
175+
TEST_CASE_11 = [
176+
{"y": torch.zeros((2, 2, 3, 3)), "y_pred": torch.zeros((2, 2, 3, 3)), "ignore_empty": False},
177+
[[1.0000, 1.0000], [1.0000, 1.0000]],
178+
]
179+
180+
TEST_CASE_12 = [
181+
{"y": torch.zeros((2, 2, 3, 3)), "y_pred": torch.ones((2, 2, 3, 3)), "ignore_empty": False},
182+
[[0.0000, 0.0000], [0.0000, 0.0000]],
183+
]
184+
185+
186+
class TestComputeMeanIoU(unittest.TestCase):
187+
@parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_9, TEST_CASE_11, TEST_CASE_12])
188+
def test_value(self, input_data, expected_value):
189+
result = compute_meaniou(**input_data)
190+
np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4)
191+
192+
@parameterized.expand([TEST_CASE_3])
193+
def test_nans(self, input_data, expected_value):
194+
result = compute_meaniou(**input_data)
195+
self.assertTrue(np.allclose(np.isnan(result.cpu().numpy()), expected_value))
196+
197+
# MeanIoU class tests
198+
@parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_10])
199+
def test_value_class(self, input_data, expected_value):
200+
201+
# same test as for compute_meaniou
202+
vals = {}
203+
vals["y_pred"] = input_data.pop("y_pred")
204+
vals["y"] = input_data.pop("y")
205+
iou_metric = MeanIoU(**input_data)
206+
iou_metric(**vals)
207+
result = iou_metric.aggregate(reduction="none")
208+
np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4)
209+
210+
@parameterized.expand([TEST_CASE_4, TEST_CASE_5, TEST_CASE_6, TEST_CASE_7, TEST_CASE_8])
211+
def test_nans_class(self, params, input_data, expected_value):
212+
213+
iou_metric = MeanIoU(**params)
214+
iou_metric(**input_data)
215+
result, _ = iou_metric.aggregate()
216+
np.testing.assert_allclose(result.cpu().numpy(), expected_value, atol=1e-4)
217+
218+
219+
if __name__ == "__main__":
220+
unittest.main()

0 commit comments

Comments
 (0)