From cf81190f9fd557b7f7f8dd64d450e420a957e57e Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 6 Jun 2024 22:14:37 +0530
Subject: [PATCH 01/27] [FEAT]: breast, fgt, and tumor seg module

---
 .../bamf_nnunet_mr_breast/config/default.yml  |  44 +++++++
 .../dockerfiles/Dockerfile                    |  39 ++++++
 .../utils/BreastPostProcessor.py              | 115 ++++++++++++++++++
 3 files changed, 198 insertions(+)
 create mode 100644 models/bamf_nnunet_mr_breast/config/default.yml
 create mode 100644 models/bamf_nnunet_mr_breast/dockerfiles/Dockerfile
 create mode 100644 models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
new file mode 100644
index 00000000..75a344b4
--- /dev/null
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -0,0 +1,44 @@
+general:
+  data_base_dir: /app/data
+  version: 1.0
+  description: Default configuration for Bamf NNUnet Breast and tumor segmentation on MR scans (dicom to dicom)
+
+execute:
+- DicomImporter
+- NiftiConverter
+- module: NNUnetRunner
+  nnunet_task: Dataset009_Breast
+  roi: BREAST,BREAST+FGT
+- module: NNUnetRunner
+  nnunet_task: Dataset011_Breast
+  roi: BREAST_TUMOR
+- BreastPostProcessor
+- DsegConverter
+- DataOrganizer
+
+modules:
+  DicomImporter:
+    source_dir: input_data
+    import_dir: sorted_data
+    sort_data: true
+    meta: 
+      mod: '%Modality'
+
+  NNUnetRunner:
+    in_data: nifti:mod=mr
+    nnunet_model: 3d_fullres
+    export_prob_maps: False
+
+  BreastPostProcessor:
+    in_breast_data: nifti:mod=seg:nnunet_task=Dataset009_Breast
+    in_tumor_data: nifti:mod=seg:nnunet_task=Dataset011_Breast
+
+  DsegConverter:
+    model_name: Bamf NNUnet MR Breast and tumor
+    target_dicom: dicom:mod=mr
+    source_segs: nifti:mod=seg:processor=bamf
+    skip_empty_slices: True
+
+  DataOrganizer:
+    targets:
+    - dicomseg-->[i:sid]/bamf_nnunet_mr_breast.seg.dcm
diff --git a/models/bamf_nnunet_mr_breast/dockerfiles/Dockerfile b/models/bamf_nnunet_mr_breast/dockerfiles/Dockerfile
new file mode 100644
index 00000000..23e605b5
--- /dev/null
+++ b/models/bamf_nnunet_mr_breast/dockerfiles/Dockerfile
@@ -0,0 +1,39 @@
+FROM mhubai/base:latest
+
+# FIXME: set this environment variable as a shortcut to avoid nnunet crashing the build
+# by pulling sklearn instead of scikit-learn
+# N.B. this is a known issue:
+# https://github.com/MIC-DKFZ/nnUNet/issues/1281 
+# https://github.com/MIC-DKFZ/nnUNet/pull/1209
+ENV SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL=True
+
+# Install nnunet and platipy
+RUN pip3 install --no-cache-dir nnunet==1.7.1
+
+# Clone the main branch of MHubAI/models
+ARG MHUB_MODELS_REPO
+RUN buildutils/import_mhub_model.sh bamf_nnunet_mr_breast ${MHUB_MODELS_REPO}
+
+# Pull nnUNet model weights into the container for Dataset009_Breast
+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
+RUN wget --directory-prefix ${WEIGHTS_DIR} ${WEIGHTS_URL}
+RUN unzip ${WEIGHTS_DIR}${WEIGHTS_FN} -d ${WEIGHTS_DIR}
+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
+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
+ENV WEIGHTS_FOLDER=$WEIGHTS_DIR
+
+# Default run script
+ENTRYPOINT ["mhub.run"]
+CMD ["--config", "/app/models/bamf_nnunet_mr_breast/config/default.yml"]
diff --git a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
new file mode 100644
index 00000000..b963b7cc
--- /dev/null
+++ b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
@@ -0,0 +1,115 @@
+"""
+-------------------------------------------------
+MHub - Run Module for ensembling nnUNet inference.
+-------------------------------------------------
+-------------------------------------------------
+Author: Rahul Soni
+Email:  rahul.soni@bamfhealth.com
+-------------------------------------------------
+"""
+
+from mhubio.core import IO
+from mhubio.core import Module, Instance, InstanceData
+import os, shutil
+import SimpleITK as sitk
+import numpy as np
+from skimage import measure, filters
+
+
+class BreastPostProcessor(Module):
+
+    def get_mask(self, ip_path):
+        # get segmentation mask
+        seg_data = sitk.GetArrayFromImage(sitk.ReadImage(ip_path))
+        seg_data[seg_data > 0] = 1
+        return seg_data
+
+    def n_connected(self, img_data):
+        img_data_mask = np.zeros(img_data.shape)
+        img_data_mask[img_data >= 1] = 1
+        img_filtered = np.zeros(img_data_mask.shape)
+        blobs_labels = measure.label(img_data_mask, background=0)
+        lbl, counts = np.unique(blobs_labels, return_counts=True)
+        lbl_dict = {}
+        for i, j in zip(lbl, counts):
+            lbl_dict[i] = j
+        sorted_dict = dict(sorted(lbl_dict.items(), key=lambda x: x[1], reverse=True))
+        count = 0
+        for key, value in sorted_dict.items():
+            if count >= 1 and count <= 2:
+                print(key, value)
+                img_filtered[blobs_labels == key] = 1
+            count += 1
+
+        img_data[img_filtered != 1] = 0
+        return img_data
+
+    def n4_bias_correction(
+        input_image,
+        shrink_factor=1,
+        mask_image=None,
+        number_of_iterations=50,
+        number_of_fitting_levels=4,
+    ):
+        if mask_image is None:
+            mask_image = sitk.OtsuThreshold(input_image, 0, 1, 200)
+
+        if shrink_factor > 1:
+            input_image = sitk.Shrink(
+                input_image, [shrink_factor] * input_image.GetDimension()
+            )
+            mask_image = sitk.Shrink(
+                mask_image, [shrink_factor] * mask_image.GetDimension()
+            )
+
+        corrector = sitk.N4BiasFieldCorrectionImageFilter()
+        corrector.SetMaximumNumberOfIterations(
+            [number_of_iterations] * number_of_fitting_levels
+        )
+
+        corrector.Execute(input_image, mask_image)
+        log_bias_field = corrector.GetLogBiasFieldAsImage(input_image)
+        corrected_image_full_resolution = input_image / sitk.Exp(log_bias_field)
+        return corrected_image_full_resolution
+
+    def get_seg_img(self, breast_and_fgt, breast_tumor, mr_path):
+        seg_data = np.zeros(breast_and_fgt.shape)
+        seg_data[breast_and_fgt == 1] = 1
+        seg_data[breast_tumor == 1] = 2
+        ref = sitk.ReadImage(mr_path)
+        seg_img = sitk.GetImageFromArray(seg_data)
+        seg_img.CopyInformation(ref)
+        return seg_img
+
+
+    @IO.Instance()
+    @IO.Input('in_breast_and_fgt_data', 'nifti:mod=seg:nnunet_task=Dataset009_Breast', the='input data from breast and fgt segmentation')
+    @IO.Input('in_breast_tumor_data', 'nifti:mod=seg:nnunet_task=Dataset011_Breast', the='input data from breast tumor segmentation')
+    @IO.Input('in_mr_data', 'nifti:mod=mr', the='input mr data')
+    @IO.Output('out_data', 'bamf_processed.nii.gz', 'nifti:mod=seg:processor=bamf:roi=BREAST,BREAST+FGT,BREAST_TUMOR',
+               the="get breast_fgt and breast segmentation file")
+    def task(self, instance: Instance, in_breast_and_fgt_data: InstanceData, in_breast_tumor_data: InstanceData,
+             in_mr_data: InstanceData, out_data: InstanceData):
+
+        print('input breast and fgt segmentation: ' + in_breast_and_fgt_data.abspath)
+        print('input breast tumor segmentation: ' + in_breast_tumor_data.abspath)
+        print('output path: ' + out_data.abspath)
+
+        breast_fgt_seg = self.get_mask(in_breast_and_fgt_data.abspath)
+        breast_fgt_seg = self.n_connected(breast_fgt_seg)
+
+        breast_tumor_seg = self.get_mask(in_breast_tumor_data.abspath)
+        breast_tumor_seg[breast_fgt_seg == 0] = 0
+
+        seg_img = self.get_seg_img(
+            breast_and_fgt = np.copy(breast_fgt_seg),
+            breast_tumor = np.copy(breast_tumor_seg),
+            mr_path = in_mr_data.abspath
+            )
+        process_dir = self.config.data.requestTempDir(label="nnunet-breast-processor")
+        process_file = os.path.join(process_dir, f'final.nii.gz')
+        sitk.WriteImage(
+            seg_img,
+            process_file,
+        )
+        shutil.copyfile(process_file, out_data.abspath)

From b972102f0e4948d0b4b34175eb8ea3fe5082f13e Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Fri, 7 Jun 2024 11:35:11 +0530
Subject: [PATCH 02/27] added meta data and refactored post processor

---
 .../bamf_nnunet_mr_breast/config/default.yml  |   2 +-
 models/bamf_nnunet_mr_breast/meta.json        | 135 ++++++++++++++++++
 .../utils/BreastPostProcessor.py              |  72 ++++------
 3 files changed, 164 insertions(+), 45 deletions(-)
 create mode 100644 models/bamf_nnunet_mr_breast/meta.json

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index 75a344b4..38d09282 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -8,7 +8,7 @@ execute:
 - NiftiConverter
 - module: NNUnetRunner
   nnunet_task: Dataset009_Breast
