parallelization ready

NiChart_DLMUSE/__main__.py

@@ -7,8 +7,11 @@
 
 import argparse
 import os
+import multiprocessing
+import threading
 
 from .dlmuse_pipeline import run_pipeline
+from .utils import split_data, remove_subfolders
 
 # VERSION = pkg_resources.require("NiChart_DLMUSE")[0].version
 VERSION = 1.0
@@ -110,8 +113,22 @@ def main() -> None:
         os.system("DLMUSE --clear_cache")
 
     # Run pipeline
-    run_pipeline(in_data, out_dir, device)
+    no_threads = 4 # 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()
+
+    remove_subfolders(in_data)
+
+    # run_pipeline(in_data, out_dir, device)
 
 if __name__ == "__main__":
     main()

NiChart_DLMUSE/utils.py

@@ -126,3 +126,46 @@ 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_*")

output/split_0/temp_working_dir/s2_dlicv/dataset.json

@@ -0,0 +1,11 @@
+{
+    "channel_names": {
+        "0": "MRI"
+    },
+    "file_ending": ".nii.gz",
+    "labels": {
+        "background": 0,
+        "class1": 1
+    },
+    "numTraining": 1806
+}