Initial commit codebase

Author: sarah-deboer

.gitignore

@@ -0,0 +1,3 @@
+*.pt
+*.pyc
+*.egg-info

Dockerfile

@@ -0,0 +1,7 @@
+FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime
+
+RUN pip3 install monai==1.4.0
+RUN pip3 install TotalSegmentator
+RUN pip3 install nnunetv2==2.4.2
+RUN pip3 install blosc2
+RUN pip3 install acvl-utils==0.2

README.md

@@ -32,24 +32,19 @@ export nnUNet_results='/path/to/.totalsegmentator/nnunet/results'
 ## Running the algorithm
 When using the container build using the Dockerfile you can mount the input path containing the CT scans to "/input", the output path to "/output", and the folder containing downloaded model weights to "/opt/ml/model". 
 
-When using a virtual or conda environment, you need to set the input, output and model directories using command line arguments:
-```bash
-    --input
-    --output
-    --model
-```
+When using a virtual or conda environment, you can either use the command line arguments to set the folder paths or set them as environment variables. 
 
 ### To start the script
-Kick-off the script with ```python main.py``` with the following input arguments:
+You can kick-off the script with:
 
 ```bash
-    --use_cropping (optionally)
+python main.py --use-cropping (optional) --input-path (optional) --output-path(optional) --model-path (optional)
 ```
 
 ## Pipeline:
 The algorithm does the following steps for each CT scan (in .mha format) in the input folder:
 1. read ct scan (.mha) using SimpleITK
-2. if flag --use_cropping is set: use TotalSegmentator to find bladder and lungs, if found, crop around that in all three directions. 
+2. if flag --use-cropping is set: use TotalSegmentator to find bladder and lungs, if found, crop around that in all three directions. 
 3. use trained nnUNet to segment the kidneys and kidney abnormalities if present
 4. postprocess the kidney abnormality masks by removing small components (<3mm) and only keeping components attached to a kidney region, except when the abnormality component is larger than 100,000 mm^3.
 

main.py

