Added segdb export, removed dseg.json, added remapped output to runne…

…r, updated Dockerfile and config

models/gc_nnunet_pancreas/config/default.yml

@@ -16,20 +16,20 @@ modules:
     import_dir: sorted_data
     sort_data: true
     meta: 
-      mod: ct
+      mod: '%Modality'
 
   MhaConverter:
     engine: panimg
+    targets: [dicom:mod=ct]
 
   DsegConverter:
     model_name: 'GC NNUnet Pancreas'
-    source_segs: ['mha:mod=seg']
+    source_segs: ['mha:mod=seg:type=remapped']
     target_dicom: dicom:mod=ct
     skip_empty_slices: True
-    json_config_path: /app/models/gc_nnunet_pancreas/config/dseg.json
 
   DataOrganizer:
     targets:
     - mha:mod=heatmap-->[i:sid]/nnunet_pancreas_heatmap.mha
-    - mha:mod=seg-->[i:sid]/nnunet_pancreas.seg.mha
+    - mha:mod=seg:type=original-->[i:sid]/nnunet_pancreas.seg.mha
     - dicomseg:mod=seg-->[i:sid]/nnunet_pancreas.seg.dcm

models/gc_nnunet_pancreas/dockerfiles/Dockerfile

@@ -3,27 +3,38 @@ FROM mhubai/base:latest
 # Specify/override authors label
 LABEL authors="[email protected]"
 
+# Install PyTorch 2.0.1 (CUDA enabled)
+RUN pip3 install --no-cache-dir torch==2.0.1+cu118 -f https://download.pytorch.org/whl/torch_stable.html
+
 # Install git-lfs (required for downloading the model weights)
