From b16a1a7ab24da9a7e2cffbb1aea2562e52af763c Mon Sep 17 00:00:00 2001 From: jaehwan Date: Tue, 28 May 2024 17:20:31 +0900 Subject: [PATCH] feat: change nnUNetPredictor --- Dockerfile | 80 +++++++++++++++++++++++------------ README.md | 40 ++++++++++++++---- build.sh | 2 +- calc_dice.py | 107 +++++++++++++++++++++++++++++++++++++++++++++++ export.sh | 5 --- process.py | 70 ++++++++++++++++--------------- requirements.txt | 2 +- test.sh | 46 ++++++++++---------- 8 files changed, 255 insertions(+), 97 deletions(-) create mode 100644 calc_dice.py delete mode 100755 export.sh diff --git a/Dockerfile b/Dockerfile index d1bf892..a2bed28 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,42 +1,68 @@ -FROM pytorch/pytorch +FROM ubuntu:20.04 -RUN groupadd -r user && useradd -m --no-log-init -r -g user user +# Set up environment variables +ENV DEBIAN_FRONTEND=noninteractive +ENV PATH="/home/user/.local/bin:${PATH}" +ENV nnUNet_results="/opt/algorithm/checkpoint/" +ENV nnUNet_raw="/opt/algorithm/nnUNet_raw_data_base" +ENV nnUNet_preprocessed="/opt/algorithm/preproc" +ENV MKL_SERVICE_FORCE_INTEL=1 +ENV OMP_NUM_THREADS=1 +ENV OPENBLAS_NUM_THREADS=1 +ENV MKL_NUM_THREADS=1 +ENV NUMEXPR_NUM_THREADS=1 -RUN mkdir -p /opt/app /input /output \ - && chown user:user /opt/app /input /output +# Install system dependencies +RUN apt-get update && apt-get install -y \ + python3.9 \ + python3.9-venv \ + python3.9-dev \ + python3-pip \ + zip \ + unzip \ + gdb \ + && rm -rf /var/lib/apt/lists/* +# Set python3.9 as the default python3 +RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1 + +# Add user +ARG UID=1000 +ARG GID=1000 +RUN groupadd -g ${GID} user && useradd -u ${UID} -g user -m --no-log-init -r -g user user + +# Create necessary directories and set permissions +RUN mkdir -p /opt/app /input /output /opt/algorithm/checkpoint/nnUNet \ + && chown -R user:user /opt/app /input /output /opt/algorithm/checkpoint/nnUNet + +# Switch to user USER user WORKDIR /opt/app -ENV PATH="/home/user/.local/bin:${PATH}" +# Install Python packages +RUN python3 -m pip install --user -U pip +RUN python3 -m pip install --user pip-tools +RUN python3 -m pip install --upgrade pip -RUN python -m pip install --user -U pip && python -m pip install --user pip-tools && python -m pip install --upgrade pip -COPY --chown=user:user nnUNet/ /opt/app/nnUNet/ -RUN python -m pip install -e nnUNet -#RUN python -m pip uninstall -y scipy -#RUN python -m pip install --user --upgrade scipy +# Install PyTorch and related packages +RUN python3 -m pip install --user torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118 -COPY --chown=user:user requirements.txt /opt/app/ -RUN python -m pip install --user -r requirements.txt +# Copy nnUNet and install +COPY --chown=user:user nnUNet/ /opt/app/nnUNet/ +RUN python3 -m pip install --user -e nnUNet +# Copy requirements and install +COPY --chown=user:user requirements.txt /opt/app/ +RUN python3 -m pip install --user -r requirements.txt -# This is the checkpoint file, uncomment the line below and modify /local/path/to/the/checkpoint to your needs +# Copy checkpoint and extract COPY --chown=user:user nnUNetTrainer__nnUNetPlans__3d_fullres.zip /opt/algorithm/checkpoint/nnUNet/ -RUN python -c "import zipfile; import os; zipfile.ZipFile('/opt/algorithm/checkpoint/nnUNet/nnUNetTrainer__nnUNetPlans__3d_fullres.zip').extractall('/opt/algorithm/checkpoint/nnUNet/')" +RUN python3 -c "import zipfile; import os; zipfile.ZipFile('/opt/algorithm/checkpoint/nnUNet/nnUNetTrainer__nnUNetPlans__3d_fullres.zip').extractall('/opt/algorithm/checkpoint/nnUNet/')" +# Copy custom scripts COPY --chown=user:user custom_algorithm.py /opt/app/ COPY --chown=user:user process.py /opt/app/ +COPY --chown=user:user calc_dice.py /opt/app/ -# COPY --chown=user:user weights /opt/algorithm/checkpoint -ENV nnUNet_results="/opt/algorithm/checkpoint/" -ENV nnUNet_raw="/opt/algorithm/nnUNet_raw_data_base" -ENV nnUNet_preprocessed="/opt/algorithm/preproc" -# ENV ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS=64(nope!) -# ENV nnUNet_def_n_proc=1 - -#ENTRYPOINT [ "python3", "-m", "process" ] - -ENV MKL_SERVICE_FORCE_INTEL=1 - -# Launches the script -ENTRYPOINT python -m process $0 $@ +# Launch the script +ENTRYPOINT ["python3", "-m", "process"] diff --git a/README.md b/README.md index e68b7fb..edadebc 100644 --- a/README.md +++ b/README.md @@ -1,13 +1,39 @@ -# DMX Solution to HaNSeg Challenge +# DMX Solution to HaNSeg -The Head and Neck oragan-at-risk CT & MR segmentation challenge. Contribution to the Grand Challenge (MICCAI 2023) +## Overview -Challenge URL: **[HaN-Seg 2023 challenge](https://han-seg2023.grand-challenge.org/)** +This repository contains scripts and tools for building a Docker algorithm, performing prediction on a test dataset, and calculating DSC (Dice Similarity Coefficient). Below are the steps to execute each of these tasks. -This solution is based on: +## Prerequisites - - [ANTsPY](https://antspy.readthedocs.io/en/latest/) - - [nnUNetv2](https://github.com/MIC-DKFZ/nnUNet/) - - [Zhack47](https://github.com/Zhack47/HaNSeg-QuantIF) +Make sure you have the following installed: +- Docker +- Python 3.9 +- Necessary Python packages (can be installed using `requirements.txt` if provided) +## Steps +### 1. Build the Docker Algorithm + +To build the Docker algorithm, run the following command in your terminal: + +```sh +sh test.sh +``` + +### 2. Prediction on Test Dataset + +To perform predictions on the test dataset, execute the following command: +```sh +python3 process.py +``` + +### 3. DSC Calculation + +To calculate the Dice Similarity Coefficient (DSC), use the following command: +```sh +python3 calc_dice.py +``` + +## Contact +Email: jaehwanhan12@gmail.com \ No newline at end of file diff --git a/build.sh b/build.sh index b9c52ff..3b868fd 100755 --- a/build.sh +++ b/build.sh @@ -1,4 +1,4 @@ #!/usr/bin/env bash SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )" # docker build --no-cache -t hanseg2023algorithm "$SCRIPTPATH" -docker build -t hanseg2023algorithm_dmx "$SCRIPTPATH" +docker build -t hanseg2023algorithm_dmx:jhhan "$SCRIPTPATH" diff --git a/calc_dice.py b/calc_dice.py new file mode 100644 index 0000000..719ecd8 --- /dev/null +++ b/calc_dice.py @@ -0,0 +1,107 @@ +import os, sys +sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__)))) + +import numpy as np +import nrrd +import pandas as pd + +LABEL_dict = { + "background": 0, + "A_Carotid_L": 1, + "A_Carotid_R": 2, + "Arytenoid": 3, + "Bone_Mandible": 4, + "Brainstem": 5, + "BuccalMucosa": 6, + "Cavity_Oral": 7, + "Cochlea_L": 8, + "Cochlea_R": 9, + "Cricopharyngeus": 10, + "Esophagus_S": 11, + "Eye_AL": 12, + "Eye_AR": 13, + "Eye_PL": 14, + "Eye_PR": 15, + "Glnd_Lacrimal_L": 16, + "Glnd_Lacrimal_R": 17, + "Glnd_Submand_L": 18, + "Glnd_Submand_R": 19, + "Glnd_Thyroid": 20, + "Glottis": 21, + "Larynx_SG": 22, + "Lips": 23, + "OpticChiasm": 24, + "OpticNrv_L": 25, + "OpticNrv_R": 26, + "Parotid_L": 27, + "Parotid_R": 28, + "Pituitary": 29, + "SpinalCord": 30, +} + +def load_nrrd(file_path): + data, _ = nrrd.read(file_path) + return data + +def dice_score(y_true, y_pred): + eps = 1e-6 + y_true = y_true.flatten() + y_pred = y_pred.flatten() + intersection = np.sum(y_true * y_pred) + return (2. * intersection) / (np.sum(y_true) + np.sum(y_pred)) + eps + + +def calculate_dice_scores(result_folder, gt_folder): + data = [] + result_files = [f for f in os.listdir(result_folder) if f.endswith('.nrrd')] + + for result_file in result_files: + case_id = result_file.split('_IMG')[0] + gt_file = f"{case_id}_all_rois.seg.nrrd" + + result_path = os.path.join(result_folder, result_file) + gt_path = os.path.join(gt_folder, gt_file) + + if not os.path.exists(gt_path): + print(f"Ground truth file not found for {result_file}") + continue + + result_data = load_nrrd(result_path) + gt_data = load_nrrd(gt_path) + + case_dice_scores = {"file_name": result_file} + + for label, label_index in LABEL_dict.items(): + result_label = (result_data == label_index).astype(np.uint8) + gt_label = (gt_data == label_index).astype(np.uint8) + + if np.sum(gt_label) == 0: + case_dice_scores[label] = None + else: + score = dice_score(gt_label, result_label) + case_dice_scores[label] = score + + # Calculate total mean DICE score for this case + valid_scores = [score for score in case_dice_scores.values() if isinstance(score, (float, int))] + total_dice_score = np.mean(valid_scores) if valid_scores else 0.0 + case_dice_scores["total"] = total_dice_score + + data.append(case_dice_scores) + print(f"Processed {result_file}") + + return data + + +if __name__ == '__main__': + result_folder = '/output/images/head_neck_oar' + gt_folder = '/input/gt' + + data = calculate_dice_scores(result_folder, gt_folder) + + # Create a DataFrame and save to CSV + df = pd.DataFrame(data) + df = df[["file_name"] + list(LABEL_dict.keys()) + ["total"]] # Ensure columns are in the correct order + csv_path = "/output/dice_scores.csv" + df.to_csv(csv_path, index=False) + + print(f"CSV file saved to {csv_path}") diff --git a/export.sh b/export.sh deleted file mode 100755 index 54abcec..0000000 --- a/export.sh +++ /dev/null @@ -1,5 +0,0 @@ -#!/usr/bin/env bash - -./build.sh - -docker save hanseg2023algorithm_dmx | gzip -c > HanSeg2023AlgorithmDMX.tar.gz diff --git a/process.py b/process.py index bbabfdd..5295555 100644 --- a/process.py +++ b/process.py @@ -1,6 +1,8 @@ import time import SimpleITK as sitk import numpy as np +import os +os.environ['OPENBLAS_NUM_THREADS'] = '1' np.lib.index_tricks.int = np.uint16 import ants from os.path import join @@ -104,39 +106,41 @@ def predict(self, *, image_ct: ants.ANTsImage, image_mrt1: ants.ANTsImage) -> si del ct_image # Shamelessly copied from nnUNet/nnunetv2/preprocessing/resampling/default_resampling.py new_shape = np.array([int(round(i / j * k)) for i, j, k in zip(fin_spacing, spacing[::-1], fin_size)]) - if new_shape.prod()< 1e8: - print(f"Image is not too large ({new_shape.prod()}), using the folds (0,1,2,3,4) with mirror") - predictor = nnUNetPredictor(tile_step_size=0.4, use_mirroring=True, perform_everything_on_gpu=True, - verbose=True, verbose_preprocessing=True, - allow_tqdm=True) - predictor.initialize_from_trained_model_folder(trained_model_path, use_folds=(0,1,2,3), - checkpoint_name="checkpoint_best.pth") - # predictor.allowed_mirroring_axes = (0, 2) - elif new_shape.prod()< 1.3e8: - print(f"Image is not too large ({new_shape.prod()}), using the folds (0,1,2,3,4)") - - predictor = nnUNetPredictor(tile_step_size=0.6, use_mirroring=True, perform_everything_on_gpu=False, - verbose=True, verbose_preprocessing=True, - allow_tqdm=True) - predictor.initialize_from_trained_model_folder(trained_model_path, use_folds=(0,1,2,3), #(0,1,2,3,4) - checkpoint_name="checkpoint_best.pth") - elif new_shape.prod()< 1.7e8: - print(f"Image is not too large ({new_shape.prod()}), using the 'all' fold with mirror") - - predictor = nnUNetPredictor(tile_step_size=0.4, use_mirroring=True, perform_everything_on_gpu=False, - verbose=True, verbose_preprocessing=True, - allow_tqdm=True) - predictor.initialize_from_trained_model_folder(trained_model_path, use_folds="0", - checkpoint_name="checkpoint_best.pth") - # predictor.allowed_mirroring_axes = (0, 2) - - else: - predictor = nnUNetPredictor(tile_step_size=0.6, use_mirroring=True, perform_everything_on_gpu=False, - verbose=True, verbose_preprocessing=True, - allow_tqdm=True) - print(f"Image is too large ({new_shape.prod()}), using the 'all' fold") - predictor.initialize_from_trained_model_folder(trained_model_path, use_folds="0", - checkpoint_name="checkpoint_best.pth") + # if new_shape.prod()< 1e8: + # print(f"Image is not too large ({new_shape.prod()}), using the folds (0,1,2,3,4) with mirror") + # predictor = nnUNetPredictor(tile_step_size=0.4, use_mirroring=True, perform_everything_on_gpu=True, + # verbose=True, verbose_preprocessing=True, + # allow_tqdm=True) + # predictor.initialize_from_trained_model_folder(trained_model_path, use_folds=(0,1,2,3), + # checkpoint_name="checkpoint_best.pth") + # # predictor.allowed_mirroring_axes = (0, 2) + # elif new_shape.prod()< 1.3e8: + # print(f"Image is not too large ({new_shape.prod()}), using the folds (0,1,2,3,4)") + + # predictor = nnUNetPredictor(tile_step_size=0.6, use_mirroring=True, perform_everything_on_gpu=False, + # verbose=True, verbose_preprocessing=True, + # allow_tqdm=True) + # predictor.initialize_from_trained_model_folder(trained_model_path, use_folds=(0,1,2,3), #(0,1,2,3,4) + # checkpoint_name="checkpoint_best.pth") + # elif new_shape.prod()< 1.7e8: + # print(f"Image is not too large ({new_shape.prod()}), using the 'all' fold with mirror") + + # predictor = nnUNetPredictor(tile_step_size=0.4, use_mirroring=True, perform_everything_on_gpu=False, + # verbose=True, verbose_preprocessing=True, + # allow_tqdm=True) + # predictor.initialize_from_trained_model_folder(trained_model_path, use_folds="0", + # checkpoint_name="checkpoint_best.pth") + # # predictor.allowed_mirroring_axes = (0, 2) + + # else: + predictor = nnUNetPredictor(tile_step_size=0.5, use_mirroring=False, perform_everything_on_gpu=True, + verbose=True, verbose_preprocessing=True, + allow_tqdm=True) + print(f"Image is too large ({new_shape.prod()}), using the 'all' fold") + predictor.initialize_from_trained_model_folder(trained_model_path, use_folds="0", + checkpoint_name="checkpoint_best.pth") + # predictor.initialize_from_trained_model_folder(trained_model_path, use_folds=(0,1,2,3), #(0,1,2,3,4) + # checkpoint_name="checkpoint_best.pth") img_temp = predictor.predict_single_npy_array(images, properties, None, None, False).astype(np.uint8) del images diff --git a/requirements.txt b/requirements.txt index 1e174c7..49cef5c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -98,7 +98,7 @@ urllib3==1.26.15 # via requests wheel==0.40.0 # via pip-tools - +pynrrd # The following packages are considered to be unsafe in a requirements file: # pip # setuptools diff --git a/test.sh b/test.sh index aa2ece8..7d4393b 100755 --- a/test.sh +++ b/test.sh @@ -1,37 +1,37 @@ #!/usr/bin/env bash - +# 현재 디렉토리 서칭 SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )" -echo $SCRIPTPATH +echo $SCRIPTPATH OUTDIR=$SCRIPTPATH/output ./build.sh +# Ensure output directory exists and set permissions +mkdir -p $OUTDIR +chmod -R 777 $OUTDIR + # Maximum is currently 30g, configurable in your algorithm image settings on grand challenge -MEM_LIMIT="12g" +MEM_LIMIT="30g" # create output dir if it does not exist if [ ! -d $OUTDIR ]; then mkdir $OUTDIR; fi echo "starting docker" -# Do not change any of the parameters to docker run, these are fixed -docker run --rm \ - --memory="${MEM_LIMIT}" \ - --memory-swap="${MEM_LIMIT}" \ - --network="none" \ - --cap-drop="ALL" \ - --security-opt="no-new-privileges" \ - --shm-size="128m" \ - --pids-limit="256" \ - --gpus="0" \ - -v $SCRIPTPATH/test/:/input/ \ - -v $SCRIPTPATH/output/:/output \ - hanseg2023algorithm_dmx -echo "docker done" +docker run --rm --name hanseg_algorithm_container \ + --memory="${MEM_LIMIT}" \ + --memory-swap="${MEM_LIMIT}" \ + --network="none" \ + --cap-drop="ALL" \ + --security-opt="no-new-privileges" \ + --shm-size="4g" \ + --pids-limit="512" \ + --gpus="all" \ + -v $SCRIPTPATH/input/:/input/ \ + -v $SCRIPTPATH/output/:/output \ + -e LOCAL_USER_ID=$(id -u) \ + -e LOCAL_GROUP_ID=$(id -g) \ + -it --entrypoint bash \ + hanseg2023algorithm_dmx:jhhan -echo -echo -echo "Compare files in $OUTDIR with the expected results to see if test is successful" -docker run --rm \ - -v $OUTDIR:/output/ \ - python:3.8-slim ls -al /output/images/head_neck_oar +echo "docker done"