test: run compute tests, format and lint loss/metric updates

Signed-off-by: Rusheel Sharma <rusheelhere@gmail.com>

Author: Rusheel86

monai/losses/focal_loss.py

@@ -174,8 +174,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
                     alpha_arg = None
                     # Move the warning INSIDE this block
                     warnings.warn(
-                        "`include_background=False`, scalar `alpha` ignored when using softmax.", 
-                        stacklevel=2
+                        "`include_background=False`, scalar `alpha` ignored when using softmax.", stacklevel=2
                     )
             loss = softmax_focal_loss(input, target, self.gamma, alpha_arg)
         else:
@@ -184,10 +183,29 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         if mask is not None:
             loss = loss * mask
 
-        if self.class_weight is not None and target.shape[1] != 1:
-            cw = torch.as_tensor(self.class_weight).to(loss)
-            broadcast_dims = [-1] + [1] * len(target.shape[2:])
-            loss = cw.view(broadcast_dims) * loss
+        if self.class_weight is not None:
+            cw = torch.as_tensor(self.class_weight, device=loss.device, dtype=loss.dtype)
+            num_classes = loss.shape[1]
+
+            if cw.ndim > 0:
+                if num_classes == 1:
+                    raise ValueError("Per-class class_weight is not supported for single-channel outputs.")
+                if cw.numel() != num_classes:
+                    raise ValueError(
+                        f"The number of class_weight ({cw.numel()}) must match the number of "
+                        f"output channels ({num_classes})."
+                    )
+                if (cw < 0).any():
+                    raise ValueError("class_weight values must be non-negative.")
+            else:
+                if cw < 0:
+                    raise ValueError("class_weight values must be non-negative.")
+
+            if cw.ndim == 0:
+                loss = loss * cw
+            else:
+                broadcast_shape = [1, num_classes] + [1] * (loss.ndim - 2)
+                loss = loss * cw.view(broadcast_shape)
 
         if self.reduction == LossReduction.SUM.value:
             loss = loss.sum()
@@ -281,10 +299,10 @@ def sigmoid_focal_loss(
                 )
             broadcast_dims = [1, -1] + [1] * len(target.shape[2:])
             alpha_t = alpha_t.view(broadcast_dims)
-            
+
             # Apply per-class weight only to positive samples
             alpha_factor = torch.where(target == 1, alpha_t, torch.ones_like(alpha_t))
 
         # This multiplication now works for both Scalar and Tensor cases
         loss = alpha_factor * loss
-    return loss
+    return loss

monai/losses/unified_focal_loss.py

@@ -69,17 +69,13 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
             raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")
 
         # Handle ignore_index:
-        # Since the wrapper already zeroed out ignore_index in y_true,
-        # we create the mask where y_true has any valid class (sum across channels > 0)
         mask = torch.ones_like(y_true)
         if self.ignore_index is not None:
-            # We identify valid pixels: where at least one channel is 1
+            # Identify valid pixels: where at least one channel is 1
             spatial_mask = (torch.sum(y_true, dim=1, keepdim=True) > 0).float()
             mask = spatial_mask.expand_as(y_true)
-            # Ensure y_pred is also masked so it doesn't create False Positives in ignored areas
             y_pred = y_pred * mask
 
-        # clip the prediction to avoid NaN
         y_pred = torch.clamp(y_pred, self.epsilon, 1.0 - self.epsilon)
         axis = list(range(2, len(y_pred.shape)))
 
@@ -91,7 +87,7 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
 
         # Calculate losses separately for each class, enhancing both classes
         back_dice = 1 - dice_class[:, 0]
-        fore_dice = (1 - dice_class[:, 1]) * torch.pow(1 - dice_class[:, 1], -self.gamma)
+        fore_dice = torch.pow(1 - dice_class[:, 1], 1 - self.gamma)
 
         # Average class scores
         loss = torch.stack([back_dice, fore_dice], dim=-1)
@@ -249,9 +245,9 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         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}")
 
-        # 2. Transform binary inputs to 2-channel space
+        # Transform binary inputs to 2-channel space
         if y_pred.shape[1] == 1:
-            y_pred = torch.cat([-y_pred, y_pred], dim=1)
+            y_pred = torch.cat([1 - y_pred, y_pred], dim=1)
 
         # Move one_hot conversion OUTSIDE the if y_pred.shape[1] == 1 block
         if self.to_onehot_y:
@@ -264,20 +260,15 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
             else:
                 y_true = one_hot(y_true, num_classes=self.num_classes)
 
-        # 3. NOW check if shapes match (They should both be [B, 2, H, W] now)
+        # Check if shapes match
+        if y_true.shape[1] == 1 and y_pred.shape[1] == 2:
+            y_true = torch.cat([1 - y_true, y_true], dim=1)
         if y_true.shape != y_pred.shape:
             raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")
 
         if torch.max(y_true) != self.num_classes - 1:
             raise ValueError(f"Please make sure the number of classes is {self.num_classes - 1}")
 
-        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)
-
         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/metrics/test_ignore_index_metrics.py

@@ -25,7 +25,6 @@
     SurfaceDiceMetric,
     SurfaceDistanceMetric,
 )
-
 from monai.utils import optional_import
 
 scipy, has_scipy = optional_import("scipy")