Merge pull request #67 from MHubAI/m-nnunet-prostate

MHub / IDC - Implementing the nnUNet Task005 Prostate MR (ADC, T2) Model

models/nnunet_prostate_zonal_task05/config/default.yml

@@ -0,0 +1,40 @@
+general:
+  data_base_dir: /app/data
+  version: 1.0
+  description: Prostate MR ADC-T2 segmentation (dicom2dicom) 
+
+execute:
+- FileStructureImporter
+- NiftiConverter
+- ProstateResampler
+- ProstateRunner
+- DsegConverter
+- DataOrganizer
+
+modules:
+  FileStructureImporter:
+    outsource_instances: True 
+    import_id: patientID/studyID
+    structures:
+    - $patientID/$studyID@instance/$part@bundle@dicom
+    - $patientID@instance:studyID=none/ADC$part@bundle@dicom
+    - $patientID@instance:studyID=none/T2$part@bundle@dicom
+
+  NiftiConverter:
+    in_datas: dicom:part=ADC|T2
+    allow_multi_input: true
+    overwrite_existing_file: true
+
+  DsegConverter:
+    model_name: nnUNet Zonal Prostate (Task05)
+    target_dicom: dicom:part=T2
+    source_segs: nifti:mod=seg:roi=*
+    body_part_examined: PROSTATE
+    skip_empty_slices: True
+    segment_id_meta_key: roi
+
+  DataOrganizer:
+    targets:
+    - DICOMSEG:mod=seg-->[i:patientID]/[i:studyID]/nnunet_prostate_zonal_task05.seg.dcm
+#   - NIFTI:mod=seg-->[i:patientID]/[i:studyID]/results.nii.gz
+#   - LOG-->[i:patientID]/[i:studyID]/logs/[d:part]/[basename]

models/nnunet_prostate_zonal_task05/dockerfiles/Dockerfile

@@ -0,0 +1,39 @@
+FROM mhubai/base:latest
+
+# Authors of the image
+LABEL authors="[email protected]"
+
+# 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
+
+# isntall additional system dependencies
+RUN apt update && apt install -y dcm2niix
+
+# install additional python dependencies
+RUN pip3 install --no-cache-dir \
+    nnunet \
+    nibabel
+
+# pull weights for platipy's nnU-Net so that the user doesn't need to every time a container is run
+ENV WEIGHTS_DIR="/root/.nnunet/nnUNet_models/nnUNet/"
+ENV WEIGHTS_URL="https://www.dropbox.com/s/igpwt45v6hlquxp/Task005_Prostate.zip"
+ENV WEIGHTS_FN="Task005_Prostate.zip"
+
+RUN wget --directory-prefix ${WEIGHTS_DIR} ${WEIGHTS_URL}
+RUN unzip ${WEIGHTS_DIR}${WEIGHTS_FN} -d ${WEIGHTS_DIR}
+RUN rm ${WEIGHTS_DIR}${WEIGHTS_FN}
+
+# Import the MHub model definiton
+ARG MHUB_MODELS_REPO
+RUN buildutils/import_mhub_model.sh nnunet_prostate_zonal_task05 ${MHUB_MODELS_REPO}
+
+# specify nnunet specific environment variables
+ENV WEIGHTS_FOLDER=$WEIGHTS_DIR
+
+# Default run script
+ENTRYPOINT ["mhub.run"]
+CMD ["--config", "/app/models/nnunet_prostate_zonal_task05/config/default.yml"]

models/nnunet_prostate_zonal_task05/meta.json

