diff --git a/models/bamf_nnunet_mr_liver/config/default.yml b/models/bamf_nnunet_mr_liver/config/default.yml new file mode 100644 index 00000000..87f6da53 --- /dev/null +++ b/models/bamf_nnunet_mr_liver/config/default.yml @@ -0,0 +1,40 @@ +general: + data_base_dir: /app/data + version: 1.0 + description: default configuration for Bamf NNUnet MR Liver 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: Task822_Liver_AMD + nnunet_model: 3d_fullres + roi: LIVER + + DsegConverter: + source_segs: nifti:mod=seg + target_dicom: dicom:mod=mr + model_name: Bamf NNUnet MR Liver + skip_empty_slices: True + + DataOrganizer: + targets: + - dicomseg-->[i:sid]/bamf_nnunet_mr_liver.seg.dcm \ No newline at end of file diff --git a/models/bamf_nnunet_mr_liver/config/slicer.yml b/models/bamf_nnunet_mr_liver/config/slicer.yml new file mode 100644 index 00000000..30ce6af6 --- /dev/null +++ b/models/bamf_nnunet_mr_liver/config/slicer.yml @@ -0,0 +1,34 @@ +general: + data_base_dir: /app/data + version: 1.0 + description: configuration for Bamf NNUnet MR Liver 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: 'Task882_MR_Liver' + nnunet_model: '3d_fullres' + roi: LIVER + + BamfProcessorRunner: + + DataOrganizer: + targets: + - nifti:mod=seg-->[basename] + - json:mod=seg-->segdef.json \ No newline at end of file diff --git a/models/bamf_nnunet_mr_liver/dockerfiles/Dockerfile b/models/bamf_nnunet_mr_liver/dockerfiles/Dockerfile new file mode 100644 index 00000000..ba8050f2 --- /dev/null +++ b/models/bamf_nnunet_mr_liver/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_liver ${MHUB_MODELS_REPO} + +# Pull weights into the container +ENV WEIGHTS_DIR=/root/.nnunet/nnUNet_models/nnUNet/ +RUN mkdir -p $WEIGHTS_DIR +ENV WEIGHTS_FN=Task882_MR_Liver.zip +ENV WEIGHTS_URL=https://zenodo.org/record/8290124/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_liver/config/default.yml"] \ No newline at end of file diff --git a/models/bamf_nnunet_mr_liver/meta.json b/models/bamf_nnunet_mr_liver/meta.json new file mode 100644 index 00000000..e8d82946 --- /dev/null +++ b/models/bamf_nnunet_mr_liver/meta.json @@ -0,0 +1,137 @@ +{ + "id": "", + "name": "bamf_nnunet_mr_liver", + "title": "AIMI MRI Liver", + "summary": { + "description": "An nnU-Net based model to segment liver from T1 weighted MRI scans", + "inputs": [ + { + "label": "Input Image", + "description": "The abdominal T1 weighted MRI scan of a patient.", + "format": "DICOM", + "modality": "MR", + "bodypartexamined": "LIVER", + "slicethickness": "2.5mm", + "non-contrast": true, + "contrast": false + } + ], + "outputs": [ + { + "label": "Segmentation", + "type": "Segmentation", + "description": "Segmentation liver", + "classes": [ + "LIVER" + ] + } + ], + "model": { + "architecture": "U-net", + "training": "supervised", + "cmpapproach": "3D" + }, + "data": { + "training": { + "vol_samples": 350 + }, + "evaluation": { + "vol_samples": 7 + }, + "public": true, + "external": true + } + }, + "details": { + "name": "AIMI MRI Liver", + "version": "1.0.0", + "devteam": "BAMF Health", + "authors": [ + "Soni, Rahul", + "McCrumb, Diana", + "Murugesan, Gowtham Krishnan", + "Van Oss, Jeff" + ], + "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-liver-mr" + }, + "info": { + "use": { + "title": "Intended Use", + "text": "This model is intended to perform liver segmentation in T1 weighted MRI 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 CT Liver segmentation task in the Medical Segmentation Decathlon challenge. The model was used to segment 67 cases from the IDC collection TCGA-LIHC [1]. Seven 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 MRI liver annotations", + "entries": { + "Dice: Radiologist": "0.91 (0.18)", + "NSD: Radiologist": "0.89 (0.20)", + "Dice: Non-expert": "0.90 (0.15)", + "NSD: Non-expert": "0.85 (0.20)" + } + } + ], + "references": [ + { + "label": "TCGA-LIHC", + "uri": "https://wiki.cancerimagingarchive.net/pages/viewpage.action?pageId=6885436" + }, + { + "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 67 cases from the IDC [1] collection TCGA-LIHC [1]. Seven 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": "Imaging Data Collections (IDC)", + "uri": "https://datacommons.cancer.gov/repository/imaging-data-commons" + }, + { + "label": "TCGA-LIHC", + "uri": "https://wiki.cancerimagingarchive.net/pages/viewpage.action?pageId=6885436" + }, + { + "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 350 MRI liver annotations taken from the AMOS [1] (N=40) and DUKE Liver Dataset V2 [2] (N=310).", + "references": [ + { + "label": "AMOS Dataset", + "uri": "https://zenodo.org/records/7262581" + }, + { + "label": "Duke Liver Dataset (MRI) v2", + "uri": "https://zenodo.org/records/7774566" + } + ] + } + } +} \ No newline at end of file diff --git a/models/bamf_nnunet_mr_liver/utils/BamfProcessorRunner.py b/models/bamf_nnunet_mr_liver/utils/BamfProcessorRunner.py new file mode 100644 index 00000000..53fc382a --- /dev/null +++ b/models/bamf_nnunet_mr_liver/utils/BamfProcessorRunner.py @@ -0,0 +1,66 @@ +""" +------------------------------------------------- +MHub - Run Module for ensembling nnUNet inference. +------------------------------------------------- +------------------------------------------------- +Author: Rahul Soni +Email: rahul.soni@bamfhealth.com +------------------------------------------------- +""" + +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 diff --git a/models/bamf_nnunet_mr_liver/utils/__init__.py b/models/bamf_nnunet_mr_liver/utils/__init__.py new file mode 100644 index 00000000..d6522730 --- /dev/null +++ b/models/bamf_nnunet_mr_liver/utils/__init__.py @@ -0,0 +1 @@ +from .BamfProcessorRunner import * \ No newline at end of file