-  roi: BREAST,BREAST+FGT
+  roi: BREAST,FGT
 - module: NNUnetRunner
   nnunet_task: Dataset011_Breast
   roi: BREAST_TUMOR
diff --git a/models/bamf_nnunet_mr_breast/meta.json b/models/bamf_nnunet_mr_breast/meta.json
new file mode 100644
index 00000000..b595f5ee
--- /dev/null
+++ b/models/bamf_nnunet_mr_breast/meta.json
@@ -0,0 +1,135 @@
+{
+    "id": "",
+    "name": "bamf_nnunet_mr_breast",
+    "title": "BAMF MR Breast Segmentation (nnU-Net)",
+    "summary": {
+      "description": "An nnU-Net based model to segment breast from MR scans",
+      "inputs": [
+        {
+          "label": "Input Image",
+          "description": "The MR scan of a patient.",
+          "format": "DICOM",
+          "modality": "MR",
+          "bodypartexamined": "BREAST,FGT,BREAST_TUMOR",
+          "slicethickness": "2.5mm",
+          "non-contrast": true,
+          "contrast": false
+        }
+      ],
+      "outputs": [
+        {
+          "label": "Segmentation",
+          "type": "Segmentation",
+          "description": "Segmentation breast, fgt, and tumor",
+          "classes": [
+            "BREAST",
+            "FGT",
+            "BREAST_TUMOR"
+          ]
+        }
+      ],
+      "model": {
+        "architecture": "U-net v2",
+        "training": "supervised",
+        "cmpapproach": "3D"
+      },
+      "data": {
+        "training": {
+          "vol_samples": "TODO. Checked the training dataset and there were 98 + 489 = 587 samples"
+        },
+        "evaluation": {
+          "vol_samples": "TODO"
+        },
+        "public": true,
+        "external": true
+      }
+    },
+    "details": {
+      "name": "AIMI Breast MR",
+      "version": "1.0.0",
+      "devteam": "BAMF Health",
+      "authors": [
+        "Murugesan, Gowtham Krishnan",
+        "Soni, Rahul",
+        "McCrumb, Diana",
+        "Van Oss, Jeff"
+      ],
+      "type": "nnU-Net v2 (U-Net structure, optimized by data-driven heuristics)",
+      "date": {
+        "code": "07.06.2024",
+        "weights": "06.06.2024",
+        "pub": "07.06.2024"
+      },
+      "cite": "Murugesan, Gowtham Krishnan, Diana McCrumb, Mariam Aboian, Tej Verma, Rahul Soni, Fatima Memon, and Jeff Van Oss. The AIMI Initiative: AI-Generated Annotations for Imaging Data Commons Collections. arXiv preprint arXiv:2310.14897 (2023).",
+      "license": {
+        "code": "MIT",
+        "weights": "CC BY-NC 4.0"
+      },
+      "publications": [
+        {
+          "title": "The AIMI Initiative: AI-Generated Annotations in IDC Collections",
+          "uri": "https://arxiv.org/abs/2310.14897"
+        }
+      ],
+      "github": "https://github.com/bamf-health/aimi-breast-mr"
+    },
+    "info": {
+      "use": {
+        "title": "Intended Use",
+        "text": "This model is intended to perform breast, fgt, and breast tumor segmentation in MR scans. The model has been trained and tested on scans aquired during clinical care of patients, so it might not be suited for a healthy population. The generalization capabilities of the model on a range of ages, genders, and ethnicities are unknown."
+      },
+      "analyses": {
+        "title": "Quantitative Analyses",
+        "text": "The model's performance was assessed using the Dice Coefficient and Normalized Surface Distance (NSD) with tolerance 7mm, as specified in the Duke-Breast-Cancer-MRI. The model was used to segment cases from the IDC collection ACRIN-Contralateral-Breast-MR [1]. Nine of those cases were reviewed and corrected by a board-certified radiologist and a non-expert. The analysis is published here [2]",
+        "tables": [
+          {
+            "label": "Label-wise metrics (mean (standard deviation)) between AI derived and manually corrected MR Breast annotations",
+            "entries": {
+              "Dice: Radiologist": "0.97 (0.07)",
+              "NSD: Radiologist": "0.96 (0.09)",
+              "Dice: Non-expert": "0.97 (0.08)",
+              "NSD: Non-expert": "0.95 (0.11)"
+            }
+          }
+        ],
+        "references": [
+          {
+            "label": "ACRIN-Contralateral-Breast-MR",
+            "uri": "https://www.cancerimagingarchive.net/collection/acrin-contralateral-breast-mr"
+          },
+          {
+            "label": "The AIMI Initiative: AI-Generated Annotations for Imaging Data Commons Collections",
+            "uri": "https://arxiv.org/abs/2310.14897"
+          }
+        ]
+      },
+      "evaluation": {
+        "title": "Evaluation Data",
+        "text": "The model was used to segment cases 89 from the ACRIN-Contralateral-Breast-MR [1]. Nine of those cases were randomly selected to be reviewed and corrected by a board-certified radiologist. The model predictions, and radiologist corrections are published on zenodo [3]",
+        "references": [
+          {
+            "label": "ACRIN-Contralateral-Breast-MR",
+            "uri": "https://www.cancerimagingarchive.net/collection/acrin-contralateral-breast-mr"
+          },
+          {
+            "label": "Image segmentations produced by the AIMI Annotations initiative",
+            "uri": "https://zenodo.org/records/10009368"
+          }
+        ]
+      },
+      "training": {
+        "title": "Training Data",
+        "text": "The training dataset consists of 587 breast, fgt, and breast tumor annotations. The breast and fgt annotations where taken from the Duke-Breast-Cancer-MRI dataset [1] (N=98) and breast tumor annotations where taken from the ACRIN-Contralateral-Breast-MR [2] (N=489).",
+        "references": [
+          {
+            "label": "Duke-Breast-Cancer-MRI Dataset",
+            "uri": "https://wiki.cancerimagingarchive.net/pages/viewpage.action?pageId=70226903"
+          },
+          {
+            "label": "ACRIN-Contralateral-Breast-MR",
+            "uri": "https://www.cancerimagingarchive.net/collection/acrin-contralateral-breast-mr"
+          }
+        ]
+      }
+    }
+  }
\ No newline at end of file
diff --git a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
index b963b7cc..ad15a8f9 100644
--- a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
+++ b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
@@ -14,6 +14,8 @@
 import SimpleITK as sitk
 import numpy as np
 from skimage import measure, filters
+from typing import Union
+from pathlib import Path
 
 
 class BreastPostProcessor(Module):
@@ -24,7 +26,10 @@ def get_mask(self, ip_path):
         seg_data[seg_data > 0] = 1
         return seg_data
 
-    def n_connected(self, img_data):
+    def n_connected(self, img_data: sitk.Image) -> sitk.Image:
+        """
+        Analyse connected components and drop smaller blobs
+        """
         img_data_mask = np.zeros(img_data.shape)
         img_data_mask[img_data >= 1] = 1
         img_filtered = np.zeros(img_data_mask.shape)
@@ -44,44 +49,20 @@ def n_connected(self, img_data):
         img_data[img_filtered != 1] = 0
         return img_data
 
-    def n4_bias_correction(
-        input_image,
-        shrink_factor=1,
-        mask_image=None,
-        number_of_iterations=50,
-        number_of_fitting_levels=4,
-    ):
-        if mask_image is None:
-            mask_image = sitk.OtsuThreshold(input_image, 0, 1, 200)
-
-        if shrink_factor > 1:
-            input_image = sitk.Shrink(
-                input_image, [shrink_factor] * input_image.GetDimension()
-            )
-            mask_image = sitk.Shrink(
-                mask_image, [shrink_factor] * mask_image.GetDimension()
-            )
-
-        corrector = sitk.N4BiasFieldCorrectionImageFilter()
-        corrector.SetMaximumNumberOfIterations(
-            [number_of_iterations] * number_of_fitting_levels
-        )
-
-        corrector.Execute(input_image, mask_image)
-        log_bias_field = corrector.GetLogBiasFieldAsImage(input_image)
-        corrected_image_full_resolution = input_image / sitk.Exp(log_bias_field)
-        return corrected_image_full_resolution
-
-    def get_seg_img(self, breast_and_fgt, breast_tumor, mr_path):
-        seg_data = np.zeros(breast_and_fgt.shape)
-        seg_data[breast_and_fgt == 1] = 1
-        seg_data[breast_tumor == 1] = 2
+    def merge_segmentations(self, breast_and_fgt_seg:sitk.Image, tumor_seg:sitk.Image, mr_path:Union[str,Path]) -> sitk.Image:
+        """
+        Merge labels of breast, fgt, and breast tumor
+        breast_and_fgt_seg: contains label=1 for breast, and label=2 for fgt
+        tumor_seg: assign label=3 for breast tumor in the merge image
+        returns: sitk image containing labels: [0,1,2,3]
+        """
+        tumor_seg[breast_and_fgt_seg == 0] = 0 # only focusing on predictions within the breast and fgt region
+        breast_and_fgt_seg[tumor_seg == 1] = 3
         ref = sitk.ReadImage(mr_path)
-        seg_img = sitk.GetImageFromArray(seg_data)
+        seg_img = sitk.GetImageFromArray(breast_and_fgt_seg)
         seg_img.CopyInformation(ref)
         return seg_img
 
