|
| 1 | +# Copyright The PyTorch Lightning team. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | +from typing import Any, Optional |
| 15 | + |
| 16 | +import torch |
| 17 | +from pytorch_lightning.metrics.classification.confusion_matrix import ConfusionMatrix |
| 18 | +from pytorch_lightning.metrics.functional.iou import _iou_from_confmat |
| 19 | + |
| 20 | + |
| 21 | +class IoU(ConfusionMatrix): |
| 22 | + r""" |
| 23 | + Computes `Intersection over union, or Jaccard index calculation <https://en.wikipedia.org/wiki/Jaccard_index>`_: |
| 24 | +
|
| 25 | + .. math:: J(A,B) = \frac{|A\cap B|}{|A\cup B|} |
| 26 | +
|
| 27 | + Where: :math:`A` and :math:`B` are both tensors of the same size, containing integer class values. |
| 28 | + They may be subject to conversion from input data (see description below). Note that it is different from box IoU. |
| 29 | +
|
| 30 | + Works with binary, multiclass and multi-label data. |
| 31 | + Accepts logits from a model output or integer class values in prediction. |
| 32 | + Works with multi-dimensional preds and target. |
| 33 | +
|
| 34 | + Forward accepts |
| 35 | +
|
| 36 | + - ``preds`` (float or long tensor): ``(N, ...)`` or ``(N, C, ...)`` where C is the number of classes |
| 37 | + - ``target`` (long tensor): ``(N, ...)`` |
| 38 | +
|
| 39 | + If preds and target are the same shape and preds is a float tensor, we use the ``self.threshold`` argument. |
| 40 | + This is the case for binary and multi-label logits. |
| 41 | +
|
| 42 | + If preds has an extra dimension as in the case of multi-class scores we perform an argmax on ``dim=1``. |
| 43 | +
|
| 44 | + Args: |
| 45 | + num_classes: Number of classes in the dataset. |
| 46 | + ignore_index: optional int specifying a target class to ignore. If given, this class index does not contribute |
| 47 | + to the returned score, regardless of reduction method. Has no effect if given an int that is not in the |
| 48 | + range [0, num_classes-1]. By default, no index is ignored, and all classes are used. |
| 49 | + absent_score: score to use for an individual class, if no instances of the class index were present in |
| 50 | + `pred` AND no instances of the class index were present in `target`. For example, if we have 3 classes, |
| 51 | + [0, 0] for `pred`, and [0, 2] for `target`, then class 1 would be assigned the `absent_score`. |
| 52 | + threshold: |
| 53 | + Threshold value for binary or multi-label logits. |
| 54 | + reduction: a method to reduce metric score over labels. |
| 55 | +
|
| 56 | + - ``'elementwise_mean'``: takes the mean (default) |
| 57 | + - ``'sum'``: takes the sum |
| 58 | + - ``'none'``: no reduction will be applied |
| 59 | +
|
| 60 | + compute_on_step: |
| 61 | + Forward only calls ``update()`` and return None if this is set to False. |
| 62 | + dist_sync_on_step: |
| 63 | + Synchronize metric state across processes at each ``forward()`` |
| 64 | + before returning the value at the step. |
| 65 | + process_group: |
| 66 | + Specify the process group on which synchronization is called. default: None (which selects the entire world) |
| 67 | +
|
| 68 | + Example: |
| 69 | + >>> from pytorch_lightning.metrics import IoU |
| 70 | + >>> target = torch.randint(0, 2, (10, 25, 25)) |
| 71 | + >>> pred = torch.tensor(target) |
| 72 | + >>> pred[2:5, 7:13, 9:15] = 1 - pred[2:5, 7:13, 9:15] |
| 73 | + >>> iou = IoU(num_classes=2) |
| 74 | + >>> iou(pred, target) |
| 75 | + tensor(0.9660) |
| 76 | +
|
| 77 | + """ |
| 78 | + |
| 79 | + def __init__( |
| 80 | + self, |
| 81 | + num_classes: int, |
| 82 | + ignore_index: Optional[int] = None, |
| 83 | + absent_score: float = 0.0, |
| 84 | + threshold: float = 0.5, |
| 85 | + reduction: str = 'elementwise_mean', |
| 86 | + compute_on_step: bool = True, |
| 87 | + dist_sync_on_step: bool = False, |
| 88 | + process_group: Optional[Any] = None, |
| 89 | + ): |
| 90 | + super().__init__( |
| 91 | + num_classes=num_classes, |
| 92 | + normalize=None, |
| 93 | + threshold=threshold, |
| 94 | + compute_on_step=compute_on_step, |
| 95 | + dist_sync_on_step=dist_sync_on_step, |
| 96 | + process_group=process_group, |
| 97 | + ) |
| 98 | + self.reduction = reduction |
| 99 | + self.ignore_index = ignore_index |
| 100 | + self.absent_score = absent_score |
| 101 | + |
| 102 | + def compute(self) -> torch.Tensor: |
| 103 | + """ |
| 104 | + Computes intersection over union (IoU) |
| 105 | + """ |
| 106 | + return _iou_from_confmat(self.confmat, self.num_classes, self.ignore_index, self.absent_score, self.reduction) |
0 commit comments