@@ -1,99 +1,43 @@
+from typing import Union
+
 import torch
-from typing import Union, Optional
-from torch.nn.modules.loss import _Loss
-from super_gradients.training.utils.segmentation_utils import to_one_hot
-from super_gradients.training.losses.loss_utils import apply_reduce, LossReduction
+
 from super_gradients.common.abstractions.abstract_logger import get_logger
+from super_gradients.training.losses.loss_utils import LossReduction
+from super_gradients.training.losses.structure_loss import AbstarctSegmentationStructureLoss
 
 logger = get_logger(__name__)
 
 
-class DiceLoss(_Loss):
+class DiceLoss(AbstarctSegmentationStructureLoss):
     """
     Compute average Dice loss between two tensors, It can support both multi-classes and binary tasks.
     Defined in the paper: "V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation"
     """
-    def __init__(self,
-                 apply_softmax: bool = True,
-                 ignore_index: int = None,
-                 smooth: float = 1.,
-                 eps: float = 1e-5,
-                 sum_over_batches: bool = False,
-                 generalized_dice: bool = False,
-                 weight: Optional[torch.Tensor] = None,
-                 reduction: Union[LossReduction, str] = "mean"):
-        """
-        :param apply_softmax: Whether to apply softmax to the predictions.
-        :param smooth: laplace smoothing, also known as additive smoothing. The larger smooth value is, closer the dice
-            coefficient is to 1, which can be used as a regularization effect.
-            As mentioned in: https://github.com/pytorch/pytorch/issues/1249#issuecomment-337999895
-        :param eps: epsilon value to avoid inf.
-        :param sum_over_batches: Whether to average dice over the batch axis if set True,
-         default is `False` to average over the classes axis.
-        :param generalized_dice: Whether to apply normalization by the volume of each class.
-        :param weight: a manual rescaling weight given to each class. If given, it has to be a Tensor of size `C`.
-        :param reduction: Specifies the reduction to apply to the output: `none` | `mean` | `sum`.
-            `none`: no reduction will be applied.
-            `mean`: the sum of the output will be divided by the number of elements in the output.
-            `sum`: the output will be summed.
-            Default: `mean`
+    def _calc_numerator_denominator(self, labels_one_hot, predict):
         """
-        super().__init__(reduction=reduction)
-        self.ignore_index = ignore_index
-        self.apply_softmax = apply_softmax
-        self.eps = eps
-        self.smooth = smooth
-        self.sum_over_batches = sum_over_batches
-        self.generalized_dice = generalized_dice
-        self.weight = weight
-        if self.generalized_dice:
-            assert self.weight is None, "Cannot use Dice Loss with weight classes and generalized dice normalization"
-            if self.eps > 1e-12:
-                logger.warning("When using GeneralizedDiceLoss, it is recommended to use eps below 1e-12, to not affect"
-                               "small values normalized terms.")
-            if self.smooth != 0:
-                logger.warning("When using GeneralizedDiceLoss, it is recommended to set smooth value as 0.")
+        Calculate dice metric's numerator and denominator.
 