-
     @IO.Instance()
     @IO.Input('in_breast_and_fgt_data', 'nifti:mod=seg:nnunet_task=Dataset009_Breast', the='input data from breast and fgt segmentation')
     @IO.Input('in_breast_tumor_data', 'nifti:mod=seg:nnunet_task=Dataset011_Breast', the='input data from breast tumor segmentation')
@@ -95,21 +76,24 @@ def task(self, instance: Instance, in_breast_and_fgt_data: InstanceData, in_brea
         print('input breast tumor segmentation: ' + in_breast_tumor_data.abspath)
         print('output path: ' + out_data.abspath)
 
-        breast_fgt_seg = self.get_mask(in_breast_and_fgt_data.abspath)
-        breast_fgt_seg = self.n_connected(breast_fgt_seg)
+        # Breast and FGT segmentation
+        breast_and_fgt_seg = self.get_mask(in_breast_and_fgt_data.abspath)
+        breast_and_fgt_seg = self.n_connected(breast_and_fgt_seg)
 
-        breast_tumor_seg = self.get_mask(in_breast_tumor_data.abspath)
-        breast_tumor_seg[breast_fgt_seg == 0] = 0
+        # Breast tumor segmentation
+        tumor_seg = self.get_mask(in_breast_tumor_data.abspath)
+        tumor_seg = self.n_connected(tumor_seg)
 
-        seg_img = self.get_seg_img(
-            breast_and_fgt = np.copy(breast_fgt_seg),
-            breast_tumor = np.copy(breast_tumor_seg),
+        # Merged segmentation masks
+        output_seg = self.merge_segmentations(
+            breast_and_fgt = np.copy(breast_and_fgt_seg),
+            breast_tumor = np.copy(tumor_seg),
             mr_path = in_mr_data.abspath
             )
         process_dir = self.config.data.requestTempDir(label="nnunet-breast-processor")
-        process_file = os.path.join(process_dir, f'final.nii.gz')
+        process_file = os.path.join(process_dir, f'bamf_nnunet_mr_breast.nii.gz')
         sitk.WriteImage(
-            seg_img,
+            output_seg,
             process_file,
         )
         shutil.copyfile(process_file, out_data.abspath)

From e4ead8b4d5cfbf982440dd240408e01da9bfbf80 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Wed, 12 Jun 2024 19:03:52 +0530
Subject: [PATCH 03/27] added nnunet v2 runner class

---
 .../bamf_nnunet_mr_breast/config/default.yml  |  14 +-
 .../utils/NNUnetRunnerV2.py                   | 210 ++++++++++++++++++
 2 files changed, 217 insertions(+), 7 deletions(-)
 create mode 100644 models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index 38d09282..07f9446a 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -6,12 +6,12 @@ general:
 execute:
 - DicomImporter
 - NiftiConverter
-- module: NNUnetRunner
-  nnunet_task: Dataset009_Breast
+- module: NNUnetRunnerV2
+  nnunet_dataset: Dataset009_Breast
   roi: BREAST,FGT
-- module: NNUnetRunner
-  nnunet_task: Dataset011_Breast
-  roi: BREAST_TUMOR
+- module: NNUnetRunnerV2
+  nnunet_dataset: Dataset011_Breast
+  roi: BREAST+TUMOR
 - BreastPostProcessor
 - DsegConverter
 - DataOrganizer
@@ -24,9 +24,9 @@ modules:
     meta: 
       mod: '%Modality'
 
-  NNUnetRunner:
+  NNUnetRunnerV2:
     in_data: nifti:mod=mr
-    nnunet_model: 3d_fullres
+    nnunet_config: 3d_fullres
     export_prob_maps: False
 
   BreastPostProcessor:
diff --git a/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
new file mode 100644
index 00000000..558daa60
--- /dev/null
+++ b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
@@ -0,0 +1,210 @@
+"""
+-------------------------------------------------
+MHub - NNU-Net Runner v2
+       Runner for pre-trained nnunet v2 models. 
+-------------------------------------------------
+
+-------------------------------------------------
+Author: Rahul Soni
+Email:  rahul.soni@bamfhealth.com
+-------------------------------------------------
+"""
+
+
+from typing import List, Optional
+import os, subprocess, shutil
+import SimpleITK as sitk, numpy as np
+from mhubio.core import Module, Instance, InstanceData, DataType, FileType, IO
+
+
+nnunet_dataset_name_regex = r"Dataset[0-9]{3}_[a-zA-Z0-9_]+"
+
+@IO.ConfigInput('in_data', 'nifti:mod=mr', the="input data to run nnunet on")
+@IO.Config('nnunet_dataset', str, None, the='nnunet dataset name')
+@IO.Config('nnunet_config', str, None, the='nnunet model name (2d, 3d_lowres, 3d_fullres, 3d_cascade_fullres)')
+@IO.Config('folds', int, None, the='number of folds to run nnunet on')
+@IO.Config('use_tta', bool, True, the='flag to enable test time augmentation')
+@IO.Config('export_prob_maps', bool, False, the='flag to export probability maps')
+@IO.Config('prob_map_segments', list, [], the='segment labels for probability maps')
+@IO.Config('roi', str, None, the='roi or comma separated list of roi the nnunet segments')
+class NNUnetRunnerV2(Module):
+
+    nnunet_dataset: str
+    nnunet_config: str
+    input_data_type: DataType
+    folds: int                          # TODO: support optional config attributes
+    use_tta: bool
+    export_prob_maps: bool
+    prob_map_segments: list
+    roi: str
+
+    def export_prob_mask(self, nnunet_out_dir: str, ref_file: InstanceData, output_dtype: str = 'float32', structure_list: Optional[List[str]] = None):
+        """
+        Convert softmax probability maps to NRRD. For simplicity, the probability maps
+        are converted by default to UInt8
+        Arguments:
+            model_output_folder : required - path to the folder where the inferred segmentation masks should be stored.
+            ref_file            : required - InstanceData object of the generated segmentation mask used as reference file.
+            output_dtype        : optional - output data type. Data type float16 is not supported by the NRRD standard,
+                                            so the choice should be between uint8, uint16 or float32. 
+            structure_list      : optional - list of the structures whose probability maps are stored in the 
+                                            first channel of the `.npz` file (output from the nnU-Net pipeline
+                                            when `export_prob_maps` is set to True). 
+        Outputs:
+            This function [...]
+        """
+
+        # initialize structure list
+        if structure_list is None:
+            if self.roi is not None:
+                structure_list = self.roi.split(',')
+            else:
+                structure_list = []
+
+        # sanity check user inputs
+        assert(output_dtype in ["uint8", "uint16", "float32"])      
+
+        # input file containing the raw information
+        pred_softmax_fn = 'VOLUME_001.npz'
+        pred_softmax_path = os.path.join(nnunet_out_dir, pred_softmax_fn)
+
+        # parse NRRD file - we will make use of if to populate the header of the
+        # NRRD mask we are going to get from the inferred segmentation mask
+        sitk_ct = sitk.ReadImage(ref_file.abspath)
+
+        # generate bundle for prob masks
+        # TODO: we really have to create folders (or add this as an option that defaults to true) automatically
+        prob_masks_bundle = ref_file.getDataBundle('prob_masks')
+        if not os.path.isdir(prob_masks_bundle.abspath):
+            os.mkdir(prob_masks_bundle.abspath)
+
+        # load softmax probability maps
+        pred_softmax_all = np.load(pred_softmax_path)["softmax"]
+
+        # iterate all channels
+        for channel in range(0, len(pred_softmax_all)):
+
+            structure = structure_list[channel] if channel < len(structure_list) else f"structure_{channel}"
+            pred_softmax_segmask = pred_softmax_all[channel].astype(dtype = np.float32)
+
+            if output_dtype == "float32":
+                # no rescale needed - the values will be between 0 and 1
+                # set SITK image dtype to Float32
+                sitk_dtype = sitk.sitkFloat32
+
+            elif output_dtype == "uint8":
+                # rescale between 0 and 255, quantize
+                pred_softmax_segmask = (255*pred_softmax_segmask).astype(np.int32)
+                # set SITK image dtype to UInt8
+                sitk_dtype = sitk.sitkUInt8
+
+            elif output_dtype == "uint16":
+                # rescale between 0 and 65536
+                pred_softmax_segmask = (65536*pred_softmax_segmask).astype(np.int32)
+                # set SITK image dtype to UInt16
+                sitk_dtype = sitk.sitkUInt16
+            else:
+                raise ValueError("Invalid output data type. Please choose between uint8, uint16 or float32.")
+                
+            pred_softmax_segmask_sitk = sitk.GetImageFromArray(pred_softmax_segmask)
+            pred_softmax_segmask_sitk.CopyInformation(sitk_ct)
+            pred_softmax_segmask_sitk = sitk.Cast(pred_softmax_segmask_sitk, sitk_dtype)
+
+            # generate data
+            prob_mask = InstanceData(f'{structure}.nrrd', DataType(FileType.NRRD, {'mod': 'prob_mask', 'structure': structure}), bundle=prob_masks_bundle)
+
+            # export file
+            writer = sitk.ImageFileWriter()
+            writer.UseCompressionOn()
+            writer.SetFileName(prob_mask.abspath)
+            writer.Execute(pred_softmax_segmask_sitk)
+
+            # check if the file was written
+            if os.path.isfile(prob_mask.abspath):
+                self.v(f" > prob mask for {structure} saved to {prob_mask.abspath}")
+                prob_mask.confirm()
+
+    @IO.Instance()
+    @IO.Input("in_data", the="input data to run nnunet on")
+    @IO.Output("out_data", 'VOLUME_001.nii.gz', 'nifti:mod=seg:model=nnunet', data='in_data', the="output data from nnunet")
+    def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData) -> None:
+        
+        # get the nnunet model to run
+        self.v("Running nnUNet_predict.")
+        self.v(f" > dataset:     {self.nnunet_dataset}")
+        self.v(f" > config:      {self.nnunet_config}")
+        self.v(f" > input data:  {in_data.abspath}")
+        self.v(f" > output data: {out_data.abspath}")
+
+        # download weights if not found
+        # NOTE: only for testing / debugging. For productiio always provide the weights in the Docker container.
+        if not os.path.isdir(os.path.join(os.environ["WEIGHTS_FOLDER"], '')):
+            print("Downloading nnUNet model weights...")
+            bash_command = ["nnUNet_download_pretrained_model", self.nnunet_dataset]
+            self.subprocess(bash_command, text=True)
+
+        # bring input data in nnunet specific format
+        # NOTE: only for nifti data as we hardcode the nnunet-formatted-filename (and extension) for now.
+        assert in_data.type.ftype == FileType.NIFTI
+        assert in_data.abspath.endswith('.nii.gz')
+        inp_dir = self.config.data.requestTempDir(label="nnunet-model-inp")
+        inp_file = f'VOLUME_001_0000.nii.gz'
+        shutil.copyfile(in_data.abspath, os.path.join(inp_dir, inp_file))
+
+        # 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
+
+        # symlink nnunet input folder to the input data with python
+        # create symlink in python
+        # NOTE: this is a workaround for the nnunet bash script that expects the input data to be in a specific folder
+        #       structure. This is not the case for the mhub data structure. So we create a symlink to the input data
+        #       in the nnunet input folder structure.
+        os.symlink(os.environ['WEIGHTS_FOLDER'], os.path.join(out_dir, 'nnUNet'))
+        
+        # NOTE: instead of running from commandline this could also be done in a pythonic way:
+        #       `nnUNet/nnunet/inference/predict.py` - but it would require
+        #       to set manually all the arguments that the user is not intended
+        #       to fiddle with; so stick with the bash executable
+
+        # construct nnunet inference command
+        bash_command  = ["nnUNetv2_predict"]
+        bash_command += ["-i", str(inp_dir)]
+        bash_command += ["-o", str(out_dir)]
+        bash_command += ["-d", self.nnunet_dataset]
+        bash_command += ["-c", self.nnunet_config]
+        
+        # add optional arguments
+        if self.folds is not None:
+            bash_command += ["--folds", str(self.folds)]
+
+        if not self.use_tta:
+            bash_command += ["--disable_tta"]
+        
+        if self.export_prob_maps:
+            bash_command += ["--save_npz"]
+
+        # run command
+        self.subprocess(bash_command, text=True)
+
+        # output meta
+        meta = {
+            "model": "nnunet",
+            "nnunet_dataset": self.nnunet_dataset,
+            "nnunet_config": self.nnunet_config,
+            "roi": self.roi
+        }
+
+        # get output data
+        out_file = f'VOLUME_001.nii.gz'
+        out_path = os.path.join(out_dir, out_file)
+
+        # copy output data to instance
+        shutil.copyfile(out_path, out_data.abspath)
+
+        # export probabiliy maps if requested as dynamic data
+        if self.export_prob_maps:
+            self.export_prob_mask(str(out_dir), out_data, 'float32', self.prob_map_segments)
+
+        # update meta dynamically
+        out_data.type.meta += meta

From c5fe136f035423c9576bd9477952bbe0efbad002 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 13 Jun 2024 13:27:40 +0530
Subject: [PATCH 04/27] updated metadata

---
 models/bamf_nnunet_mr_breast/meta.json | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/models/bamf_nnunet_mr_breast/meta.json b/models/bamf_nnunet_mr_breast/meta.json
index b595f5ee..52920e10 100644
--- a/models/bamf_nnunet_mr_breast/meta.json
+++ b/models/bamf_nnunet_mr_breast/meta.json
@@ -10,7 +10,7 @@
           "description": "The MR scan of a patient.",
           "format": "DICOM",
           "modality": "MR",
-          "bodypartexamined": "BREAST,FGT,BREAST_TUMOR",
+          "bodypartexamined": "BREAST,FGT,BREAST+TUMOR",
           "slicethickness": "2.5mm",
           "non-contrast": true,
           "contrast": false
@@ -20,11 +20,11 @@
         {
           "label": "Segmentation",
           "type": "Segmentation",
-          "description": "Segmentation breast, fgt, and tumor",
+          "description": "Segmentation for breast, fgt, and tumor",
           "classes": [
             "BREAST",
             "FGT",
-            "BREAST_TUMOR"
+            "BREAST+TUMOR"
           ]
         }
       ],
