Add nnUNet segmentation application and dependencies

Signed-off-by: Simone Bendazzoli <[email protected]>

Author: SimoneBendazzoli93

examples/apps/ai_spleen_nnunet_seg_app/app.py

@@ -22,28 +22,32 @@
 from monai.deploy.operators.dicom_seg_writer_operator import DICOMSegmentationWriterOperator, SegmentDescription
 from monai.deploy.operators.dicom_series_selector_operator import DICOMSeriesSelectorOperator
 from monai.deploy.operators.dicom_series_to_volume_operator import DICOMSeriesToVolumeOperator
-from monai.deploy.operators.monai_bundle_inference_operator import BundleConfigNames, IOMapping
-from monai.deploy.operators.monet_bundle_inference_operator import MONetBundleInferenceOperator
+from monai.deploy.operators.monai_bundle_inference_operator import (
+    BundleConfigNames,
+    IOMapping,
+
+)
+from monai.deploy.operators.monai_nnunet_bundle_inference_operator import MonainnUNetBundleInferenceOperator
 from monai.deploy.operators.stl_conversion_operator import STLConversionOperator
 
 
 # @resource(cpu=1, gpu=1, memory="7Gi")
 # pip_packages can be a string that is a path(str) to requirements.txt file or a list of packages.
 # The monai pkg is not required by this class, instead by the included operators.
 class AISpleennnUNetSegApp(Application):
-    """Demonstrates inference with built-in MONet Bundle inference operator with DICOM files as input/output
+    """Demonstrates inference with built-in MONAI nnUNet Bundle inference operator with DICOM files as input/output
 
     This application loads a set of DICOM instances, select the appropriate series, converts the series to
-    3D volume image, performs inference with the built-in MONet Bundle inference operator, including nnUNet resampling, pre-processing
+    3D volume image, performs inference with the built-in MONAI nnUNet Bundle inference operator, including pre-processing
     and post-processing, save the segmentation image in a DICOM Seg OID in an instance file, and optionally the
     surface mesh in STL format.
 
     Pertinent nnUNet MONAI Bundle:
       <Upload to the MONAI Model Zoo>
 
     Execution Time Estimate:
-      With a Nvidia RTXA600 48GB GPU, for an input DICOM Series of size 106x415x415 and patches of size 64x192x160, the execution time is around
-      50 seconds with saving both DICOM Seg and surface mesh STL file.
+      With a Nvidia RTXA600 48GB GPU, for an input DICOM Series of 139 instances, the execution time is around
+      75 seconds with saving both DICOM Seg and surface mesh STL file.
     """
 
     def __init__(self, *args, **kwargs):
@@ -83,7 +87,7 @@ def compose(self):
 
         config_names = BundleConfigNames(config_names=["inference"])  # Same as the default
 
-        bundle_spleen_seg_op = MONetBundleInferenceOperator(
+        bundle_spleen_seg_op = MonainnUNetBundleInferenceOperator(
             self,
             input_mapping=[IOMapping("image", Image, IOType.IN_MEMORY)],
             output_mapping=[IOMapping("pred", Image, IOType.IN_MEMORY)],
@@ -158,6 +162,16 @@ def compose(self):
 """
 
 if __name__ == "__main__":
+    # Creates the app and test it standalone. When running is this mode, please note the following:
+    #     -m <model file>, for model file path
+    #     -i <DICOM folder>, for input DICOM CT series folder
+    #     -o <output folder>, for the output folder, default $PWD/output
+    # e.g.
+    #     monai-deploy exec app.py -i input -m model/model.ts
+    #
+    # export HOLOSCAN_INPUT_PATH=dcm
+    # export HOLOSCAN_MODEL_PATH=spleen_model/model.ts
+    # export HOLOSCAN_OUTPUT_PATH="output"
     logging.info(f"Begin {__name__}")
     AISpleennnUNetSegApp().run()
     logging.info(f"End {__name__}")

examples/apps/ai_spleen_nnunet_seg_app/app.yaml

@@ -15,7 +15,7 @@
 # limitations under the License.
 ---
 application:
-  title: MONAI Deploy App Package - Spleen MONet Seg Inference
+  title: MONAI Deploy App Package - Spleen nnUNet Seg Inference
   version: 1.0
   inputFormats: ["file"]
   outputFormats: ["file"]

examples/apps/ai_spleen_nnunet_seg_app/requirements.txt

@@ -7,7 +7,5 @@ numpy-stl>=2.12.0
 trimesh>=3.8.11
 nibabel>=3.2.1
 torch>=1.12.0
-nvflare
-git+https://github.com/SimoneBendazzoli93/dynamic-network-architectures.git
-git+https://github.com/SimoneBendazzoli93/MONAI.git@dev
-git+https://github.com/SimoneBendazzoli93/nnUNet.git
+monai>=1.0.0
+nnunetv2>=2.5.1

