Small changes and documentation

Author: schilling40

flamingo_tools/segmentation/__init__.py

@@ -1,2 +1,2 @@
-from .unet_prediction import run_unet_prediction, run_unet_prediction_slurm
+from .unet_prediction import run_unet_prediction
 from .postprocessing import filter_isolated_objects

flamingo_tools/segmentation/unet_prediction.py

@@ -19,6 +19,11 @@
 from torch_em.util.prediction import predict_with_halo
 from tqdm import tqdm
 
+"""
+Prediction using distance U-Net.
+Parallelization using multiple GPUs is currently only possible by calling functions directly.
+Functions for the parallelization end with '_slurm' and divide the process into preprocessing, prediction, and segmentation.
+"""
 
 class SelectChannel(SimpleTransformationWrapper):
     def __init__(self, volume, channel):
@@ -281,7 +286,9 @@ def run_unet_prediction(
     pmap_out = os.path.join(output_folder, "predictions.zarr")
     segmentation_impl(pmap_out, output_folder, min_size=min_size, original_shape=original_shape)
 
-def run_unet_prediction_slurm_preprocess(
+#---Workflow for parallel prediction using slurm---
+
+def run_unet_prediction_preprocess_slurm(
         input_path, input_key, output_folder,
 ):
     """

scripts/convert_tif_to_n5.py

@@ -5,29 +5,36 @@
 
 
 def main(input_path, output_path):
-	if not os.path.isfile(input_path):
-		sys.exit("Input file does not exist.")
+    """
+    Convert tif file to n5 format.
+    If no output_path is supplied, the output file is created in the same directory as the input.
 
-	if input_path.split(".")[-1] not in ["TIFF", "TIF", "tiff", "tif"]:
-		sys.exit("Input file must be in tif format.")
+    :param str input_path: Input tif
+    :param str output_path: Output path for n5 format
+    """
+    if not os.path.isfile(input_path):
+        sys.exit("Input file does not exist.")
 
-	basename = "".join(input_path.split("/")[-1].split(".")[:-1])
-	input_dir = input_path.split(basename)[0]
-	input_dir = os.path.abspath(input_dir)
+    if input_path.split(".")[-1] not in ["TIFF", "TIF", "tiff", "tif"]:
+        sys.exit("Input file must be in tif format.")
 
-	if "" == output_path:
-		output_path = os.path.join(input_dir, basename + ".n5")
-	img = imageio.imread(input_path)
-	pybdv.make_bdv(img, output_path)
+    basename = "".join(input_path.split("/")[-1].split(".")[:-1])
+    input_dir = input_path.split(basename)[0]
+    input_dir = os.path.abspath(input_dir)
+
+    if "" == output_path:
+        output_path = os.path.join(input_dir, basename + ".n5")
+    img = imageio.imread(input_path)
+    pybdv.make_bdv(img, output_path)
 
 if __name__ == "__main__":
 
-	parser = argparse.ArgumentParser(
-		description="Script to transform file from tif into n5 format.")
+    parser = argparse.ArgumentParser(
+        description="Script to transform file from tif into n5 format.")
 
-	parser.add_argument('input', type=str, help="Input file")
-	parser.add_argument('-o', "--output", type=str, default="", help="Output file")
+    parser.add_argument('input', type=str, help="Input file")
+    parser.add_argument('-o', "--output", type=str, default="", help="Output file. Default: <basename>.n5")
 
-	args = parser.parse_args()
+    args = parser.parse_args()
 
-	main(args.input, args.output)
+    main(args.input, args.output)

scripts/extract_block.py

@@ -1,7 +1,6 @@
 import os
 import argparse
 import numpy as np
-import h5py
 import z5py
 
 """
@@ -31,8 +30,6 @@ def main(input_file, output_dir, input_key, resolution, coords, roi_halo):
     if "" == output_dir:
         output_dir = input_dir
 
-    input_key = "setup0/timepoint0/s0"
-
     output_file = os.path.join(output_dir, basename + "_crop" + ".n5")
 
     #M_LR_000167_R, coords = '806,1042,1334', coords = (z, y, x) compared to MoBIE view

scripts/prediction/run_prediction_distance_unet.py

@@ -6,17 +6,21 @@
 
 sys.path.append("../..")
 
+"""
+Prediction using distance U-Net.
+Parallelization using multiple GPUs is currently only possible by calling functions located in segmentation/unet_prediction.py directly.
+Functions for the parallelization end with '_slurm' and divide the process into preprocessing, prediction, and segmentation.
+"""
 
 def main():
-    from flamingo_tools.segmentation import run_unet_prediction, run_unet_prediction_slurm
+    from flamingo_tools.segmentation import run_unet_prediction
 
     parser = argparse.ArgumentParser()
     parser.add_argument("-i", "--input", required=True)
     parser.add_argument("-o", "--output_folder", required=True)
     parser.add_argument("-m", "--model", required=True)
     parser.add_argument("-k", "--input_key", default=None)
     parser.add_argument("-s", "--scale", default=None, type=float, help="Downscale the image by the given factor.")
-    parser.add_argument("-n", "--number_gpu", default=1, type=int, help="Number of GPUs to use in parallel.")
 
     args = parser.parse_args()
 
@@ -37,24 +41,11 @@ def main():
         block_shape = tuple([2 * ch for ch in chunks]) if have_cuda else tuple(chunks)
         halo = (16, 64, 64) if have_cuda else (8, 32, 32)
 
-    prediction_instances = args.number_gpu if have_cuda else 1
-
-    if 1 > prediction_instances:
-        # FIXME: only does prediction part, no segmentation yet
-        # FIXME: implement array job
-        run_unet_prediction_slurm(
-            args.input, args.input_key, args.output_folder, args.model,
-            scale=scale,
-            block_shape=block_shape, halo=halo,
-            prediction_instances=prediction_instances,
-        )
-    else:
-
-        run_unet_prediction(
-            args.input, args.input_key, args.output_folder, args.model,
-            scale=scale, min_size=min_size,
-            block_shape=block_shape, halo=halo,
-        )
+    run_unet_prediction(
+        args.input, args.input_key, args.output_folder, args.model,
+        scale=scale, min_size=min_size,
+        block_shape=block_shape, halo=halo,
+    )
 
 
 if __name__ == "__main__":