-    def forward(self, predict, target):
-        if self.apply_softmax:
-            predict = torch.softmax(predict, dim=1)
-        # target to one hot format
-        if target.size() == predict.size():
-            labels_one_hot = target
-        elif len(target.size()) == 3:       # if target tensor is in class indexes format.
-            if predict.size(1) == 1 and self.ignore_index is None:    # if one class prediction task
-                labels_one_hot = target.unsqueeze(1)
-            else:
-                labels_one_hot = to_one_hot(target, num_classes=predict.shape[1], ignore_index=self.ignore_index)
-        else:
-            raise AssertionError(f"Mismatch of target shape: {target.size()} and prediction shape: {predict.size()},"
-                                 f" target must be [NxWxH] tensor for to_one_hot conversion"
-                                 f" or to have the same num of channels like prediction tensor")
-
-        reduce_spatial_dims = list(range(2, len(predict.shape)))
-        reduce_dims = [1] + reduce_spatial_dims if self.sum_over_batches else [0] + reduce_spatial_dims
-
-        intersection = torch.sum(labels_one_hot * predict, dim=reduce_dims)
+        :param labels_one_hot: target in one hot format.   shape: [BS, num_classes, img_width, img_height]
+        :param predict: predictions tensor.                shape: [BS, num_classes, img_width, img_height]
+        :return:
+            numerator = intersection between predictions and target. shape: [BS, num_classes, img_width, img_height]
+            denominator = sum of predictions and target areas.       shape: [BS, num_classes, img_width, img_height]
+        """
+        numerator = labels_one_hot * predict
         denominator = labels_one_hot + predict
-        # exclude ignore labels from denominator, false positive predicted on ignore samples are not included in
-        # total denominator.
-        if self.ignore_index is not None:
-            valid_mask = target.ne(self.ignore_index).unsqueeze(1).expand_as(denominator)
-            denominator *= valid_mask
-        denominator = torch.sum(denominator, dim=reduce_dims)
+        return numerator, denominator
 
-        if self.generalized_dice:
-            weights = 1. / (torch.sum(labels_one_hot, dim=reduce_dims) ** 2)
-            # if some classes are not in batch, weights will be inf.
-            infs = torch.isinf(weights)
-            weights[infs] = 0.0
-            intersection *= weights
-            denominator *= weights
+    def _calc_loss(self, numerator, denominator):
+        """
+        Calculate dice loss.
+        All tensors are of shape [BS] if self.reduce_over_batches else [num_classes].
 
-        dices = 1. - ((2. * intersection + self.smooth) / (denominator + self.eps + self.smooth))
-        if self.weight is not None:
-            dices *= self.weight
-        return apply_reduce(dices, reduction=self.reduction)
+        :param numerator: intersection between predictions and target.
+        :param denominator: total number of pixels of prediction and target.
+        """
+        loss = 1. - ((2. * numerator + self.smooth) / (denominator + self.eps + self.smooth))
+        return loss
 
 
 class BinaryDiceLoss(DiceLoss):
@@ -112,7 +56,7 @@ class BinaryDiceLoss(DiceLoss):
             As mentioned in: https://github.com/pytorch/pytorch/issues/1249#issuecomment-337999895
         :param eps: epsilon value to avoid inf.
         """
-        super().__init__(apply_softmax=False, ignore_index=None, smooth=smooth, eps=eps, sum_over_batches=False)
+        super().__init__(apply_softmax=False, ignore_index=None, smooth=smooth, eps=eps, reduce_over_batches=False)
         self.apply_sigmoid = apply_sigmoid
 
     def forward(self, predict, target):
@@ -138,7 +82,7 @@ class GeneralizedDiceLoss(DiceLoss):
                  ignore_index: int = None,
                  smooth: float = 0.0,
                  eps: float = 1e-17,
-                 sum_over_batches: bool = False,
+                 reduce_over_batches: bool = False,
                  reduction: Union[LossReduction, str] = "mean"
                  ):
         """
@@ -147,7 +91,7 @@ class GeneralizedDiceLoss(DiceLoss):
             coefficient is to 1, which can be used as a regularization effect.
             As mentioned in: https://github.com/pytorch/pytorch/issues/1249#issuecomment-337999895
         :param eps: epsilon value to avoid inf.
-        :param sum_over_batches: Whether to average dice over the batch axis if set True,
+        :param reduce_over_batches: Whether to apply reduction over the batch axis if set True,
          default is `False` to average over the classes axis.
         :param reduction: Specifies the reduction to apply to the output: `none` | `mean` | `sum`.
             `none`: no reduction will be applied.
@@ -156,4 +100,5 @@ class GeneralizedDiceLoss(DiceLoss):
             Default: `mean`
         """
         super().__init__(apply_softmax=apply_softmax, ignore_index=ignore_index, smooth=smooth, eps=eps,
-                         sum_over_batches=sum_over_batches, generalized_dice=True, weight=None, reduction=reduction)
+                         reduce_over_batches=reduce_over_batches, generalized_metric=True, weight=None,
+                         reduction=reduction)

