Added Focal tversky loss

segmentation_models_pytorch/losses/__init__.py

@@ -8,6 +8,7 @@
 from .soft_ce import SoftCrossEntropyLoss
 from .tversky import TverskyLoss
 from .mcc import MCCLoss
+from .focal_tversky import FocalTverskyLoss
 
 __all__ = [
     "BINARY_MODE",
@@ -21,4 +22,5 @@
     "SoftCrossEntropyLoss",
     "TverskyLoss",
     "MCCLoss",
+    "FocalTverskyLoss",
 ]

segmentation_models_pytorch/losses/_functional.py

@@ -205,6 +205,37 @@ def soft_tversky_score(
     ).clamp_min(eps)
     return tversky_score
 
+def focal_tversky_loss(
+    output: torch.Tensor,
+    target: torch.Tensor,
+    alpha: float,
+    beta: float,
+    gamma: float,
+    smooth: float = 0.0,
+    eps: float = 1e-7,
+    dims=None,
+) -> torch.Tensor:
+    """Focal Tversky loss
+
+    References:
+        https://arxiv.org/pdf/1810.07842
+
+    """
+    assert output.size() == target.size()
+
+    output_sum = torch.sum(output, dim=dims)
+    target_sum = torch.sum(target, dim=dims)
+    difference = LA.vector_norm(output - target, ord=1, dim=dims)
+
+    intersection = (output_sum + target_sum - difference) / 2  # TP
+    fp = output_sum - intersection
+    fn = target_sum - intersection
+
+    tversky_score = (intersection + smooth) / (
+        intersection + alpha * fp + beta * fn + smooth
+    ).clamp_min(eps)
+    tversky_loss = (1 - tversky_score) ** gamma
+    return tversky_loss
 
 def wing_loss(
     output: torch.Tensor, target: torch.Tensor, width=5, curvature=0.5, reduction="mean"

segmentation_models_pytorch/losses/focal_tversky.py

@@ -0,0 +1,69 @@
+from typing import List, Optional
+
+import torch
+from ._functional import focal_tversky_loss
+from .constants import BINARY_MODE, MULTICLASS_MODE, MULTILABEL_MODE
+from .dice import DiceLoss
+
+__all__ = ["TverskyLoss"]
+
+
+class FocalTverskyLoss(DiceLoss):
+    """Focal Tversky loss for image segmentation tasks.
+    FP and FN are weighted by alpha and beta parameters, respectively.
+    With alpha == beta == 0.5, this loss becomes equivalent to DiceLoss.
+    The gamma parameter focuses the loss on hard-to-classify examples.
+    If gamma > 1, the function focuses more on misclassified examples.
+    If gamma = 1, it is equivalent to Tversky Loss.
+    This loss supports binary, multiclass, and multilabel cases.
+
+    Args:
+        mode: Metric mode {'binary', 'multiclass', 'multilabel'}
+        classes: Optional list of classes that contribute to loss computation;
+        By default, all channels are included.
+        log_loss: If True, computes loss as ``-log(tversky)``; otherwise as ``1 - tversky``
+        from_logits: If True, assumes input is raw logits
+        smooth: Smoothing factor to avoid division by zero
+        ignore_index: Label indicating ignored pixels (not contributing to loss)
+        eps: Small epsilon for numerical stability
+        alpha: Weight constant that penalizes False Positives (FPs)
+        beta: Weight constant that penalizes False Negatives (FNs)
+        gamma: Focal factor to adjust the focus on harder examples (defaults to 1.0)
+
+
+    Return:
+        loss: torch.Tensor
+
+    """
+
+    def __init__(
+        self,
+        mode: str,
+        classes: List[int] = None,
+        log_loss: bool = False,
+        from_logits: bool = True,
+        smooth: float = 0.0,
+        ignore_index: Optional[int] = None,
+        eps: float = 1e-7,
+        alpha: float = 0.5,
+        beta: float = 0.5,
+        gamma: float = 1.0,
+    ):
+        assert mode in {BINARY_MODE, MULTILABEL_MODE, MULTICLASS_MODE}
+        super().__init__(
+            mode, classes, log_loss, from_logits, smooth, ignore_index, eps
+        )
+        self.alpha = alpha
+        self.beta = beta
+        self.gamma = gamma
+
+    def aggregate_loss(self, loss):
+        return loss.mean() ** self.gamma
+
+    def compute_score(
+        self, output, target, smooth=0.0, eps=1e-7, dims=None
+    ) -> torch.Tensor:
+        tversky_score = focal_tversky_loss(
+            output, target, self.alpha, self.beta, self.gamma, smooth, eps, dims
+        )
+        return tversky_score

tests/test_losses.py

@@ -9,6 +9,7 @@
     SoftCrossEntropyLoss,
     TverskyLoss,
     MCCLoss,
+    focal_tversky,
 )
 
 
@@ -92,6 +93,23 @@ def test_soft_tversky_score(y_true, y_pred, expected, eps, alpha, beta):
     actual = F.soft_tversky_score(y_pred, y_true, eps=eps, alpha=alpha, beta=beta)
     assert float(actual) == pytest.approx(expected, eps)
 
+@pytest.mark.parametrize(
+    ["y_true", "y_pred", "expected", "eps", "alpha", "beta", "gamma"],
+    [
+        [[1, 1, 1, 1], [1, 1, 1, 1], 0.0, 1e-5, 0.5, 0.5, 1.0], 
+        [[0, 1, 1, 0], [0, 1, 1, 0], 0.0, 1e-5, 0.5, 0.5, 1.0], 
+        [[1, 1, 1, 1], [0, 0, 0, 0], 1.0, 1e-5, 0.5, 0.5, 1.0], 
+        [[1, 0, 1, 0], [0, 1, 0, 0], 1.0, 1e-5, 0.5, 0.5, 1.0],  
+        [[1, 0, 1, 0], [1, 1, 0, 0], 0.5, 1e-5, 0.5, 0.5, 1.0],  
+    ],
+)
+
+def test_focal_tversky_score(y_true, y_pred, expected, eps, alpha, beta, gamma):
+    y_true = torch.tensor(y_true, dtype=torch.float32)
+    y_pred = torch.tensor(y_pred, dtype=torch.float32)
+    actual = F.focal_tversky_loss(y_pred, y_true, eps=eps, alpha=alpha, beta=beta, gamma=gamma)
+    assert float(actual) == pytest.approx(expected, eps)
+
 
 @torch.no_grad()
 def test_dice_loss_binary():