@@ -0,0 +1,178 @@
+{
+    "id": "f2eb536b-448a-4e9a-8981-3efc51301f62",
+    "name": "nnunet_prostate_zonal_task05",
+    "title": "nnU-Net (Prostate transitional zone and peripheral zone segmentation)",
+    "summary": {
+        "description": "nnU-Net's zonal prostate segmentation model is a multi-modality input AI-based pipeline for the automated segmentation of the peripheral and transition zone of the prostate on MRI 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.6 mm",
+                "non-contrast": true,
+                "contrast": false
+            },
+            {
+                "label": "ADC Input Image",
+                "description": "The ADC axial sequence being one of the two input image",
+                "format": "DICOM",
+                "modality": "MR",
+                "bodypartexamined": "Prostate",
+                "slicethickness": "3.6 mm",
+                "non-contrast": true,
+                "contrast": false
+            }
+        ],
+        "outputs": [
+            {
+                "type": "Segmentation",
+                "classes": [
+                    "PROSTATE_PERIPHERAL_ZONE",
+                    "PROSTATE_TRANSITION_ZONE"
+                ]
+            }
+        ],
+        "model": {
+            "architecture": "U-net",
+            "training": "supervised",
+            "cmpapproach": "3D"
+        },
+        "data": {
+            "training": {
+                "vol_samples": 32
+            },
+            "evaluation": {
+                "vol_samples": 16
+            },
+            "public": true,
+            "external": false
+        }
+    },
+    "details": {
+        "name": "nnU-Net Zonal prostate regions Segmentation Model",
+        "version": "1.0.0",
+        "devteam": "MIC-DKFZ (Helmholtz Imaging Applied Computer Vision Lab)",
+        "type": "nnU-Net (U-Net structure, optimized by data-driven heuristics)",
+        "date": {
+            "weights": "2020",
+            "code": "2020",
+            "pub": "2020"
+        },
+        "cite": "Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2020). nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation. Nature Methods, 1-9.",
+        "license": {
+            "code": "Apache 2.0",
+            "weights": "CC BY-NC 4.0"
+        },
+        "publications": [
+            {
+                "title": "nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation",
+                "uri": "https://www.nature.com/articles/s41592-020-01008-z"
+            }
+        ],
+        "github": "https://github.com/MIC-DKFZ/nnUNet/tree/nnunetv1",
+        "zenodo": "https://zenodo.org/record/4485926"
+    },
+    "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 3.6mm. Input ADC and T2 modalities are co-registered during training. No endorectal coil was present during training."
+        },
+        "analyses": {
+            "title": "Quantitative Analyses",
+            "text": "The model's performance was assessed using the Dice Coefficient, in the context of the Medical Segmentation Decathlon challenge. The complete breakdown of the metrics can be consulted on GrandChallenge [1] and is reported in the supplementary material to the publication [2].",
+            "references": [
+                {
+                    "label": "Medical Segmentation Decathlon on GrandChallenge",
+                    "uri": "https://decathlon-10.grand-challenge.org/evaluation/challenge/leaderboard"
+                },
+                {
+                    "label": "nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation",
+                    "uri": "https://www.nature.com/articles/s41592-020-01008-z"
+                }
+            ]
+        },
+        "evaluation": {
+            "title": "Evaluation Data",
+            "text": "The evaluation dataset consists of 16 validation samples coming from the Medical Decathlon collection.",
+            "tables": [{
+                "label": "mean DSC peripheral zone results on internal training data, using five fold cross-validation",
+                "entries": {
+                    "2D": "0.6285",
+                    "3D_fullres": "0.6663",
+                    "Best ensemble (2D + 3D_fullres)": "0.6611",
+                    "Postprocessed": "0.6611"
+                }
+            },
+            {
+                "label": "mean DSC transition zone results on internal training data, using five fold cross-validation",
+                "entries": {
+                    "2D": "0.8380",
+                    "3D_fullres": "0.8410",
+                    "Best ensemble (2D + 3D_fullres)": "0.8575",
+                    "Postprocessed": "0.8577"
+                }
+            },
+            {
+                "label": "mean DSC prostate zonal regions results on internal test data",
+                "entries": {
+                    "mean DSC for PZ": "0.77",
+                    "mean DSC for TZ": "0.90"
+                }
+            }],
+            "references": [
+                {
+                    "label": "Medical Segmentation Decathlon",
+                    "uri": "https://www.nature.com/articles/s41467-022-30695-9"
+                },
+                {
+                    "label": "Medical Decathlon Prostate dataset (direct download)",
+                    "uri": "https://drive.google.com/drive/folders/1HqEgzS8BV2c7xYNrZdEAnrHk7osJJ--2"
+                }
+            ]
+        },
+        "training": {
+            "title": "Training Data",
+            "text": "The training dataset consists of 32 MRI cases containing the prostate, from the Medical Segmentation Decathlon. The authors report the following characteristics for the portal venous phase CT scans of the training dataset:",
+            "tables": [
+                {
+                    "label": "Medical Image Decathlon dataset (training)",
+                    "entries": {
+                        "Slice Thickness": "3.6 mm",
+                        "In-Plane Resolution": "0.62 mm"
+                    }
+                }
+            ],
+            "references": [
+                {
+                    "label": "Medical Segmentation Decathlon",
+                    "uri": "https://www.nature.com/articles/s41467-022-30695-9"
+                },
+                {
+                    "label": "Medical Decathlon Prostate dataset (direct download)",
+                    "uri": "https://drive.google.com/drive/folders/1HqEgzS8BV2c7xYNrZdEAnrHk7osJJ--2"
+                }
+            ]
+        },
+        "limitations":{
+            "title": "Dealing with multi-modality input",
+            "text": "Authors recommend co-registration of ADC and T2 input sequences, as applied during training. At the very least, the ADC and T2 sequence need to have identical geometry for nnUNet to run. Since evaluated ADC and T2 sequences during evaluation might more often that not fail this requirement, we apply resampling of the ADC sequence to the T2 sequence, since T2 tends to have a higher resolution. Below are some references regarding nnUnet recommendations for multi-modality input, alongside the paper describing the registration process of Medical Image Decathlon dataset for the ADC and T2 sequences.",
+            "references": [
+                {
+                    "label": "Litjens et al., A pattern recognition approach to zonal segmentation of the prostate on MRI",
+                    "uri": "https://pubmed.ncbi.nlm.nih.gov/23286075/"
+                },
+                {
+                    "label": "Alignment of multi channel inputs for nnunet #502",
+                    "uri": "https://github.com/MIC-DKFZ/nnUNet/issues/502"
+                },
+                {
+                    "label": "Multi-modality dataset conversion issue #306",
+                    "uri": "https://github.com/MIC-DKFZ/nnUNet/issues/306"
+                }
+            ]
+        }
+    }
+}

