BAMF - NNUnet Prostate MR

models/bamf_nnunet_mr_prostate/config/default.yml

@@ -0,0 +1,42 @@
+general:
+  data_base_dir: /app/data
+  version: 1.0
+  description: default configuration for Bamf NNUnet MR Prostate segmentation (dicom to dicom)
+
+execute:
+- DicomImporter
+- NiftiConverter
+- NNUnetRunner
+- BamfProcessorRunner
+- DsegConverter
+- DataOrganizer
+
+modules:
+  DicomImporter:
+    source_dir: input_data
+    import_dir: sorted_data
+    sort_data: true
+    meta: 
+      mod: '%Modality'
+
+  NiftiConverter:
+    in_datas: dicom:mod=mr
+    engine: plastimatch
+
+  NNUnetRunner:
+    in_data: nifti:mod=mr
+    nnunet_task: Task788_ProstateX
+    nnunet_model: 3d_fullres
+    roi: PROSTATE
+
+  BamfProcessorRunner:
+
+  DsegConverter:
+    source_segs: nifti:mod=seg
+    target_dicom: dicom:mod=mr
+    model_name: Bamf NNUnet MR Prostate
+    skip_empty_slices: True
+
+  DataOrganizer:
+    targets:
+    - dicomseg-->[i:sid]/bamf_nnunet_mr_prostate.seg.dcm

models/bamf_nnunet_mr_prostate/config/slicer.yml

@@ -0,0 +1,33 @@
+general:
+  data_base_dir: /app/data
+  version: 1.0
+  description: configuration for Bamf NNUnet MR Prostate segmentation  in 3D Slicer (nrrd to nifti)
+
+execute:
+- NrrdImporter
+- NiftiConverter
+- NNUnetRunner
+- BamfProcessorRunner
+- JsonSegExporter
+- DataOrganizer
+
+modules:
+  NrrdImporter:
+    input_dir: input_data
+    input_file_name: image.nrrd
+
+  JsonSegExporter:
+    segment_id_meta_key: roi
+    targets:
+    - nifti:mod=seg-->[basename]
+
+  NNUnetRunner:
+    nnunet_task: Task778_MR_Prostate
+    nnunet_model: 3d_fullres
+
+  BamfProcessorRunner:
+
+  DataOrganizer:
+    targets:
+    - nifti:mod=seg-->[basename]
+    - json:mod=seg-->segdef.json

models/bamf_nnunet_mr_prostate/dockerfiles/Dockerfile

@@ -0,0 +1,32 @@
+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 
+
+# Clone the main branch of MHubAI/models
+ARG MHUB_MODELS_REPO
+RUN buildutils/import_mhub_model.sh bamf_nnunet_mr_prostate ${MHUB_MODELS_REPO}
+
+# Pull weights into the container
+ENV WEIGHTS_DIR=/root/.nnunet/nnUNet_models/nnUNet/
+RUN mkdir -p $WEIGHTS_DIR
+ENV WEIGHTS_FN=Task788_Prostate.zip
+ENV WEIGHTS_URL=https://zenodo.org/record/8290093/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_prostate/config/default.yml"]

models/bamf_nnunet_mr_prostate/meta.json

