fix preprocessing

.gitignore

@@ -172,3 +172,6 @@ cython_debug/
 
 # PyPI configuration file
 .pypirc
+
+# Project-specific files
+test.py

exaonepath.py

@@ -1,5 +1,6 @@
 import math
 import typing as t
+from functools import partial
 
 import torch
 import torch.nn as nn
@@ -22,7 +23,13 @@ if t.TYPE_CHECKING:
     from _typeshed import StrPath
 
 
-class PadToDivisible(T.Transform):
+class Transform(T.Transform):
+    # For compatibility with torchvision <= 0.20
+    def _transform(self, inpt, params):
+        return self.transform(inpt, params)
+
+
+class PadToDivisible(Transform):
     def __init__(self, size: int, pad_value: float | None = None):
         super().__init__()
         self.size = size
@@ -44,32 +51,6 @@ class PadToDivisible(T.Transform):
         return inpt
 
 
-class Preprocessing(T.Transform):
-    def __init__(
-        self, small_tile_size_with_this_mpp: int, small_tile_size_with_target_mpp: int
-    ):
-        self.small_tile_size_with_this_mpp = small_tile_size_with_this_mpp
-        self.small_tile_size_with_target_mpp = small_tile_size_with_target_mpp
-
-    def transform(self, inpt, params):
-        assert isinstance(inpt, torch.Tensor) and inpt.ndim >= 3
-
-        # Scale the input tensor to the target MPP
-        if self.small_tile_size_with_this_mpp != self.small_tile_size_with_target_mpp:
-            inpt = TF.resize(
-                inpt,
-                [
-                    self.small_tile_size_with_target_mpp,
-                    self.small_tile_size_with_target_mpp,
-                ],
-            )
-
-        # Normalize the input tensor
-        inpt = scale_and_normalize(inpt)
-
-        return inpt
-
-
 class EXAONEPathV20(nn.Module, PyTorchModelHubMixin):
     def __init__(
         self,
@@ -103,7 +84,8 @@ class EXAONEPathV20(nn.Module, PyTorchModelHubMixin):
                     self.model_first_stg,
                     small_tiles,
                     batch_size_on_gpu=first_stg_batch_size,
-                    preproc_fn=Preprocessing(
+                    preproc_fn=partial(
+                        _preproc,
                         small_tile_size_with_this_mpp=small_tile_size,
                         small_tile_size_with_target_mpp=self.small_tile_size,
                     ),
@@ -111,14 +93,14 @@ class EXAONEPathV20(nn.Module, PyTorchModelHubMixin):
                     out_device=self.device,
                     dtype=torch.bfloat16,
                 )
-                act1 = format_first_stg_act_as_second_stg_inp(
+                act1_formatted = format_first_stg_act_as_second_stg_inp(
                     act1,
                     height=height,
                     width=width,
                     small_tile_size=small_tile_size,
                     large_tile_size=large_tile_size,
                 )
-                act2: torch.Tensor = self.model_second_stg(act1)
+                act2: torch.Tensor = self.model_second_stg(act1_formatted)
                 act2_formatted = format_second_stg_act_as_third_stg_inp(
                     act2,
                     height=height,
@@ -126,6 +108,7 @@ class EXAONEPathV20(nn.Module, PyTorchModelHubMixin):
                     large_tile_size=large_tile_size,
                 )
                 act3: torch.Tensor = self.model_third_stg(act2_formatted)
+
         return act1[is_tile_valid], act2, act3
 
     def _load_wsi(self, svs_path: "StrPath", target_mpp: float):
@@ -163,3 +146,21 @@ class EXAONEPathV20(nn.Module, PyTorchModelHubMixin):
         is_tile_valid = mask_tile.sum(dim=(1, 2)) > 0
 
         return x, is_tile_valid, padded_size, small_tile_size, large_tile_size
+
+
+def _preproc(
+    x: torch.Tensor,
+    small_tile_size_with_this_mpp: int,
+    small_tile_size_with_target_mpp: int,
+):
+    # Scale the input tensor to the target MPP
+    if small_tile_size_with_this_mpp != small_tile_size_with_target_mpp:
+        x = TF.resize(
+            x,
+            [small_tile_size_with_target_mpp, small_tile_size_with_target_mpp],
+        )
+
+    # Normalize the input tensor
+    x = scale_and_normalize(x)
+
+    return x