@@ -35,10 +35,10 @@
       },
       "data": {
         "training": {
-          "vol_samples": "TODO. Checked the training dataset and there were 98 + 489 = 587 samples"
+          "vol_samples": 587
         },
         "evaluation": {
-          "vol_samples": "TODO"
+          "vol_samples": 92
         },
         "public": true,
         "external": true
@@ -80,15 +80,14 @@
       },
       "analyses": {
         "title": "Quantitative Analyses",
-        "text": "The model's performance was assessed using the Dice Coefficient and Normalized Surface Distance (NSD) with tolerance 7mm, as specified in the Duke-Breast-Cancer-MRI. The model was used to segment cases from the IDC collection ACRIN-Contralateral-Breast-MR [1]. Nine of those cases were reviewed and corrected by a board-certified radiologist and a non-expert. The analysis is published here [2]",
+        "text": "The model's performance was assessed using the Dice Coefficient with tolerance 7mm, as specified in the Duke-Breast-Cancer-MRI. The model was used to segment cases from the IDC collection ACRIN-Contralateral-Breast-MR [1]. Nine of those cases were reviewed and corrected by a board-certified radiologist. The analysis is published here [2]",
         "tables": [
           {
-            "label": "Label-wise metrics (mean (standard deviation)) between AI derived and manually corrected MR Breast annotations",
+            "label": "Label-wise metrics (mean) between AI derived and manually corrected MR Breast annotations",
             "entries": {
-              "Dice: Radiologist": "0.97 (0.07)",
-              "NSD: Radiologist": "0.96 (0.09)",
-              "Dice: Non-expert": "0.97 (0.08)",
-              "NSD: Non-expert": "0.95 (0.11)"
+              "Dice: BREAST": "0.98",
+              "Dice: FGT": "0.80",
+              "Dice: BREAST+TUMOR": "0.56"
             }
           }
         ],
@@ -105,7 +104,7 @@
       },
       "evaluation": {
         "title": "Evaluation Data",
-        "text": "The model was used to segment cases 89 from the ACRIN-Contralateral-Breast-MR [1]. Nine of those cases were randomly selected to be reviewed and corrected by a board-certified radiologist. The model predictions, and radiologist corrections are published on zenodo [3]",
+        "text": "The model was used to segment 92 cases from the ACRIN-Contralateral-Breast-MR [1]. All of those cases were reviewed and corrected by a board-certified radiologist. The model predictions, and radiologist corrections would be published on zenodo soon",
         "references": [
           {
             "label": "ACRIN-Contralateral-Breast-MR",
@@ -132,4 +131,4 @@
         ]
       }
     }
-  }
\ No newline at end of file
+  }

From 07adc87301ee0b6c86a25b41c4e76766b024ba20 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 13 Jun 2024 13:29:15 +0530
Subject: [PATCH 05/27] updated metadata

---
 models/bamf_nnunet_mr_breast/meta.json | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/bamf_nnunet_mr_breast/meta.json b/models/bamf_nnunet_mr_breast/meta.json