@@ -0,0 +1,137 @@
+{
+    "id": "",
+    "name": "bamf_nnunet_mr_prostate",
+    "title": "Bamf NNunet MR Prostate Model",
+    "summary": {
+        "description": "bamf_nnunet_mr_prostate model is a semantic segmentation model where an NNUnet model is finetuned for MR images",
+        "inputs": [
+            {
+                "label": "Input Image",
+                "description": "The MR scan of a patient",
+                "format": "DICOM",
+                "modality": "MR",
+                "bodypartexamined": "PROSTATE",
+                "slicethickness": "2.5mm",
+                "non-contrast": true,
+                "contrast": true
+            }
+        ],
+        "outputs": [
+            {
+                "label": "Segmentation",
+                "type": "Segmentation",
+                "description": "Segmentation Prostate",
+                "classes": [
+                    "Prostate"
+                ]
+            }
+        ],
+        "model": {
+            "architecture": "U-net",
+            "training": "supervised",
+            "cmpapproach": "3D"
+        },
+        "data": {
+            "training": {
+                "vol_samples": 439
+            },
+            "evaluation": {
+                "vol_samples": 81
+            },
+            "public": true,
+            "external": true
+        }        
+    },
+    "details": {
+        "name": "AIMI MR Prostate",
+        "version": "1.0.0",
+        "devteam": "BAMF Health",
+        "type": "nnU-Net (U-Net structure, optimized by data-driven heuristics)",
+        "date": {
+            "code": "17.10.2023",
+            "weights": "28.08.2023",
+            "pub": "23.10.2023"
+        },
+        "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-prostate-mr"
+    },
+    "info": {
+        "use": {
+            "title": "Intended Use",
+            "text": "This model is intended to perform segmentations of prostate region 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 MR prostate segmentation task in the Medical Segmentation Decathlon challenge. The model was used to segment cases from the IDC collection ProstateX [1]. Thirty four 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 prostate annotations",
+                    "entries": {
+                        "Dice: Radiologist": "1.00 (0.00)",
+                        "NSD: Radiologist": "0.00 (0.00)"
+                    }
+                }
+            ],
+            "references": [
+                {
+                    "label": "PROSTATEx",
+                    "uri": "https://wiki.cancerimagingarchive.net/pages/viewpage.action?pageId=23691656"
+                },
+                {
+                    "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 from the PROSTATEx dataset. A test/holdout validation split of 81/34 was created from the 115 prostates. These were reviewed and corrected by a board-certified radiologist. The model predictions, and radiologist corrections are published on zenodo [2]",
+            "references": [
+                {
+                    "label": "Imaging Data Collections (IDC)",
+                    "uri": "https://datacommons.cancer.gov/repository/imaging-data-commons"
+                },
+                {
+                    "label": "PROSTATEx",
+                    "uri": "https://wiki.cancerimagingarchive.net/pages/viewpage.action?pageId=23691656"
+                },
+                {
+                    "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 439 T2W MRI prostate annotations taken from IDC [1] (N=98), ProstateX dataset found publicly [2] (N=134), T2W MRI scans from Prostate158 [3] (N=138), and ISBI-MR-Prostate-2013 dataset [4] (N=69)",
+            "references": [
+                {
+                    "label": "Imaging Data Collections (IDC)",
+                    "uri": "https://datacommons.cancer.gov/repository/imaging-data-commons"
+                },
+                {
+                    "label": "ProstateX dataset",
+                    "uri": "https://wiki.cancerimagingarchive.net/pages/viewpage.action?pageId=23691656"
+                },
+                {
+                    "label": "Prostate158 dataset",
+                    "uri": "https://zenodo.org/records/6481141"
+                },
+                {
+                    "label": "ISBI-MR-Prostate-2013 dataset",
+                    "uri": "https://wiki.cancerimagingarchive.net/display/public/nci-isbi+2013+challenge+-+automated+segmentation+of+prostate+structures"
+                }
+            ]
+        }
+    }
+}

models/bamf_nnunet_mr_prostate/utils/BamfProcessorRunner.py

@@ -0,0 +1,66 @@
+"""
+-------------------------------------------------
+MHub - Run Module for ensembling nnUNet inference.
+-------------------------------------------------
+-------------------------------------------------
+Author: Rahul Soni
+Email:  [email protected]
+-------------------------------------------------
+"""
+
+from mhubio.core import Instance, InstanceData
+from mhubio.core import Module, IO
+import numpy as np
+import SimpleITK as sitk
+from skimage import measure
+import numpy as np
+
+
+
+class BamfProcessorRunner(Module):
+
+    @IO.Instance
+    @IO.Input('in_data', 'nifti:mod=ct|mr', the='input data to run nnunet on')
+    @IO.Output('out_data', 'bamf_processed.nrrd', 'nrrd:mod=seg:processor=bamf', data='in_data', the="keep the two largest connected components of the segmentation and remove all other ones")
+    def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData) -> None:
+
+        # Log bamf runner info
+        self.log("Running BamfProcessor on....")
+        self.log(f" > input data:  {in_data.abspath}")
+        self.log(f" > output data: {out_data.abspath}")
+
+        # read image
+        self.log(f"Reading image from {in_data.abspath}")
+        img_itk = sitk.ReadImage(in_data.abspath)
+        img_np = sitk.GetArrayFromImage(img_itk)
+
+        # apply post-processing
+        img_bamf_processed = self.n_connected(img_np)
+
+        # store image temporarily
+        self.log(f"Writing tmp image to {out_data.abspath}")
+        img_bamf_processed_itk = sitk.GetImageFromArray(img_bamf_processed)
+        img_bamf_processed_itk.CopyInformation(img_itk)
+        sitk.WriteImage(img_bamf_processed_itk, out_data.abspath)
+
+
+    def n_connected(self, img_data):
+        img_data_mask = np.zeros(img_data.shape)
+        img_data_mask[img_data > 0] = 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:
+                print(key, value)
+                img_filtered[blobs_labels == key] = 1
+            count += 1
+
+        img_data[img_filtered != 1] = 0
+        return img_data

models/bamf_nnunet_mr_prostate/utils/__init__.py

@@ -0,0 +1 @@
+from .BamfProcessorRunner import *