monai/deploy/operators/__init__.py

@@ -25,7 +25,7 @@
     IOMapping
     ModelInfo
     MonaiBundleInferenceOperator
-    MONetBundleInferenceOperator
+    MonainnUNetBundleInferenceOperator
     MonaiSegInferenceOperator
     PNGConverterOperator
     PublisherOperator
@@ -48,8 +48,8 @@
 from .dicom_utils import EquipmentInfo, ModelInfo, random_with_n_digits, save_dcm_file, write_common_modules
 from .inference_operator import InferenceOperator
 from .monai_bundle_inference_operator import BundleConfigNames, IOMapping, MonaiBundleInferenceOperator
+from .monai_nnunet_bundle_inference_operator import MonainnUNetBundleInferenceOperator
 from .monai_seg_inference_operator import MonaiSegInferenceOperator
-from .monet_bundle_inference_operator import MONetBundleInferenceOperator
 from .nii_data_loader_operator import NiftiDataLoader
 from .png_converter_operator import PNGConverterOperator
 from .publisher_operator import PublisherOperator

monai/deploy/operators/monai_nnunet_bundle_inference_operator.py

@@ -0,0 +1,88 @@
+# Copyright 2002 MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from monai.deploy.operators.monai_bundle_inference_operator import MonaiBundleInferenceOperator, get_bundle_config
+from monai.deploy.utils.importutil import optional_import
+from typing import Any, Dict, Tuple, Union
+from monai.deploy.core import Image
+from pathlib import Path
+MONAI_UTILS = "monai.utils"
+nibabel, _ = optional_import("nibabel", "3.2.1")
+torch, _ = optional_import("torch", "1.10.2")
+
+NdarrayOrTensor, _ = optional_import("monai.config", name="NdarrayOrTensor")
+MetaTensor, _ = optional_import("monai.data.meta_tensor", name="MetaTensor")
+PostFix, _ = optional_import("monai.utils.enums", name="PostFix")  # For the default meta_key_postfix
+first, _ = optional_import("monai.utils.misc", name="first")
+ensure_tuple, _ = optional_import(MONAI_UTILS, name="ensure_tuple")
+convert_to_dst_type, _ = optional_import(MONAI_UTILS, name="convert_to_dst_type")
+Key, _ = optional_import(MONAI_UTILS, name="ImageMetaKey")
+MetaKeys, _ = optional_import(MONAI_UTILS, name="MetaKeys")
+SpaceKeys, _ = optional_import(MONAI_UTILS, name="SpaceKeys")
+Compose_, _ = optional_import("monai.transforms", name="Compose")
+ConfigParser_, _ = optional_import("monai.bundle", name="ConfigParser")
+MapTransform_, _ = optional_import("monai.transforms", name="MapTransform")
+SimpleInferer, _ = optional_import("monai.inferers", name="SimpleInferer")
+
+Compose: Any = Compose_
+MapTransform: Any = MapTransform_
+ConfigParser: Any = ConfigParser_
+__all__ = ["MonainnUNetBundleInferenceOperator"]
+
+
+class MonainnUNetBundleInferenceOperator(MonaiBundleInferenceOperator):
+    """
+    A specialized operator for performing inference using the MONAI nnUNet bundle.
+    This operator extends the `MonaiBundleInferenceOperator` to support nnUNet-specific
+    configurations and prediction logic. It initializes the nnUNet predictor and provides
+    a method for performing inference on input data.
+    
+    Attributes
+    ----------
+    _nnunet_predictor : torch.nn.Module
+        The nnUNet predictor module used for inference.
+    
+    Methods
+    -------
+    _init_config(config_names)
+        Initializes the configuration for the nnUNet bundle, including parsing the bundle
+        configuration and setting up the nnUNet predictor.
+    predict(data, *args, **kwargs)
+        Performs inference on the input data using the nnUNet predictor.
+    """
+
+    def __init__(
+        self,
+        *args,
+        **kwargs,
+    ):
+
+
+        super().__init__(*args, **kwargs)
+        
+        self._nnunet_predictor : torch.nn.Module = None
+        
+       
+    def _init_config(self, config_names):   
+
+        super()._init_config(config_names)
+        parser = get_bundle_config(str(self._bundle_path), config_names)       
+        parser['bundle_root'] = str(Path(self._bundle_path).parent.parent.parent)
+        self._parser = parser
+
+        self._nnunet_predictor = parser.get_parsed_content("network_def")
+
+    def predict(self, data: Any, *args, **kwargs) -> Union[Image, Any, Tuple[Any, ...], Dict[Any, Any]]:
+        """Predicts output using the inferer."""
+
+        self._nnunet_predictor.predictor.network = self._model_network
+        #os.environ['nnUNet_def_n_proc'] = "1"
+        return self._nnunet_predictor(torch.unsqueeze(data, 0))