Make 'y' optional in ignore_background

Signed-off-by: ytl0623 <david89062388@gmail.com>

Author: ytl0623

monai/metrics/active_learning_metrics.py

@@ -130,7 +130,6 @@ def compute_variance(
 
     if not include_background:
         y = y_pred
-        # TODO If this utils is made to be optional for 'y' it would be nice
         y_pred, y = ignore_background(y_pred=y_pred, y=y)
 
     # Set any values below 0 to threshold

monai/metrics/utils.py

@@ -15,7 +15,7 @@
 from collections.abc import Iterable, Sequence
 from functools import cache, partial
 from types import ModuleType
-from typing import Any
+from typing import Any, overload
 
 import numpy as np
 import torch
@@ -51,20 +51,35 @@
 ]
 
 
-def ignore_background(y_pred: NdarrayTensor, y: NdarrayTensor) -> tuple[NdarrayTensor, NdarrayTensor]:
+@overload
+def ignore_background(y_pred: NdarrayTensor, y: NdarrayTensor) -> tuple[NdarrayTensor, NdarrayTensor]: ...
+
+
+@overload
+def ignore_background(y_pred: NdarrayTensor, y: None = ...) -> tuple[NdarrayTensor, None]: ...
+
+
+def ignore_background(
+    y_pred: NdarrayTensor, y: NdarrayTensor | None = None
+) -> tuple[NdarrayTensor, NdarrayTensor | None]:
     """
     This function is used to remove background (the first channel) for `y_pred` and `y`.
 
     Args:
         y_pred: predictions. As for classification tasks,
             `y_pred` should has the shape [BN] where N is larger than 1. As for segmentation tasks,
             the shape should be [BNHW] or [BNHWD].
-        y: ground truth, the first dim is batch.
+        y: ground truth, the first dim is batch. (Optional)
+
+    Returns:
+        Tuple of background-removed predictions and ground truth. `y` is None if not provided.
 
     """
 
-    y = y[:, 1:] if y.shape[1] > 1 else y  # type: ignore[assignment]
+    if y is not None:
+        y = y[:, 1:] if y.shape[1] > 1 else y  # type: ignore[assignment]
     y_pred = y_pred[:, 1:] if y_pred.shape[1] > 1 else y_pred  # type: ignore[assignment]
+
     return y_pred, y