diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..01a602e --- /dev/null +++ b/Dockerfile @@ -0,0 +1,30 @@ +FROM pytorch/pytorch + + +RUN groupadd -r algorithm && useradd -m --no-log-init -r -g algorithm algorithm + +RUN mkdir -p /opt/algorithm /input /output \ + && chown algorithm:algorithm /opt/algorithm /input /output + +USER algorithm + +WORKDIR /opt/algorithm + +ENV PATH="/home/algorithm/.local/bin:${PATH}" + +RUN python -m pip install --user -U pip + + + +COPY --chown=algorithm:algorithm requirements.txt /opt/algorithm/ +RUN python -m pip install --user -r requirements.txt + +COPY --chown=algorithm:algorithm process.py /opt/algorithm/ +COPY --chown=algorithm:algorithm best_metric_model_segmentation2d_dict.pth /opt/algorithm/ + +ENTRYPOINT python -m process $0 $@ + +## ALGORITHM LABELS ## + +# These labels are required +LABEL nl.diagnijmegen.rse.algorithm.name=VesselSegmentation diff --git a/best_metric_model_segmentation2d_dict.pth b/best_metric_model_segmentation2d_dict.pth new file mode 100644 index 0000000..554badc Binary files /dev/null and b/best_metric_model_segmentation2d_dict.pth differ diff --git a/build.sh b/build.sh new file mode 100755 index 0000000..8a579a9 --- /dev/null +++ b/build.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )" + +docker build -t vesselsegmentation "$SCRIPTPATH" diff --git a/process.py b/process.py new file mode 100644 index 0000000..6a43f5d --- /dev/null +++ b/process.py @@ -0,0 +1,83 @@ + +import SimpleITK +import numpy as np +import torch +import monai +from scipy.special import expit +from skimage import transform + +from evalutils import SegmentationAlgorithm +from evalutils.validators import ( + UniquePathIndicesValidator, + UniqueImagesValidator, +) + + +class Vesselsegmentation(SegmentationAlgorithm): + def __init__(self): + super().__init__( + validators=dict( + input_image=( + UniqueImagesValidator(), + UniquePathIndicesValidator(), + ) + ), + ) + + # use GPU if available, otherwise use the CPU + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + print("==> Using ", self.device) + print("==> Initializing model") + + # define the U-Net with monai + self.model = monai.networks.nets.UNet( + dimensions=2, + in_channels=3, + out_channels=1, + channels=(16, 32, 64, 128, 256), + strides=(2, 2, 2, 2), + num_res_units=2, + ).to(self.device) + print(f'==> Model created') + + res = self.model.eval() # set to inference mode + print(f'==> Model evaluated') + + self.model.load_state_dict( + torch.load( + "./best_metric_model_segmentation2d_dict.pth", + map_location=self.device, + ) + ) + + print("==> Weights loaded") + + def predict(self, *, input_image: SimpleITK.Image) -> SimpleITK.Image: + + image = SimpleITK.GetArrayFromImage(input_image) + image = np.array(image) + shape = image.shape + + # Pre-process the image + image = transform.resize(image, (512, 512), order=3) # resize all images to 512 x 512 in shape + image = image.astype(np.float32) / 255. # normalize + image = image.transpose((2, 0, 1)) # flip the axes and bring color to the first channel + image = torch.from_numpy(image).to(self.device).reshape(1, 3, 512, 512) + + # Do the forward pass + out = self.model(image).squeeze().data.cpu().numpy() + + # Post-process the image + out = transform.resize(out, shape[:-1], order=3) + out = (expit(out) > 0.99) # apply the sigmoid filter and binarize the predictions + out = (out * 255).astype(np.uint8) + out = SimpleITK.GetImageFromArray(out) # convert numpy array to SimpleITK image for grand-challenge.org + + print("==> Forward pass done") + + return out + + +if __name__ == "__main__": + Vesselsegmentation().process() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..2fb54c7 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,6 @@ + +evalutils==0.3.0 +monai==0.4.0 +scikit-learn==1.0 +scipy==1.6.3 +scikit-image==0.18.1 diff --git a/test.sh b/test.sh new file mode 100755 index 0000000..3ff95ee --- /dev/null +++ b/test.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash + +SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )" + +./build.sh + +VOLUME_SUFFIX=$(cat /dev/urandom | tr -dc 'a-zA-Z0-9' | fold -w 8 | head -n 1) + +docker volume create vesselsegmentation-output-$VOLUME_SUFFIX + +# run the forward pass and store the outputs in a temporary Docker volume +docker run --rm \ + --gpus=all \ + -v $SCRIPTPATH/test/input/:/input/ \ + -v vesselsegmentation-output-$VOLUME_SUFFIX:/output/ \ + vesselsegmentation + +# compare the outputs in the Docker volume with the outputs in ./test/expected_output/ +docker run --rm \ + -v vesselsegmentation-output-$VOLUME_SUFFIX:/output/ \ + -v $SCRIPTPATH/test/expected_output/:/expected_output/ \ + biocontainers/simpleitk:v1.0.1-3-deb-py3_cv1 python3 -c """ +import SimpleITK as sitk +import os + +print(os.listdir('/output/images/')) +output = sitk.ReadImage('/output/images/01_test.tif') +expected_output = sitk.ReadImage('/expected_output/images/01_test.tif') + +label_filter = sitk.LabelOverlapMeasuresImageFilter() +label_filter.Execute(output, expected_output) +dice_score = label_filter.GetDiceCoefficient() + +if dice_score == 1.0: + print('Test passed!') +else: + print('Test failed!') +""" + +docker volume rm vesselsegmentation-output-$VOLUME_SUFFIX diff --git a/test/expected_output/images/01_test.tif b/test/expected_output/images/01_test.tif new file mode 100644 index 0000000..e7cb719 Binary files /dev/null and b/test/expected_output/images/01_test.tif differ diff --git a/test/expected_output/results.json b/test/expected_output/results.json new file mode 100644 index 0000000..e23fe17 --- /dev/null +++ b/test/expected_output/results.json @@ -0,0 +1 @@ +[{"outputs": [{"type": "metaio_image", "filename": "01_test.tif"}], "inputs": [{"type": "metaio_image", "filename": "01_test.tif"}], "error_messages": []}] \ No newline at end of file diff --git a/test/input/01_test.tif b/test/input/01_test.tif new file mode 100644 index 0000000..bd3e2dc Binary files /dev/null and b/test/input/01_test.tif differ