From f48bf52fd548417e4894895f78237120b10ad1b4 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Tue, 18 Jun 2024 01:17:03 +0530
Subject: [PATCH] update nnunet version and refactor env vars

---
 models/bamf_nnunet_mr_breast/dockerfiles/Dockerfile  | 6 +++---
 models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/models/bamf_nnunet_mr_breast/dockerfiles/Dockerfile b/models/bamf_nnunet_mr_breast/dockerfiles/Dockerfile
index 23e605b5..cd8a0b7b 100644
--- a/models/bamf_nnunet_mr_breast/dockerfiles/Dockerfile
+++ b/models/bamf_nnunet_mr_breast/dockerfiles/Dockerfile
@@ -8,7 +8,7 @@ FROM mhubai/base:latest
 ENV SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL=True
 
 # Install nnunet and platipy
-RUN pip3 install --no-cache-dir nnunet==1.7.1
+RUN pip3 install --no-cache-dir nnunetv2
 
 # Clone the main branch of MHubAI/models
 ARG MHUB_MODELS_REPO
@@ -18,7 +18,7 @@ RUN buildutils/import_mhub_model.sh bamf_nnunet_mr_breast ${MHUB_MODELS_REPO}
 ENV WEIGHTS_DIR=/root/.nnunet/nnUNet_models/nnUNet/
 RUN mkdir -p $WEIGHTS_DIR
 ENV WEIGHTS_FN=Dataset009_Breast.zip
-ENV WEIGHTS_URL=https://zenodo.org/record/11503782/files/$WEIGHTS_FN
+ENV WEIGHTS_URL=https://zenodo.org/record/11998679/files/$WEIGHTS_FN
 RUN wget --directory-prefix ${WEIGHTS_DIR} ${WEIGHTS_URL}
 RUN unzip ${WEIGHTS_DIR}${WEIGHTS_FN} -d ${WEIGHTS_DIR}
 RUN rm ${WEIGHTS_DIR}${WEIGHTS_FN}
@@ -26,7 +26,7 @@ RUN rm ${WEIGHTS_DIR}${WEIGHTS_FN}
 # Pull nnUNet model weights into the container for Task775_CT_NSCLC_RG
 ENV TASK_NAME_NSCLC_RG=Task775_CT_NSCLC_RG
 ENV WEIGHTS_FN=Dataset011_Breast.zip
-ENV WEIGHTS_URL=https://zenodo.org/record/11504675/files/$WEIGHTS_FN
+ENV WEIGHTS_URL=https://zenodo.org/record/11998632/files/$WEIGHTS_FN
 RUN wget --directory-prefix ${WEIGHTS_DIR} ${WEIGHTS_URL}
 RUN unzip ${WEIGHTS_DIR}${WEIGHTS_FN} -d ${WEIGHTS_DIR}
 RUN rm ${WEIGHTS_DIR}${WEIGHTS_FN}
diff --git a/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
index 558daa60..e0ea53b4 100644
--- a/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
+++ b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
@@ -153,7 +153,7 @@ def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData
 
         # define output folder (temp dir) and also override environment variable for nnunet
         out_dir = self.config.data.requestTempDir(label="nnunet-model-out")
-        os.environ['nnUNet_results'] = out_dir
+        os.environ['nnUNet_results'] = os.environ['WEIGHTS_FOLDER'] # set model weights path
 
         # symlink nnunet input folder to the input data with python
         # create symlink in python