index 52920e10..16d1669f 100644
--- a/models/bamf_nnunet_mr_breast/meta.json
+++ b/models/bamf_nnunet_mr_breast/meta.json
@@ -80,7 +80,7 @@
       },
       "analyses": {
         "title": "Quantitative Analyses",
-        "text": "The model's performance was assessed using the Dice Coefficient with tolerance 7mm, as specified in the Duke-Breast-Cancer-MRI. The model was used to segment cases from the IDC collection ACRIN-Contralateral-Breast-MR [1]. Nine of those cases were reviewed and corrected by a board-certified radiologist. The analysis is published here [2]",
+        "text": "The model's performance was assessed using the Dice Coefficient with tolerance 7mm, as specified in the Duke-Breast-Cancer-MRI. The model was used to segment cases from the IDC collection ACRIN-Contralateral-Breast-MR [1]. All validation cases (=92) were reviewed and corrected by board-certified radiologists. The analysis is published here [2]",
         "tables": [
           {
             "label": "Label-wise metrics (mean) between AI derived and manually corrected MR Breast annotations",

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 06/27] 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

From 6dc7a289dbd7f46fe749ef1b7d1c95f5a7be9da4 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Tue, 18 Jun 2024 17:47:10 +0530
Subject: [PATCH 07/27] fix rois

---
 .../bamf_nnunet_mr_breast/utils/BreastPostProcessor.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
index ad15a8f9..cdf0a3a8 100644
--- a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
+++ b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
@@ -64,10 +64,10 @@ def merge_segmentations(self, breast_and_fgt_seg:sitk.Image, tumor_seg:sitk.Imag
         return seg_img
 
     @IO.Instance()
-    @IO.Input('in_breast_and_fgt_data', 'nifti:mod=seg:nnunet_task=Dataset009_Breast', the='input data from breast and fgt segmentation')
-    @IO.Input('in_breast_tumor_data', 'nifti:mod=seg:nnunet_task=Dataset011_Breast', the='input data from breast tumor segmentation')
+    @IO.Input('in_breast_and_fgt_data', 'nifti:mod=seg:nnunet_dataset=Dataset009_Breast', the='input data from breast and fgt segmentation')
+    @IO.Input('in_breast_tumor_data', 'nifti:mod=seg:nnunet_dataset=Dataset011_Breast', the='input data from breast tumor segmentation')
     @IO.Input('in_mr_data', 'nifti:mod=mr', the='input mr data')
-    @IO.Output('out_data', 'bamf_processed.nii.gz', 'nifti:mod=seg:processor=bamf:roi=BREAST,BREAST+FGT,BREAST_TUMOR',
+    @IO.Output('out_data', 'bamf_processed.nii.gz', 'nifti:mod=seg:processor=bamf:roi=BREAST,FGT,BREAST+TUMOR',
                the="get breast_fgt and breast segmentation file")
     def task(self, instance: Instance, in_breast_and_fgt_data: InstanceData, in_breast_tumor_data: InstanceData,
              in_mr_data: InstanceData, out_data: InstanceData):
@@ -86,8 +86,8 @@ def task(self, instance: Instance, in_breast_and_fgt_data: InstanceData, in_brea
 
         # Merged segmentation masks
         output_seg = self.merge_segmentations(
-            breast_and_fgt = np.copy(breast_and_fgt_seg),
-            breast_tumor = np.copy(tumor_seg),
+            breast_and_fgt_seg = np.copy(breast_and_fgt_seg),
+            tumor_seg = np.copy(tumor_seg),
             mr_path = in_mr_data.abspath
             )
         process_dir = self.config.data.requestTempDir(label="nnunet-breast-processor")

From b47061f9cb7a51c54a2b6574497813185a2d31b9 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Tue, 18 Jun 2024 21:52:10 +0530
Subject: [PATCH 08/27] minor description update

---
 models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
index cdf0a3a8..f0075a1f 100644
--- a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
+++ b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
@@ -68,7 +68,7 @@ def merge_segmentations(self, breast_and_fgt_seg:sitk.Image, tumor_seg:sitk.Imag
     @IO.Input('in_breast_tumor_data', 'nifti:mod=seg:nnunet_dataset=Dataset011_Breast', the='input data from breast tumor segmentation')
     @IO.Input('in_mr_data', 'nifti:mod=mr', the='input mr data')
     @IO.Output('out_data', 'bamf_processed.nii.gz', 'nifti:mod=seg:processor=bamf:roi=BREAST,FGT,BREAST+TUMOR',
-               the="get breast_fgt and breast segmentation file")
+               the="get breast,fgt, and tumor segmentation masks")
     def task(self, instance: Instance, in_breast_and_fgt_data: InstanceData, in_breast_tumor_data: InstanceData,
              in_mr_data: InstanceData, out_data: InstanceData):
 

From 204238fe5058bb9b62366c326f8cf648b8db3038 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Wed, 10 Jul 2024 00:20:25 +0530
Subject: [PATCH 09/27] custom SegDB keys

---
 .../bamf_nnunet_mr_breast/config/default.yml  | 22 +++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index 07f9446a..ac98d191 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -3,6 +3,28 @@ general:
   version: 1.0
   description: Default configuration for Bamf NNUnet Breast and tumor segmentation on MR scans (dicom to dicom)
 
+segdb:
+  triplets:
+    BREAST_FGT_TRIPLET:
+      code: 12
+      meaning: breast and fgt triplets
+  segmentations:
+    FGT_SEGMENTATION:
+      name: breast and fgt segmentation
+      category: C_BODY_STRUCTURE
+      type: BREAST_FGT_TRIPLET
+      color: [255, 255, 255]
+    BREAST:
+      name: custom label and color for breast
+      category: C_BODY_STRUCTURE
+      type: T_BREAST
+      color: [255, 0, 0]
+    FGT:
+      name: custom label and color for fibrograndular tissue
+      category: C_BODY_STRUCTURE
+      type: T_FGT
+      color: [0, 255, 0]
+
 execute:
 - DicomImporter
 - NiftiConverter

From 10f587d111edec5a73838c3d3f83449efd04af12 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Wed, 10 Jul 2024 01:13:02 +0530
Subject: [PATCH 10/27] custom SegDB keys

---
 models/bamf_nnunet_mr_breast/config/default.yml | 15 +++++----------
 1 file changed, 5 insertions(+), 10 deletions(-)

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index ac98d191..5bda3c95 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -6,23 +6,18 @@ general:
 segdb:
   triplets:
     BREAST_FGT_TRIPLET:
-      code: 12
+      code: 123
       meaning: breast and fgt triplets
   segmentations:
-    FGT_SEGMENTATION:
-      name: breast and fgt segmentation
-      category: C_BODY_STRUCTURE
-      type: BREAST_FGT_TRIPLET
-      color: [255, 255, 255]
     BREAST:
-      name: custom label and color for breast
+      name: custom label and color for breast segmentation
       category: C_BODY_STRUCTURE
-      type: T_BREAST
+      type: T_BREAST_STRUCTURE
       color: [255, 0, 0]
     FGT:
-      name: custom label and color for fibrograndular tissue
+      name: custom label and color for fibrograndular tissue segmentation
       category: C_BODY_STRUCTURE
-      type: T_FGT
+      type: T_BREAST_STRUCTURE
       color: [0, 255, 0]
 
 execute:

From c7f2007315c8e5ef714933aa9178892776853d02 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 11 Jul 2024 22:22:27 +0530
Subject: [PATCH 11/27] custom SegDB module

---
 .../bamf_nnunet_mr_breast/config/default.yml  | 17 -------
 .../utils/BreastPostProcessor.py              | 47 ++++++++++++++++++-
 2 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index 5bda3c95..07f9446a 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -3,23 +3,6 @@ general:
   version: 1.0
   description: Default configuration for Bamf NNUnet Breast and tumor segmentation on MR scans (dicom to dicom)
 
-segdb:
-  triplets:
-    BREAST_FGT_TRIPLET:
-      code: 123
-      meaning: breast and fgt triplets
-  segmentations:
-    BREAST:
-      name: custom label and color for breast segmentation
-      category: C_BODY_STRUCTURE
-      type: T_BREAST_STRUCTURE
-      color: [255, 0, 0]
-    FGT:
-      name: custom label and color for fibrograndular tissue segmentation
-      category: C_BODY_STRUCTURE
-      type: T_BREAST_STRUCTURE
-      color: [0, 255, 0]
-
 execute:
 - DicomImporter
 - NiftiConverter
diff --git a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
index f0075a1f..a932d6cd 100644
--- a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
+++ b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
@@ -16,6 +16,49 @@
 from skimage import measure, filters
 from typing import Union
 from pathlib import Path
+import yaml
+import segdb
+from segdb.classes import Segment
+from segdb.classes.Triplet import Triplet
+
+
+custom_seg_config = """
+segdb:
+    triplets:
+        T_BREAST_FGT_STRUCTURE:
+            code: C114467
+            meaning: Breast Fibroglandular Tissue
+            scheme_designator: NCIt
+        T_BREAST_CARCINOMA_MASS:
+            code: 4147007
+            meaning: Mass
+            scheme_designator: SCT
+    segmentations:
+        FGT:
+            name: FGT
+            category: C_RADIOLOGIC_FINDING
+            type: T_BREAST_FGT_STRUCTURE
+            color: [128, 174, 128]
+        BREAST_CARCINOMA:
+            name: Breast_Carcinoma
+            category: C_MORPHOLOGICALLY_ABNORMAL_STRUCTURE
+            type: T_BREAST_CARCINOMA_MASS
+            modifier: M_RIGHT_AND_LEFT
+            color: [144, 238, 144]
+"""
+parsed_config = yaml.safe_load(custom_seg_config)
+
+if 'segdb' in parsed_config:
+    if 'segmentations' in parsed_config['segdb'] and isinstance(parsed_config['segdb']['segmentations'], dict):
+        from segdb.classes.Segment import Segment
+        for seg_id, seg_data in parsed_config['segdb']['segmentations'].items():
+            print("added segment", seg_id, seg_data )
+            Segment.register(seg_id, **seg_data)
+    if 'triplets' in parsed_config['segdb'] and isinstance(parsed_config['segdb']['triplets'], dict):
+        from segdb.classes.Triplet import Triplet
+        for trp_id, trp_data in parsed_config['segdb']['triplets'].items():
+            print("added triplet", trp_id, trp_data )
+            Triplet.register(trp_id, overwrite=True, **trp_data)
 
 
 class BreastPostProcessor(Module):
@@ -67,8 +110,8 @@ def merge_segmentations(self, breast_and_fgt_seg:sitk.Image, tumor_seg:sitk.Imag
     @IO.Input('in_breast_and_fgt_data', 'nifti:mod=seg:nnunet_dataset=Dataset009_Breast', the='input data from breast and fgt segmentation')
     @IO.Input('in_breast_tumor_data', 'nifti:mod=seg:nnunet_dataset=Dataset011_Breast', the='input data from breast tumor segmentation')
     @IO.Input('in_mr_data', 'nifti:mod=mr', the='input mr data')
-    @IO.Output('out_data', 'bamf_processed.nii.gz', 'nifti:mod=seg:processor=bamf:roi=BREAST,FGT,BREAST+TUMOR',
-               the="get breast,fgt, and tumor segmentation masks")
+    @IO.Output('out_data', 'bamf_processed.nii.gz', 'nifti:mod=seg:processor=bamf:roi=BREAST,FGT,BREAST+BREAST_CARCINOMA',
+               the="get breast, fgt, and breast-tumor segmentation masks")
     def task(self, instance: Instance, in_breast_and_fgt_data: InstanceData, in_breast_tumor_data: InstanceData,
              in_mr_data: InstanceData, out_data: InstanceData):
 

From 5d7e268f5e1e2388d3328fff91ddda4414e8f102 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Tue, 30 Jul 2024 22:30:15 +0530
Subject: [PATCH 12/27] move custom segDB triplets to config

---
 .../bamf_nnunet_mr_breast/config/default.yml  | 25 ++++++++++-
 .../utils/BreastPostProcessor.py              | 41 -------------------
 2 files changed, 24 insertions(+), 42 deletions(-)

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index 07f9446a..cae0c151 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -3,6 +3,29 @@ general:
   version: 1.0
   description: Default configuration for Bamf NNUnet Breast and tumor segmentation on MR scans (dicom to dicom)
 
+segdb:
+  triplets:
+    T_BREAST_FGT_STRUCTURE:
+      code: C114467
+      meaning: Breast Fibroglandular Tissue
+      scheme_designator: NCIt
+    T_BREAST_CARCINOMA_MASS:
+      code: 4147007
+      meaning: Mass
+      scheme_designator: SCT
+  segments:
+    FGT:
+      name: FGT
+      category: C_RADIOLOGIC_FINDING
+      type: T_BREAST_FGT_STRUCTURE
+      color: [128, 174, 128]
+    BREAST_CARCINOMA:
+      name: Breast_Carcinoma
+      category: C_MORPHOLOGICALLY_ABNORMAL_STRUCTURE
+      type: T_BREAST_CARCINOMA_MASS
+      modifier: M_RIGHT_AND_LEFT
+      color: [144, 238, 144]
+
 execute:
 - DicomImporter
 - NiftiConverter
@@ -11,7 +34,7 @@ execute:
   roi: BREAST,FGT
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset011_Breast
-  roi: BREAST+TUMOR
+  roi: BREAST+BREAST_CARCINOMA
 - BreastPostProcessor
 - DsegConverter
 - DataOrganizer
diff --git a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
index a932d6cd..bc03da52 100644
--- a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
+++ b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
@@ -17,49 +17,8 @@
 from typing import Union
 from pathlib import Path
 import yaml
-import segdb
-from segdb.classes import Segment
-from segdb.classes.Triplet import Triplet
 
 
-custom_seg_config = """
-segdb:
-    triplets:
-        T_BREAST_FGT_STRUCTURE:
-            code: C114467
-            meaning: Breast Fibroglandular Tissue
-            scheme_designator: NCIt
-        T_BREAST_CARCINOMA_MASS:
-            code: 4147007
-            meaning: Mass
-            scheme_designator: SCT
-    segmentations:
-        FGT:
-            name: FGT
-            category: C_RADIOLOGIC_FINDING
-            type: T_BREAST_FGT_STRUCTURE
-            color: [128, 174, 128]
-        BREAST_CARCINOMA:
-            name: Breast_Carcinoma
-            category: C_MORPHOLOGICALLY_ABNORMAL_STRUCTURE
-            type: T_BREAST_CARCINOMA_MASS
-            modifier: M_RIGHT_AND_LEFT
-            color: [144, 238, 144]
-"""
-parsed_config = yaml.safe_load(custom_seg_config)
-
-if 'segdb' in parsed_config:
-    if 'segmentations' in parsed_config['segdb'] and isinstance(parsed_config['segdb']['segmentations'], dict):
-        from segdb.classes.Segment import Segment
-        for seg_id, seg_data in parsed_config['segdb']['segmentations'].items():
-            print("added segment", seg_id, seg_data )
-            Segment.register(seg_id, **seg_data)
-    if 'triplets' in parsed_config['segdb'] and isinstance(parsed_config['segdb']['triplets'], dict):
-        from segdb.classes.Triplet import Triplet
-        for trp_id, trp_data in parsed_config['segdb']['triplets'].items():
-            print("added triplet", trp_id, trp_data )
-            Triplet.register(trp_id, overwrite=True, **trp_data)
-
 
 class BreastPostProcessor(Module):
 

From 312ae557fcebe2c1e6ee33cbabd1ce556d8b905e Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Mon, 12 Aug 2024 17:33:40 +0400
Subject: [PATCH 13/27] update documentation

---
 models/bamf_nnunet_mr_breast/meta.json | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/models/bamf_nnunet_mr_breast/meta.json b/models/bamf_nnunet_mr_breast/meta.json
index 16d1669f..756f7c16 100644
--- a/models/bamf_nnunet_mr_breast/meta.json
+++ b/models/bamf_nnunet_mr_breast/meta.json
@@ -67,8 +67,8 @@
       },
       "publications": [
         {
-          "title": "The AIMI Initiative: AI-Generated Annotations in IDC Collections",
-          "uri": "https://arxiv.org/abs/2310.14897"
+          "title": "",
+          "uri": ""
         }
       ],
       "github": "https://github.com/bamf-health/aimi-breast-mr"
@@ -112,7 +112,7 @@
           },
           {
             "label": "Image segmentations produced by the AIMI Annotations initiative",
-            "uri": "https://zenodo.org/records/10009368"
+            "uri": "https://zenodo.org/records/13244892"
           }
         ]
       },

From 5997cda5ed591825eff550016ddc6192872c4bea Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Wed, 21 Aug 2024 20:00:03 +0400
Subject: [PATCH 14/27] minor refactoring for PR-82

---
 .../bamf_nnunet_mr_breast/config/default.yml  |   2 +-
 models/bamf_nnunet_mr_breast/meta.json        |   9 +-
 .../utils/BreastPostProcessor.py              |   6 +-
 .../utils/NNUnetRunnerV2.py                   | 102 ------------------
 4 files changed, 7 insertions(+), 112 deletions(-)

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index cae0c151..23867b4b 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -31,7 +31,7 @@ execute:
 - NiftiConverter
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset009_Breast
-  roi: BREAST,FGT
+  roi: BREAST+FGT
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset011_Breast
   roi: BREAST+BREAST_CARCINOMA
diff --git a/models/bamf_nnunet_mr_breast/meta.json b/models/bamf_nnunet_mr_breast/meta.json
index 756f7c16..9774a843 100644
--- a/models/bamf_nnunet_mr_breast/meta.json
+++ b/models/bamf_nnunet_mr_breast/meta.json
@@ -10,7 +10,7 @@
           "description": "The MR scan of a patient.",
           "format": "DICOM",
           "modality": "MR",
-          "bodypartexamined": "BREAST,FGT,BREAST+TUMOR",
+          "bodypartexamined": "BREAST",
           "slicethickness": "2.5mm",
           "non-contrast": true,
           "contrast": false
@@ -20,11 +20,10 @@
         {
           "label": "Segmentation",
           "type": "Segmentation",
-          "description": "Segmentation for breast, fgt, and tumor",
+          "description": "Segmentation for breast, fgt, and breast carcinoma",
           "classes": [
-            "BREAST",
-            "FGT",
-            "BREAST+TUMOR"
+            "BREAST+FGT",
+            "BREAST+BREAST_CARCINOMA"
           ]
         }
       ],
diff --git a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
index bc03da52..3d55f469 100644
--- a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
+++ b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
@@ -92,10 +92,8 @@ def task(self, instance: Instance, in_breast_and_fgt_data: InstanceData, in_brea
             tumor_seg = np.copy(tumor_seg),
             mr_path = in_mr_data.abspath
             )
-        process_dir = self.config.data.requestTempDir(label="nnunet-breast-processor")
-        process_file = os.path.join(process_dir, f'bamf_nnunet_mr_breast.nii.gz')
         sitk.WriteImage(
             output_seg,
-            process_file,
+            out_data.abspath,
         )
-        shutil.copyfile(process_file, out_data.abspath)
+
diff --git a/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
index e0ea53b4..0aab9b1a 100644
--- a/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
+++ b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
@@ -23,9 +23,6 @@
 @IO.Config('nnunet_dataset', str, None, the='nnunet dataset name')
 @IO.Config('nnunet_config', str, None, the='nnunet model name (2d, 3d_lowres, 3d_fullres, 3d_cascade_fullres)')
 @IO.Config('folds', int, None, the='number of folds to run nnunet on')
-@IO.Config('use_tta', bool, True, the='flag to enable test time augmentation')
-@IO.Config('export_prob_maps', bool, False, the='flag to export probability maps')
-@IO.Config('prob_map_segments', list, [], the='segment labels for probability maps')
 @IO.Config('roi', str, None, the='roi or comma separated list of roi the nnunet segments')
 class NNUnetRunnerV2(Module):
 
@@ -33,97 +30,8 @@ class NNUnetRunnerV2(Module):
     nnunet_config: str
     input_data_type: DataType
     folds: int                          # TODO: support optional config attributes
-    use_tta: bool
-    export_prob_maps: bool
-    prob_map_segments: list
     roi: str
 
-    def export_prob_mask(self, nnunet_out_dir: str, ref_file: InstanceData, output_dtype: str = 'float32', structure_list: Optional[List[str]] = None):
-        """
-        Convert softmax probability maps to NRRD. For simplicity, the probability maps
-        are converted by default to UInt8
-        Arguments:
-            model_output_folder : required - path to the folder where the inferred segmentation masks should be stored.
-            ref_file            : required - InstanceData object of the generated segmentation mask used as reference file.
-            output_dtype        : optional - output data type. Data type float16 is not supported by the NRRD standard,
-                                            so the choice should be between uint8, uint16 or float32. 
-            structure_list      : optional - list of the structures whose probability maps are stored in the 
-                                            first channel of the `.npz` file (output from the nnU-Net pipeline
-                                            when `export_prob_maps` is set to True). 
-        Outputs:
-            This function [...]
-        """
-
-        # initialize structure list
-        if structure_list is None:
-            if self.roi is not None:
-                structure_list = self.roi.split(',')
-            else:
-                structure_list = []
-
-        # sanity check user inputs
-        assert(output_dtype in ["uint8", "uint16", "float32"])      
-
-        # input file containing the raw information
-        pred_softmax_fn = 'VOLUME_001.npz'
-        pred_softmax_path = os.path.join(nnunet_out_dir, pred_softmax_fn)
-
-        # parse NRRD file - we will make use of if to populate the header of the
-        # NRRD mask we are going to get from the inferred segmentation mask
-        sitk_ct = sitk.ReadImage(ref_file.abspath)
-
-        # generate bundle for prob masks
-        # TODO: we really have to create folders (or add this as an option that defaults to true) automatically
-        prob_masks_bundle = ref_file.getDataBundle('prob_masks')
-        if not os.path.isdir(prob_masks_bundle.abspath):
-            os.mkdir(prob_masks_bundle.abspath)
-
-        # load softmax probability maps
-        pred_softmax_all = np.load(pred_softmax_path)["softmax"]
-
-        # iterate all channels
-        for channel in range(0, len(pred_softmax_all)):
-
-            structure = structure_list[channel] if channel < len(structure_list) else f"structure_{channel}"
-            pred_softmax_segmask = pred_softmax_all[channel].astype(dtype = np.float32)
-
-            if output_dtype == "float32":
-                # no rescale needed - the values will be between 0 and 1
-                # set SITK image dtype to Float32
-                sitk_dtype = sitk.sitkFloat32
-
-            elif output_dtype == "uint8":
-                # rescale between 0 and 255, quantize
-                pred_softmax_segmask = (255*pred_softmax_segmask).astype(np.int32)
-                # set SITK image dtype to UInt8
-                sitk_dtype = sitk.sitkUInt8
-
-            elif output_dtype == "uint16":
-                # rescale between 0 and 65536
-                pred_softmax_segmask = (65536*pred_softmax_segmask).astype(np.int32)
-                # set SITK image dtype to UInt16
-                sitk_dtype = sitk.sitkUInt16
-            else:
-                raise ValueError("Invalid output data type. Please choose between uint8, uint16 or float32.")
-                
-            pred_softmax_segmask_sitk = sitk.GetImageFromArray(pred_softmax_segmask)
-            pred_softmax_segmask_sitk.CopyInformation(sitk_ct)
-            pred_softmax_segmask_sitk = sitk.Cast(pred_softmax_segmask_sitk, sitk_dtype)
-
-            # generate data
-            prob_mask = InstanceData(f'{structure}.nrrd', DataType(FileType.NRRD, {'mod': 'prob_mask', 'structure': structure}), bundle=prob_masks_bundle)
-
-            # export file
-            writer = sitk.ImageFileWriter()
-            writer.UseCompressionOn()
-            writer.SetFileName(prob_mask.abspath)
-            writer.Execute(pred_softmax_segmask_sitk)
-
-            # check if the file was written
-            if os.path.isfile(prob_mask.abspath):
-                self.v(f" > prob mask for {structure} saved to {prob_mask.abspath}")
-                prob_mask.confirm()
-
     @IO.Instance()
     @IO.Input("in_data", the="input data to run nnunet on")
     @IO.Output("out_data", 'VOLUME_001.nii.gz', 'nifti:mod=seg:model=nnunet', data='in_data', the="output data from nnunet")
@@ -177,13 +85,7 @@ def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData
         # add optional arguments
         if self.folds is not None:
             bash_command += ["--folds", str(self.folds)]
-
-        if not self.use_tta:
-            bash_command += ["--disable_tta"]
         
-        if self.export_prob_maps:
-            bash_command += ["--save_npz"]
-
         # run command
         self.subprocess(bash_command, text=True)
 
@@ -202,9 +104,5 @@ def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData
         # copy output data to instance
         shutil.copyfile(out_path, out_data.abspath)
 
-        # export probabiliy maps if requested as dynamic data
-        if self.export_prob_maps:
-            self.export_prob_mask(str(out_dir), out_data, 'float32', self.prob_map_segments)
-
         # update meta dynamically
         out_data.type.meta += meta

From 24d3b6b277d1050d93021576133935cf6f43f6a9 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 22 Aug 2024 00:31:26 +0400
Subject: [PATCH 15/27] minor refactoring for PR-82

---
 models/bamf_nnunet_mr_breast/config/default.yml      | 4 +---
 models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py | 9 +++------
 2 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index 23867b4b..3b0a704e 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -31,7 +31,7 @@ execute:
 - NiftiConverter
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset009_Breast
-  roi: BREAST+FGT
+  roi: BREAST,FGT
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset011_Breast
   roi: BREAST+BREAST_CARCINOMA
@@ -49,8 +49,6 @@ modules:
 
   NNUnetRunnerV2:
     in_data: nifti:mod=mr
-    nnunet_config: 3d_fullres
-    export_prob_maps: False
 
   BreastPostProcessor:
     in_breast_data: nifti:mod=seg:nnunet_task=Dataset009_Breast
diff --git a/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
index 0aab9b1a..f8dc85bb 100644
--- a/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
+++ b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
@@ -21,13 +21,11 @@
 
 @IO.ConfigInput('in_data', 'nifti:mod=mr', the="input data to run nnunet on")
 @IO.Config('nnunet_dataset', str, None, the='nnunet dataset name')
-@IO.Config('nnunet_config', str, None, the='nnunet model name (2d, 3d_lowres, 3d_fullres, 3d_cascade_fullres)')
 @IO.Config('folds', int, None, the='number of folds to run nnunet on')
 @IO.Config('roi', str, None, the='roi or comma separated list of roi the nnunet segments')
 class NNUnetRunnerV2(Module):
 
     nnunet_dataset: str
-    nnunet_config: str
     input_data_type: DataType
     folds: int                          # TODO: support optional config attributes
     roi: str
@@ -40,7 +38,6 @@ def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData
         # get the nnunet model to run
         self.v("Running nnUNet_predict.")
         self.v(f" > dataset:     {self.nnunet_dataset}")
-        self.v(f" > config:      {self.nnunet_config}")
         self.v(f" > input data:  {in_data.abspath}")
         self.v(f" > output data: {out_data.abspath}")
 
@@ -80,7 +77,7 @@ def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData
         bash_command += ["-i", str(inp_dir)]
         bash_command += ["-o", str(out_dir)]
         bash_command += ["-d", self.nnunet_dataset]
-        bash_command += ["-c", self.nnunet_config]
+        bash_command += ["-c", "3d_fullres"]
         
         # add optional arguments
         if self.folds is not None:
@@ -91,9 +88,9 @@ def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData
 
         # output meta
         meta = {
-            "model": "nnunet",
+            "model": "nnunet-v2",
             "nnunet_dataset": self.nnunet_dataset,
-            "nnunet_config": self.nnunet_config,
+            "nnunet_config": "3d_fullres",
             "roi": self.roi
         }
 

From 5217425fad37f44eb323167fbd9be3db795e011c Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 22 Aug 2024 00:33:13 +0400
Subject: [PATCH 16/27] [PR-82] update output classes in meta

---
 models/bamf_nnunet_mr_breast/meta.json | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/models/bamf_nnunet_mr_breast/meta.json b/models/bamf_nnunet_mr_breast/meta.json
index 9774a843..1bd86309 100644
--- a/models/bamf_nnunet_mr_breast/meta.json
+++ b/models/bamf_nnunet_mr_breast/meta.json
@@ -22,7 +22,8 @@
           "type": "Segmentation",
           "description": "Segmentation for breast, fgt, and breast carcinoma",
           "classes": [
-            "BREAST+FGT",
+            "BREAST",
+            "FGT",
             "BREAST+BREAST_CARCINOMA"
           ]
         }

