[FEAT]: BAMF nnUNet Breast MR Model

models/bamf_nnunet_mr_breast/config/default.yml

@@ -0,0 +1,65 @@
+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)
+
+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
+- module: NNUnetRunnerV2
+  nnunet_dataset: Dataset011_Breast
+  roi: BREAST+BREAST_CARCINOMA
+- BreastPostProcessor
+- DsegConverter
+- DataOrganizer
+
+modules:
+  DicomImporter:
+    source_dir: input_data
+    import_dir: sorted_data
+    sort_data: true
+    meta: 
+      mod: '%Modality'
+
+  NNUnetRunnerV2:
+    in_data: nifti:mod=mr
+
+  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

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 nnunetv2
+
+# 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/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}
+
+# 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/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}
+
+# 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"]

models/bamf_nnunet_mr_breast/meta.json

@@ -0,0 +1,134 @@
+{
+    "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",
+          "slicethickness": "2.5mm",
+          "non-contrast": true,
+          "contrast": false
+        }
+      ],
+      "outputs": [
+        {
+          "label": "Segmentation",
+          "type": "Segmentation",
+          "description": "Segmentation for breast, fgt, and breast carcinoma",
+          "classes": [
+            "BREAST",
+            "FGT",
+            "BREAST+BREAST_CARCINOMA"
+          ]
+        }
+      ],
+      "model": {
+        "architecture": "U-net v2",
+        "training": "supervised",
+        "cmpapproach": "3D"
+      },
+      "data": {
+        "training": {
+          "vol_samples": 587
+        },
+        "evaluation": {
+          "vol_samples": 92
+        },
+        "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": "",
+          "uri": ""
+        }
+      ],
+      "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 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",
+            "entries": {
+              "Dice: BREAST": "0.98",
+              "Dice: FGT": "0.80",
+              "Dice: BREAST+TUMOR": "0.56"
+            }
+          }
+        ],
+        "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 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",
+            "uri": "https://www.cancerimagingarchive.net/collection/acrin-contralateral-breast-mr"
+          },
+          {
+            "label": "Image segmentations produced by the AIMI Annotations initiative",
+            "uri": "https://zenodo.org/records/13244892"
+          }
+        ]
+      },
+      "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"
+          }
+        ]
+      }
+    }
+  }

models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py

@@ -0,0 +1,99 @@
+"""
+-------------------------------------------------
+MHub - Run Module for ensembling nnUNet inference.
+-------------------------------------------------
+-------------------------------------------------
+Author: Rahul Soni
+Email:  [email protected]
+-------------------------------------------------
+"""
+
+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
+from typing import Union
+from pathlib import Path
+import yaml
+
+
+
+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: 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)
+        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 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(breast_and_fgt_seg)
+        seg_img.CopyInformation(ref)
+        return seg_img
+
+    @IO.Instance()
+    @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',
+               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):
+
+        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 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 segmentation
+        tumor_seg = self.get_mask(in_breast_tumor_data.abspath)
+        tumor_seg = self.n_connected(tumor_seg)
+
+        # Merged segmentation masks
+        output_seg = self.merge_segmentations(
+            breast_and_fgt_seg = np.copy(breast_and_fgt_seg),
+            tumor_seg = np.copy(tumor_seg),
+            mr_path = in_mr_data.abspath
+            )
+        sitk.WriteImage(
+            output_seg,
+            out_data.abspath,
+        )
+