BAMF Liver and Tumor segmentation

models/bamf_ct_liver_tumor/config/default.yml

@@ -24,10 +24,6 @@ modules:
 
   NNUnetRunnerV2:
     in_data: nifti:mod=ct
-    nnunet_dataset: Dataset006_Liver
-    nnunet_config: 3d_fullres
-    roi: SPLEEN,KIDNEY,GALLBLADDER,DUODENUM,PANCREAS,SMALL_INTESTINE,LUNG,LIVER,LIVER+NEOPLASM_MALIGNANT
-    export_prob_maps: False
 
   DsegConverter:
     model_name: bamf_ct_liver_tumor
@@ -38,16 +34,3 @@ modules:
   DataOrganizer:
     targets:
     - dicomseg-->[i:sid]/bamf_ct_liver_tumor.seg.dcm
-
-segdb:
-    triplets:
-        T_LIVER_LESION:
-            code: C159516
-            meaning: Liver lesion
-            scheme_designator: NCIt
-    segments:
-        NEOPLASM_MALIGNANT:
-            name: Neoplasm Malignant
-            category: C_RADIOLOGIC_FINDING
-            type: T_LIVER_LESION
-            color: [255, 0, 0]

models/bamf_ct_liver_tumor/dockerfiles/Dockerfile

@@ -8,8 +8,7 @@ FROM mhubai/base:latest
 ENV SKLEARN_ALLOW_DEPRECATED_SKLEARN_PACKAGE_INSTALL=True
 
 # Install nnunet and platipy
-RUN pip3 install --no-cache-dir nnunetv2==2.0 \
-        dcm2niix==1.0.20220715
+RUN pip3 install --no-cache-dir nnunetv2==2.0
 
 # Clone the main branch of MHubAI/models
 ARG MHUB_MODELS_REPO

models/bamf_ct_liver_tumor/meta.json

@@ -23,7 +23,7 @@
         "description": "Segmentation liver",
         "classes": [
           "LIVER",
-          "LIVER+NEOPLASM_MALIGNANT"
+          "LIVER+NEOPLASM"
         ]
       }
     ],
@@ -70,12 +70,12 @@
         "uri": "https://arxiv.org/abs/2310.14897"
       }
     ],
-    "github": "https://github.com/MHubAI/models"
+    "github": "https://github.com/bamf-health/aimi-liver-tumor-mr"
   },
   "info": {
     "use": {
       "title": "Intended Use",
-      "text": "This model is intended to perform liver and tumor segmentation in CT scans. The liver is a common site of primary (i.e. originating in the liver like hepatocellular carcinoma, HCC) or secondary (i.e. spreading to the liver like colorectal cancer) tumor development.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."
+      "text": "This model is intended to perform liver and tumor segmentation in CT scans. The liver is a common site of primary (i.e. originating in the liver like hepatocellular carcinoma, HCC) or secondary (i.e. spreading to the liver like colorectal cancer) tumor development."
     },
     "analyses": {
       "title": "Quantitative Analyses",
@@ -120,17 +120,25 @@
     },
     "training": {
       "title": "Training Data",
-      "text": "CT images from the LiTS 2017 and Medical Decathlon datasets were used to train the nnU-Net model for liver and tumor segmentation. We utilized selected totalsegmentator outputs to develop anatomically informed model. The liver-CT model was trained to predict liver and liver tumors, as well as other abdominal organs including the duodenum, gallbladder, intestines, kidneys, lungs, pancreas, and spleen. ",
+      "text": "131 CT images from the LiTS 2017 and 131 CT images from the Medical Segmentation Decathlon datasets were used to train the nnU-Net model for liver and tumor segmentation. We utilized selected totalsegmentator outputs to develop anatomically informed model. The liver-CT model was trained to predict liver and liver tumors, as well as other abdominal organs including the duodenum, gallbladder, intestines, kidneys, lungs, pancreas, and spleen. ",
       "references": [
         {
           "label": "LiTS - Liver Tumor Segmentation Challenge",
           "uri": "https://competitions.codalab.org/competitions/17094"
         },
+        {
+          "label": "Medical Segmentation Decathlon",
+          "uri": "https://doi.org/10.1038/s41467-022-30695-9"
+        },
         {
           "label": "TotalSegmentator",
           "uri": "https://mhub.ai/models/totalsegmentator"
         }
       ]
+    },
+    "limitations": {
+      "title": "Limitations",
+      "text": "The model has been trained and tested on scans acquired 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."
     }
   }
 }

models/bamf_ct_liver_tumor/utils/NNUnetRunnerV2.py

@@ -1,46 +1,35 @@
 """
 -------------------------------------------------
 MHub - NNU-Net Runner v2
-       Runner for pre-trained nnunet v2 models. 
+       Custom Runner for pre-trained nnunet v2 models.
 -------------------------------------------------
 
 -------------------------------------------------
-Author: Rahul Soni
-Email:  rahul.soni@bamfhealth.com
+Author: Jithendra kumar
+Email:  jithendra.kumar@bamfhealth.com
 -------------------------------------------------
 """
 