from typing import Union

import torch

from super_gradients.common.abstractions.abstract_logger import get_logger
from super_gradients.training.losses.loss_utils import LossReduction
from super_gradients.training.losses.structure_loss import AbstarctSegmentationStructureLoss

logger = get_logger(__name__)


class IoULoss(AbstarctSegmentationStructureLoss):
    """
    Compute average IoU loss between two tensors, It can support both multi-classes and binary tasks.
    """
    def _calc_numerator_denominator(self, labels_one_hot, predict):
        """
        Calculate iou metric's numerator and denominator.

        :param labels_one_hot: target in one hot format.   shape: [BS, num_classes, img_width, img_height]
        :param predict: predictions tensor.                shape: [BS, num_classes, img_width, img_height]
        :return:
            numerator = intersection between predictions and target.    shape: [BS, num_classes, img_width, img_height]
            denominator = area of union between predictions and target. shape: [BS, num_classes, img_width, img_height]
        """
        numerator = labels_one_hot * predict
        denominator = labels_one_hot + predict - numerator
        return numerator, denominator

    def _calc_loss(self, numerator, denominator):
        """
        Calculate iou loss.
        All tensors are of shape [BS] if self.reduce_over_batches else [num_classes]

        :param numerator: intersection between predictions and target.
        :param denominator: area of union between prediction pixels and target pixels.
        """
        loss = 1. - ((numerator + self.smooth) / (denominator + self.eps + self.smooth))
        return loss


class BinaryIoULoss(IoULoss):
    """
    Compute IoU Loss for binary class tasks (1 class only).
    Except target to be a binary map with 0 and 1 values.
    """
    def __init__(self,
                 apply_sigmoid: bool = True,
                 smooth: float = 1.,
                 eps: float = 1e-5):
        """
        :param apply_sigmoid: Whether to apply sigmoid to the predictions.
        :param smooth: laplace smoothing, also known as additive smoothing. The larger smooth value is, closer the IoU
            coefficient is to 1, which can be used as a regularization effect.
            As mentioned in: https://github.com/pytorch/pytorch/issues/1249#issuecomment-337999895
        :param eps: epsilon value to avoid inf.
        """
        super().__init__(apply_softmax=False, ignore_index=None, smooth=smooth, eps=eps, reduce_over_batches=False)
        self.apply_sigmoid = apply_sigmoid

    def forward(self, predict, target):
        if self.apply_sigmoid:
            predict = torch.sigmoid(predict)
        return super().forward(predict=predict, target=target)


class GeneralizedIoULoss(IoULoss):
    """
    Compute the Generalised IoU loss, contribution of each label is normalized by the inverse of its volume, in order
     to deal with class imbalance.
    Args:
        smooth (float): default value is 0, smooth laplacian is not recommended to be used with GeneralizedIoULoss.
         because the weighted values to be added are very small.
        eps (float): default value is 1e-17, must be a very small value, because weighted `intersection` and
        `denominator` are very small after multiplication with `1 / counts ** 2`
    """
    def __init__(self,
                 apply_softmax: bool = True,
                 ignore_index: int = None,
                 smooth: float = 0.0,
                 eps: float = 1e-17,
                 reduce_over_batches: bool = False,
                 reduction: Union[LossReduction, str] = "mean"
                 ):
        """
        :param apply_softmax: Whether to apply softmax to the predictions.
        :param smooth: laplace smoothing, also known as additive smoothing. The larger smooth value is, closer the iou
            coefficient is to 1, which can be used as a regularization effect.
            As mentioned in: https://github.com/pytorch/pytorch/issues/1249#issuecomment-337999895
        :param eps: epsilon value to avoid inf.
        :param reduce_over_batches: Whether to apply reduction over the batch axis if set True,
         default is `False` to average over the classes axis.
        :param reduction: Specifies the reduction to apply to the output: `none` | `mean` | `sum`.
            `none`: no reduction will be applied.
            `mean`: the sum of the output will be divided by the number of elements in the output.
            `sum`: the output will be summed.
            Default: `mean`
        """
        super().__init__(apply_softmax=apply_softmax, ignore_index=ignore_index, smooth=smooth, eps=eps,
                         reduce_over_batches=reduce_over_batches, generalized_metric=True, weight=None,
                         reduction=reduction)