@@ -0,0 +1,147 @@
+#  Copyright 2022 Diagnostic Image Analysis Group, Radboudumc, Nijmegen, The Netherlands
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import gc
+import sys
+from argparse import ArgumentParser
+from pathlib import Path
+
+import SimpleITK
+
+from src.kidney_abnormality_segmentation.config import (
+    default_input_path,
+    default_model_path,
+    default_output_path,
+)
+from src.kidney_abnormality_segmentation.postprocessing.postprocess_segmentation_mask import (
+    postprocess_segmentation_mask,
+)
+from src.kidney_abnormality_segmentation.preprocessing.extract_roi import extract_roi
+from src.kidney_abnormality_segmentation.segmentation.segment_ct_image import (
+    segment_ct_image,
+)
+from src.kidney_abnormality_segmentation.utils import resample_volume, stem
+
+
+def build_parser() -> ArgumentParser:
+    """Create and return the CLI argument parser."""
+    p = ArgumentParser(description="Process images and perform segmentation")
+    p.add_argument(
+        "--use-cropping",
+        action="store_true",
+        default=False,
+        help="Enable ROI cropping based on TotalSegmentator (default: disabled).",
+    )
+    # Use Path typing and safer argument names (avoid shadowing builtins like `input`)
+    p.add_argument(
+        "--input-path",
+        type=Path,
+        default=default_input_path(),
+        help="Input folder path (defaults to /input or $INPUT_PATH).",
+    )
+    p.add_argument(
+        "--output-path",
+        type=Path,
+        default=default_output_path(),
+        help="Output folder path (defaults to /output or $OUTPUT_PATH).",
+    )
+    p.add_argument(
+        "--model-path",
+        type=Path,
+        default=default_model_path(),
+        help="Model weights folder path (defaults to /opt/ml/model or $MODEL_PATH).",
+    )
+    return p
+
+
+def run():
+    parser = build_parser()
+    args = parser.parse_args()
+
+    # List all CT files under /input
+    ct_folder = args.input_path
+    if not ct_folder.exists():
+        raise FileNotFoundError(f"Input folder does not exist: {ct_folder}")
+    if not ct_folder.is_dir():
+        raise NotADirectoryError(f"Input path is not a directory: {ct_folder}")
+
+    try:
+        all_cts = list(ct_folder.rglob("*.mha"))
+    except PermissionError as e:
+        raise PermissionError(f"Cannot access {args.input_path}: {e}") from e
+
+    if not all_cts:
+        print(f"No CT files found under {ct_folder}")
+        sys.exit(1)
+
+    print(f"[run] Found {len(all_cts)} input CTs to process")
+    crop_roi = args.use_cropping
+
+    for input_ct_image_path in all_cts:
+        print(f"[run] Processing {input_ct_image_path.name}")
+        if input_ct_image_path.name.startswith("."):
+            print(f"[run] Skipping {input_ct_image_path.name} because not an image.")
+            continue
+        image_name = stem(str(input_ct_image_path))
+        out_folder = args.output_path
+        out_folder.mkdir(parents=True, exist_ok=True)
+        out_path = out_folder / f"{image_name}.mha"
+        if out_path.is_file():
+            print(
+                f"[run] Skipping {input_ct_image_path.name} because output segmentation already exists for this image."
+            )
+            continue
+        # 3) Decide what to hand to segment_ct_image:
+        #    - If cropping: read into memory, crop, then pass the cropped SITK.Image.
+        #    - If no cropping: NEVER read the full CT. Pass the filepath string instead.
+        orig_spacing = SimpleITK.ReadImage(str(input_ct_image_path)).GetSpacing()
+        if crop_roi:
+            print("[run] Cropping ROI; will read full CT into memory.")
+            full_ct = SimpleITK.ReadImage(str(input_ct_image_path))
+            input_for_seg = extract_roi(full_ct)
+            # free the full CT
+            del full_ct
+            gc.collect()
+        else:
+            print(
+                "[run] No cropping requested; will segment from disk without reading full CT."
+            )
+            input_for_seg = str(input_ct_image_path)
+
+        # 4) Segment (this now never loads the full on-disk CT into RAM)
+        print("[run] Calling segment_ct_image() …")
+        segmentation_sitk = segment_ct_image(input_for_seg, str(args.model_path))
+
+        # 5) Free any remaining cropped image if it was in RAM
+        if isinstance(input_for_seg, SimpleITK.Image):
+            del input_for_seg
+            gc.collect()
+
+        # 6) Postprocess & write out
+        post_sitk = postprocess_segmentation_mask(segmentation_sitk)
+        final_image = resample_volume(
+            post_sitk,
+            new_spacing=orig_spacing,
+            interpolator=SimpleITK.sitkNearestNeighbor,
+        )
+        print(f"[run] Writing final mask to: {out_path}")
+
+        SimpleITK.WriteImage(final_image, str(out_path))
+
+    print("[run] Done.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(run())

pyproject.toml

@@ -0,0 +1,22 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "kidney_abnormality_segmentation"
+version = "1.0.0"
+description = "Kidney Abnormality Segmentation using nnUNet."
+readme = "README.md"
+requires-python = ">=3.9"
+license = {text = "Apache License, Version 2.0 (the 'License)"}
+authors = [{name = "Sarah de Boer"}]
+dependencies = [
+  "monai==1.4.0",
+  "TotalSegmentator",
+  "nnunetv2==2.4.2",
+  "blosc2",
+  "acvl-utils==0.2"
+]
+
+[project.optional-dependencies]
+dev = ["pytest", "black", "ruff"]

requirements.yml

@@ -0,0 +1,16 @@
+name: kidney-abnormality-segmentation
+channels:
+  - pytorch
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.11
+  - pip
+  - pytorch=2.2.2
+  - monai=1.4.0
+  - pip:
+      - TotalSegmentator
+      - nnunetv2==2.4.2
+      - blosc2
+      - acvl-utils==0.2
+      - -e .

src/kidney_abnormality_segmentation/__init__.py

@@ -0,0 +1,22 @@
+import os
+from pathlib import Path
+
+# Defaults match docker mounts
+DEFAULT_INPUT_PATH = Path("/input")
+DEFAULT_OUTPUT_PATH = Path("/output")
+DEFAULT_MODEL_PATH = Path("/opt/ml/model")
+
+
+def default_input_path() -> Path:
+    """Return default input path, allowing override via env var."""
+    return Path(os.getenv("INPUT_PATH", DEFAULT_INPUT_PATH))
+
+
+def default_output_path() -> Path:
+    """Return default output path, allowing override via env var."""
+    return Path(os.getenv("OUTPUT_PATH", DEFAULT_OUTPUT_PATH))
+
+
+def default_model_path() -> Path:
+    """Return default model path, allowing override via env var."""
+    return Path(os.getenv("MODEL_PATH", DEFAULT_MODEL_PATH))