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..70a110e1 --- /dev/null +++ b/models/bamf_nnunet_mr_breast/config/default.yml @@ -0,0 +1,42 @@ +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: NNUnetRunnerV2 + nnunet_dataset: Dataset009_Breast + roi: BREAST,BREAST_FIBROGLANDULAR_TISSUE +- 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 Breast and Tumor AI Segmentation + 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..cd8a0b7b --- /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 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"] diff --git a/models/bamf_nnunet_mr_breast/meta.json b/models/bamf_nnunet_mr_breast/meta.json new file mode 100644 index 00000000..edc0d486 --- /dev/null +++ b/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, 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" + }, + "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" + } + ] + } + } + } diff --git a/models/bamf_nnunet_mr_breast/mhub.toml b/models/bamf_nnunet_mr_breast/mhub.toml new file mode 100644 index 00000000..9960faae --- /dev/null +++ b/models/bamf_nnunet_mr_breast/mhub.toml @@ -0,0 +1,2 @@ +[model.deployment] +test = "https://zenodo.org/records/13843612/files/bamf_nnunet_mr_breast.test.zip?download=1" 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..4d580916 --- /dev/null +++ b/models/bamf_nnunet_mr_breast/utils/BreastPostProcessor.py @@ -0,0 +1,99 @@ +""" +------------------------------------------------- +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 +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,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): + + 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, + ) + 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..583d90db --- /dev/null +++ b/models/bamf_nnunet_mr_breast/utils/NNUnetRunnerV2.py @@ -0,0 +1,105 @@ +""" +------------------------------------------------- +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('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 + input_data_type: DataType + folds: int # TODO: support optional config attributes + roi: str + + @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 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}") + + # 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'] = os.environ['WEIGHTS_FOLDER'] # set model weights path + + # 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", "3d_fullres"] + + # add optional arguments + if self.folds is not None: + bash_command += ["--folds", str(self.folds)] + + # run command + self.subprocess(bash_command, text=True) + + # output meta + meta = { + "model": "nnunet-v2", + "nnunet_dataset": self.nnunet_dataset, + "nnunet_config": "3d_fullres", + "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) + + # update meta dynamically + out_data.type.meta += meta