Reformat code

Author: spirosmaggioros

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()

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,18 +75,17 @@ 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):
         # 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)
 
@@ -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(

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

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"]},
 )