From 4f2c0f992743ff267153ae22ab4fc376357f0b9d Mon Sep 17 00:00:00 2001 From: Spiros Maggioros Date: Thu, 31 Oct 2024 00:14:29 +0200 Subject: [PATCH] [BIDS] Accept BIDS format input TODO: Generate BIDS output format if the user specify it --- NiChart_DLMUSE/__main__.py | 104 +++++++++++++++++++++++++------------ NiChart_DLMUSE/utils.py | 42 +++++++++++++++ 2 files changed, 114 insertions(+), 32 deletions(-) diff --git a/NiChart_DLMUSE/__main__.py b/NiChart_DLMUSE/__main__.py index b012cc0..2971cf2 100644 --- a/NiChart_DLMUSE/__main__.py +++ b/NiChart_DLMUSE/__main__.py @@ -10,10 +10,10 @@ import threading from .dlmuse_pipeline import run_pipeline -from .utils import merge_output_data, remove_subfolders, split_data +from .utils import merge_output_data, remove_subfolders, split_data, collect_T1 # VERSION = pkg_resources.require("NiChart_DLMUSE")[0].version -VERSION = 1.0 +VERSION = "1.0.5" def main() -> None: @@ -24,7 +24,7 @@ def main() -> None: ICV calculation, brain segmentation, and ROI extraction pipelines for structural MRI data. required arguments: - [-i, --in_data] Input images. The input should be: + [-i, --in_dir] Input images. The input should be: - a single image file (.nii.gz or .nii), or - a directory containing image files, or - a list with the full path for each input image (one in each row) @@ -35,7 +35,7 @@ def main() -> None: [-h, --help] Show this help message and exit. [-V, --version] Show program's version number and exit. EXAMPLE USAGE: - NiChart_DLMUSE --in_data /path/to/input \ + NiChart_DLMUSE --in_dir /path/to/input \ --out_dir /path/to/output \ """.format( VERSION=VERSION @@ -48,7 +48,7 @@ def main() -> None: # INDIR argument parser.add_argument( "-i", - "--in_data", + "--in_dir", type=str, help="Input images.", default=None, @@ -84,6 +84,14 @@ def main() -> None: required=False, ) + parser.add_argument( + "--bids", + type=bool, + help="Specify if the provided dataset is BIDS type", + default=False, + required=False + ) + # VERSION argument help = "Show the version and exit" parser.add_argument( @@ -122,7 +130,7 @@ def main() -> None: args = parser.parse_args() - in_data = args.in_data + in_dir = args.in_dir out_dir = args.out_dir device = args.device dlicv_extra_args = args.dlicv_args @@ -141,33 +149,65 @@ def main() -> None: os.system("DLICV -i ./dummy -o ./dummy --clear_cache") os.system("DLMUSE -i ./dummy -o ./dummy --clear_cache") + # Run pipeline - no_threads = int(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, - dlmuse_extra_args, - dlicv_extra_args, - i, - ), - ) - curr_thread.start() - threads.append(curr_thread) - - for t in threads: - t.join() - - merge_output_data(out_dir) - remove_subfolders(in_data) - remove_subfolders(out_dir) + if args.bids == True: + collect_T1(in_dir) + + no_threads = int(args.cores) + subfolders = split_data("raw_temp_T1", 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, + dlmuse_extra_args, + dlicv_extra_args, + i, + ), + ) + curr_thread.start() + threads.append(curr_thread) + + for t in threads: + t.join() + + merge_output_data(out_dir) + remove_subfolders("raw_temp_T1") + remove_subfolders(out_dir) + + + else: + no_threads = int(args.cores) + subfolders = split_data(in_dir, 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, + dlmuse_extra_args, + dlicv_extra_args, + i, + ), + ) + curr_thread.start() + threads.append(curr_thread) + + for t in threads: + t.join() + + merge_output_data(out_dir) + remove_subfolders(in_dir) + remove_subfolders(out_dir) if __name__ == "__main__": diff --git a/NiChart_DLMUSE/utils.py b/NiChart_DLMUSE/utils.py index f23825c..4028b15 100644 --- a/NiChart_DLMUSE/utils.py +++ b/NiChart_DLMUSE/utils.py @@ -147,6 +147,44 @@ def dir_size(in_dir: str) -> int: return size +def dir_foldercount(in_dir: str) -> int: + """ + Returns the number of subfolders that the input directory has + """ + + size = 0 + for path in os.listdir(in_dir): + if os.path.isdir(os.path.join(in_dir, path)): + size += 1 + + return size + + +def collect_T1(in_dir: str) -> None: + """ + This function collects all the raw T1 images from the passed BIDS input dir and + it creates a temporary folder that will act as a generic dataset with only T1 images + """ + if os.path.isdir("raw_temp_T1"): + os.system("rm -r raw_temp_T1/*") + else: + # create the raw_temp_T1 folder that will host all the T1 images + os.system("mkdir raw_temp_T1") + + total_subs = dir_foldercount(in_dir) + accepted_subfolders = [] + for i in range(1, total_subs): + if i < 10: + accepted_subfolders.append(f"sub-0{i}") + else: + accepted_subfolders.append(f"sub-{i}") + + for root, subs, files in os.walk(in_dir): + for sub in subs: + if(sub in accepted_subfolders): + os.system(f"cp {os.path.join(in_dir, sub)}/anat/* raw_temp_T1") + + def split_data(in_dir: str, N: int) -> list: """ @@ -176,6 +214,10 @@ def split_data(in_dir: str, N: int) -> list: os.system(f"cp {file} {in_dir}/split_{current_folder}") current_file += 1 + if current_file < no_files_in_folders: + # Don't forget the last split if it has less files than the maximum files in a subfolder + subfolders.append(f"{in_dir}/split_{current_folder}") + return subfolders