Fixed usage bug of multitarget normalizer

Author: andreped

torchstain/base/normalizers/__init__.py

@@ -1,3 +1,4 @@
 from .he_normalizer import HENormalizer
 from .macenko import MacenkoNormalizer
+from .multitarget import MultiMacenkoNormalizer
 from .reinhard import ReinhardNormalizer

torchstain/torch/normalizers/__init__.py

@@ -1,3 +1,3 @@
 from torchstain.torch.normalizers.macenko import TorchMacenkoNormalizer
-from torchstain.torch.normalizers.multitarget import MultiMacenkoNormalizer
+from torchstain.torch.normalizers.multitarget import TorchMultiMacenkoNormalizer
 from torchstain.torch.normalizers.reinhard import TorchReinhardNormalizer

torchstain/torch/normalizers/multitarget.py

@@ -1,16 +1,17 @@
 import torch
 from torchstain.torch.utils import cov, percentile
+
 """
 Implementation of the multi-target normalizer from the paper: https://arxiv.org/pdf/2406.02077
 """
-class MultiMacenkoNormalizer:
-    def __init__(self, norm_mode='avg-post'):
+class TorchMultiMacenkoNormalizer:
+    def __init__(self, norm_mode="avg-post"):
         self.norm_mode = norm_mode
         self.HERef = torch.tensor([[0.5626, 0.2159],
                                    [0.7201, 0.8012],
                                    [0.4062, 0.5581]])
         self.maxCRef = torch.tensor([1.9705, 1.0308])
-        self.updated_lstsq = hasattr(torch.linalg, 'lstsq')
+        self.updated_lstsq = hasattr(torch.linalg, "lstsq")
 
     def __convert_rgb2od(self, I, Io, beta):
         I = I.permute(1, 2, 0)
@@ -59,15 +60,15 @@ def __compute_matrices_single(self, I, Io, alpha, beta):
         return HE, C, maxC
 
     def fit(self, Is, Io=240, alpha=1, beta=0.15):
-        if self.norm_mode == 'avg-post':
+        if self.norm_mode == "avg-post":
             HEs, _, maxCs = zip(*(
                 self.__compute_matrices_single(I, Io, alpha, beta)
                 for I in Is
             ))
 
             self.HERef = torch.stack(HEs).mean(dim=0)
             self.maxCRef = torch.stack(maxCs).mean(dim=0)
-        elif self.norm_mode == 'concat':
+        elif self.norm_mode == "concat":
             ODs, ODhats = zip(*(
                 self.__convert_rgb2od(I, Io, beta)
                 for I in Is
@@ -83,7 +84,7 @@ def fit(self, Is, Io=240, alpha=1, beta=0.15):
             maxCs = torch.stack([percentile(C[0, :], 99), percentile(C[1, :], 99)])
             self.HERef = HE
             self.maxCRef = maxCs
-        elif self.norm_mode == 'avg-pre':
+        elif self.norm_mode == "avg-pre":
             ODs, ODhats = zip(*(
                 self.__convert_rgb2od(I, Io, beta)
                 for I in Is
@@ -100,7 +101,7 @@ def fit(self, Is, Io=240, alpha=1, beta=0.15):
             maxCs = torch.stack([percentile(C[0, :], 99), percentile(C[1, :], 99)])
             self.HERef = HE
             self.maxCRef = maxCs
-        elif self.norm_mode == 'fixed-single' or self.norm_mode == 'stochastic-single':
+        elif self.norm_mode == "fixed-single" or self.norm_mode == "stochastic-single":
             # single img
             self.HERef, _, self.maxCRef = self.__compute_matrices_single(Is[0], Io, alpha, beta)
         else:
@@ -127,4 +128,4 @@ def normalize(self, I, Io=240, alpha=1, beta=0.15, stains=True):
             E[E > 255] = 255
             E = E.T.reshape(h, w, c).int()
 
-        return Inorm, H, E
+        return Inorm, H, E