Update compartment segmentation code

Author: constantinpape

scripts/cooper/ground_truth/compartments/run_prediction_04.py

@@ -4,11 +4,11 @@
 import h5py
 from tqdm import tqdm
 
-from synaptic_reconstruction.inference.utils import _Scaler
+from synaptic_reconstruction.inference.util import _Scaler
 from synaptic_reconstruction.inference.compartments import segment_compartments
 
-INPUT_ROOT = ""
-MODEL_PATH = ""
+INPUT_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/ground_truth/04Dataset_for_vesicle_eval"  # noqa
+MODEL_PATH = "/mnt/lustre-emmy-hdd/projects/nim00007/compartment_models/compartment_model_3d.pt"
 OUTPUT = "./predictions"
 
 
@@ -17,10 +17,12 @@ def segment_volume(input_path, model_path):
         raw = f["raw"][:]
 
     scale = (0.25, 0.25, 0.25)
-    scaler = _Scaler(scale)
+    scaler = _Scaler(scale, verbose=False)
     raw = scaler.scale_input(raw)
 
-    seg = segment_compartments(raw, model_path, verbose=False)
+    n_slices_exclude = 4
+    seg = segment_compartments(raw, model_path, verbose=False, n_slices_exclude=n_slices_exclude)
+    raw, seg = raw[n_slices_exclude:-n_slices_exclude], seg[n_slices_exclude:-n_slices_exclude]
 
     return raw, seg
 

synaptic_reconstruction/inference/compartments.py

@@ -117,6 +117,7 @@ def segment_compartments(
     return_predictions: bool = False,
     scale: Optional[List[float]] = None,
     mask: Optional[np.ndarray] = None,
+    n_slices_exclude: int = 5,
 ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
     """
     Segment synaptic compartments in an input volume.
@@ -129,6 +130,7 @@ def segment_compartments(
         verbose: Whether to print timing information.
         return_predictions: Whether to return the predictions (foreground, boundaries) alongside the segmentation.
         scale: The scale factor to use for rescaling the input volume before prediction.
+        n_slices_exclude:
 
     Returns:
         The segmentation mask as a numpy array, or a tuple containing the segmentation mask
@@ -156,7 +158,7 @@ def segment_compartments(
     if input_volume.ndim == 2:
         seg = _segment_compartments_2d(pred)
     else:
-        seg = _segment_compartments_3d(pred)
+        seg = _segment_compartments_3d(pred, n_slices_exclude=n_slices_exclude)
     if verbose:
         print("Run segmentation in", time.time() - t0, "s")
     seg = scaler.rescale_output(seg, is_segmentation=True)

synaptic_reconstruction/inference/util.py

@@ -388,7 +388,7 @@ def get_default_tiling():
             tile = {"x": 352, "y": 352, "z": 48}
         else:
             # TODO determine tilings for smaller VRAM
-            raise NotImplementedError
+            raise NotImplementedError(f"Estimating the tile size for a GPU with {vram} GB is not yet supported.")
 
         print(f"Determined tile size: {tile}")
         tiling = {"tile": tile, "halo": halo}