Update UNETR inference application to publish images for Clara Render Server (#182)

* Added the transform to save the input image in nii.gz, for RS publishing

Signed-off-by: mmelqin <mingmelvinq@nvidia.com>

* Add publishing input and seg image for rendering

Signed-off-by: mmelqin <mingmelvinq@nvidia.com>

* Make use of the new publish operator for saving dataset for Render Server

Signed-off-by: mmelqin <mingmelvinq@nvidia.com>

* Update by run check --autofix

Signed-off-by: mmelqin <mingmelvinq@nvidia.com>

* Fix unused from os import path

Signed-off-by: mmelqin <mingmelvinq@nvidia.com>

* Change to use Path.mkdir

Signed-off-by: mmelqin <mingmelvinq@nvidia.com>

Author: MMelQin

examples/apps/ai_unetr_seg_app/app.py

@@ -18,6 +18,7 @@
 from monai.deploy.operators.dicom_seg_writer_operator import DICOMSegmentationWriterOperator
 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.publisher_operator import PublisherOperator
 
 
 @resource(cpu=1, gpu=1, memory="7Gi")
@@ -46,6 +47,9 @@ def compose(self):
         series_to_vol_op = DICOMSeriesToVolumeOperator()
         # Model specific inference operator, supporting MONAI transforms.
         unetr_seg_op = UnetrSegOperator()
+        # Create the publisher operator
+        publisher_op = PublisherOperator()
+
         # Creates DICOM Seg writer with segment label name in a string list
         dicom_seg_writer = DICOMSegmentationWriterOperator(
             seg_labels=[
@@ -64,15 +68,19 @@ def compose(self):
                 "lad",
             ]
         )
-
         # Create the processing pipeline, by specifying the upstream and downstream operators, and
         # ensuring the output from the former matches the input of the latter, in both name and type.
         self.add_flow(study_loader_op, series_selector_op, {"dicom_study_list": "dicom_study_list"})
         self.add_flow(series_selector_op, series_to_vol_op, {"dicom_series": "dicom_series"})
         self.add_flow(series_to_vol_op, unetr_seg_op, {"image": "image"})
         # Note below the dicom_seg_writer requires two inputs, each coming from a upstream operator.
+        # Also note that the DICOMSegmentationWriterOperator may throw exception with some inputs.
+        #   Bug has been created to track the issue.
         self.add_flow(series_selector_op, dicom_seg_writer, {"dicom_series": "dicom_series"})
         self.add_flow(unetr_seg_op, dicom_seg_writer, {"seg_image": "seg_image"})
+        # Add the publishing operator to save the input and seg images for Render Server.
+        # Note the PublisherOperator has temp impl till a proper rendering module is created.
+        self.add_flow(unetr_seg_op, publisher_op, {"saved_images_folder": "saved_images_folder"})
 
         self._logger.debug(f"End {self.compose.__name__}")
 

examples/apps/ai_unetr_seg_app/unetr_seg_operator.py

@@ -10,12 +10,11 @@
 # limitations under the License.
 
 import logging
-from os import path
 
 from numpy import uint8
 
 import monai.deploy.core as md
-from monai.deploy.core import ExecutionContext, Image, InputContext, IOType, Operator, OutputContext
+from monai.deploy.core import DataPath, ExecutionContext, Image, InputContext, IOType, Operator, OutputContext
 from monai.deploy.operators.monai_seg_inference_operator import InMemImageReader, MonaiSegInferenceOperator
 from monai.transforms import (
     Activationsd,
@@ -35,6 +34,7 @@
 
 @md.input("image", Image, IOType.IN_MEMORY)
 @md.output("seg_image", Image, IOType.IN_MEMORY)
+@md.output("saved_images_folder", DataPath, IOType.DISK)
 @md.env(pip_packages=["monai==0.6.0", "torch>=1.5", "numpy>=1.17", "nibabel"])
 class UnetrSegOperator(Operator):
     """Performs multi-organ segmentation using UNETR model with an image converted from a DICOM CT series.
@@ -62,12 +62,16 @@ def compute(self, op_input: InputContext, op_output: OutputContext, context: Exe
         # Get the output path from the execution context for saving file(s) to app output.
         # Without using this path, operator would be saving files to its designated path, e.g.
         # $PWD/.monai_workdir/operators/6048d75a-5de1-45b9-8bd1-2252f88827f2/0/output
-        output_path = context.output.get().path
+        op_output_folder_name = DataPath("saved_images_folder")
+        op_output.set(op_output_folder_name, "saved_images_folder")
+        op_output_folder_path = op_output.get("saved_images_folder").path
+        op_output_folder_path.mkdir(parents=True, exist_ok=True)
+        print(f"Operator output folder path: {op_output_folder_path}")
 
         # This operator gets an in-memory Image object, so a specialized ImageReader is needed.
         _reader = InMemImageReader(input_image)
         pre_transforms = self.pre_process(_reader)
-        post_transforms = self.post_process(pre_transforms, path.join(output_path, "prediction_output"))
+        post_transforms = self.post_process(pre_transforms, op_output_folder_path)
 
         # Delegates inference and saving output to the built-in operator.
         infer_operator = MonaiSegInferenceOperator(
@@ -115,5 +119,12 @@ def post_process(self, pre_transforms: Compose, out_dir: str = "./prediction_out
                     keys=pred_key, transform=pre_transforms, orig_keys=self._input_dataset_key, nearest_interp=True
                 ),
                 SaveImaged(keys=pred_key, output_dir=out_dir, output_postfix="seg", output_dtype=uint8, resample=False),
+                SaveImaged(
+                    keys=self._input_dataset_key,
+                    output_dir=out_dir,
+                    output_postfix="",
+                    output_dtype=uint8,
+                    resample=False,
+                ),
             ]
         )