Merge pull request #62 from MHubAI/patch-models

Patch models
MHubAI · Nov 13, 2023 · 51aad4d · 51aad4d
2 parents 204abd8 + 426e065
commit 51aad4d
Show file tree

Hide file tree

Showing 10 changed files with 70 additions and 47 deletions.
diff --git a/base/utils/download_weights.sh → base/buildutils/download_weights.sh b/base/utils/download_weights.sh → base/buildutils/download_weights.sh
diff --git a/base/buildutils/import_mhub_model.sh b/base/buildutils/import_mhub_model.sh
@@ -0,0 +1,42 @@
+#!/bin/bash
+
+# Script to import the MHub model definition from GitHub.
+# provide the name of the model as a parameter.
+# Usage: utils/import_mhub_model.sh <model_name> <(repo_url=https://github.com/MHubAI/models.git::main)> <(branch=main)>
+
+# parameters extraction
+MODEL_NAME=$1
+REPO_AND_BRANCH=${2:-https://github.com/MHubAI/models.git::main}
+REPO_URL=$(echo $REPO_AND_BRANCH | awk -F :: '{print $1}')
+REPO_BRANCH=$(echo $REPO_AND_BRANCH | awk -F :: '{print $2}')
+REPO_BRANCH=${REPO_BRANCH:-$3}
+REPO_BRANCH=${REPO_BRANCH:-main}
+
+# printout paramaters (this happens during the docker build...)
+echo "Importing model definition from MHub models repository."
+echo "├── MODEL NAME ..... ${MODEL_NAME}"
+echo "├── REPOSITORY ..... ${REPO_URL}"
+echo "└── BRANCH ......... ${REPO_BRANCH}"
+echo 
+
+# fail if model name is empty
+if [ -z "$MODEL_NAME" ]; then
+    echo "Error: no model name provided."
+    exit 1
+fi
+
+# print a warning that the model definition is not from the 
+# the official MHub Models repository and therefore only 
+# suitable for development
+if [ "$REPO_URL@$REPO_BRANCH" != "https://github.com/MHubAI/models.git@main" ]; then
+    echo 
+    echo "Warning: the model definition is not from the official MHub Models repository and therefore only suitable for development."
+    echo 
+fi
+
+# perform a sparse checkout of the model definition folder 
+# (models/<model_name>) from the referenced repository and branch
+git stash
+git fetch ${REPO_URL} ${REPO_BRANCH}
+git merge FETCH_HEAD
+git sparse-checkout set "models/${MODEL_NAME}"
diff --git a/base/dockerfiles/Dockerfile b/base/dockerfiles/Dockerfile
@@ -58,10 +58,10 @@ RUN pip3 install --upgrade pip && pip3 install --no-cache-dir \
 RUN pip3 install git+https://github.com/MHubAI/mhubio.git \
   && pip3 install git+https://github.com/MHubAI/segdb.git \
   && git init \
-  && git sparse-checkout set "base/utils" "base/bin" \
+  && git sparse-checkout set "base/buildutils" "base/bin" \
   && git fetch https://github.com/MHubAI/models.git main \
   && git merge FETCH_HEAD \
-  && mv base/utils . \
+  && mv base/buildutils . \
   && chmod +x base/bin/* \
   && cp base/bin/* /usr/bin/ \
   && rm -r base

diff --git a/models/casust/dockerfiles/Dockerfile b/models/casust/dockerfiles/Dockerfile
@@ -14,16 +14,10 @@ RUN pip3 install --no-cache-dir \
   torchvision \
   torchio
 
-# Clone the main branch of MHubAI/models
-RUN git stash \
- && git sparse-checkout set "models/casust" \
- && git fetch https://github.com/MHubAI/models.git main \
- && git merge FETCH_HEAD
-
 # Clone the casust model
 RUN git clone https://github.com/LennyN95/CaSuSt /app/models/casust/src
 
-# pull weights for casust so that the user doesn't need to every time a container is run
+# Pull weights for casust so that the user doesn't need to every time a container is run
 ENV WEIGHTS_DIR="/app/models/casust/src/weights/"
 ENV WEIGHTS_URL="https://zenodo.org/record/7836696/files/casust_weights_7roi.zip"
 ENV WEIGHTS_FN="casust_weights_7roi.zip"
@@ -32,7 +26,7 @@ RUN wget --directory-prefix ${WEIGHTS_DIR} ${WEIGHTS_URL}
 RUN unzip ${WEIGHTS_DIR}${WEIGHTS_FN} -d ${WEIGHTS_DIR}
 RUN rm ${WEIGHTS_DIR}${WEIGHTS_FN}
 
-# pull weights for platipy's nnU-Net so that the user doesn't need to every time a container is run
+# Pull weights for platipy's nnU-Net so that the user doesn't need to every time a container is run
 ENV WEIGHTS_DIR="/root/.platipy/nnUNet_models/nnUNet/"
 ENV WEIGHTS_URL="https://zenodo.org/record/6585664/files/Task400_OPEN_HEART_3d_lowres.zip"
 ENV WEIGHTS_FN="Task400_OPEN_HEART_3d_lowres.zip"
@@ -41,9 +35,13 @@ RUN wget --directory-prefix ${WEIGHTS_DIR} ${WEIGHTS_URL}
 RUN unzip ${WEIGHTS_DIR}${WEIGHTS_FN} -d ${WEIGHTS_DIR}
 RUN rm ${WEIGHTS_DIR}${WEIGHTS_FN}
 
-# specify nnunet specific environment variables
+# Specify nnunet specific environment variables
 ENV WEIGHTS_FOLDER=$WEIGHTS_DIR
 
+# Import the MHub model definiton
+ARG MHUB_MODELS_REPO
+RUN buildutils/import_mhub_model.sh casust ${MHUB_MODELS_REPO}
+
 # Default run script
 ENTRYPOINT ["python3", "-m", "mhubio.run"]
 CMD ["--config", "/app/models/casust/config/default.yml"]
diff --git a/models/gc_lunglobes/dockerfiles/Dockerfile b/models/gc_lunglobes/dockerfiles/Dockerfile
@@ -15,13 +15,9 @@ RUN pip3 install --no-cache-dir \
 # SimpleITK downgrade required for legacy Resample::Execute operation
 RUN pip3 install --no-cache-dir --force-reinstall SimpleITK==1.2.4 
 
-# FIXME: temporary fix waiting for the latest base image update
-# Clone the main branch of MHubAI/models
-RUN git stash \
-    && git fetch https://github.com/MHubAI/models.git main \
-    && git merge FETCH_HEAD \
-    && git sparse-checkout set "models/gc_lunglobes" \
-    && git fetch https://github.com/MHubAI/models.git main
+# Import the MHub model definiton
+ARG MHUB_MODELS_REPO
+RUN buildutils/import_mhub_model.sh gc_lunglobes ${MHUB_MODELS_REPO}
 
 # Install Xie's pulmonary lobe segmentation algorithm and model weights
 RUN git clone https://github.com/DIAGNijmegen/bodyct-pulmonary-lobe-segmentation.git src && \

diff --git a/models/lungmask/dockerfiles/Dockerfile b/models/lungmask/dockerfiles/Dockerfile
@@ -4,14 +4,9 @@ FROM mhubai/base:latest
 RUN pip3 install --no-cache-dir \
     lungmask==0.2.16
 
-# FIXME: temporary fix waiting for the latest base image update
-# Clone the main branch of MHubAI/models
-RUN git stash \
-  && git fetch https://github.com/MHubAI/models.git main \
-  && git merge FETCH_HEAD \
-  && git sparse-checkout set "models/lungmask" \
-  && git fetch https://github.com/MHubAI/models.git main \
-  && git merge FETCH_HEAD
+# Import the MHub model definiton
+ARG MHUB_MODELS_REPO
+RUN buildutils/import_mhub_model.sh lungmask ${MHUB_MODELS_REPO}
 
 # pull the weights for the lung segmentation 2D U-Net model
 ENV WEIGHTS_DIR="/root/.cache/torch/hub/checkpoints/"

diff --git a/models/nnunet_liver/dockerfiles/Dockerfile b/models/nnunet_liver/dockerfiles/Dockerfile
@@ -11,11 +11,9 @@ ENV SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL=True
 RUN pip3 install --no-cache-dir \
   nnunet==1.7.1 
 
-# Clone the main branch of MHubAI/models
-RUN git stash \
- && git sparse-checkout set "models/nnunet_liver" \
- && git fetch https://github.com/MHubAI/models.git main \
- && git merge FETCH_HEAD
+# Import the MHub model definiton
+ARG MHUB_MODELS_REPO
+RUN buildutils/import_mhub_model.sh nnunet_liver ${MHUB_MODELS_REPO}
 
 # pull weights for platipy's nnU-Net so that the user doesn't need to every time a container is run
 ENV WEIGHTS_DIR="/root/.nnunet/nnUNet_models/nnUNet/"

diff --git a/models/nnunet_pancreas/dockerfiles/Dockerfile b/models/nnunet_pancreas/dockerfiles/Dockerfile
@@ -11,11 +11,9 @@ ENV SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL=True
 RUN pip3 install --no-cache-dir \
   nnunet==1.7.1 
 
-# Clone the main branch of MHubAI/models
-RUN git stash \
- && git sparse-checkout set "models/nnunet_pancreas" \
- && git fetch https://github.com/MHubAI/models.git main \
- && git merge FETCH_HEAD
+# Import the MHub model definiton
+ARG MHUB_MODELS_REPO
+RUN buildutils/import_mhub_model.sh nnunet_pancreas ${MHUB_MODELS_REPO}
 
 # pull weights for platipy's nnU-Net so that the user doesn't need to every time a container is run
 ENV WEIGHTS_DIR="/root/.nnunet/nnUNet_models/nnUNet/"

diff --git a/models/platipy/dockerfiles/Dockerfile b/models/platipy/dockerfiles/Dockerfile
@@ -16,11 +16,9 @@ RUN pip3 install --no-cache-dir \
   nnunet \
   platipy[cardiac]==0.7.0
 
-# Clone the main branch of MHubAI/models
-RUN git stash \
- && git sparse-checkout set "models/platipy" \
- && git fetch https://github.com/MHubAI/models.git main \
- && git merge FETCH_HEAD
+# Import the MHub model definiton
+ARG MHUB_MODELS_REPO
+RUN buildutils/import_mhub_model.sh platipy ${MHUB_MODELS_REPO}
 
 # pull weights for platipy's nnU-Net so that the user doesn't need to every time a container is run
 ENV WEIGHTS_DIR="/root/.platipy/nnUNet_models/nnUNet/"

diff --git a/models/totalsegmentator/dockerfiles/Dockerfile b/models/totalsegmentator/dockerfiles/Dockerfile
@@ -10,15 +10,13 @@ ENV SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL=True
 # Install TotalSegmentator
 RUN pip3 install --no-cache-dir totalsegmentator==1.5.6 
 
-# Clone the main branch of MHubAI/models
-RUN git stash \
- && git sparse-checkout set "models/totalsegmentator" \
- && git fetch https://github.com/MHubAI/models.git main \
- && git merge FETCH_HEAD
+# Import the MHub model definiton
+ARG MHUB_MODELS_REPO
+RUN buildutils/import_mhub_model.sh totalsegmentator ${MHUB_MODELS_REPO}
 
 # Download weights using the script in utils
 # Usage: utils/download_weights.sh <weights_dir> <url1> <url2> ...
-RUN utils/download_weights.sh \
+RUN buildutils/download_weights.sh \
   /root/.totalsegmentator/nnunet/results/nnUNet/3d_fullres/ \
   https://zenodo.org/record/6802052/files/Task256_TotalSegmentator_3mm_1139subj.zip \
   https://zenodo.org/record/6802342/files/Task251_TotalSegmentator_part1_organs_1139subj.zip \