added preprocess to predcition function so that input for predict_with_halo can be prepocessed per tile

Author: lufre1

synapse_net/inference/mitochondria.py

@@ -1,5 +1,5 @@
 import time
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, Optional, Tuple, Union
 
 import elf.parallel as parallel
 import numpy as np
@@ -65,6 +65,7 @@ def segment_mitochondria(
     ws_halo: Tuple[int, ...] = (48, 48, 48),
     boundary_threshold: float = 0.25,
     area_threshold: int = 5000,
+    preprocess: Callable = None,
 ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
     """Segment mitochondria in an input volume.
 
@@ -99,7 +100,10 @@ def segment_mitochondria(
     # Rescale the mask if it was given and run prediction.
     if mask is not None:
         mask = scaler.scale_input(mask, is_segmentation=True)
-    pred = get_prediction(input_volume, model_path=model_path, model=model, tiling=tiling, mask=mask, verbose=verbose)
+    pred = get_prediction(
+        input_volume, model_path=model_path, model=model, tiling=tiling, mask=mask, verbose=verbose,
+        preprocess=preprocess
+    )
 
     # Run segmentation and rescale the result if necessary.
     foreground, boundaries = pred[:2]

synapse_net/inference/util.py

@@ -126,6 +126,7 @@ def get_prediction(
     mask: Optional[ArrayLike] = None,
     prediction: Optional[ArrayLike] = None,
     devices: Optional[List[str]] = None,
+    preprocess: Optional[callable] = None,
 ) -> ArrayLike:
     """Run prediction on a given volume.
 
@@ -192,7 +193,7 @@ def get_prediction(
         # print(f"updated_tiling {updated_tiling}")
         prediction = get_prediction_torch_em(
             input_volume, updated_tiling, model_path, model, verbose, with_channels,
-            mask=mask, prediction=prediction, devices=devices,
+            mask=mask, prediction=prediction, devices=devices, preprocess=preprocess,
         )
 
     return prediction
@@ -208,6 +209,7 @@ def get_prediction_torch_em(
     mask: Optional[ArrayLike] = None,
     prediction: Optional[ArrayLike] = None,
     devices: Optional[List[str]] = None,
+    preprocess: Optional[callable] = None,
 ) -> np.ndarray:
     """Run prediction using torch-em on a given volume.
 
@@ -258,7 +260,10 @@ def get_prediction_torch_em(
                 print("Run prediction with mask.")
             mask = mask.astype("bool")
 
-        preprocess = None if isinstance(input_volume, np.ndarray) else torch_em.transform.raw.standardize
+        if preprocess is None:
+            preprocess = None if isinstance(input_volume, np.ndarray) else torch_em.transform.raw.standardize
+        else:
+            preprocess = preprocess
         prediction = predict_with_halo(
             input_volume, model, gpu_ids=devices,
             block_shape=block_shape, halo=halo,