[BIDS] Accept BIDS format input

TODO: Generate BIDS output format if the user specify it

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__":

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