from abc import ABC, abstractmethod
from typing import Union, Optional

import torch
from torch.nn.modules.loss import _Loss

from super_gradients.common.abstractions.abstract_logger import get_logger
from super_gradients.training.losses.loss_utils import apply_reduce, LossReduction
from super_gradients.training.utils.segmentation_utils import to_one_hot

logger = get_logger(__name__)


class AbstarctSegmentationStructureLoss(_Loss, ABC):
    """
    Abstract computation of structure loss between two tensors, It can support both multi-classes and binary tasks.
    """
    def __init__(self,
                 apply_softmax: bool = True,
                 ignore_index: int = None,
                 smooth: float = 1.,
                 eps: float = 1e-5,
                 reduce_over_batches: bool = False,
                 generalized_metric: bool = False,
                 weight: Optional[torch.Tensor] = None,
                 reduction: Union[LossReduction, str] = "mean"):
        """
        :param apply_softmax: Whether to apply softmax to the predictions.
        :param smooth: laplace smoothing, also known as additive smoothing. The larger smooth value is, closer the metric
            coefficient is to 1, which can be used as a regularization effect.
            As mentioned in: https://github.com/pytorch/pytorch/issues/1249#issuecomment-337999895
        :param eps: epsilon value to avoid inf.
        :param reduce_over_batches: Whether to average metric over the batch axis if set True,
         default is `False` to average over the classes axis.
        :param generalized_metric: Whether to apply normalization by the volume of each class.
        :param weight: a manual rescaling weight given to each class. If given, it has to be a Tensor of size `C`.
        :param reduction: Specifies the reduction to apply to the output: `none` | `mean` | `sum`.
            `none`: no reduction will be applied.
            `mean`: the sum of the output will be divided by the number of elements in the output.
            `sum`: the output will be summed.
            Default: `mean`
        """
        super().__init__(reduction=reduction)
        self.ignore_index = ignore_index
        self.apply_softmax = apply_softmax
        self.eps = eps
        self.smooth = smooth
        self.reduce_over_batches = reduce_over_batches
        self.generalized_metric = generalized_metric
        self.weight = weight
        if self.generalized_metric:
            assert self.weight is None, "Cannot use structured Loss with weight classes and generalized normalization"
            if self.eps > 1e-12:
                logger.warning("When using GeneralizedLoss, it is recommended to use eps below 1e-12, to not affect"
                               "small values normalized terms.")
            if self.smooth != 0:
                logger.warning("When using GeneralizedLoss, it is recommended to set smooth value as 0.")

    @abstractmethod
    def _calc_numerator_denominator(self, labels_one_hot, predict) -> (torch.Tensor, torch.Tensor):
        """
        All base classes must implement this function.
        Return: 2 tensor of shape [BS, num_classes, img_width, img_height].
        """
        raise NotImplementedError()

    @abstractmethod
    def _calc_loss(self, numerator, denominator) -> torch.Tensor:
        """
        All base classes must implement this function.
        Return a tensors of shape [BS] if self.reduce_over_batches else [num_classes].
        """
        raise NotImplementedError()

    def forward(self, predict, target):
        if self.apply_softmax:
            predict = torch.softmax(predict, dim=1)
        # target to one hot format
        if target.size() == predict.size():
            labels_one_hot = target
        elif target.dim() == 3:       # if target tensor is in class indexes format.
            if predict.size(1) == 1 and self.ignore_index is None:    # if one class prediction task
                labels_one_hot = target.unsqueeze(1)
            else:
                labels_one_hot = to_one_hot(target, num_classes=predict.shape[1], ignore_index=self.ignore_index)
        else:
            raise AssertionError(f"Mismatch of target shape: {target.size()} and prediction shape: {predict.size()},"
                                 f" target must be [NxWxH] tensor for to_one_hot conversion"
                                 f" or to have the same num of channels like prediction tensor")

        reduce_spatial_dims = list(range(2, len(predict.shape)))
        reduce_dims = [1] + reduce_spatial_dims if self.reduce_over_batches else [0] + reduce_spatial_dims

        # Calculate the numerator and denominator of the chosen metric
        numerator, denominator = self._calc_numerator_denominator(labels_one_hot, predict)

        # exclude ignore labels from numerator and denominator, false positive predicted on ignore samples
        # are not included in the total calculation.
        if self.ignore_index is not None:
            valid_mask = target.ne(self.ignore_index).unsqueeze(1).expand_as(denominator)
            numerator *= valid_mask
            denominator *= valid_mask

        numerator = torch.sum(numerator, dim=reduce_dims)
        denominator = torch.sum(denominator, dim=reduce_dims)

        if self.generalized_metric:
            weights = 1. / (torch.sum(labels_one_hot, dim=reduce_dims) ** 2)
            # if some classes are not in batch, weights will be inf.
            infs = torch.isinf(weights)
            weights[infs] = 0.0
            numerator *= weights
            denominator *= weights

        # Calculate the loss of the chosen metric
        losses = self._calc_loss(numerator, denominator)
        if self.weight is not None:
            losses *= self.weight
        return apply_reduce(losses, reduction=self.reduction)