From a3a34743785f3f035704087c416ef2f23b819525 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 22 Aug 2024 13:54:37 +0400
Subject: [PATCH 17/27] [PR-82] indicate nnUNetv2 in debug output

---
 models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
index f8dc85bb..583d90db 100644
--- a/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
+++ b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py
@@ -36,7 +36,7 @@ class NNUnetRunnerV2(Module):
     def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData) -> None:
         
         # get the nnunet model to run
-        self.v("Running nnUNet_predict.")
+        self.v("Running nnUNetv2_predict.")
         self.v(f" > dataset:     {self.nnunet_dataset}")
         self.v(f" > input data:  {in_data.abspath}")
         self.v(f" > output data: {out_data.abspath}")

From e23f9c43b7b1a38a8bb33a52d5d2d06157730390 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 22 Aug 2024 22:57:30 +0400
Subject: [PATCH 18/27] BREAST+FGT label correction

---
 models/bamf_nnunet_mr_breast/config/default.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index 3b0a704e..1de2a6df 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -31,7 +31,7 @@ execute:
 - NiftiConverter
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset009_Breast
-  roi: BREAST,FGT
+  roi: BREAST,BREAST+FGT
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset011_Breast
   roi: BREAST+BREAST_CARCINOMA

From 3eb991e08be6801387ebf525d6c6bf00d2e37556 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 22 Aug 2024 23:02:09 +0400
Subject: [PATCH 19/27] Switch FGT to SCT scheme designator

