Refactor nnUNet bundle inference operator to MONetBundleInferenceOperator and add implementation

Signed-off-by: Simone Bendazzoli <simben@kth.se>

Author: SimoneBendazzoli93

examples/apps/ai_spleen_nnunet_seg_app/app.py

@@ -22,49 +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
-<<<<<<< HEAD
-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
->>>>>>> dacfaa9 (Add nnUNet segmentation application and dependencies)
+from monai.deploy.operators.monai_nnunet_bundle_inference_operator import MONetBundleInferenceOperator
 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):
-<<<<<<< HEAD
-    """Demonstrates inference with built-in MONet 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
-=======
     """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 MONAI nnUNet Bundle inference operator, including pre-processing
->>>>>>> dacfaa9 (Add nnUNet segmentation application and dependencies)
     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:
-<<<<<<< HEAD
-      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.
->>>>>>> dacfaa9 (Add nnUNet segmentation application and dependencies)
     """
 
     def __init__(self, *args, **kwargs):
@@ -104,11 +87,7 @@ def compose(self):
 
         config_names = BundleConfigNames(config_names=["inference"])  # Same as the default
 
-<<<<<<< HEAD
         bundle_spleen_seg_op = MONetBundleInferenceOperator(
-=======
-        bundle_spleen_seg_op = MonainnUNetBundleInferenceOperator(
->>>>>>> dacfaa9 (Add nnUNet segmentation application and dependencies)
             self,
             input_mapping=[IOMapping("image", Image, IOType.IN_MEMORY)],
             output_mapping=[IOMapping("pred", Image, IOType.IN_MEMORY)],
@@ -183,8 +162,6 @@ def compose(self):
 """
 
 if __name__ == "__main__":
-<<<<<<< HEAD
-=======
     # 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
@@ -195,7 +172,6 @@ def compose(self):
     # export HOLOSCAN_INPUT_PATH=dcm
     # export HOLOSCAN_MODEL_PATH=spleen_model/model.ts
     # export HOLOSCAN_OUTPUT_PATH="output"
->>>>>>> dacfaa9 (Add nnUNet segmentation application and dependencies)
     logging.info(f"Begin {__name__}")
     AISpleennnUNetSegApp().run()
     logging.info(f"End {__name__}")

monai/deploy/operators/__init__.py

@@ -52,7 +52,7 @@
 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 .monet_bundle_inference_operator import MONetBundleInferenceOperator
 from .monai_seg_inference_operator import MonaiSegInferenceOperator
 from .monet_bundle_inference_operator import MONetBundleInferenceOperator
 from .nii_data_loader_operator import NiftiDataLoader

monai/deploy/operators/monai_nnunet_bundle_inference_operator.py

@@ -9,30 +9,47 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Tuple, Union
-
-from monai.deploy.core import Image
 from monai.deploy.operators.monai_bundle_inference_operator import MonaiBundleInferenceOperator, get_bundle_config
 from monai.deploy.utils.importutil import optional_import
-from monai.transforms import ConcatItemsd, ResampleToMatch
+from typing import Any, Dict, Tuple, Union
+from monai.deploy.core import Image
 
+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__ = ["MONetBundleInferenceOperator"]
 
 
 class MONetBundleInferenceOperator(MonaiBundleInferenceOperator):
     """
-    A specialized operator for performing inference using the MONet bundle.
+    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)
@@ -48,31 +65,25 @@ def __init__(
         **kwargs,
     ):
 
-        super().__init__(*args, **kwargs)
-
-        self._nnunet_predictor: torch.nn.Module = None
 
-    def _init_config(self, config_names):
+        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 = get_bundle_config(str(self._bundle_path), config_names)       
         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. If multimodal data is provided as keyword arguments,
-        it concatenates the data with the main input data."""
+        """Predicts output using the inferer."""
 
         self._nnunet_predictor.predictor.network = self._model_network
-
-        if len(kwargs) > 0:
-            multimodal_data = {"image": data}
-            for key in kwargs.keys():
-                if isinstance(kwargs[key], MetaTensor):
-                    multimodal_data[key] = ResampleToMatch(mode="bilinear")(kwargs[key], img_dst=data
-                                                         )
-            data = ConcatItemsd(keys=list(multimodal_data.keys()),name="image")(multimodal_data)["image"]
+        #os.environ['nnUNet_def_n_proc'] = "1"
         if len(data.shape) == 4:
             data = data[None]
         return self._nnunet_predictor(data)