Skip to content

Commit ef742c0

Browse files
FJDorfnerFJDorfner
authored
[PW41] Add MRSegmentator Model (#90)
Pull Request to add the MRSegmentator model to mhub.
1 parent abb66b1 commit ef742c0

File tree

8 files changed

+258
-0
lines changed

8 files changed

+258
-0
lines changed

‎models/mrsegmentator/README.md

+3
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
# MRSegmentator @ MHub.ai
2+
3+
For details on the model or how to run the end-to-end pipeline on your data in a single command, visit [mhub.ai/models/mrsegmentator](https://mhub.ai/models/mrsegmentator)

‎models/mrsegmentator/__init__.py

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
from .utils import *

‎models/mrsegmentator/config/default.yml

+32
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
general:
2+
data_base_dir: /app/data
3+
version: 1.0.0
4+
description: MRSegmentator default config (dicom to dicom)
5+
6+
execute:
7+
- DicomImporter
8+
- NiftiConverter
9+
- MRSegmentatorMLRunner
10+
- DsegConverter
11+
- DataOrganizer
12+
13+
modules:
14+
DicomImporter:
15+
source_dir: input_data
16+
import_dir: sorted_data
17+
sort_data: true
18+
meta:
19+
mod: '%Modality'
20+
21+
MRSegmentatorMLRunner:
22+
23+
DsegConverter:
24+
model_name: MRSegmentator
25+
body_part_examined: WHOLEBODY
26+
source_segs: nifti:mod=seg
27+
skip_empty_slices: True
28+
target_dicom: dicom:mod=ct|mr
29+
30+
DataOrganizer:
31+
targets:
32+
- dicomseg:mod=seg-->[i:sid]/MRSegmentator.seg.dcm

‎models/mrsegmentator/dockerfiles/Dockerfile

+28
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
FROM mhubai/base:latest
2+
3+
# clone mhub implementation
4+
ARG MHUB_MODELS_REPO
5+
RUN buildutils/import_mhub_model.sh mrsegmentator ${MHUB_MODELS_REPO}
6+
7+
8+
# Note: There currently is a dependency issue for nnunet because of batchgenerators. This solves the issue in the description file
9+
RUN git clone https://github.com/MIC-DKFZ/batchgenerators.git /tmp/batchgenerators && \
10+
cd /tmp/batchgenerators && \
11+
git checkout 1185d57bbc002f4b88c03fdea885dc28537ad8e7
12+
13+
# Fix the setup.cfg file - change description-file to description_file
14+
RUN sed -i 's/description-file = README.md/description_file = README.md/' /tmp/batchgenerators/setup.cfg
15+
# Install batchgenerators
16+
RUN uv pip install -e /tmp/batchgenerators
17+
18+
# Install MRSegmentator
19+
RUN uv pip install mrsegmentator==1.2.3
20+
21+
# Execute mrsegmentator once to download the weights, this will fail to run but download regardless
22+
RUN touch .temp_image.nii.gz
23+
RUN uv run mrsegmentator -i .temp_image.nii.gz; exit 0
24+
RUN rm .temp_image.nii.gz
25+
26+
# Default run script
27+
ENTRYPOINT ["mhub.run"]
28+
CMD ["--config", "/app/models/mrsegmentator/config/default.yml"]

‎models/mrsegmentator/meta.json

+140
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
{
2+
"id": "3b8e2f1d-4c5a-4b2a-8f3e-6a7d9e2c1b0d",
3+
"name": "mrsegmentator",
4+
"title": "MRSegmentator",
5+
"summary": {
6+
"description": "MRSegmentator is an AI-based pipeline for the segmentation of 40 anatomical structures in MR images (with and without contrast).",
7+
"inputs": [
8+
{
9+
"label": "Input Image",
10+
"description": "The MR or CT scan of a patient (Thorax, Abdomen and Pelvis).",
11+
"format": "DICOM",
12+
"modality": "MRI|CT",
13+
"bodypartexamined": "WHOLEBODY",
14+
"slicethickness": "n/a",
15+
"non-contrast": true,
16+
"contrast": true
17+
}
18+
],
19+
"outputs": [
20+
{
21+
"type": "Segmentation",
22+
"classes": [
23+
"SPLEEN",
24+
"RIGHT_KIDNEY",
25+
"LEFT_KIDNEY",
26+
"GALLBLADDER",
27+
"LIVER",
28+
"STOMACH",
29+
"PANCREAS",
30+
"RIGHT_ADRENAL_GLAND",
31+
"LEFT_ADRENAL_GLAND",
32+
"LEFT_LUNG",
33+
"RIGHT_LUNG",
34+
"HEART",
35+
"AORTA",
36+
"INFERIOR_VENA_CAVA",
37+
"PORTAL_AND_SPLENIC_VEIN",
38+
"LEFT_ILIAC_ARTERY",
39+
"RIGHT_ILIAC_ARTERY",
40+
"ESOPHAGUS",
41+
"SMALL_INTESTINE",
42+
"DUODENUM",
43+
"COLON",
44+
"URINARY_BLADDER",
45+
"SPINE",
46+
"SACRUM",
47+
"LEFT_HIP",
48+
"RIGHT_HIP",
49+
"LEFT_FEMUR",
50+
"RIGHT_FEMUR",
51+
"LEFT_AUTOCHTHONOUS_BACK_MUSCLE",
52+
"RIGHT_AUTOCHTHONOUS_BACK_MUSCLE",
53+
"LEFT_ILIOPSOAS",
54+
"RIGHT_ILIOPSOAS",
55+
"LEFT_GLUTEUS_MAXIMUS",
56+
"RIGHT_GLUTEUS_MAXIMUS",
57+
"LEFT_GLUTEUS_MEDIUS",
58+
"RIGHT_GLUTEUS_MEDIUS",
59+
"LEFT_GLUTEUS_MINIMUS",
60+
"RIGHT_GLUTEUS_MINIMUS"
61+
]
62+
}
63+
],
64+
"model": {
65+
"architecture": "U-net",
66+
"training": "supervised",
67+
"cmpapproach": "ensemble"
68+
},
69+
"data": {
70+
"training": {
71+
"vol_samples": 2649
72+
},
73+
"evaluation": {
74+
"vol_samples": 960
75+
},
76+
"public": true,
77+
"external": true
78+
}
79+
},
80+
"details": {
81+
"name": "MRSegmentator",
82+
"version": "1.1.0",
83+
"devteam": "AI-Assisted Healthcare Lab at Technical University Munich and Charité Universitätsmedizin Berlin",
84+
"type": "nnU-Net (U-Net structure, optimized by data-driven heuristics)",
85+
"date": {
86+
"weights": "05/18/24",
87+
"code": "05/18/24",
88+
"pub": "2024"
89+
},
90+
"cite": "Hartmut Häntze, Lina Xu, Felix J. Dorfner, Leonhard Donle, Daniel Truhn, Hugo Aerts, Mathias Prokop, Bram van Ginneken, Alessa Hering, Lisa C. Adams, and Keno K. Bressem. MRSegmentator: Robust multi-modality segmentation of 40 classes in MRI and CT sequences. arXiv, 2024.",
91+
"license": {
92+
"code": "Apache 2.0",
93+
"weights": "Apache 2.0"
94+
},
95+
"publications": [
96+
{
97+
"title": "MRSegmentator: Robust multi-modality segmentation of 40 classes in MRI and CT sequences",
98+
"uri": "https://arxiv.org/pdf/2405.06463"
99+
}
100+
],
101+
"github": "https://github.com/hhaentze/MRSegmentator",
102+
"slicer": false
103+
},
104+
"info": {
105+
"use": {
106+
"title": "Intended Use",
107+
"text": "Contrary to CT scans, where tools for automatic multi-structure segmentation are quite mature, segmentation tasks in MRI scans are often either focused on the brain region or on a subset of few organs in other body regions. MRSegmentator aims to extend this and accurately segment 40 organs and structures in human MRI scans of the abdominal, pelvic and thorax regions. The segmentation works well on different sequence types, including T1- and T2-weighted, Dixon sequences and even CT images."
108+
},
109+
"analyses": {
110+
"title": "Quantitative Analyses",
111+
"text": "The model's performance was assessed using the Dice Coefficient on three different external datasets. For more information, please refer to the model's publication [1].",
112+
"references": [
113+
{
114+
"label": "MRSegmentator: Robust multi-modality segmentation of 40 classes in MRI and CT sequences",
115+
"uri": "https://arxiv.org/pdf/2405.06463"
116+
}
117+
]
118+
},
119+
"evaluation": {
120+
"title": "Evaluation Data",
121+
"text": "The model was evaluated on test data from three different dataset. The NAKO dataset included 600 MRI examinationsm, the AMOS MRI dataset included 60 MRI examinations and the AMOS CT dataset which included a total of 300 CT examinations.",
122+
"references": [
123+
{
124+
"label": "MRSegmentator: Robust multi-modality segmentation of 40 classes in MRI and CT sequences",
125+
"uri": "https://arxiv.org/pdf/2405.06463"
126+
}
127+
]
128+
},
129+
"training": {
130+
"title": "Training Data",
131+
"text": "Three different datasets were used for training. The in-house dataset included 221 MRI scans with T1, T2 and T1fs post contrast sequences. The UK Biobank dataset used consisted of 1200 MRI examinations acquired using IN, OPP, W, F sequences. Furthermore 1228 CT examinations from the TotalSegmentator dataset were used. All segmentations were performed by a radiologist, using a human-in-the-loop annotation approach to efficiently create high-quality segmentations for the training data.",
132+
"references": [
133+
{
134+
"label": "MRSegmentator: Robust multi-modality segmentation of 40 classes in MRI and CT sequences",
135+
"uri": "https://arxiv.org/pdf/2405.06463"
136+
}
137+
]
138+
}
139+
}
140+
}

‎models/mrsegmentator/mhub.toml

+2
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
[model.deployment]
2+
test = "https://zenodo.org/records/15079716/files/mrsegmentator.test.zip"

‎models/mrsegmentator/utils/MRSegmentatorMLRunner.py

+52
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
"""
2+
-------------------------------------------------
3+
MHub - Run Module for MRSegmentator.
4+
-------------------------------------------------
5+
6+
-------------------------------------------------
7+
Author: Felix Dorfner
8+
Email: [email protected]
9+
-------------------------------------------------
10+
"""
11+
12+
from mhubio.core import Module, Instance, InstanceData, DataType, FileType, CT, SEG, IO, DataTypeQuery
13+
import os, subprocess
14+
from segdb.classes.Segment import Segment
15+
import shutil
16+
17+
18+
19+
# register custom segmentation before class definition
20+
Segment.register("SPINE", name="Spine")
21+
22+
class MRSegmentatorMLRunner(Module):
23+
24+
@IO.Instance()
25+
@IO.Input('in_data', 'nifti:mod=ct|mr', the="input whole body mr/ct scan")
26+
@IO.Output('out_data', 'segmentations.nii.gz', 'nifti:mod=seg:model=MRSegmentator:roi=SPLEEN,RIGHT_KIDNEY,LEFT_KIDNEY,GALLBLADDER,LIVER,STOMACH,PANCREAS,RIGHT_ADRENAL_GLAND,LEFT_ADRENAL_GLAND,LEFT_LUNG,RIGHT_LUNG,HEART,AORTA,INFERIOR_VENA_CAVA,PORTAL_AND_SPLENIC_VEIN,LEFT_ILIAC_ARTERY,RIGHT_ILIAC_ARTERY,ESOPHAGUS,SMALL_INTESTINE,DUODENUM,COLON,URINARY_BLADDER,SPINE,SACRUM,LEFT_HIP,RIGHT_HIP,LEFT_FEMUR,RIGHT_FEMUR,LEFT_AUTOCHTHONOUS_BACK_MUSCLE,RIGHT_AUTOCHTHONOUS_BACK_MUSCLE,LEFT_ILIOPSOAS,RIGHT_ILIOPSOAS,LEFT_GLUTEUS_MAXIMUS,RIGHT_GLUTEUS_MAXIMUS,LEFT_GLUTEUS_MEDIUS,RIGHT_GLUTEUS_MEDIUS,LEFT_GLUTEUS_MINIMUS,RIGHT_GLUTEUS_MINIMUS', data='in_data', the="output segmentation mask containing all labels")
27+
def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData) -> None:
28+
29+
tmp_dir = self.config.data.requestTempDir("mr_segmentator")
30+
31+
bash_command = ["mrsegmentator"]
32+
bash_command += ["-i", in_data.abspath]
33+
bash_command += ["--outdir", tmp_dir]
34+
bash_command += ["--split_level", "1"]
35+
36+
self.v(">> run: ", " ".join(bash_command))
37+
38+
# run the model
39+
self.subprocess(bash_command, text=True)
40+
41+
# Find the output file in the temporary directory
42+
output_file = None
43+
for filename in os.listdir(tmp_dir):
44+
if filename.endswith('.nii.gz'):
45+
output_file = filename
46+
break
47+
48+
if output_file is None:
49+
raise FileNotFoundError("No output segmentation files found in the temporary directory.")
50+
51+
# copy data
52+
shutil.copy(os.path.join(tmp_dir, output_file), out_data.abspath)

‎models/mrsegmentator/utils/__init__.py

Whitespace-only changes.

0 commit comments

Comments
 (0)