fix bugs

Signed-off-by: ytl0623 <[email protected]>

Author: ytl0623

monai/losses/unified_focal_loss.py

@@ -66,12 +66,11 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
             y_true: ground truth labels. Shape should match y_pred.
         """
 
-        # Auto-handle single channel input (binary segmentation case)
-        if y_pred.shape[1] == 1 and not self.use_softmax:
+        if y_pred.shape[1] == 1:
             y_pred = torch.sigmoid(y_pred)
             y_pred = torch.cat([1 - y_pred, y_pred], dim=1)
             is_already_prob = True
-            # Expand y_true to match if it's single channel
+
             if y_true.shape[1] == 1:
                 y_true = one_hot(y_true, num_classes=2)
         else:
@@ -122,12 +121,12 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         # Apply reduction
         if self.reduction == LossReduction.MEAN.value:
             return torch.mean(all_losses)
-        if self.reduction == LossReduction.SUM.value:
+        elif self.reduction == LossReduction.SUM.value:
             return torch.sum(all_losses)
-        if self.reduction == LossReduction.NONE.value:
+        elif self.reduction == LossReduction.NONE.value:
             return all_losses
-
-        return torch.mean(all_losses)
+        else:
+            return torch.mean(all_losses)
 
 
 class AsymmetricFocalLoss(_Loss):
@@ -253,6 +252,14 @@ def __init__(
             delta: background/foreground balancing weight. Defaults to 0.7.
             reduction: specifies the reduction to apply to the output. Defaults to "mean".
             use_softmax: whether to use softmax for probability conversion. Defaults to False.
+
+        Example:
+            >>> import torch
+            >>> from monai.losses import AsymmetricUnifiedFocalLoss
+            >>> pred = torch.ones((1,1,32,32), dtype=torch.float32)
+            >>> grnd = torch.ones((1,1,32,32), dtype=torch.int64)
+            >>> fl = AsymmetricUnifiedFocalLoss(to_onehot_y=True)
+            >>> fl(pred, grnd)
         """
         super().__init__(reduction=LossReduction(reduction).value)
         self.to_onehot_y = to_onehot_y
@@ -283,30 +290,36 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
             y_pred: Prediction logits. Shape: (B, C, H, W, [D]).
                     Supports binary (C=1 or C=2) and multi-class (C>2) segmentation.
             y_true: Ground truth labels. Shape should match y_pred (or be indices if to_onehot_y is True).
