♻️ Move argmax postprocessing to utils.

Author: shaneahmed

tiatoolbox/models/architecture/unet.py

@@ -2,18 +2,24 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
-import numpy as np
 import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import nn
 from torchvision.models.resnet import Bottleneck as ResNetBottleneck
 from torchvision.models.resnet import ResNet
 
-from tiatoolbox.models.architecture.utils import UpSample2x, centre_crop
+from tiatoolbox.models.architecture.utils import (
+    UpSample2x,
+    argmax_last_axis,
+    centre_crop,
+)
 from tiatoolbox.models.models_abc import ModelABC
 
+if TYPE_CHECKING:  # pragma: no cover
+    import numpy as np
+
 
 class ResNetEncoder(ResNet):
     """A subclass of ResNet defined in torch.
@@ -463,9 +469,9 @@ def infer_batch(
 
     @staticmethod
     def postproc(image: np.ndarray) -> np.ndarray:
-        """Define the post-processing of this class of model.
+        """Define post-processing of this class of model.
 
         This simply applies argmax along last axis of the input.
 
         """
-        return np.argmax(image, axis=-1)
+        return argmax_last_axis(image=image)

tiatoolbox/models/architecture/utils.py

@@ -233,3 +233,20 @@ def forward(self: UpSample2x, x: torch.Tensor) -> torch.Tensor:
         ret = torch.tensordot(x, mat, dims=1)  # bxcxhxwxshxsw
         ret = ret.permute(0, 1, 2, 4, 3, 5)
         return ret.reshape((-1, input_shape[1], input_shape[2] * 2, input_shape[3] * 2))
+
+
+def argmax_last_axis(image: np.ndarray) -> np.ndarray:
+    """Define the post-processing of this class of model.
+
+    This simply applies argmax along last axis of the input.
+
+    Args:
+        image (np.ndarray):
+            The input image array.
+
+    Returns:
+        np.ndarray:
+            The post-processed image array.
+
+    """
+    return np.argmax(image, axis=-1)

tiatoolbox/models/architecture/vanilla.py

@@ -4,16 +4,17 @@
 
 from typing import TYPE_CHECKING
 
-import numpy as np
 import timm
 import torch
 import torchvision.models as torch_models
 from timm.layers import SwiGLUPacked
 from torch import nn
 
+from tiatoolbox.models.architecture.utils import argmax_last_axis
 from tiatoolbox.models.models_abc import ModelABC
 
 if TYPE_CHECKING:  # pragma: no cover
+    import numpy as np
     from torchvision.models import WeightsEnum
 
 
@@ -205,23 +206,6 @@ def _get_timm_architecture(
     raise ValueError(msg)
 
 
-def _postproc(image: np.ndarray) -> np.ndarray:
-    """Define the post-processing of this class of model.
-
-    This simply applies argmax along last axis of the input.
-
-    Args:
-        image (np.ndarray):
-            The input image array.
-
-    Returns:
-        np.ndarray:
-            The post-processed image array.
-
-    """
-    return np.argmax(image, axis=-1)
-
-
 def _infer_batch(
     model: nn.Module,
     batch_data: torch.Tensor,
@@ -339,7 +323,7 @@ def postproc(image: np.ndarray) -> np.ndarray:
                 The post-processed image array.
 
         """
-        return _postproc(image=image)
+        return argmax_last_axis(image=image)
 
     @staticmethod
     def infer_batch(
@@ -463,7 +447,7 @@ def postproc(image: np.ndarray) -> np.ndarray:
                 The post-processed image array.
 
         """
-        return _postproc(image=image)
+        return argmax_last_axis(image=image)
 
     @staticmethod
     def infer_batch(