added a local variable is_converted_to_prob that tracks whether we've already converted to probabilities

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

Author: ytl0623

monai/losses/unified_focal_loss.py

@@ -71,20 +71,26 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
             y_pred: prediction logits or probabilities. Shape should be (B, C, spatial_dims).
             y_true: ground truth labels. Shape should match y_pred.
         """
-
+        
+        # Track whether we've already converted to probabilities
+        is_converted_to_prob = self.is_prob_input
+        
         # Handle single-channel binary case
         if y_pred.shape[1] == 1:
             # For binary segmentation, enforce num_classes == 2
             if self.num_classes != 2:
-                raise ValueError(f"Single-channel input requires num_classes=2, but got num_classes={self.num_classes}")
-
+                raise ValueError(
+                    f"Single-channel input requires num_classes=2, but got num_classes={self.num_classes}"
+                )
+            
             # Convert to probability if needed
-            if not self.is_prob_input:
+            if not is_converted_to_prob:
                 y_pred = torch.sigmoid(y_pred)
-
+                is_converted_to_prob = True
+            
             # Expand to 2 channels
             y_pred = torch.cat([1 - y_pred, y_pred], dim=1)
-
+            
             # Convert y_true to one-hot with 2 classes (matching expanded y_pred)
             if y_true.shape[1] == 1:
                 y_true = one_hot(y_true, num_classes=2)
@@ -107,14 +113,16 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
                 if y_true.shape[1] == 1:
                     # Validate class indices
                     if torch.max(y_true) >= self.num_classes:
-                        raise ValueError(f"Ground truth contains class indices >= {self.num_classes}")
+                        raise ValueError(
+                            f"Ground truth contains class indices >= {self.num_classes}"
+                        )
                     y_true = one_hot(y_true, num_classes=self.num_classes)
 
         if y_true.shape != y_pred.shape:
             raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")
 
         # Convert logits to probabilities if not already done
-        if not self.is_prob_input and y_pred.shape[1] != 1:
+        if not is_converted_to_prob:
             if self.use_softmax:
                 y_pred = torch.softmax(y_pred, dim=1)
             else:
@@ -203,19 +211,25 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
             y_true: ground truth labels.
         """
 
+        # Track whether we've already converted to probabilities
+        is_converted_to_prob = self.is_prob_input
+        
         # Handle single-channel binary case
         if y_pred.shape[1] == 1:
             # For binary segmentation, enforce num_classes == 2
             if self.num_classes != 2:
-                raise ValueError(f"Single-channel input requires num_classes=2, but got num_classes={self.num_classes}")
-
+                raise ValueError(
+                    f"Single-channel input requires num_classes=2, but got num_classes={self.num_classes}"
+                )
+            
             # Convert to probability if needed
-            if not self.is_prob_input:
+            if not is_converted_to_prob:
                 y_pred = torch.sigmoid(y_pred)
-
+                is_converted_to_prob = True
+            
             # Expand to 2 channels
             y_pred = torch.cat([1 - y_pred, y_pred], dim=1)
-
+            
             # Convert y_true to one-hot with 2 classes (matching expanded y_pred)
             if y_true.shape[1] == 1:
                 y_true = one_hot(y_true, num_classes=2)
@@ -238,14 +252,16 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
                 if y_true.shape[1] == 1:
                     # Validate class indices
                     if torch.max(y_true) >= self.num_classes:
-                        raise ValueError(f"Ground truth contains class indices >= {self.num_classes}")
+                        raise ValueError(
+                            f"Ground truth contains class indices >= {self.num_classes}"
+                        )
                     y_true = one_hot(y_true, num_classes=self.num_classes)
 
         if y_true.shape != y_pred.shape:
             raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")
 
         # Convert logits to probabilities if not already done
-        if not self.is_prob_input and y_pred.shape[1] != 1:
+        if not is_converted_to_prob:
             if self.use_softmax:
                 y_pred = torch.softmax(y_pred, dim=1)
             else:
@@ -267,10 +283,12 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
 
         # Apply reduction
         if self.reduction == LossReduction.MEAN.value:
+            # For MEAN: average over all elements (batch and classes and spatial)
             return torch.mean(all_ce)
         if self.reduction == LossReduction.SUM.value:
             return torch.sum(all_ce)
         if self.reduction == LossReduction.NONE.value:
+            # For NONE: return full tensor with all spatial dimensions
             return all_ce
 
         return torch.mean(all_ce)
@@ -316,11 +334,13 @@ def __init__(
             >>> fl(pred, grnd)
         """
         super().__init__(reduction=LossReduction(reduction).value)
-
+        
         # Validate configuration
         if use_softmax and num_classes < 2:
-            raise ValueError(f"use_softmax=True requires num_classes >= 2, but got num_classes={num_classes}")
-
+            raise ValueError(
+                f"use_softmax=True requires num_classes >= 2, but got num_classes={num_classes}"
+            )
+        
         self.to_onehot_y = to_onehot_y
         self.num_classes = num_classes
         self.gamma = gamma
@@ -362,18 +382,18 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
             ValueError: When the number of classes in ground truth exceeds the configured `num_classes`.
             ValueError: When use_softmax=True with single-channel input.
         """
-
+        
         # Validate input dimensions
         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}")
-
+        
         # Validate use_softmax with single channel
         if self.use_softmax and y_pred.shape[1] == 1:
             raise ValueError(
                 "use_softmax=True is invalid with single-channel input (C=1). "
                 "Use use_softmax=False for binary segmentation with sigmoid activation."
             )
-
+        
         # Shape validation
         if y_pred.shape != y_true.shape:
             # Allow mismatch only for valid cases
@@ -386,13 +406,14 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
                     "and this mismatch cannot be resolved by `to_onehot_y` or binary expansion."
                 )
 
+        # Pre-process y_true if needed (will be done inside sub-losses, but validate here)
         if self.to_onehot_y and 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."
                 )
-
+        
         # Get losses from sub-losses
         # AsymmetricFocalLoss with NONE reduction returns (B, C, H, W, [D])
         # AsymmetricFocalTverskyLoss with NONE reduction returns (B, C)
@@ -421,4 +442,4 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         else:
             loss = torch.mean(combined_loss)
 
-        return loss
+        return loss