fix: raise clear ValueError when one-hot tensors passed in multiclass mode

Author: lapertor

segmentation_models_pytorch/metrics/functional.py

@@ -140,26 +140,27 @@ def get_stats(
 
     if output.shape != target.shape:
         # Check if user accidentally passed a one-hot / logits tensor in multiclass mode
-        if mode == "multiclass" and output.ndim == target.ndim + 1:
-            raise ValueError(
-                f"In 'multiclass' mode, ``output`` should contain class indices of shape "
-                f"(N, ...), but got shape {output.shape}. "
-                f"It looks like you passed a one-hot or logits tensor. "
-                f"Please convert it first with ``output.argmax(dim=1)``."
-            )
-        if mode == "multiclass" and target.ndim == output.ndim + 1:
-            raise ValueError(
-                f"In 'multiclass' mode, ``target`` should contain class indices of shape "
-                f"(N, ...), but got shape {target.shape}. "
-                f"It looks like you passed a one-hot tensor. "
-                f"Please convert it first with ``target.argmax(dim=1)``."
-            )
+        if mode == "multiclass":
+            if output.ndim >= 3 and output.shape[1] == num_classes:
+                raise ValueError(
+                    f"In 'multiclass' mode, ``output`` should contain class indices of shape "
+                    f"(N, H, W) or (N,), but got shape {tuple(output.shape)}. "
+                    f"It looks like you passed a one-hot or logits tensor of shape (N, C, ...). "
+                    f"For that use case, please use mode='multilabel' instead, "
+                    f"or convert your tensor with ``output.argmax(dim=1)`` first."
+                )
+            if target.ndim >= 3 and target.shape[1] == num_classes:
+                raise ValueError(
+                    f"In 'multiclass' mode, ``target`` should contain class indices of shape "
+                    f"(N, H, W) or (N,), but got shape {tuple(target.shape)}. "
+                    f"It looks like you passed a one-hot encoded tensor of shape (N, C, ...). "
+                    f"Convert it with ``target.argmax(dim=1)`` first."
+                )
         raise ValueError(
             "Dimensions should match, but ``output`` shape is not equal to ``target`` "
             + f"shape, {output.shape} != {target.shape}"
         )
 
-
     if mode != "multiclass" and ignore_index is not None:
         raise ValueError(
             f"``ignore_index`` parameter is not supported for '{mode}' mode"
@@ -179,21 +180,37 @@ def get_stats(
         )
 
     if mode == "multiclass":
-        if output.ndim > 1 and output.shape[1] == num_classes and output.ndim >= 3:
+        if output.ndim >= 3 and output.shape[1] == num_classes:
             raise ValueError(
                 f"In 'multiclass' mode, ``output`` should contain class indices of shape "
                 f"(N, H, W) or (N,), but got shape {tuple(output.shape)}. "
                 f"It looks like you passed a one-hot or logits tensor of shape (N, C, ...). "
                 f"For that use case, please use mode='multilabel' instead, "
                 f"or convert your tensor with ``output.argmax(dim=1)`` first."
             )
-        if target.ndim > 1 and target.shape[1] == num_classes and target.ndim >= 3:
+        if target.ndim >= 3 and target.shape[1] == num_classes:
             raise ValueError(
                 f"In 'multiclass' mode, ``target`` should contain class indices of shape "
                 f"(N, H, W) or (N,), but got shape {tuple(target.shape)}. "
                 f"It looks like you passed a one-hot encoded tensor of shape (N, C, ...). "
                 f"Convert it with ``target.argmax(dim=1)`` first."
             )
+        # Additional robustness: a 4D tensor (N, C, H, W) is always incorrect in
+        # 'multiclass' mode, regardless of whether shape[1] matches num_classes
+        if output.ndim == 4:
+            raise ValueError(
+                f"In 'multiclass' mode, ``output`` should contain class indices of shape "
+                f"(N, H, W), but got shape {tuple(output.shape)}. "
+                f"A 4D tensor is always incorrect in 'multiclass' mode. "
+                f"Convert it with ``output.argmax(dim=1)`` first."
+            )
+        if target.ndim == 4:
+            raise ValueError(
+                f"In 'multiclass' mode, ``target`` should contain class indices of shape "
+                f"(N, H, W), but got shape {tuple(target.shape)}. "
+                f"A 4D tensor is always incorrect in 'multiclass' mode. "
+                f"Convert it with ``target.argmax(dim=1)`` first."
+            )
         tp, fp, fn, tn = _get_stats_multiclass(
             output, target, num_classes, ignore_index
         )