+
+        Raises:
+            ValueError: When ground truth shape does not match input shape.
+            ValueError: When input tensor shape is not 4D or 5D.
+            ValueError: When the number of classes in ground truth exceeds the configured `num_classes`.
         """
         if y_pred.shape != y_true.shape:
-            is_binary_logits = y_pred.shape[1] == 1 and not self.use_softmax
-            if not self.to_onehot_y and not is_binary_logits:
-                raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")
-            raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")
+            is_binary_logits = (y_pred.shape[1] == 1) and (not self.use_softmax)
+            is_target_needs_onehot = self.to_onehot_y and (y_true.shape[1] == 1)
 
-        if len(y_pred.shape) != 4 and len(y_pred.shape) != 5:
-            raise ValueError(f"input shape must be 4 or 5, but got {y_pred.shape}")
-
-        if y_pred.shape[1] == 1:
-            y_pred = one_hot(y_pred, num_classes=self.num_classes)
-            y_true = one_hot(y_true, num_classes=self.num_classes)
+            if not is_binary_logits and not is_target_needs_onehot:
+                raise ValueError(
+                    f"Ground truth has different shape ({y_true.shape}) from input ({y_pred.shape}), "
+                    "and this mismatch cannot be resolved by `to_onehot_y` or binary expansion."
+                )
 
-        if torch.max(y_true) != self.num_classes - 1:
-            raise ValueError(f"Please make sure the number of classes is {self.num_classes - 1}")
+        if len(y_pred.shape) not in [4, 5]:
+            raise ValueError(f"Input shape must be 4 (2D) or 5 (3D), but got {y_pred.shape}")
 
-        n_pred_ch = y_pred.shape[1]
         if self.to_onehot_y:
-            if n_pred_ch == 1:
-                warnings.warn("single channel prediction, `to_onehot_y=True` ignored.")
-            else:
-                y_true = one_hot(y_true, num_classes=n_pred_ch)
-
+            # Only convert if y_true is single channel (Indices)
+            if y_true.shape[1] == 1:
+                # Check indices validity before conversion
+                if torch.max(y_true) >= self.num_classes:
+                    raise ValueError(
+                        f"Ground truth contains class indices >= {self.num_classes}, which exceeds num_classes."
+                    )
+
+                # Convert to One-hot
+                y_true = one_hot(y_true, num_classes=self.num_classes)
         asy_focal_loss = self.asy_focal_loss(y_pred, y_true)
         asy_focal_tversky_loss = self.asy_focal_tversky_loss(y_pred, y_true)
 

tests/losses/test_unified_focal_loss.py

@@ -20,24 +20,20 @@
 from monai.losses import AsymmetricUnifiedFocalLoss
 
 # 1. Binary Case (Logits input): Prediction matches GT perfectly
-# Input Shape: (B, 1, H, W) -> Auto expanded internally
 TEST_CASE_BINARY_LOGITS = [
     {"y_pred": torch.tensor([[[[10.0, -10.0], [-10.0, 10.0]]]]), "y_true": torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]])},
     0.0,
-    {"use_softmax": False, "to_onehot_y": False},
+    {"use_softmax": False, "to_onehot_y": False, "num_classes": 2},
 ]
 
 # 2. Binary Case (2 Channels input): Prediction matches GT perfectly
-# Input Shape: (B, 2, H, W)
 TEST_CASE_BINARY_2CH = [
     {
-        "y_pred": torch.tensor(
-            [[[[-10.0, 10.0], [10.0, -10.0]], [[10.0, -10.0], [-10.0, 10.0]]]]  # Ch0 (Background): Low, High, High, Low
-        ),  # Ch1 (Foreground): High, Low, Low, High
+        "y_pred": torch.tensor([[[[-10.0, 10.0], [10.0, -10.0]], [[10.0, -10.0], [-10.0, 10.0]]]]),
         "y_true": torch.tensor([[[[1, 0], [0, 1]]]]),
     },
     0.0,
-    {"use_softmax": True, "to_onehot_y": True},
+    {"use_softmax": True, "to_onehot_y": True, "num_classes": 2},
 ]
 
 # 3. Multi-Class Case (3 Channels): Prediction matches GT perfectly
@@ -46,16 +42,16 @@
         "y_pred": torch.tensor(
             [
                 [
-                    [[10.0, -10.0], [-10.0, 10.0]],  # Class 0 Logits
-                    [[-10.0, 10.0], [-10.0, -10.0]],  # Class 1 Logits
-                    [[-10.0, -10.0], [10.0, -10.0]],
+                    [[10.0, -10.0], [-10.0, 10.0]],  # Class 0
+                    [[-10.0, 10.0], [-10.0, -10.0]],  # Class 1
+                    [[-10.0, -10.0], [10.0, -10.0]],  # Class 2
                 ]
             ]
-        ),  # Class 2 Logits
-        "y_true": torch.tensor([[[[0, 1], [2, 0]]]]),  # Indices
+        ),
+        "y_true": torch.tensor([[[[0, 1], [2, 0]]]]),
     },
     0.0,
-    {"use_softmax": True, "to_onehot_y": True},
+    {"use_softmax": True, "to_onehot_y": True, "num_classes": 3},
 ]
 
 # 4. Multi-Class Case: Wrong Prediction
@@ -64,10 +60,10 @@
         "y_pred": torch.tensor(
             [[[[-10.0, -10.0], [-10.0, -10.0]], [[10.0, 10.0], [10.0, 10.0]], [[-10.0, -10.0], [-10.0, -10.0]]]]
         ),
-        "y_true": torch.tensor([[[[0, 0], [0, 0]]]]),  # GT is class 0, but Pred is class 1
+        "y_true": torch.tensor([[[[0, 0], [0, 0]]]]),
     },
     None,
-    {"use_softmax": True, "to_onehot_y": True},
+    {"use_softmax": True, "to_onehot_y": True, "num_classes": 3},
 ]
 
 
@@ -77,11 +73,11 @@ class TestAsymmetricUnifiedFocalLoss(unittest.TestCase):
     def test_perfect_prediction(self, input_data, expected_val, args):
         loss_func = AsymmetricUnifiedFocalLoss(**args)
         result = loss_func(**input_data)
-        # We use a small tolerance because 10.0 logits is not exactly probability 1.0
+        # Using a relaxed tolerance for logits -> probability conversion
         np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-3, rtol=1e-3)
 
     @parameterized.expand([TEST_CASE_MULTICLASS_WRONG])
-    def test_wrong_prediction(self, input_data, expected_val, args):
+    def test_wrong_prediction(self, input_data, _, args):
         loss_func = AsymmetricUnifiedFocalLoss(**args)
         result = loss_func(**input_data)
         self.assertGreater(result.item(), 1.0, "Loss should be high for wrong predictions")
@@ -93,7 +89,6 @@ def test_ill_shape(self):
 
     def test_with_cuda(self):
         if not torch.cuda.is_available():
-            print("CUDA not available, skipping test_with_cuda")
             return
 
         loss = AsymmetricUnifiedFocalLoss(use_softmax=False, to_onehot_y=False)
@@ -102,7 +97,6 @@ def test_with_cuda(self):
         j = torch.tensor([[[[1.0, 0], [0, 1.0]]], [[[1.0, 0], [0, 1.0]]]]).cuda()
 
         output = loss(i, j)
-        print(f"CUDA Output: {output.item()}")
         self.assertTrue(output.is_cuda)
         self.assertLess(output.item(), 1.0)