Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

BAMF - NNUnet Liver MR #45

Merged
merged 20 commits into from
Mar 29, 2024
Merged
Show file tree
Hide file tree
Changes from 19 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 40 additions & 0 deletions models/bamf_nnunet_mr_liver/config/default.yml
Original file line number Diff line number Diff line change
@@ -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
34 changes: 34 additions & 0 deletions models/bamf_nnunet_mr_liver/config/slicer.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
general:
rahul99 marked this conversation as resolved.
Show resolved Hide resolved
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'
rahul99 marked this conversation as resolved.
Show resolved Hide resolved
input_file_name: 'image.nrrd'
rahul99 marked this conversation as resolved.
Show resolved Hide resolved

JsonSegExporter:
segment_id_meta_key: roi
targets:
- nifti:mod=seg-->[basename]

NNUnetRunner:
nnunet_task: 'Task882_MR_Liver'
rahul99 marked this conversation as resolved.
Show resolved Hide resolved
nnunet_model: '3d_fullres'
rahul99 marked this conversation as resolved.
Show resolved Hide resolved
roi: LIVER

BamfProcessorRunner:
rahul99 marked this conversation as resolved.
Show resolved Hide resolved

DataOrganizer:
targets:
- nifti:mod=seg-->[basename]
- json:mod=seg-->segdef.json
32 changes: 32 additions & 0 deletions models/bamf_nnunet_mr_liver/dockerfiles/Dockerfile
Original file line number Diff line number Diff line change
@@ -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"]
137 changes: 137 additions & 0 deletions models/bamf_nnunet_mr_liver/meta.json
Original file line number Diff line number Diff line change
@@ -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",
LennyN95 marked this conversation as resolved.
Show resolved Hide resolved
"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": {
LennyN95 marked this conversation as resolved.
Show resolved Hide resolved
"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 cases from the IDC collection TCGA-LIHC [1]. Nine of those cases were reviewed and corrected by a board-certified radiologist and a non-expert. The analysis is published here [2]",
LennyN95 marked this conversation as resolved.
Show resolved Hide resolved
"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 cases from the IDC [1] collection TCGA-LIHC [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]",
LennyN95 marked this conversation as resolved.
Show resolved Hide resolved
"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"
}
]
}
}
}
66 changes: 66 additions & 0 deletions models/bamf_nnunet_mr_liver/utils/BamfProcessorRunner.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
"""
rahul99 marked this conversation as resolved.
Show resolved Hide resolved
-------------------------------------------------
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
1 change: 1 addition & 0 deletions models/bamf_nnunet_mr_liver/utils/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
from .BamfProcessorRunner import *
Loading