[Parallelization] Performed batch splitting and parallelization of the pipeline

NiChart_DLMUSE/__main__.py

@@ -7,8 +7,10 @@
 
 import argparse
 import os
+import threading
 
 from .dlmuse_pipeline import run_pipeline
+from .utils import merge_output_data, remove_subfolders, split_data
 
 # VERSION = pkg_resources.require("NiChart_DLMUSE")[0].version
 VERSION = 1.0
@@ -73,6 +75,15 @@ def main() -> None:
         required=True,
     )
 
+    parser.add_argument(
+        "-c",
+        "--cores",
+        type=str,
+        help="Number of cores",
+        default=4,
+        required=False,
+    )
+
     # VERSION argument
     help = "Show the version and exit"
     parser.add_argument(
@@ -110,7 +121,23 @@ def main() -> None:
         os.system("DLMUSE --clear_cache")
 
     # Run pipeline
-    run_pipeline(in_data, out_dir, device)
+    no_threads = args.cores  # for now
+    subfolders = split_data(in_data, no_threads)
+
+    threads = []
+    for i in range(len(subfolders)):
+        curr_out_dir = out_dir + f"/split_{i}"
+        curr_thread = threading.Thread(
+            target=run_pipeline, args=(subfolders[i], curr_out_dir, device)
+        )
+        curr_thread.start()
+        threads.append(curr_thread)
+
+    for t in threads:
+        t.join()
+
+    merge_output_data(out_dir)
+    remove_subfolders(in_data)
 
 
 if __name__ == "__main__":

NiChart_DLMUSE/utils.py

@@ -126,3 +126,83 @@ def make_img_list(in_data: str) -> pd.DataFrame:
 
     # Return out dataframe
     return df_out
+
+
+def dir_size(in_dir: str) -> int:
+    """
+    Returns the number of images the user passed
+    """
+    size = 0
+    for path in os.listdir(in_dir):
+        if os.path.isfile(os.path.join(in_dir, path)):
+            size += 1
+
+    return size
+
+
+def split_data(in_dir: str, N: int) -> list:
+    """
+    Splits the input data directory into subfolders of size N
+    """
+    assert N > 0
+    data_size = dir_size(in_dir)
+    no_files_in_folders = data_size / N if (data_size % N == 0) else (data_size / N) + 1
+    assert no_files_in_folders > 0
+    subfolders = []
+
+    current_folder = 1
+    current_file = 0
+    os.system(f"mkdir {in_dir}/split_{current_folder}")
+    for img in os.listdir(in_dir):
+        if current_file >= no_files_in_folders:
+            subfolders.append(f"{in_dir}/split_{current_folder}")
+            current_folder += 1
+            os.system(f"mkdir {in_dir}/split_{current_folder}")
+            current_file = 0
+
+        file = os.path.join(in_dir, img)
+        if os.path.isfile(file):
+            os.system(f"cp {file} {in_dir}/split_{current_folder}")
+            current_file += 1
+
+    return subfolders
+
+
+def remove_subfolders(in_dir: str) -> None:
+    os.system(f"rm -r {in_dir}/split_*")
+
+
+def merge_output_data(in_dir: str) -> None:
+    os.system(f"mkdir {in_dir}/results")
+    os.system(f"mkdir {in_dir}/results/s1_reorient_lps")
+    os.system(f"mkdir {in_dir}/results/s2_dlicv")
+    os.system(f"mkdir {in_dir}/results/s3_masked")
+    os.system(f"mkdir {in_dir}/results/s4_dlmuse")
+    os.system(f"mkdir {in_dir}/results/s5_relabeled")
+    os.system(f"mkdir {in_dir}/results/s6_combined")
+
+    for dir in os.listdir(in_dir):
+        if dir == "results":
+            continue
+
+        os.system(
+            f"mv {in_dir}/{dir}/temp_working_dir/s1_reorient_lps/* {in_dir}/results/s1_reorient_lps/"
+        )
+        os.system(
+            f"mv {in_dir}/{dir}/temp_working_dir/s2_dlicv/* {in_dir}/results/s2_dlicv/"
+        )
+        os.system(
+            f"mv {in_dir}/{dir}/temp_working_dir/s3_masked/* {in_dir}/results/s3_masked/"
+        )
+        os.system(
+            f"mv {in_dir}/{dir}/temp_working_dir/s4_dlmuse/* {in_dir}/results/s4_dlmuse/"
+        )
+        os.system(
+            f"mv {in_dir}/{dir}/temp_working_dir/s5_relabeled/* {in_dir}/results/s5_relabeled/"
+        )
+        os.system(
+            f"mv {in_dir}/{dir}/temp_working_dir/s6_combined/* {in_dir}/results/s6_combined/"
+        )
+        os.system(f"mv {in_dir}/{dir}/*.nii.gz {in_dir}/results/")
+
+    os.system(f"rm -r {in_dir}/split_*")