Enforce tensor range [0,1]

Author: C-Achard

napari_cellseg3d/code_models/worker_inference.py

@@ -34,6 +34,7 @@
     ONNXModelWrapper,
     QuantileNormalization,
     QuantileNormalizationd,
+    RemapTensor,
     WeightsDownloader,
 )
 
@@ -715,11 +716,20 @@ def inference(self):
             # )
 
             if not post_process_config.thresholding.enabled:
-                post_process_transforms = EnsureType()
+                post_process_transforms = Compose(
+                    [
+                        RemapTensor(new_max=1.0, new_min=0.0),
+                        EnsureType(),
+                    ]
+                )
             else:
                 t = post_process_config.thresholding.threshold_value
                 post_process_transforms = Compose(
-                    AsDiscrete(threshold=t), EnsureType()
+                    [
+                        RemapTensor(new_max=1.0, new_min=0.0),
+                        AsDiscrete(threshold=t),
+                        EnsureType(),
+                    ]
                 )
 
             is_folder = self.config.images_filepaths is not None

napari_cellseg3d/code_models/worker_training.py

@@ -44,6 +44,7 @@
     PRETRAINED_WEIGHTS_DIR,
     LogSignal,
     QuantileNormalizationd,
+    RemapTensor,
     TrainingReport,
     WeightsDownloader,
 )
@@ -638,12 +639,15 @@ def get_loader_func(num_samples):
 
                             # TODO : more parameters/flexibility
                             post_pred = Compose(
-                                AsDiscrete(threshold=0.5),  # needed ?
-                                EnsureType(),
+                                [
+                                    RemapTensor(new_max=1, new_min=0),
+                                    AsDiscrete(threshold=0.5),  # needed ?
+                                    EnsureType(),
+                                ]
                             )  #
                             post_label = EnsureType()
 
-                            output_raw = [t for t in pred]
+                            output_raw = [RemapTensor(0, 1)(t) for t in pred]
                             val_outputs = [
                                 post_pred(res_tensor) for res_tensor in pred
                             ]

napari_cellseg3d/code_models/workers_utils.py

@@ -203,6 +203,16 @@ def __call__(self, img):
         return utils.quantile_normalization(img)
 
 
+class RemapTensor(Transform):
+    def __init__(self, new_max, new_min):
+        super().__init__()
+        self.max = new_max
+        self.min = new_min
+
+    def __call__(self, img):
+        return utils.remap_image(img, new_max=self.max, new_min=self.min)
+
+
 @dataclass
 class InferenceResult:
     """Class to record results of a segmentation job"""