Merge pull request #71 from ccosmin97/monai_prostate158

MHUB/IDC - Implementing the Prostate158 whole prostate gland segmentation model available in MONAI (T2 only)

models/monai_prostate158/config/default.yml

@@ -0,0 +1,36 @@
+general:
+  data_base_dir: /app/data
+  version: 1.0
+  description: default configuration for MONAI Prostate158 MR Prostate zonal regions segmentation (dicom to dicom)
+
+execute:
+- DicomImporter
+- NiftiConverter
+- Prostate158Runner
+- 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: dcm2niix
+
+  Prostate158Runner:
+    in_data: nifti:mod=mr
+
+  DsegConverter:
+    source_segs: nifti:mod=seg:roi=PROSTATE_TRANSITION_ZONE,PROSTATE_PERIPHERAL_ZONE
+    target_dicom: dicom:mod=mr
+    model_name: 'MONAI Prostate158'
+    skip_empty_slices: True
+
+  DataOrganizer:
+    targets:
+    - dicomseg-->[i:sid]/monai_prostate158.seg.dcm

models/monai_prostate158/dockerfiles/Dockerfile

@@ -0,0 +1,30 @@
+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
+
+ARG MONAI_BUNDLE_DIR='https://github.com/Project-MONAI/model-zoo/releases/download/hosting_storage_v1/prostate_mri_anatomy_v0.1.0.zip'
+ARG MONAI_MODEL_NAME='prostate_mri_anatomy'
+
+# Install nnunet and platipy
+RUN python3 -m pip install --upgrade pip && pip3 install --no-cache-dir "monai[ignite]" fire nibabel simpleITK
+
+# Clone the main branch of MHubAI/models
+ARG MHUB_MODELS_REPO
+RUN buildutils/import_mhub_model.sh monai_prostate158 ${MHUB_MODELS_REPO}
+
+# Pull weights into the container
+ENV WEIGHTS_DIR=/app/models/monai_prostate158/bundle
+RUN mkdir -p $WEIGHTS_DIR
+RUN python3 -m monai.bundle download "prostate_mri_anatomy" --bundle_dir ${WEIGHTS_DIR}
+
+#define path to bundle root
+ENV BUNDLE_ROOT=/app/models/monai_prostate158/bundle/prostate_mri_anatomy
+
+# Default run script
+ENTRYPOINT ["mhub.run"]
+CMD ["--config", "/app/models/monai_prostate158/config/default.yml"]

models/monai_prostate158/meta.json

