vessel segmentation

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

build.sh

@@ -0,0 +1,4 @@
+#!/usr/bin/env bash
+SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )"
+
+docker build -t vesselsegmentation "$SCRIPTPATH"

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

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

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

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": []}]