update

Author: tfunck

brainbuilder/align/align_2d.py

@@ -20,16 +20,17 @@
     shell,
     threshold,
 )
-from joblib import Parallel, cpu_count, delayed
+from joblib import Parallel, cpu_count, delayed, parallel_backend
 from skimage.transform import resize
 
 logger = utils.get_logger(__name__)
 
+
 def get_base_from_raw(x):
     """Get base from raw filename."""
-    
     return os.path.basename(x).split(".")[0]
-    
+
+
 def resample_reference_to_sections(
     resolution: float,
     input_fn: str,
@@ -239,7 +240,8 @@ def ants_registration_2d_section(
             else:
                 init_str = f"--initial-moving-transform [{fx_fn},{mv_fn},1]"
 
-            command_str = f"antsRegistration -v 1 -d 2 --write-composite-transform {write_composite_transform} {init_str} -o [{prefix}_{transform}_{metric}_,{mv_rsl_fn},/tmp/out_inv.nii.gz] -t {transform}[{step}]  -m {metric}[{fx_fn},{mv_fn},1,{bins},Random,{sampling}] -s {s_str} -f {f_str} -c {itr_str} "
+            inv_out = f"{prefix}_{transform}_{metric}_inv.nii.gz"
+            command_str = f"antsRegistration -v 1 -d 2 --write-composite-transform {write_composite_transform} {init_str} -o [{prefix}_{transform}_{metric}_,{mv_rsl_fn},{inv_out}] -t {transform}[{step}]  -m {metric}[{fx_fn},{mv_fn},1,{bins},Random,{sampling}] -s {s_str} -f {f_str} -c {itr_str} "
 
             if int(verbose) <= 0:
                 command_str += f" &> {output_log_fname}"
@@ -379,7 +381,6 @@ def align_2d_parallel(
     y = int(row["sample"])
     base = get_base_from_raw(row["raw"])
 
-
     prefix = f"{tfm_dir}/{base}_y-{y}"
 
     fx_fn = row["fx"]
@@ -490,12 +491,14 @@ def apply_transforms_parallel(
 
     tfm_fn = row["2d_tfm"]
 
-    if not isinstance(tfm_fn, str) : # assume identity transform
-        print(f"Warning: transform file does not exist, assuming identity transform in:\n{base}")
-        # copy img_rsl_fn to out_fn 
+    if not isinstance(tfm_fn, str):  # assume identity transform
+        print(
+            f"Warning: transform file does not exist, assuming identity transform in:\n{base}"
+        )
+        # copy img_rsl_fn to out_fn
         shutil.copy(img_rsl_fn, out_fn)
         return out_fn
-    
+
     cmd = f"antsApplyTransforms -v {int(verbose)} -d 2 -n {interpolation} -i {img_rsl_fn} -r {fx_fn} -t {tfm_fn} -o {out_fn} "
 
     shell(cmd, True)
@@ -610,44 +613,43 @@ def align_sections(
 
     tfm_dir = output_dir + os.sep + "tfm"
 
-
-
     sect_info = get_align_filenames(tfm_dir, sect_info)
 
-
     # get lists of files that need to be aligned and resampled
     to_do_sect_info, to_do_resample_sect_info = get_align_2d_to_do(sect_info)
 
     if len(to_do_sect_info) > 0:
-        Parallel(n_jobs=num_cores, backend="multiprocessing")(
-            delayed(align_2d_parallel)(
-                tfm_dir,
-                resolution,
-                resolution_list,
-                row,
-                base_lin_itr=base_lin_itr,
-                base_nl_itr=base_nl_itr,
-                file_to_align=file_to_align,
-                use_syn=use_syn,
-                verbose=verbose,
-            )
-            for row in to_do_sect_info
-        )
-
-    if len(to_do_resample_sect_info) > 0:
-        for tissue_str in ["", "_cls"]:
-            Parallel(n_jobs=num_cores, backend="multiprocessing")(
-                delayed(apply_transforms_parallel)(
+        with parallel_backend("loky", inner_max_num_threads=1):
+            Parallel(n_jobs=num_cores, backend="loky", prefer="processes")(
+                delayed(align_2d_parallel)(
                     tfm_dir,
                     resolution,
+                    resolution_list,
                     row,
-                    tissue_str=tissue_str,
-                    interpolation=interpolation,
+                    base_lin_itr=base_lin_itr,
+                    base_nl_itr=base_nl_itr,
+                    file_to_align=file_to_align,
+                    use_syn=use_syn,
                     verbose=verbose,
                 )
-                for row in to_do_resample_sect_info
+                for row in to_do_sect_info
             )
 
+    if len(to_do_resample_sect_info) > 0:
+        for tissue_str in ["", "_cls"]:
+            with parallel_backend("loky", inner_max_num_threads=1):
+                Parallel(n_jobs=num_cores, backend="loky", prefer="processes")(
+                    delayed(apply_transforms_parallel)(
+                        tfm_dir,
+                        resolution,
+                        row,
+                        tissue_str=tissue_str,
+                        interpolation=interpolation,
+                        verbose=verbose,
+                    )
+                    for row in to_do_resample_sect_info
+                )
+
     sect_info["2d_align"] = sect_info["2d_align_out"]
     sect_info["2d_align_cls"] = sect_info["2d_align_cls_out"]
 
@@ -764,9 +766,7 @@ def align_2d(
     )
 
     logger.info("\t\tStep 4: 2d nl alignment")
-    sect_info["base"] = sect_info["raw"].apply(
-        lambda x: get_base_from_raw(x)
-    )
+    sect_info["base"] = sect_info["raw"].apply(lambda x: get_base_from_raw(x))
 
     # Align 2D sections to sections from reference volume using ANTs
     sect_info = align_sections(