@@ -0,0 +1,131 @@
+{
+    "id": "...",
+    "name": "monai_prostate158",
+    "title": "Prostate158 (Prostate transitional zone and peripheral zone segmentation)",
+    "summary": {
+        "description": "Prostate158 is a zonal prostate segmentation model, a multi-modality input AI-based pipeline for the automated segmentation of the peripheral and central gland of the prostate on MRI T2 axial scans.",
+        "inputs": [
+            {
+                "label": "T2 input image",
+                "description": "The T2 axial sequence being one of the two input image",
+                "format": "DICOM",
+                "modality": "MR",
+                "bodypartexamined": "Prostate",
+                "slicethickness": "3 mm",
+                "non-contrast": true,
+                "contrast": false
+            }
+        ],
+        "outputs": [
+            {
+                "type": "Segmentation",
+                "classes": [
+                    "PROSTATE_TRANSITION_ZONE",
+                    "PROSTATE_PERIPHERAL_ZONE"
+                ]
+            }
+        ],
+        "model": {
+            "architecture": "U-net",
+            "training": "supervised",
+            "cmpapproach": "3D"
+        },
+        "data": {
+            "training": {
+                "vol_samples": 139
+            },
+            "evaluation": {
+                "vol_samples": 20
+            },
+            "public": true,
+            "external": false
+        }
+    },
+    "details": {
+        "name": "Prostate158 - An expert-annotated 3T MRI dataset and algorithm for prostate cancer detection",
+        "version": "1.0.0",
+        "devteam": "Lisa C. Adams, Keno K. Bressem",
+        "type": "Prostate158 (U-Net structure for prostate segmentation)",
+        "date": {
+            "weights": "March 2022",
+            "code": "April 2022",
+            "pub": "September 2022"
+        },
+        "cite": "Lisa C. Adams and Marcus R. Makowski and Günther Engel and Maximilian Rattunde and Felix Busch and Patrick Asbach and Stefan M. Niehues and Shankeeth Vinayahalingam and Bram {van Ginneken} and Geert Litjens and Keno K. Bressem, Prostate158 - An expert-annotated 3T MRI dataset and algorithm for prostate cancer detection",
+        "license": {
+            "code": "MIT",
+            "weights": "CC BY-NC 4.0"
+        },
+        "publications": [
+            {
+                "title": "Prostate158 - An expert-annotated 3T MRI dataset and algorithm for prostate cancer detection",
+                "uri": "https://doi.org/10.1016/j.compbiomed.2022.105817"
+            }
+        ],
+        "github": "https://github.com/Project-MONAI/model-zoo/tree/dev/models/prostate_mri_anatomy",
+        "zenodo": "https://zenodo.org/records/6481141"
+    },
+    "info": {
+        "use": {
+            "title": "Intended Use",
+            "text": "This model is intended to perform prostate regions anatomy segmentation in MR ADC and T2 scans. The slice thickness of the training data is 3mm. T2 input modality is used during training. To align with the model training pre-processing scheme, center-cropping of the input T2 image is recommended. No endorectal coil was present during training."
+        },
+        "analyses": {
+            "title": "Quantitative Analyses",
+            "text": "The model's performance was assessed using the Dice Coefficient, on an internal test set and ProstateX collection. The complete breakdown of the metrics can be consulted in the publication.",
+            "references": [
+                {
+                    "label": "Prostate158 - An expert-annotated 3T MRI dataset and algorithm for prostate cancer detection",
+                    "uri": "https://doi.org/10.1016/j.compbiomed.2022.105817"
+                }
+            ]
+        },
+        "evaluation": {
+            "title": "External Evaluation Data",
+            "text": "The evaluation datasets consist of 186 ProstateX samples and 32 prostate MRI Medical Decathlon dataset samples.",
+            "tables": [
+                {
+                    "label": "Medical Decathlon mean DSC for the segmentation of the central gland and peripheral zone",
+                    "entries": {
+                        "Central gland": "0.82",
+                        "Peripheral zone": "0.64"
+                    }
+                },
+                {
+                    "label": "ProstateX mean DSC for the segmentation of the central gland and peripheral zone",
+                    "entries": {
+                        "Central gland": "0.86",
+                        "Peripheral zone": "0.71"
+                    }
+                }
+            ],
+            "references": [{
+                    "label": "Medical Segmentation Decathlon",
+                    "uri": "https://www.nature.com/articles/s41467-022-30695-9"
+                },
+                {
+                    "label": "Quality control and whole-gland, zonal and lesion annotations for the PROSTATEx challenge public dataset",
+                    "uri": "https://www.sciencedirect.com/science/article/abs/pii/S0720048X21001273"
+                }]
+        },
+        "training": {
+            "title": "Training Data",
+            "text": "The training dataset consists of 139 MRI cases containing the prostate, from the Prostate158 collection. The authors report the following characteristics for the T2 imaging sequeneces:",
+            "tables": [
+                {
+                    "label": "Prostate158 dataset (training)",
+                    "entries": {
+                        "Slice Thickness": "3 mm",
+                        "In-Plane Resolution": "0.47 mm"
+                    }
+                }
+            ],
+            "references": [
+                {
+                    "label": "Prostate158 dataset (Zenodo access)",
+                    "uri": "https://zenodo.org/records/6481141"
+                }
+            ]
+        }
+    }
+}

models/monai_prostate158/utils/Prostate158Runner.py

