Implement an inverse for the identity.

ahcore/lit_module.py

@@ -89,7 +89,6 @@ def __init__(
         self._use_test_time_augmentation = True
         self._tta_steps = len(self._tta_augmentations)
 
-
     @property
     def wsi_metrics(self):
         return self._wsi_metrics
@@ -170,34 +169,24 @@ def do_step(self, batch, batch_idx: int, stage: TrainerFn):
     def _get_inference_prediction(self, _input: torch.Tensor) -> dict[str, torch.Tensor]:
         output = {}
 
-        if self._use_test_time_augmentation:
-            _predictions = torch.zeros([self._tta_steps, *_input.size()], device=self.device)
-            _collected_features = None
-
-            with ExtractFeaturesHook(self._model, layer_names=self._attach_feature_layers) as hook:
-                for idx, augmentation in enumerate(self._tta_augmentations):
-                    model_prediction = self._model(augmentation(_input))
-                    _predictions[idx] = augmentation.inverse(model_prediction)
+        _predictions = torch.zeros([self._tta_steps, *_input.size()], device=self.device)
+        _collected_features = None
 
-                    if self._attach_feature_layers:
-                        _features = hook.features
-                        if _collected_features is None:
-                            _collected_features = torch.zeros([self._tta_steps, *_features.size()], device=self.device)
-                        _features[idx] = _collected_features
+        with ExtractFeaturesHook(self._model, layer_names=self._attach_feature_layers) as hook:
+            for idx, augmentation in enumerate(self._tta_augmentations):
+                model_prediction = self._model(augmentation(_input))
+                _predictions[idx] = augmentation.inverse(model_prediction)
 
-                output["prediction"] = _predictions.mean(dim=0)
+                if self._attach_feature_layers:
+                    _features = hook.features
+                    if _collected_features is None:
+                        _collected_features = torch.zeros([self._tta_steps, *_features.size()], device=self.device)
+                    _features[idx] = _collected_features
 
-            if self._attach_feature_layers:
-                if _collected_features is None:
-                    output["features"] = hook.features.unsqueeze(0)
-                else:
-                    output["features"] = _collected_features
+            output["prediction"] = _predictions.mean(dim=0)
 
-        else:
-            with ExtractFeaturesHook(self._model, layer_names=self._attach_feature_layers) as hook:
-                _prediction = self._model(_input)
-                if self._attach_feature_layers:
-                    output["features"] = hook.features.unsqueeze(0)
+        if self._attach_feature_layers:
+            output["features"] = _collected_features
 
         return output
 

ahcore/transforms/augmentations.py

@@ -15,7 +15,7 @@
 
 from ahcore.utils.data import DataDescription
 from ahcore.utils.io import get_logger
-
+import kornia as K
 logger = get_logger(__name__)
 
 
@@ -84,6 +84,7 @@ def forward(self, *args: torch.Tensor, **kwargs):
 
         return output
 
+
 class Identity(nn.Module):
     """
     Identity transform.
@@ -95,6 +96,8 @@ def __init__(self):
     def forward(self, *args: torch.Tensor, **kwargs):
         return args
 
+    def inverse(self, *args: torch.Tensor, **kwargs):
+        return args
 
 class CenterCrop(nn.Module):
     """Perform a center crop of the image and target"""