@@ -16,6 +16,7 @@ from tests.unit_tests.lr_warmup_test import LRWarmupTest
 from tests.unit_tests.kd_ema_test import KDEMATest
 from tests.unit_tests.kd_model_test import KDModelTest
 from tests.unit_tests.dice_loss_test import DiceLossTest
+from tests.unit_tests.iou_loss_test import IoULossTest
 from tests.unit_tests.update_param_groups_unit_test import UpdateParamGroupsTest
 from tests.unit_tests.vit_unit_test import TestViT
 from tests.unit_tests.yolox_unit_test import TestYOLOX
@@ -68,6 +69,7 @@ class CoreUnitTestSuiteRunner:
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(ContextMethodsTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(UpdateParamGroupsTest))
         self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(MaskAttentionLossTest))
+        self.unit_tests_suite.addTest(self.test_loader.loadTestsFromModule(IoULossTest))
 
     def _add_modules_to_end_to_end_tests_suite(self):
         """

import torch
import unittest

from super_gradients.training.losses.iou_loss import IoULoss, GeneralizedIoULoss, BinaryIoULoss


class IoULossTest(unittest.TestCase):

    def setUp(self) -> None:
        self.img_size = 32
        self.eps = 1e-5
        self.num_classes = 2

    def _get_default_predictions_tensor(self, fill_value: float):
        return torch.empty(3, self.num_classes, self.img_size, self.img_size).fill_(fill_value)

    def _get_default_target_zeroes_tensor(self):
        return torch.zeros((3, self.img_size, self.img_size)).long()

    def _assertion_iou_torch_values(self, expected_value: torch.Tensor, found_value: torch.Tensor, rtol: float = 1e-5):
        self.assertTrue(
            torch.allclose(found_value, expected_value, rtol=rtol),
            msg=f"Unequal iou loss: excepted: {expected_value}, found: {found_value}"
        )

    def test_iou(self):
        predictions = self._get_default_predictions_tensor(0.)
        # only label 0 is predicted as positive.
        predictions[:, 0] = 1.
        target = self._get_default_target_zeroes_tensor()
        # half target with label 0, the other half with 1.
        target[:, :self.img_size // 2] = 1

        intersection = torch.tensor([0.5, 0.])
        union = torch.tensor([1., 0.5])

        expected_iou_loss = 1. - (intersection / (union + self.eps))
        expected_iou_loss = expected_iou_loss.mean()

        criterion = IoULoss(smooth=0, eps=self.eps, apply_softmax=False)
        iou_loss = criterion(predictions, target)

        self._assertion_iou_torch_values(expected_iou_loss, iou_loss)

    def test_iou_binary(self):
        # all predictions are 0.6
        predictions = torch.ones((1, 1, self.img_size, self.img_size)) * 0.6
        target = self._get_default_target_zeroes_tensor()
        # half target with label 0, the other half with 1.
        target[:, :self.img_size // 2] = 1

        intersection = torch.tensor([0.6 * 0.5])
        union = torch.tensor([0.6 + 0.5 - 0.6 * 0.5])

        expected_iou_loss = 1. - (intersection / (union + self.eps))
        expected_iou_loss = expected_iou_loss.mean()

        criterion = BinaryIoULoss(smooth=0, eps=self.eps, apply_sigmoid=False)
        iou_loss = criterion(predictions, target)

        self._assertion_iou_torch_values(expected_iou_loss, iou_loss, rtol=1e-3)

    def test_iou_weight_classes(self):
        weight = torch.tensor([0.25, 0.66])
        predictions = self._get_default_predictions_tensor(0.)
        # only label 0 is predicted as positive.
        predictions[:, 0] = 1.
        target = self._get_default_target_zeroes_tensor()
        # half target with label 0, the other half with 1.
        target[:, :self.img_size // 2] = 1

        intersection = torch.tensor([0.5, 0.])
        union = torch.tensor([1., 0.5])

        expected_iou_loss = 1. - (intersection / (union + self.eps))
        expected_iou_loss *= weight
        expected_iou_loss = expected_iou_loss.mean()

        criterion = IoULoss(smooth=0, eps=self.eps, apply_softmax=False, weight=weight)
        iou_loss = criterion(predictions, target)

        self._assertion_iou_torch_values(expected_iou_loss, iou_loss)

    def test_iou_with_ignore(self):
        ignore_index = 2
        predictions = self._get_default_predictions_tensor(0.)
        # only label 0 is predicted as positive.
        predictions[:, 0] = 1.
        target = self._get_default_target_zeroes_tensor()
        # half target with label 0, quarter with 1 and quarter with ignore.
        target[:, :self.img_size // 2, :self.img_size // 2] = 1
        target[:, :self.img_size // 2, self.img_size // 2:] = ignore_index
        # ignore samples are excluded in both intersection and union.
        intersection = torch.tensor([0.5, 0.])
        union = torch.tensor([0.75, 0.25])

        expected_iou_loss = 1. - (intersection / (union + self.eps))
        expected_iou_loss = expected_iou_loss.mean()

        criterion = IoULoss(smooth=0, eps=self.eps, apply_softmax=False, ignore_index=ignore_index)
        iou_loss = criterion(predictions, target)

        self._assertion_iou_torch_values(expected_iou_loss, iou_loss)

    def test_generalized_iou(self):
        predictions = self._get_default_predictions_tensor(0.)
        # half prediction are 0 class, the other half 1 class.
        predictions[:, 0, :self.img_size // 2] = 1.
        predictions[:, 1, self.img_size // 2:] = 1.
        # only 0 class in target.
        target = self._get_default_target_zeroes_tensor()

        intersection = torch.tensor([0.5, 0.])
        union = torch.tensor([1., 0.5])
        counts = torch.tensor([target.numel(), 0.])
        weights = 1 / (counts ** 2)
        weights[1] = 0.0        # instead of inf

        eps = 1e-17
        expected_iou_loss = 1. - ((weights * intersection) / (weights * union + eps))
        expected_iou_loss = expected_iou_loss.mean()

        criterion = GeneralizedIoULoss(smooth=0, eps=eps, apply_softmax=False)
        iou_loss = criterion(predictions, target)

        self._assertion_iou_torch_values(expected_iou_loss, iou_loss)


if __name__ == '__main__':
    unittest.main()