models/nnunet_prostate_zonal_task05/utils/ProstateResampler.py

@@ -0,0 +1,36 @@
+import os
+import pyplastimatch as pypla
+
+from mhubio.core import Module, Instance, DataType, InstanceData, FileType, IO
+
+# for specific use case, resample ADC to match T2 (T2 is his 'sesired_grid' property value)
+# TODO: add reference to colab notebook?
+class ProstateResampler(Module):
+
+    @IO.Instance()
+    @IO.Input('in_data', 'nifti:part=ADC', the="ADC image")
+    @IO.Input('fixed_data', 'nifti:part=T2', the="T2 image")
+    @IO.Output('out_data', 'resampled.nii.gz', 'nifti:part=ADC:resampled_to=T2', data='in_data', the="ADC image resampled to T2")
+    def task(self, instance: Instance, in_data: InstanceData, fixed_data: InstanceData, out_data: InstanceData):
+
+        # log data
+        log_data = InstanceData('_pypla.log', DataType(FileType.LOG, in_data.type.meta + {
+            "log-origin": "plastimatch",
+            "log-task": "resampling",
+            "log-caller": "Resampler",
+            "log-instance": str(instance)
+        }), data=in_data, auto_increment=True)
+
+        # process
+        resample_args = {
+            'input': in_data.abspath,
+            'output': out_data.abspath,
+            'fixed': fixed_data.abspath,
+        }
+
+        # TODO add log file
+        pypla.resample(
+            verbose=self.config.verbose,   
+            path_to_log_file=log_data.abspath,
+            **resample_args # type: ignore
+        )

models/nnunet_prostate_zonal_task05/utils/ProstateRunner.py

@@ -0,0 +1,49 @@
+import os, shutil
+from mhubio.core import Module, Instance, InstanceData, IO
+
+@IO.Config('use_tta', bool, False, the='flag to enable test time augmentation')
+@IO.Config('nnunet_model', str, '3d_fullres', the='nnunet model name (2d, 3d_lowres, 3d_fullres, 3d_cascade_fullres)')
+class ProstateRunner(Module):
+
+    use_tta: bool
+    nnunet_model: str
+
+    @IO.Instance()
+    @IO.Input('T2', 'nifti:part=T2', the="T2 image")
+    @IO.Input('ADC', 'nifti:part=ADC:resampled_to=T2', the="ADC image resampled to T2")
+    @IO.Output('P', 'VOLUME_001.nii.gz', 'nifti:mod=seg:model=nnunet_t005_prostate:roi=PROSTATE_PERIPHERAL_ZONE,PROSTATE_TRANSITION_ZONE', bundle='nnunet-out', the="Prostate segmentation")
+    def task(self, instance: Instance, T2: InstanceData, ADC: InstanceData, P: InstanceData) -> None:
+
+        # copy input files to align with the nnunet input folder and file name format
+        # T2:  0000 
+        # ADC: 0001
+        inp_dir = self.config.data.requestTempDir(label="nnunet-model-inp")
+        inp_file_T2 =  f'VOLUME_001_0000.nii.gz'
+        inp_file_ADC = f'VOLUME_001_0001.nii.gz'
+        shutil.copyfile(T2.abspath, os.path.join(inp_dir, inp_file_T2))
+        shutil.copyfile(ADC.abspath, os.path.join(inp_dir, inp_file_ADC))
+
+        # define output folder (temp dir) and also override environment variable for nnunet
+        assert P.bundle is not None, f"Output bundle is required: {str(P)}"
+        os.environ['RESULTS_FOLDER'] = P.bundle.abspath
+
+        # 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(P.bundle.abspath, 'nnUNet'))
+
+        # construct nnunet inference command
+        bash_command  = ["nnUNet_predict"]
+        bash_command += ["--input_folder", str(inp_dir)]
+        bash_command += ["--output_folder", str(P.bundle.abspath)]
+        bash_command += ["--task_name", 'Task005_Prostate']
+        bash_command += ["--model", self.nnunet_model]
+
+        # optional / customizable arguments
+        if not self.use_tta:
+            bash_command += ["--disable_tta"]
+
+        # run command
+        self.subprocess(bash_command, text=True)