Fixed requirement of SLURM_ARRAY_TASK_ID

Author: schilling40

flamingo_tools/segmentation/unet_prediction.py

@@ -80,7 +80,7 @@ def prediction_impl(input_path, input_key, output_folder, model_path, scale, blo
         halo = (16, 32, 32)
 
     # Compute the global mean and standard deviation.
-    n_threads = min(2, mp.cpu_count())
+    n_threads = min(16, mp.cpu_count())
     mean, std = parallel.mean_and_std(
         input_, block_shape=tuple([2* i for i in input_.chunks]), n_threads=n_threads, verbose=True,
         mask=image_mask
@@ -243,25 +243,17 @@ def run_unet_prediction(
     output_folder, model_path,
     min_size, scale=None,
     block_shape=None, halo=None,
-    prediction_instances=1,
 ):
-    if prediction_instances > 1:
-        run_unet_prediction_slurm(
-            input_path, input_key, output_folder, model_path,
-            scale=scale, block_shape=block_shape, halo=halo,
-            prediction_instances=prediction_instances,
-        )
-    else:
-        os.makedirs(output_folder, exist_ok=True)
+    os.makedirs(output_folder, exist_ok=True)
 
-        find_mask(input_path, input_key, output_folder)
+    find_mask(input_path, input_key, output_folder)
 
-        original_shape = prediction_impl(
-            input_path, input_key, output_folder, model_path, scale, block_shape, halo
-        )
+    original_shape = prediction_impl(
+        input_path, input_key, output_folder, model_path, scale, block_shape, halo
+    )
 
-        pmap_out = os.path.join(output_folder, "predictions.zarr")
-        segmentation_impl(pmap_out, output_folder, min_size=min_size, original_shape=original_shape)
+    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(
     input_path, input_key, output_folder, model_path,
@@ -271,8 +263,11 @@ def run_unet_prediction_slurm(
     os.makedirs(output_folder, exist_ok=True)
     prediction_instances = int(prediction_instances)
     slurm_task_id = os.environ.get("SLURM_ARRAY_TASK_ID")
+
     if slurm_task_id is not None:
         slurm_task_id = int(slurm_task_id)
+    else:
+        raise ValueError("The SLURM_ARRAY_TASK_ID is not set. Ensure that you are using the '-a' option with SBATCH.")
 
     find_mask(input_path, input_key, output_folder)
 

scripts/prediction/run_prediction_distance_unet.py

@@ -8,7 +8,7 @@
 
 
 def main():
-    from flamingo_tools.segmentation import run_unet_prediction
+    from flamingo_tools.segmentation import run_unet_prediction, run_unet_prediction_slurm
 
     parser = argparse.ArgumentParser()
     parser.add_argument("-i", "--input", required=True)
@@ -39,12 +39,22 @@ def main():
 
     prediction_instances = args.number_gpu if have_cuda else 1
 
-    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,
-        prediction_instances=prediction_instances,
-    )
+    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,
+        )
 
 
 if __name__ == "__main__":