---
 models/bamf_nnunet_mr_breast/config/default.yml | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index 1de2a6df..355ce422 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -6,9 +6,9 @@ general:
 segdb:
   triplets:
     T_BREAST_FGT_STRUCTURE:
-      code: C114467
-      meaning: Breast Fibroglandular Tissue
-      scheme_designator: NCIt
+      code: 1342375003
+      meaning: Structure of fibroglandular tissue of breast
+      scheme_designator: SCT
     T_BREAST_CARCINOMA_MASS:
       code: 4147007
       meaning: Mass

From 66f8ce819df9d75af98a834b0ca0d6102166e446 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Fri, 23 Aug 2024 18:10:06 +0400
Subject: [PATCH 20/27] rollback custom SegDB triplet: keys now available in
 SegDB

---
 .../bamf_nnunet_mr_breast/config/default.yml  | 25 +------------------
 1 file changed, 1 insertion(+), 24 deletions(-)

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index 355ce422..dfb92b1a 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -3,35 +3,12 @@ general:
   version: 1.0
   description: Default configuration for Bamf NNUnet Breast and tumor segmentation on MR scans (dicom to dicom)
 
-segdb:
-  triplets:
-    T_BREAST_FGT_STRUCTURE:
-      code: 1342375003
-      meaning: Structure of fibroglandular tissue of breast
-      scheme_designator: SCT
-    T_BREAST_CARCINOMA_MASS:
-      code: 4147007
-      meaning: Mass
-      scheme_designator: SCT
-  segments:
-    FGT:
-      name: FGT
-      category: C_RADIOLOGIC_FINDING
-      type: T_BREAST_FGT_STRUCTURE
-      color: [128, 174, 128]
-    BREAST_CARCINOMA:
-      name: Breast_Carcinoma
-      category: C_MORPHOLOGICALLY_ABNORMAL_STRUCTURE
-      type: T_BREAST_CARCINOMA_MASS
-      modifier: M_RIGHT_AND_LEFT
-      color: [144, 238, 144]
-
 execute:
 - DicomImporter
 - NiftiConverter
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset009_Breast
-  roi: BREAST,BREAST+FGT
+  roi: BREAST,BREAST+BREAST_FIBROGLANDULAR_TISSUE
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset011_Breast
   roi: BREAST+BREAST_CARCINOMA

