Skip to content

Commit

Permalink
parallelization ready
Browse files Browse the repository at this point in the history
  • Loading branch information
spirosmaggioros committed Oct 9, 2024
1 parent 6fa4c8d commit 6e475da
Show file tree
Hide file tree
Showing 6 changed files with 527 additions and 1 deletion.
19 changes: 18 additions & 1 deletion NiChart_DLMUSE/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
43 changes: 43 additions & 0 deletions NiChart_DLMUSE/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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_*")
11 changes: 11 additions & 0 deletions output/split_0/temp_working_dir/s2_dlicv/dataset.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
{
"channel_names": {
"0": "MRI"
},
"file_ending": ".nii.gz",
"labels": {
"background": 0,
"class1": 1
},
"numTraining": 1806
}
Loading

0 comments on commit 6e475da

Please sign in to comment.