-RUN apt update && apt install -y --no-install-recommends git-lfs && rm -rf /var/lib/apt/lists/*
+RUN apt update && \
+    apt install -y --no-install-recommends git-lfs && \
+    rm -rf /var/lib/apt/lists/*
 
 # Install the model weights and the algorithm files
 #  * Pull algorithm from repo into /opt/algorithm (main branch, commit e4f4008c6e18e60a79f693448562a340a9252aa8)
 #  * Remove .git folder to keep docker layer small
 #  * Replace input images path in process.py with an existing folder to avoid errors
+#  * Add specific data types and compression options to output data structures in process.py to reduce generated output footprint
 RUN git clone https://github.com/DIAGNijmegen/CE-CT_PDAC_AutomaticDetection_nnUnet.git /opt/algorithm && \
     cd /opt/algorithm && \
     git reset --hard e4f4008c6e18e60a79f693448562a340a9252aa8 && \
     rm -rf /opt/algorithm/.git && \
-    sed -i 's/Path("\/input\/images\/")/Path("\/app")/g' /opt/algorithm/process.py
-
-# FIXME: set this environment variable as a shortcut to avoid nnunet crashing the build
+    sed -i 's/Path("\/input\/images\/")/Path("\/app")/g' /opt/algorithm/process.py && \
+    sed -i 's/pred_2_np  = sitk\.GetArrayFromImage(pred_2_nii)/pred_2_np  = sitk\.GetArrayFromImage(pred_2_nii)\.astype(np\.uint8)/g' /opt/algorithm/process.py && \
+    sed -i 's/pm_image = np\.zeros(image_np\.shape)/pm_image = np\.zeros(image_np\.shape, dtype=np\.float32)/g' /opt/algorithm/process.py && \
+    sed -i 's/segmentation_np = np\.zeros(image_np\.shape)/segmentation_np = np\.zeros(image_np\.shape, dtype=np\.uint8)/g' /opt/algorithm/process.py && \
+    sed -i 's/sitk\.WriteImage(segmentation_image, str(self\.segmentation))/sitk\.WriteImage(segmentation_image, str(self\.segmentation), True)/g' /opt/algorithm/process.py && \
+    sed -i 's/sitk\.WriteImage(pred_itk_resampled, str(self\.heatmap))/sitk\.WriteImage(pred_itk_resampled, str(self\.heatmap), True)/g' /opt/algorithm/process.py
+
+# Set this environment variable as a shortcut to avoid nnunet 1.7.0 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 other requirements (should install PyTorch as well...)
+# Install nnUNet 1.7.0 and other requirements
 RUN pip3 install --no-cache-dir -r /opt/algorithm/requirements.txt
 
 # Extend the nnUNet installation with custom trainers

models/gc_nnunet_pancreas/utils/GCNNUnetPancreasRunner.py

@@ -12,26 +12,45 @@
 from mhubio.core import Module, Instance, InstanceData, DataType, Meta, IO
 
 from pathlib import Path
+import SimpleITK
+import numpy as np
 
 # Import the algorithm pipeline class from the CE-CT_PDAC_AutomaticDetection_nnUnet repository
 from process import PDACDetectionContainer
 
-# TODO should move to MHubio/core/templates.py
-HEATMAP = Meta(mod="heatmap")
-
 
 class GCNNUnetPancreasRunner(Module):
     @IO.Instance()
     @IO.Input('in_data', 'mha:mod=ct', the="input data")
     @IO.Output('heatmap', 'heatmap.mha', 'mha:mod=heatmap:model=GCNNUnetPancreas', data="in_data",
                the="heatmap of the pancreatic tumor likelihood")
-    @IO.Output('segmentation', 'segmentation.mha', 'mha:mod=seg:model=GCNNUnetPancreas', data="in_data",
-               the="segmentation of the pancreas, with the following classes: "
-                   "1-veins, 2-arteries, 3-pancreas, 4-pancreatic duct, 5-bile duct, 6-cysts, 7-renal vein")
-    def task(self, instance: Instance, in_data: InstanceData, heatmap: InstanceData, segmentation: InstanceData, **kwargs) -> None:
+    @IO.Output('segmentation', 'segmentation.mha', 'mha:mod=seg:type=original:model=GCNNUnetPancreas', data="in_data",
+               the="original segmentation of the pancreas, with the following classes: "
+                   "0-background, 1-veins, 2-arteries, 3-pancreas, 4-pancreatic duct, 5-bile duct, 6-cysts, 7-renal vein")
+    @IO.Output('segmentation_remapped', 'segmentation_remapped.mha', 'mha:mod=seg:type=remapped:model=GCNNUnetPancreas:roi=PANCREAS,PANCREATIC_DUCT,BILE_DUCT,PANCREAS+CYST,RENAL_VEIN', data="in_data",
+               the="remapped segmentation of the pancreas (without the veins and arteries), with the following classes: "
+                   "0-background, 1-pancreas, 2-pancreatic duct, 3-bile duct, 4-cysts, 5-renal vein")
+    def task(self, instance: Instance, in_data: InstanceData, heatmap: InstanceData, segmentation: InstanceData, segmentation_remapped: InstanceData, **kwargs) -> None:
         # Configure the algorithm pipeline class and run it
         algorithm = PDACDetectionContainer()
         algorithm.ct_image = in_data.abspath  # set as str not Path
         algorithm.heatmap = Path(heatmap.abspath)
         algorithm.segmentation = Path(segmentation.abspath)
         algorithm.process()
+
+        # Generate remapped segmentation
+        self.remap_segementation(
+            segmentation=segmentation,
+            segmentation_remapped=segmentation_remapped
+        )
+
+    def remap_segementation(self, segmentation: InstanceData, segmentation_remapped: InstanceData):
+        mapping = {0:0, 1:0, 2:0, 3:1, 4:2, 5:3, 6:4, 7:5}
+        mapping_numpy = np.array(list(mapping.values()), dtype=np.uint8)
+        self.log("Creating remapped segmentation", level="NOTICE")
+        seg_sitk = SimpleITK.ReadImage(segmentation.abspath)
+        seg_numpy = SimpleITK.GetArrayFromImage(seg_sitk)
+        remapped_numpy = mapping_numpy[seg_numpy]
+        remapped_sitk = SimpleITK.GetImageFromArray(remapped_numpy)
+        remapped_sitk.CopyInformation(seg_sitk)
+        SimpleITK.WriteImage(remapped_sitk, segmentation_remapped.abspath, True)