From c04ab11a69514a8713f60051f2c2961cfb43a6ac Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Mon, 26 Aug 2024 13:17:01 +0400
Subject: [PATCH 21/27] refactor roi key name

---
 models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
index 3d55f469..4d580916 100644
--- a/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
+++ b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py
@@ -69,7 +69,7 @@ def merge_segmentations(self, breast_and_fgt_seg:sitk.Image, tumor_seg:sitk.Imag
     @IO.Input('in_breast_and_fgt_data', 'nifti:mod=seg:nnunet_dataset=Dataset009_Breast', the='input data from breast and fgt segmentation')
     @IO.Input('in_breast_tumor_data', 'nifti:mod=seg:nnunet_dataset=Dataset011_Breast', the='input data from breast tumor segmentation')
     @IO.Input('in_mr_data', 'nifti:mod=mr', the='input mr data')
-    @IO.Output('out_data', 'bamf_processed.nii.gz', 'nifti:mod=seg:processor=bamf:roi=BREAST,FGT,BREAST+BREAST_CARCINOMA',
+    @IO.Output('out_data', 'bamf_processed.nii.gz', 'nifti:mod=seg:processor=bamf:roi=BREAST,BREAST+BREAST_FIBROGLANDULAR_TISSUE,BREAST+BREAST_CARCINOMA',
                the="get breast, fgt, and breast-tumor segmentation masks")
     def task(self, instance: Instance, in_breast_and_fgt_data: InstanceData, in_breast_tumor_data: InstanceData,
              in_mr_data: InstanceData, out_data: InstanceData):

From 833aed7af16bf317cb7bb879f7ee59a09e76ba72 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Mon, 26 Aug 2024 22:57:45 +0400
Subject: [PATCH 22/27] refactor model display name

---
 models/bamf_nnunet_mr_breast/config/default.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index dfb92b1a..c827222b 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -32,7 +32,7 @@ modules:
     in_tumor_data: nifti:mod=seg:nnunet_task=Dataset011_Breast
 
   DsegConverter:
-    model_name: Bamf NNUnet MR Breast and tumor
+    model_name: BAMF Breast and Tumor AI Segmentation
     target_dicom: dicom:mod=mr
     source_segs: nifti:mod=seg:processor=bamf
     skip_empty_slices: True

From f44935db3ee2f4203b554224a4ef45c78857b38a Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Wed, 28 Aug 2024 16:29:33 +0400
Subject: [PATCH 23/27] omit context for tissue roi

---
 models/bamf_nnunet_mr_breast/config/default.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/bamf_nnunet_mr_breast/config/default.yml b/models/bamf_nnunet_mr_breast/config/default.yml
index c827222b..70a110e1 100644
--- a/models/bamf_nnunet_mr_breast/config/default.yml
+++ b/models/bamf_nnunet_mr_breast/config/default.yml
@@ -8,7 +8,7 @@ execute:
 - NiftiConverter
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset009_Breast
-  roi: BREAST,BREAST+BREAST_FIBROGLANDULAR_TISSUE
+  roi: BREAST,BREAST_FIBROGLANDULAR_TISSUE
 - module: NNUnetRunnerV2
   nnunet_dataset: Dataset011_Breast
   roi: BREAST+BREAST_CARCINOMA

From 745ed8203fedd9d1bbdf908c548af79cb1d64528 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 26 Sep 2024 21:16:41 +0800
Subject: [PATCH 24/27] Add mhub deployment instance

---
 models/bamf_nnunet_mr_breast/mhub.toml | 2 ++
 1 file changed, 2 insertions(+)
 create mode 100644 models/bamf_nnunet_mr_breast/mhub.toml

diff --git a/models/bamf_nnunet_mr_breast/mhub.toml b/models/bamf_nnunet_mr_breast/mhub.toml
new file mode 100644
index 00000000..820d585b
--- /dev/null
+++ b/models/bamf_nnunet_mr_breast/mhub.toml
@@ -0,0 +1,2 @@
+[model.deployment]
+test = "https://zenodo.org/records/13843404"
\ No newline at end of file

From 4ae4c9acb0e06df52bbc79df9918ad0115771790 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 26 Sep 2024 21:37:35 +0800
Subject: [PATCH 25/27] Add mhub deployment instance v2

---
 models/bamf_nnunet_mr_breast/mhub.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/bamf_nnunet_mr_breast/mhub.toml b/models/bamf_nnunet_mr_breast/mhub.toml
index 820d585b..888b4387 100644
--- a/models/bamf_nnunet_mr_breast/mhub.toml
+++ b/models/bamf_nnunet_mr_breast/mhub.toml
@@ -1,2 +1,2 @@
 [model.deployment]
-test = "https://zenodo.org/records/13843404"
\ No newline at end of file
+test = "https://zenodo.org/records/13843612"
\ No newline at end of file

From 9bdce2f48659b0a8a38feb5976ef0c15542ca51a Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Mon, 30 Sep 2024 15:54:22 +0700
Subject: [PATCH 26/27] update down url

---
 models/bamf_nnunet_mr_breast/mhub.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/bamf_nnunet_mr_breast/mhub.toml b/models/bamf_nnunet_mr_breast/mhub.toml
index 888b4387..9960faae 100644
--- a/models/bamf_nnunet_mr_breast/mhub.toml
+++ b/models/bamf_nnunet_mr_breast/mhub.toml
@@ -1,2 +1,2 @@
 [model.deployment]
-test = "https://zenodo.org/records/13843612"
\ No newline at end of file
+test = "https://zenodo.org/records/13843612/files/bamf_nnunet_mr_breast.test.zip?download=1"

From d7e5d042869d77b45a5392468c5b16c3d644f685 Mon Sep 17 00:00:00 2001
From: Rahul Soni <rahul.soni@bamfhealth.com>
Date: Thu, 3 Oct 2024 20:10:29 +0700
Subject: [PATCH 27/27] add publication in meta.json

---
 models/bamf_nnunet_mr_breast/meta.json | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/models/bamf_nnunet_mr_breast/meta.json b/models/bamf_nnunet_mr_breast/meta.json
index 1bd86309..edc0d486 100644
--- a/models/bamf_nnunet_mr_breast/meta.json
+++ b/models/bamf_nnunet_mr_breast/meta.json
@@ -60,7 +60,7 @@
         "weights": "06.06.2024",
         "pub": "07.06.2024"
       },
-      "cite": "Murugesan, Gowtham Krishnan, Diana McCrumb, Mariam Aboian, Tej Verma, Rahul Soni, Fatima Memon, and Jeff Van Oss. The AIMI Initiative: AI-Generated Annotations for Imaging Data Commons Collections. arXiv preprint arXiv:2310.14897 (2023).",
+      "cite": "Murugesan, Gowtham Krishnan, Diana McCrumb, Rahul Soni, Jithendra Kumar, Leonard Nuernberg, Linmin Pei, Ulrike Wagner, Sutton Granger, Andrey Y. Fedorov, Stephen Moore, Jeff Van Oss. AI generated annotations for Breast, Brain, Liver, Lungs and Prostate cancer collections in National Cancer Institute Imaging Data Commons. arXiv preprint arXiv:2409.20342 (2024).",
       "license": {
         "code": "MIT",
         "weights": "CC BY-NC 4.0"