From c7e2065d2ae1ddd981d57616c955c95c4b53f2b7 Mon Sep 17 00:00:00 2001 From: spirosmaggioros Date: Fri, 10 Jan 2025 17:20:21 +0000 Subject: [PATCH 1/8] Updating DLMUSE with pipeline function --- DLMUSE/__init__.py | 1 + DLMUSE/__main__.py | 191 ++++++-------------------------------- DLMUSE/dlmuse_pipeline.py | 183 ++++++++++++++++++++++++++++++++++++ DLMUSE/utils.py | 16 +++- requirements.txt | 2 +- setup.py | 2 +- 6 files changed, 228 insertions(+), 167 deletions(-) create mode 100644 DLMUSE/dlmuse_pipeline.py diff --git a/DLMUSE/__init__.py b/DLMUSE/__init__.py index e69de29..553788f 100644 --- a/DLMUSE/__init__.py +++ b/DLMUSE/__init__.py @@ -0,0 +1 @@ +from dlmuse_pipeline import dlmuse_pipeline diff --git a/DLMUSE/__main__.py b/DLMUSE/__main__.py index 396a086..fc00b20 100644 --- a/DLMUSE/__main__.py +++ b/DLMUSE/__main__.py @@ -1,19 +1,7 @@ import argparse -import json -import os -from pathlib import Path -import shutil -import sys -import warnings +from DLMUSE.dlmuse_pipeline import dlmuse_pipeline -import torch - -from .utils import prepare_data_folder, rename_and_copy_files - -warnings.simplefilter(action="ignore", category=FutureWarning) -warnings.simplefilter(action="ignore", category=UserWarning) - -VERSION = 1.0 +VERSION = "1.0.3" def main() -> None: prog="DLMUSE" @@ -23,7 +11,7 @@ def main() -> None: usage=""" DLMUSE v{VERSION} Segmentation of the brain into MUSE ROIs from the Nifti image (.nii.gz) of the LPS oriented Intra Cranial Volume (ICV - see DLICV method). - + Required arguments: [-i, --in_dir] The filepath of the input directory [-o, --out_dir] The filepath of the output directory @@ -43,17 +31,19 @@ def main() -> None: # Required Arguments parser.add_argument( "-i", + "--in_dir", type=str, required=True, help="[REQUIRED] Input folder with LPS oriented T1 sMRI Intra Cranial Volumes (ICV) in Nifti format (nii.gz).", ) parser.add_argument( "-o", + "--out_dir", type=str, required=True, help="[REQUIRED] Output folder for the segmentation results in Nifti format (nii.gz).", ) - + # Optional Arguments parser.add_argument( "-device", @@ -208,158 +198,35 @@ def main() -> None: required=False, default=0, help="[nnUnet Arg] If multiple nnUNetv2_predict exist, which one is this? IDs start with 0 " - "can end with num_parts - 1. So when you submit 5 nnUNetv2_predict calls you need to set " + "can end with num_parts - 1. So when you submit 5 nnUNetv2_predict calls you need to set " "-num_parts 5 and use -part_id 0, 1, 2, 3 and 4. Note: You are yourself responsible to make these run on separate GPUs! " "Use CUDA_VISIBLE_DEVICES.", ) - - args = parser.parse_args() - args.f = [args.f] - - if args.clear_cache: - shutil.rmtree(os.path.join( - Path(__file__).parent, - "nnunet_results" - )) - shutil.rmtree(os.path.join( - Path(__file__).parent, - ".cache" - )) - if not args.i or not args.o: - print("Cache cleared and missing either -i / -o. Exiting.") - sys.exit(0) - - if not args.i or not args.o: - parser.error("The following arguments are required: -i, -o") - - # data conversion - src_folder = args.i # input folder - - if not os.path.exists(args.o): # create output folder if it does not exist - os.makedirs(args.o) - - des_folder = os.path.join(args.o, "renamed_image") - - # check if -i argument is a folder, list (csv), or a single file (nii.gz) - if os.path.isdir(args.i): # if args.i is a directory - src_folder = args.i - prepare_data_folder(des_folder) - rename_dic, rename_back_dict = rename_and_copy_files(src_folder, des_folder) - datalist_file = os.path.join(des_folder, "renaming.json") - with open(datalist_file, "w", encoding="utf-8") as f: - json.dump(rename_dic, f, ensure_ascii=False, indent=4) - print(f"Renaming dic is saved to {datalist_file}") - else: - print("Input directory not found. Exiting DLMUSE.") - sys.exit() - - model_folder = os.path.join( - Path(__file__).parent, - "nnunet_results", - "Dataset%s_Task%s_DLMUSEV2/nnUNetTrainer__nnUNetPlans__%s/" - % (args.d, args.d, args.c), - ) - - - # Check if model exists. If not exist, download using HuggingFace - print(f"Using model folder: {model_folder}") - if not os.path.exists(model_folder): - # HF download model - print("DLMUSE model not found, downloading...") - - from huggingface_hub import snapshot_download - local_src = Path(__file__).parent - snapshot_download(repo_id="nichart/DLMUSE", local_dir=local_src) - - print("DLMUSE model has been successfully downloaded!") - else: - print("Loading the model...") - - prepare_data_folder(des_folder) - - assert ( - args.part_id < args.num_parts - ), "part_id < num_parts. Please see nnUNetv2_predict -h." - - assert args.device in [ - "cpu", - "cuda", - "mps", - ], f"-device must be either cpu, mps or cuda. Got: {args.device}." - - if args.device == "cpu": - import multiprocessing - # use half of the available threads in the system. - torch.set_num_threads(multiprocessing.cpu_count() // 2) - device = torch.device("cpu") - print("Running in CPU mode.") - elif args.device == "cuda": - # multithreading in torch doesn't help nnU-Net if run on GPU - torch.set_num_threads(1) - torch.set_num_interop_threads(1) - device = torch.device("cuda") - print("Running in CUDA mode.") - else: - device = torch.device("mps") - print("Running in MPS mode.") - - # exports for nnunetv2 purposes - os.environ["nnUNet_raw"] = "/nnunet_raw/" - os.environ["nnUNet_preprocessed"] = "/nnunet_preprocessed" - os.environ["nnUNet_results"] = ( - "/nnunet_results" # where model will be located (fetched from HF) - ) - - from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor - - # Initialize nnUnetPredictor - predictor = nnUNetPredictor( - tile_step_size=args.step_size, - use_gaussian=True, - use_mirroring=not args.disable_tta, - perform_everything_on_device=True, - device=device, - verbose=args.verbose, - verbose_preprocessing=args.verbose, - allow_tqdm=not args.disable_progress_bar, + dlmuse_pipeline( + args.in_dir, + args.out_dir, + args.device, + args.verbose, + args.save_probabilities, + args.continue_prediction, + args.disable_progress_bar, + args.clear_cache, + args.disable_tta, + args.d, + args.p, + args.tr, + args.c, + [args.f], + args.step_size, + args.chk, + args.npp, + args.nps, + args.prev_stage_predictions, + args.num_parts, + args.part_id ) - # Retrieve the model and it's weight - predictor.initialize_from_trained_model_folder( - model_folder, args.f, checkpoint_name=args.chk - ) - - # Final prediction - predictor.predict_from_files( - des_folder, - args.o, - save_probabilities=args.save_probabilities, - overwrite=not args.continue_prediction, - num_processes_preprocessing=args.npp, - num_processes_segmentation_export=args.nps, - folder_with_segs_from_prev_stage=args.prev_stage_predictions, - num_parts=args.num_parts, - part_id=args.part_id, - ) - - # After prediction, convert the image name back to original - files_folder = args.o - - for filename in os.listdir(files_folder): - if filename.endswith(".nii.gz"): - original_name = rename_back_dict[filename] - os.rename( - os.path.join(files_folder, filename), - os.path.join(files_folder, original_name), - ) - # Remove the (temporary) des_folder directory - if os.path.exists(des_folder): - shutil.rmtree(des_folder) - - print("Inference Process Done!") - - if __name__ == "__main__": main() diff --git a/DLMUSE/dlmuse_pipeline.py b/DLMUSE/dlmuse_pipeline.py new file mode 100644 index 0000000..f10170e --- /dev/null +++ b/DLMUSE/dlmuse_pipeline.py @@ -0,0 +1,183 @@ +from typing import Optional +import logging +import json +import os +from pathlib import Path +import shutil +import sys +import warnings + +import torch + +from DLMUSE.utils import prepare_data_folder, rename_and_copy_files + +warnings.simplefilter(action="ignore", category=FutureWarning) +warnings.simplefilter(action="ignore", category=UserWarning) + +def dlmuse_pipeline( + in_dir: str, + out_dir: str, + device: str, + verbose: bool = False, + save_probabilities: bool = False, + continue_prediction: bool = False, + disable_progress_bar: bool = False, + clear_cache: bool = False, + disable_tta: bool = False, + d: str = "903", + p: str = "nnUNetPlans", + tr: str = "nnUNetTrainer", + c: str = "3d_fullres", + f: list = [0], + step_size: float = 0.5, + chk: str = "checkpoint_final.pth", + npp: int = 2, + nps: int = 2, + prev_stage_predictions: Optional[str] = None, + num_parts: int = 1, + part_id: int = 0, +): + + if clear_cache: + shutil.rmtree(os.path.join( + Path(__file__).parent, + "nnunet_results" + )) + shutil.rmtree(os.path.join( + Path(__file__).parent, + ".cache" + )) + if not in_dir or not out_dir: + logging.error("Cache cleared and missing either -i / -o. Exiting.") + sys.exit(0) + + if not in_dir or not out_dir: + logging.error("The following arguments are required: -i, -o") + sys.exit(0) + + # data conversion + src_folder = in_dir # input folder + + if not os.path.exists(out_dir): # create output folder if it does not exist + logging.info(f"Can't find {out_dir}, creating it...") + os.makedirs(out_dir) + + des_folder = os.path.join(out_dir, "renamed_image") + + # check if -i argument is a folder, list (csv), or a single file (nii.gz) + if os.path.isdir(in_dir): # if in_dir is a directory + src_folder = in_dir + prepare_data_folder(des_folder) + rename_dic, rename_back_dict = rename_and_copy_files(src_folder, des_folder) + datalist_file = os.path.join(des_folder, "renaming.json") + with open(datalist_file, "w", encoding="utf-8") as _f: + json.dump(rename_dic, _f, ensure_ascii=False, indent=4) + logging.info(f"Renaming dic is saved to {datalist_file}") + else: + logging.error("Input directory not found. Exiting DLMUSE.") + sys.exit(0) + + model_folder = os.path.join( + Path(__file__).parent, + "nnunet_results", + "Dataset%s_Task%s_DLMUSEV2/nnUNetTrainer__nnUNetPlans__%s/" + % (d, d, c), + ) + + + # Check if model exists. If not exist, download using HuggingFace + logging.info(f"Using model folder: {model_folder}") + if not os.path.exists(model_folder): + # HF download model + logging.info("DLMUSE model not found, downloading...") + + from huggingface_hub import snapshot_download + local_src = Path(__file__).parent + snapshot_download(repo_id="nichart/DLMUSE", local_dir=local_src) + + logging.info("DLMUSE model has been successfully downloaded!") + else: + logging.info("Loading the model...") + + prepare_data_folder(des_folder) + + assert ( + part_id < num_parts + ), "part_id < num_parts. Please see nnUNetv2_predict -h." + + assert device in [ + "cpu", + "cuda", + "mps", + ], f"-device must be either cpu, mps or cuda. Got: {device}." + + if device == "cpu": + import multiprocessing + # use half of the available threads in the system. + torch.set_num_threads(multiprocessing.cpu_count() // 2) + device = torch.device("cpu") + logging.info("Running in CPU mode.") + elif device == "cuda": + # multithreading in torch doesn't help nnU-Net if run on GPU + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + device = torch.device("cuda") + logging.info("Running in CUDA mode.") + else: + device = torch.device("mps") + logging.info("Running in MPS mode.") + + # exports for nnunetv2 purposes + os.environ["nnUNet_raw"] = "/nnunet_raw/" + os.environ["nnUNet_preprocessed"] = "/nnunet_preprocessed" + os.environ["nnUNet_results"] = ( + "/nnunet_results" # where model will be located (fetched from HF) + ) + + from nnunetv2.inference.predict_from_raw_data import nnUNetPredictor + + # Initialize nnUnetPredictor + predictor = nnUNetPredictor( + tile_step_size=step_size, + use_gaussian=True, + use_mirroring=not disable_tta, + perform_everything_on_device=True, + device=device, + verbose=verbose, + verbose_preprocessing=verbose, + allow_tqdm=not disable_progress_bar, + ) + + # Retrieve the model and it's weight + predictor.initialize_from_trained_model_folder( + model_folder, f, checkpoint_name=chk + ) + + # Final prediction + predictor.predict_from_files( + des_folder, + out_dir, + save_probabilities=save_probabilities, + overwrite=not continue_prediction, + num_processes_preprocessing=npp, + num_processes_segmentation_export=nps, + folder_with_segs_from_prev_stage=prev_stage_predictions, + num_parts=num_parts, + part_id=part_id, + ) + + # After prediction, convert the image name back to original + files_folder = out_dir + + for filename in os.listdir(files_folder): + if filename.endswith(".nii.gz"): + original_name = rename_back_dict[filename] + os.rename( + os.path.join(files_folder, filename), + os.path.join(files_folder, original_name), + ) + # Remove the (temporary) des_folder directory + if os.path.exists(des_folder): + shutil.rmtree(des_folder) + + logging.info("Inference Process Done!") diff --git a/DLMUSE/utils.py b/DLMUSE/utils.py index 00705a3..9605a1d 100644 --- a/DLMUSE/utils.py +++ b/DLMUSE/utils.py @@ -25,17 +25,27 @@ def rename_and_copy_files(src_folder: str, des_folder: str) -> Tuple[dict, dict] rename_back_dict: a dictionary will be use to mapping backto the original name """ + if not os.path.exists(src_folder): + raise FileNotFoundError(f"Source folder '{src_folder}' does not exist.") + if not os.path.exists(des_folder): + raise FileNotFoundError(f"Source folder '{des_folder}' does not exist.") + files = os.listdir(src_folder) rename_dict = {} rename_back_dict = {} for idx, filename in enumerate(files): old_name = os.path.join(src_folder, filename) + if not os.path.isfile(old_name): # We only want files! + continue rename_file = f"case_{idx: 04d}_0000.nii.gz" rename_back = f"case_{idx: 04d}.nii.gz" new_name = os.path.join(des_folder, rename_file) - shutil.copy2(old_name, new_name) - rename_dict[filename] = rename_file - rename_back_dict[rename_back] = "DLMUSE_mask_" + filename + try: + shutil.copy2(old_name, new_name) + rename_dict[filename] = rename_file + rename_back_dict[rename_back] = "DLMUSE_mask_" + filename + except Exception as e: + print(F"Error copying file '{filename}' to '{new_name}': {e}") return rename_dict, rename_back_dict diff --git a/requirements.txt b/requirements.txt index e976eb9..18e6d12 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -torch>=2.2.1 +torch==2.3.1 nnunetv2==2.5.1 argparse huggingface_hub diff --git a/setup.py b/setup.py index f6310d7..2fe83e1 100644 --- a/setup.py +++ b/setup.py @@ -24,7 +24,7 @@ download_url="https://github.com/CBICA/DLMUSE/", url="https://github.com/CBICA/DLMUSE/", packages=find_packages(exclude=["tests", ".github"]), - python_requires=">=3.8", + python_requires=">=3.9", install_requires=required, entry_points={"console_scripts": ["DLMUSE = DLMUSE.__main__:main"]}, classifiers=[ From 9ec4a5e83886ffc7d925018b704acec3dc5a9aad Mon Sep 17 00:00:00 2001 From: spirosmaggioros Date: Fri, 10 Jan 2025 17:22:50 +0000 Subject: [PATCH 2/8] Added new workflows and badges --- .github/workflows/macos-build.yml | 34 ++++++++++++++++++++++++++++++ .github/workflows/ubuntu-build.yml | 34 ++++++++++++++++++++++++++++++ README.md | 8 +++++-- 3 files changed, 74 insertions(+), 2 deletions(-) create mode 100644 .github/workflows/macos-build.yml create mode 100644 .github/workflows/ubuntu-build.yml diff --git a/.github/workflows/macos-build.yml b/.github/workflows/macos-build.yml new file mode 100644 index 0000000..871310d --- /dev/null +++ b/.github/workflows/macos-build.yml @@ -0,0 +1,34 @@ +name: Macos Build + +on: + push: + branches: + - main + - spiros-dev + - spiros-dev-2 + pull_request: + branches: + - main + - spiros-dev + - spiros-dev-2 + +jobs: + test: + runs-on: macos-latest + strategy: + matrix: + python-version: ["3.9", "3.10", "3.11", "3.12"] + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Setup python version ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + pip install -r requirements.txt + python3 -m pip install -e . diff --git a/.github/workflows/ubuntu-build.yml b/.github/workflows/ubuntu-build.yml new file mode 100644 index 0000000..c7a7ff5 --- /dev/null +++ b/.github/workflows/ubuntu-build.yml @@ -0,0 +1,34 @@ +name: Ubuntu Build + +on: + push: + branches: + - main + - spiros-dev + - spiros-dev-2 + pull_request: + branches: + - main + - spiros-dev + - spiros-dev-2 + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.9", "3.10", "3.11", "3.12"] + + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Setup python version ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + pip install -r requirements.txt + python3 -m pip install -e . diff --git a/README.md b/README.md index 6114746..745c72b 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,7 @@ +![macos tests](https://github.com/CBICA/DLMUSE/actions/workflows/macos-build.yml/badge.svg) +![ubuntu tests](https://github.com/CBICA/DLMUSE/actions/workflows/ubuntu-build.yml/badge.svg) +![PyPI Stable](https://img.shields.io/pypi/v/DLMUSE) + # DLMUSE - Deep Learning MUlti-atlas region Segmentation utilizing Ensembles of registration algorithms and parameters ## Overview @@ -21,7 +25,7 @@ pip install -e . ``` ### Installing PyTorch -Depending on your system configuration and supported CUDA version, you may need to follow the [PyTorch Installation Instructions](https://pytorch.org/get-started/locally/). +Depending on your system configuration and supported CUDA version, you may need to follow the [PyTorch Installation Instructions](https://pytorch.org/get-started/locally/). ## Usage @@ -38,7 +42,7 @@ DLMUSE -h ``` ## \[Windows Users\] Troubleshooting model download failures -Our model download process creates several deep directory structures. If you are on Windows and your model download process fails, it may be due to Windows file path limitations. +Our model download process creates several deep directory structures. If you are on Windows and your model download process fails, it may be due to Windows file path limitations. To enable long path support in Windows 10, version 1607, and later, the registry key `HKEY_LOCAL_MACHINE\SYSTEM\CurrentControlSet\Control\FileSystem LongPathsEnabled (Type: REG_DWORD)` must exist and be set to 1. From 4a520417aa9409ae200d6fbcd29d96eb2e1e1a6d Mon Sep 17 00:00:00 2001 From: spirosmaggioros Date: Fri, 10 Jan 2025 21:21:47 +0000 Subject: [PATCH 3/8] Fix pipeline function name --- DLMUSE/__init__.py | 2 +- DLMUSE/__main__.py | 4 ++-- DLMUSE/dlmuse_pipeline.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/DLMUSE/__init__.py b/DLMUSE/__init__.py index 553788f..3589bea 100644 --- a/DLMUSE/__init__.py +++ b/DLMUSE/__init__.py @@ -1 +1 @@ -from dlmuse_pipeline import dlmuse_pipeline +from dlmuse_pipeline import run_dlmuse_pipeline diff --git a/DLMUSE/__main__.py b/DLMUSE/__main__.py index fc00b20..80bc15d 100644 --- a/DLMUSE/__main__.py +++ b/DLMUSE/__main__.py @@ -1,5 +1,5 @@ import argparse -from DLMUSE.dlmuse_pipeline import dlmuse_pipeline +from DLMUSE.dlmuse_pipeline import run_dlmuse_pipeline VERSION = "1.0.3" @@ -204,7 +204,7 @@ def main() -> None: ) args = parser.parse_args() - dlmuse_pipeline( + run_dlmuse_pipeline( args.in_dir, args.out_dir, args.device, diff --git a/DLMUSE/dlmuse_pipeline.py b/DLMUSE/dlmuse_pipeline.py index f10170e..e7dfae0 100644 --- a/DLMUSE/dlmuse_pipeline.py +++ b/DLMUSE/dlmuse_pipeline.py @@ -14,7 +14,7 @@ warnings.simplefilter(action="ignore", category=FutureWarning) warnings.simplefilter(action="ignore", category=UserWarning) -def dlmuse_pipeline( +def run_dlmuse_pipeline( in_dir: str, out_dir: str, device: str, From 9d6e9ad66aeb84c1533ff52bdd609b5a6ffbc775 Mon Sep 17 00:00:00 2001 From: spirosmaggioros Date: Fri, 10 Jan 2025 22:24:51 +0000 Subject: [PATCH 4/8] Make version update automatically --- DLMUSE/__main__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/DLMUSE/__main__.py b/DLMUSE/__main__.py index 80bc15d..0bce125 100644 --- a/DLMUSE/__main__.py +++ b/DLMUSE/__main__.py @@ -1,7 +1,9 @@ import argparse +import pkg_resources + from DLMUSE.dlmuse_pipeline import run_dlmuse_pipeline -VERSION = "1.0.3" +VERSION = pkg_resources.require("DLMUSE")[0].version def main() -> None: prog="DLMUSE" From 320834dc96c8531a66b6b959a4a6be2f075471d9 Mon Sep 17 00:00:00 2001 From: spirosmaggioros Date: Fri, 10 Jan 2025 22:25:57 +0000 Subject: [PATCH 5/8] Reformat code --- DLMUSE/__main__.py | 19 ++++++++++++------- DLMUSE/dlmuse_pipeline.py | 33 ++++++++++++--------------------- DLMUSE/utils.py | 6 +++--- setup.py | 6 ++---- 4 files changed, 29 insertions(+), 35 deletions(-) diff --git a/DLMUSE/__main__.py b/DLMUSE/__main__.py index 0bce125..fac8799 100644 --- a/DLMUSE/__main__.py +++ b/DLMUSE/__main__.py @@ -1,12 +1,14 @@ import argparse -import pkg_resources + +import pkg_resources # type: ignore from DLMUSE.dlmuse_pipeline import run_dlmuse_pipeline VERSION = pkg_resources.require("DLMUSE")[0].version + def main() -> None: - prog="DLMUSE" + prog = "DLMUSE" parser = argparse.ArgumentParser( prog=prog, description="DLMUSE - MUlti-atlas region Segmentation utilizing Ensembles of registration algorithms and parameters.", @@ -26,8 +28,10 @@ def main() -> None: -o /path/to/output \ -device cpu|cuda|mps - """.format(VERSION=VERSION), - add_help=False + """.format( + VERSION=VERSION + ), + add_help=False, ) # Required Arguments @@ -91,7 +95,7 @@ def main() -> None: action="store_true", required=False, default=False, - help="Set this flag to clear any cached models before running. This is recommended if a previous download failed." + help="Set this flag to clear any cached models before running. This is recommended if a previous download failed.", ) parser.add_argument( "--disable_tta", @@ -101,7 +105,7 @@ def main() -> None: help="[nnUnet Arg] Set this flag to disable test time data augmentation in the form of mirroring. " "Faster, but less accurate inference. Not recommended.", ) - ### DEPRECIATED #### + # DEPRECIATED ## # parser.add_argument( # "-m", # type=str, @@ -227,8 +231,9 @@ def main() -> None: args.nps, args.prev_stage_predictions, args.num_parts, - args.part_id + args.part_id, ) + if __name__ == "__main__": main() diff --git a/DLMUSE/dlmuse_pipeline.py b/DLMUSE/dlmuse_pipeline.py index e7dfae0..e326725 100644 --- a/DLMUSE/dlmuse_pipeline.py +++ b/DLMUSE/dlmuse_pipeline.py @@ -1,11 +1,11 @@ -from typing import Optional -import logging import json +import logging import os -from pathlib import Path import shutil import sys import warnings +from pathlib import Path +from typing import Optional import torch @@ -14,6 +14,7 @@ warnings.simplefilter(action="ignore", category=FutureWarning) warnings.simplefilter(action="ignore", category=UserWarning) + def run_dlmuse_pipeline( in_dir: str, out_dir: str, @@ -36,17 +37,11 @@ def run_dlmuse_pipeline( prev_stage_predictions: Optional[str] = None, num_parts: int = 1, part_id: int = 0, -): +) -> None: if clear_cache: - shutil.rmtree(os.path.join( - Path(__file__).parent, - "nnunet_results" - )) - shutil.rmtree(os.path.join( - Path(__file__).parent, - ".cache" - )) + shutil.rmtree(os.path.join(Path(__file__).parent, "nnunet_results")) + shutil.rmtree(os.path.join(Path(__file__).parent, ".cache")) if not in_dir or not out_dir: logging.error("Cache cleared and missing either -i / -o. Exiting.") sys.exit(0) @@ -80,11 +75,9 @@ def run_dlmuse_pipeline( model_folder = os.path.join( Path(__file__).parent, "nnunet_results", - "Dataset%s_Task%s_DLMUSEV2/nnUNetTrainer__nnUNetPlans__%s/" - % (d, d, c), + "Dataset%s_Task%s_DLMUSEV2/nnUNetTrainer__nnUNetPlans__%s/" % (d, d, c), ) - # Check if model exists. If not exist, download using HuggingFace logging.info(f"Using model folder: {model_folder}") if not os.path.exists(model_folder): @@ -92,6 +85,7 @@ def run_dlmuse_pipeline( logging.info("DLMUSE model not found, downloading...") from huggingface_hub import snapshot_download + local_src = Path(__file__).parent snapshot_download(repo_id="nichart/DLMUSE", local_dir=local_src) @@ -101,9 +95,7 @@ def run_dlmuse_pipeline( prepare_data_folder(des_folder) - assert ( - part_id < num_parts - ), "part_id < num_parts. Please see nnUNetv2_predict -h." + assert part_id < num_parts, "part_id < num_parts. Please see nnUNetv2_predict -h." assert device in [ "cpu", @@ -113,6 +105,7 @@ def run_dlmuse_pipeline( if device == "cpu": import multiprocessing + # use half of the available threads in the system. torch.set_num_threads(multiprocessing.cpu_count() // 2) device = torch.device("cpu") @@ -149,9 +142,7 @@ def run_dlmuse_pipeline( ) # Retrieve the model and it's weight - predictor.initialize_from_trained_model_folder( - model_folder, f, checkpoint_name=chk - ) + predictor.initialize_from_trained_model_folder(model_folder, f, checkpoint_name=chk) # Final prediction predictor.predict_from_files( diff --git a/DLMUSE/utils.py b/DLMUSE/utils.py index 9605a1d..972eb78 100644 --- a/DLMUSE/utils.py +++ b/DLMUSE/utils.py @@ -26,7 +26,7 @@ def rename_and_copy_files(src_folder: str, des_folder: str) -> Tuple[dict, dict] """ if not os.path.exists(src_folder): - raise FileNotFoundError(f"Source folder '{src_folder}' does not exist.") + raise FileNotFoundError(f"Source folder '{src_folder}' does not exist.") if not os.path.exists(des_folder): raise FileNotFoundError(f"Source folder '{des_folder}' does not exist.") @@ -36,7 +36,7 @@ def rename_and_copy_files(src_folder: str, des_folder: str) -> Tuple[dict, dict] for idx, filename in enumerate(files): old_name = os.path.join(src_folder, filename) - if not os.path.isfile(old_name): # We only want files! + if not os.path.isfile(old_name): # We only want files! continue rename_file = f"case_{idx: 04d}_0000.nii.gz" rename_back = f"case_{idx: 04d}.nii.gz" @@ -46,6 +46,6 @@ def rename_and_copy_files(src_folder: str, des_folder: str) -> Tuple[dict, dict] rename_dict[filename] = rename_file rename_back_dict[rename_back] = "DLMUSE_mask_" + filename except Exception as e: - print(F"Error copying file '{filename}' to '{new_name}': {e}") + print(f"Error copying file '{filename}' to '{new_name}': {e}") return rename_dict, rename_back_dict diff --git a/setup.py b/setup.py index 2fe83e1..97fb934 100644 --- a/setup.py +++ b/setup.py @@ -7,7 +7,7 @@ this_directory = Path(__file__).parent long_description = (this_directory / "README.md").read_text() -with open('requirements.txt') as f: +with open("requirements.txt") as f: required = f.read().splitlines() @@ -46,7 +46,5 @@ "nnU-Net", "nnunet", ], - package_data={ - "DLMUSE": ["VERSION"] - }, + package_data={"DLMUSE": ["VERSION"]}, ) From fb57fb3cc5aa7299af10ac29037c24deb4d5720f Mon Sep 17 00:00:00 2001 From: spirosmaggioros Date: Fri, 10 Jan 2025 23:31:37 +0000 Subject: [PATCH 6/8] Fixing setup.py --- setup.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 97fb934..73c3e16 100644 --- a/setup.py +++ b/setup.py @@ -25,7 +25,12 @@ url="https://github.com/CBICA/DLMUSE/", packages=find_packages(exclude=["tests", ".github"]), python_requires=">=3.9", - install_requires=required, + install_requires=[ + "torch<=2.3.1", + "nnunetv2<=2.5.1", + "argparse", + "huggingface_hub" + ], entry_points={"console_scripts": ["DLMUSE = DLMUSE.__main__:main"]}, classifiers=[ "Intended Audience :: Developers", From 84fc6626a3f039c5db21f5f004fd39e4d5593155 Mon Sep 17 00:00:00 2001 From: spirosmaggioros Date: Fri, 10 Jan 2025 23:31:50 +0000 Subject: [PATCH 7/8] Reformat code --- setup.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/setup.py b/setup.py index 73c3e16..7fdab56 100644 --- a/setup.py +++ b/setup.py @@ -25,12 +25,7 @@ url="https://github.com/CBICA/DLMUSE/", packages=find_packages(exclude=["tests", ".github"]), python_requires=">=3.9", - install_requires=[ - "torch<=2.3.1", - "nnunetv2<=2.5.1", - "argparse", - "huggingface_hub" - ], + install_requires=["torch<=2.3.1", "nnunetv2<=2.5.1", "argparse", "huggingface_hub"], entry_points={"console_scripts": ["DLMUSE = DLMUSE.__main__:main"]}, classifiers=[ "Intended Audience :: Developers", From 4a34e8b1b304d0b33b79b14096bde63d292a4c1d Mon Sep 17 00:00:00 2001 From: spirosmaggioros Date: Fri, 10 Jan 2025 23:32:16 +0000 Subject: [PATCH 8/8] Remove unused code --- setup.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/setup.py b/setup.py index 7fdab56..5f1ebd3 100644 --- a/setup.py +++ b/setup.py @@ -7,10 +7,6 @@ this_directory = Path(__file__).parent long_description = (this_directory / "README.md").read_text() -with open("requirements.txt") as f: - required = f.read().splitlines() - - setup( name="DLMUSE", version="1.0.3",