-
-from typing import List, Optional
-import os, subprocess, shutil
-import SimpleITK as sitk, numpy as np
+import os, shutil
 from mhubio.core import Module, Instance, InstanceData, DataType, FileType, IO
 
 
-nnunet_dataset_name_regex = r"Dataset[0-9]{3}_[a-zA-Z0-9_]+"
-
-@IO.ConfigInput('in_data', 'nifti:mod=mr', the="input data to run nnunet on")
-@IO.Config('nnunet_dataset', str, None, the='nnunet dataset name')
-@IO.Config('nnunet_config', str, None, the='nnunet model name (2d, 3d_lowres, 3d_fullres, 3d_cascade_fullres)')
-@IO.Config('folds', int, None, the='number of folds to run nnunet on')
-@IO.Config('use_tta', bool, True, the='flag to enable test time augmentation')
-@IO.Config('roi', str, None, the='roi or comma separated list of roi the nnunet segments')
+@IO.ConfigInput('in_data', 'nifti', the="input data to run nnunet on")
 class NNUnetRunnerV2(Module):
 
-    nnunet_dataset: str
-    nnunet_config: str
+    nnunet_dataset: str = 'Dataset006_Liver'
+    nnunet_config: str = '3d_fullres'
     input_data_type: DataType
-    folds: int                          # TODO: support optional config attributes
-    use_tta: bool
-    roi: str
 
     @IO.Instance()
     @IO.Input("in_data", the="input data to run nnunet on")
-    @IO.Output("out_data", 'VOLUME_001.nii.gz', 'nifti:mod=seg:model=nnunet', data='in_data', the="output data from nnunet")
+    @IO.Output("out_data", 'VOLUME_001.nii.gz',
+               'nifti:mod=seg:model=nnunet:roi=SPLEEN,KIDNEY,GALLBLADDER,DUODENUM,PANCREAS,SMALL_INTESTINE,LUNG,LIVER,LIVER+NEOPLASM',
+               data='in_data', the="output data from nnunet")
     def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData) -> None:
 
         # get the nnunet model to run
-        self.v("Running nnUNet_predict.")
+        self.v("Running nnUNetv2_predict.")
         self.v(f" > dataset:     {self.nnunet_dataset}")
         self.v(f" > config:      {self.nnunet_config}")
         self.v(f" > input data:  {in_data.abspath}")
@@ -65,51 +54,26 @@ def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData
         out_dir = self.config.data.requestTempDir(label="nnunet-model-out")
         os.environ['nnUNet_results'] = out_dir
 
-        # 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
+        # NOTE: this is a workaround for the nnunet bash script that expects the model files to be in a output folder
+        #       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(out_dir, 'nnUNet'))
         os.symlink(os.path.join(os.environ['WEIGHTS_FOLDER'], self.nnunet_dataset), os.path.join(out_dir, self.nnunet_dataset))
-
-        # NOTE: instead of running from commandline this could also be done in a pythonic way:
-        #       `nnUNet/nnunet/inference/predict.py` - but it would require
-        #       to set manually all the arguments that the user is not intended
-        #       to fiddle with; so stick with the bash executable
 
         # construct nnunet inference command
         bash_command  = ["nnUNetv2_predict"]
         bash_command += ["-i", str(inp_dir)]
         bash_command += ["-o", str(out_dir)]
         bash_command += ["-d", self.nnunet_dataset]
         bash_command += ["-c", self.nnunet_config]
-
-        # add optional arguments
-        if self.folds is not None:
-            bash_command += ["-f", str(self.folds)]
-
-        if not self.use_tta:
-            bash_command += ["--disable_tta"]
 
         self.v(f" > bash_command:     {bash_command}")
         # run command
         self.subprocess(bash_command, text=True)
 
-        # output meta
-        meta = {
-            "model": "nnunet",
-            "nnunet_dataset": self.nnunet_dataset,
-            "nnunet_config": self.nnunet_config,
-            "roi": self.roi
-        }
-
         # get output data
         out_file = f'VOLUME_001.nii.gz'
         out_path = os.path.join(out_dir, out_file)
 
         # copy output data to instance
         shutil.copyfile(out_path, out_data.abspath)
-
-        # update meta dynamically
-        out_data.type.meta += meta

models/bamf_ct_liver_tumor/utils/PostProcessor.py

@@ -48,7 +48,7 @@ def n_connected(self, img_data: np.ndarray) -> np.ndarray:
 
     @IO.Instance()
     @IO.Input('in_data', 'nifti:mod=seg:model=nnunet', the='input segmentations')
-    @IO.Output('out_data', 'bamf_processed.nii.gz', 'nifti:mod=seg:processor=bamf:roi=LIVER,LIVER+NEOPLASM_MALIGNANT', data='in_data', the="filtered Liver and tumor segmentation")
+    @IO.Output('out_data', 'bamf_processed.nii.gz', 'nifti:mod=seg:processor=bamf:roi=LIVER,LIVER+NEOPLASM', data='in_data', the="filtered Liver and tumor segmentation")
     def task(self, instance: Instance, in_data: InstanceData, out_data: InstanceData) -> None:
 
         # Log bamf runner info