@@ -0,0 +1,83 @@
+"""
+-------------------------------------------------
+MHub - MONAI Prostate158 Runner
+-------------------------------------------------
+
+-------------------------------------------------
+Author: Cosmin Ciausu
+Email:  [email protected]
+-------------------------------------------------
+"""
+# TODO: support multi-i/o and batch processing on multiple instances
+
+from typing import List, Optional
+import os, subprocess, shutil, glob, sys
+import SimpleITK as sitk, numpy as np
+from mhubio.core import Module, Instance, InstanceData, DataType, FileType, IO
+from mhubio.modules.runner.ModelRunner import ModelRunner
+import json
+
+@IO.Config('apply_center_crop', bool, True, the='flag to apply center cropping to input_image')
+class Prostate158Runner(Module):
+
+    apply_center_crop : bool
+
+    @IO.Instance()
+    @IO.Input("in_data", 'nifti:mod=mr', the="input T2 sequence data to run prostate158 on")
+    @IO.Output('out_data', 'monai_prostate158.nii.gz', 
+               'nifti:mod=seg:model=MonaiProstate158:roi=PROSTATE_TRANSITION_ZONE,PROSTATE_PERIPHERAL_ZONE', 
+               data='in_data', bundle='model', the="predicted segmentation model")
+    def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData) -> None:
+
+        # 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')
+        datalist = [in_data.abspath]
+
+        if self.apply_center_crop:
+            in_dir_cropped = self.config.data.requestTempDir(label="monai-crop-in")
+            in_data_processed = os.path.join(in_dir_cropped, "image_cropped.nii.gz")
+            self.subprocess([sys.executable, f"{os.path.join(os.environ['BUNDLE_ROOT'], 'scripts', 'center_crop.py')}",
+            "--file_name", in_data.abspath, "--out_name",in_data_processed], text=True)
+            datalist = [in_data_processed]
+
+        # define output folder (temp dir) and also override environment variable for nnunet
+        out_dir = self.config.data.requestTempDir(label="monai-model-out")
+
+        bash_command = [sys.executable, 
+        "-m", "monai.bundle", "run", "evaluating"]
+        bash_command += ["--meta_file", os.path.join(os.environ['BUNDLE_ROOT'], "configs", "metadata.json")]
+        bash_command += ["--config_file", os.path.join(os.environ['BUNDLE_ROOT'], "configs", "inference.json")]
+        bash_command += ["--datalist", str(datalist)]
+        bash_command += ["--output_dir", out_dir]
+        bash_command += ["--bundle_root", os.environ['BUNDLE_ROOT']]
+        bash_command += ["--dataloader#num_workers", "0"]
+        print(bash_command)
+        self.subprocess(bash_command, text=True)
+
+        # get output data
+        out_path = glob.glob(os.path.join(out_dir, "**", 
+        "*.nii.gz"), recursive=True)[0]
+
+        if self.apply_center_crop:
+            out_dir_padded = self.config.data.requestTempDir(label="monai-padded-out")
+            out_data_padded = os.path.join(out_dir_padded, "seg_padded.nii.gz")
+            paddedFilter = sitk.ConstantPadImageFilter()
+            seg_image = sitk.ReadImage(out_path)
+            t2_image = sitk.ReadImage(in_data.abspath)
+            out_seg_shape = sitk.GetArrayFromImage(seg_image).shape
+            t2_image_shape = sitk.GetArrayFromImage(t2_image).shape
+            x_bound_lower =  int((t2_image_shape[2] - out_seg_shape[2])/2)
+            x_bound_upper =  int(int(t2_image_shape[2] - out_seg_shape[2])/2 + ((t2_image_shape[2] - out_seg_shape[2]) % 2))
+            y_bound_lower = int((t2_image_shape[1] - out_seg_shape[1])/2)
+            y_bound_upper = int(int(t2_image_shape[1] - out_seg_shape[1])/2 + ((t2_image_shape[1] - out_seg_shape[1]) % 2))
+            paddedFilter.SetConstant(0)
+            paddedFilter.SetPadLowerBound([x_bound_lower, y_bound_lower, 0])
+            paddedFilter.SetPadUpperBound([x_bound_upper, y_bound_upper, 0])
+            padded_img = paddedFilter.Execute(seg_image)
+            sitk.WriteImage(padded_img, out_data_padded)
+            out_path = out_data_padded
+
+        # copy output data to instance
+        shutil.copyfile(out_path, out_data.abspath)

models/monai_prostate158/utils/__init__.py

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