Migrated the UNetR app

MMelQin · MMelQin · commit b697b0e6775c · 2023-05-08T23:13:58.000-07:00
Signed-off-by: M Q &lt;mingmelvinq@nvidia.com&gt;
diff --git a/examples/apps/ai_unetr_seg_app/__main__.py b/examples/apps/ai_unetr_seg_app/__main__.py
@@ -1,55 +1,11 @@
 import logging
-import shutil
-import traceback
-from pathlib import Path
-from typing import List
 
 from app import AIUnetrSegApp
 
-if __name__ == "__main__":
-    logging.basicConfig(level=logging.DEBUG)
-    # This main function is an example to show how a batch of input can be processed.
-    # It assumes that in the app input folder there are a number of subfolders, each
-    # containing a discrete input to be processed. Each discrete payload can have
-    # multiple DICOM instances file, optionally organized in its own folder structure.
-    # The application object is first created, and on its init the model network is
-    # loaded as well as pre and post processing transforms. This app object is then
-    # run multiple times, each time with a single discrete payload.
-
-    app = AIUnetrSegApp(do_run=False)
-
-    # Preserve the application top level input and output folder path, as the path
-    # in the context may change on each run if the I/O arguments are passed in.
-    app_input_path = Path(app.context.input_path)
-    app_output_path = Path(app.context.output_path)
-
-    # Get subfolders in the input path, assume each one contains a discrete payload
-    input_dirs = [path for path in app_input_path.iterdir() if path.is_dir()]
-
-    # Set the output path for each run under the app's output path, and do run
-    work_dirs: List[str] = []  # strings resprenting folder path
-    for idx, dir in enumerate(input_dirs):
-        try:
-            output_path = app_output_path / f"{dir.name}_output"
-            # Note: the work_dir should be mapped to the host drive when used in
-            #       a container for better performance.
-            work_dir = f".unetr_app_workdir{idx}"
-            work_dirs.extend(work_dir)
+from monai.deploy.logger import load_env_log_level
 
-            logging.info(f"Start processing input in: {dir} with results in: {output_path}")
-
-            # Run app with specific input and output path.
-            # Passing in the input and output do have the side effect of changing
-            # app context. This side effect will likely be eliminated in later releases.
-            app.run(input=dir, output=output_path, workdir=work_dir)
-
-            logging.info(f"Completed processing input in: {dir} with results in: {output_path}")
-        except Exception as ex:
-            logging.error(f"Failed processing input in {dir}, due to: {ex}\n")
-            traceback.print_exc()
-        finally:
-            # Remove the workdir; alternatively do this later, if storage space is not a concern.
-            shutil.rmtree(work_dir, ignore_errors=True)
-
-    # Alternative. Explicitly remove the working dirs at the end of main.
-    # [shutil.rmtree(work_dir, ignore_errors=True) for work_dir in work_dirs]
+if __name__ == "__main__":
+    load_env_log_level()
+    logging.info(f"Begin {__name__}")
+    AIUnetrSegApp().run()
+    logging.info(f"End {__name__}")
diff --git a/examples/apps/ai_unetr_seg_app/app.py b/examples/apps/ai_unetr_seg_app/app.py
@@ -1,4 +1,4 @@
-# Copyright 2021 MONAI Consortium
+# Copyright 2021-2023 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
@@ -10,24 +10,45 @@
 # limitations under the License.
 
 import logging
+from pathlib import Path
 from typing import List
 
-# Required for setting SegmentDescription attributes. Direct import as this is not part of App SDK package.
-from pydicom.sr.codedict import codes
+from pydicom.sr.codedict import codes  # Required for setting SegmentDescription attributes.
 from unetr_seg_operator import UnetrSegOperator
 
-from monai.deploy.core import Application, resource
+from monai.deploy.conditions import CountCondition
+from monai.deploy.core import AppContext, Application
+from monai.deploy.logger import load_env_log_level
 from monai.deploy.operators.dicom_data_loader_operator import DICOMDataLoaderOperator
 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.publisher_operator import PublisherOperator
 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.
+# This is a sample series selection rule in JSON, simply selecting CT series.
+# If the study has more than 1 CT series, then all of them will be selected.
+# Please see more detail in DICOMSeriesSelectorOperator.
+# For list of string values, e.g. "ImageType": ["PRIMARY", "ORIGINAL"], it is a match if all elements
+# are all in the multi-value attribute of the DICOM series.
+
+Sample_Rules_Text = """
+{
+    "selections": [
+        {
+            "name": "CT Series",
+            "conditions": {
+                "Modality": "(?i)CT",
+                "ImageType": ["PRIMARY", "ORIGINAL"],
+                "PhotometricInterpretation": "MONOCHROME2"
+            }
+        }
+    ]
+}
+"""
+
+
+# @resource(cpu=1, gpu=1, memory="7Gi")
 class AIUnetrSegApp(Application):
     def __init__(self, *args, **kwargs):
         """Creates an application instance."""
@@ -37,25 +58,36 @@ def __init__(self, *args, **kwargs):
 
     def run(self, *args, **kwargs):
         # This method calls the base class to run. Can be omitted if simply calling through.
-        self._logger.debug(f"Begin {self.run.__name__}")
+        self._logger.info(f"Begin {self.run.__name__}")
         super().run(*args, **kwargs)
-        self._logger.debug(f"End {self.run.__name__}")
+        self._logger.info(f"End {self.run.__name__}")
 
     def compose(self):
         """Creates the app specific operators and chain them up in the processing DAG."""
 
-        self._logger.debug(f"Begin {self.compose.__name__}")
+        self._logger.info(f"Begin {self.compose.__name__}")
+        app_context = AppContext({})  # Let it figure out all the attributes without overriding
+        app_input_path = Path(app_context.input_path)
+        app_output_path = Path(app_context.output_path)
+        model_path = Path(app_context.model_path)
+
+        self._logger.info(f"App input and output path: {app_input_path}, {app_output_path}")
+
         # Creates the custom operator(s) as well as SDK built-in operator(s).
-        study_loader_op = DICOMDataLoaderOperator()
-        series_selector_op = DICOMSeriesSelectorOperator()
-        series_to_vol_op = DICOMSeriesToVolumeOperator()
+        study_loader_op = DICOMDataLoaderOperator(
+            self, CountCondition(self, 1), input_folder=app_input_path, name="dcm_loader_op"
+        )
+        series_selector_op = DICOMSeriesSelectorOperator(self, rules=Sample_Rules_Text, name="series_selector_op")
+        series_to_vol_op = DICOMSeriesToVolumeOperator(self, name="series_to_vol_op")
         # Model specific inference operator, supporting MONAI transforms.
-        unetr_seg_op = UnetrSegOperator()
-        # Create the publisher operator
-        publisher_op = PublisherOperator()
+        seg_op = UnetrSegOperator(self, model_path=model_path, name="seg_op")
+
         # Create the surface mesh STL conversion operator, for all segments
         stl_conversion_op = STLConversionOperator(
-            output_file="stl/multi-organs.stl", keep_largest_connected_component=False
+            self,
+            output_file=app_output_path.joinpath("stl/mesh.stl"),
+            keep_largest_connected_component=False,
+            name="stl_op",
         )
 
         # Create DICOM Seg writer providing the required segment description for each segment with
@@ -100,28 +132,26 @@ def compose(self):
             for organ in organs
         ]
 
-        dicom_seg_writer = DICOMSegmentationWriterOperator(segment_descriptions)
+        dicom_seg_writer = DICOMSegmentationWriterOperator(
+            self, segment_descriptions=segment_descriptions, output_folder=app_output_path, name="dcm_seg_writer_op"
+        )
 
         # Create the processing pipeline, by specifying the source and destination 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(study_loader_op, series_selector_op, {("dicom_study_list", "dicom_study_list")})
         self.add_flow(
-            series_selector_op, series_to_vol_op, {"study_selected_series_list": "study_selected_series_list"}
+            series_selector_op, series_to_vol_op, {("study_selected_series_list", "study_selected_series_list")}
         )
-        self.add_flow(series_to_vol_op, unetr_seg_op, {"image": "image"})
-        self.add_flow(unetr_seg_op, stl_conversion_op, {"seg_image": "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.add_flow(series_to_vol_op, seg_op, {("image", "image")})
+        self.add_flow(seg_op, stl_conversion_op, {("seg_image", "image")})
 
         # Note below the dicom_seg_writer requires two inputs, each coming from a source operator.
         self.add_flow(
-            series_selector_op, dicom_seg_writer, {"study_selected_series_list": "study_selected_series_list"}
+            series_selector_op, dicom_seg_writer, {("study_selected_series_list", "study_selected_series_list")}
         )
-        self.add_flow(unetr_seg_op, dicom_seg_writer, {"seg_image": "seg_image"})
+        self.add_flow(seg_op, dicom_seg_writer, {("seg_image", "seg_image")})
 
-        self._logger.debug(f"End {self.compose.__name__}")
+        self._logger.info(f"End {self.compose.__name__}")
 
 
 if __name__ == "__main__":
@@ -132,53 +162,8 @@ def compose(self):
     # e.g.
     #     python3 app.py -i input -m model/model.ts
     #
-    import shutil
-    import traceback
-    from pathlib import Path
-
-    logging.basicConfig(level=logging.DEBUG)
-    # This main function is an example to show how a batch of input can be processed.
-    # It assumes that in the app input folder there are a number of subfolders, each
-    # containing a discrete input to be processed. Each discrete payload can have
-    # multiple DICOM instances file, optionally organized in its own folder structure.
-    # The application object is first created, and on its init the model network is
-    # loaded as well as pre and post processing transforms. This app object is then
-    # run multiple times, each time with a single discrete payload.
-
-    app = AIUnetrSegApp(do_run=False)
-
-    # Preserve the application top level input and output folder path, as the path
-    # in the context may change on each run if the I/O arguments are passed in.
-    app_input_path = Path(app.context.input_path)
-    app_output_path = Path(app.context.output_path)
-
-    # Get subfolders in the input path, assume each one contains a discrete payload
-    input_dirs = [path for path in app_input_path.iterdir() if path.is_dir()]
-
-    # Set the output path for each run under the app's output path, and do run
-    work_dirs: List[str] = []  # strings resprenting folder path
-    for idx, dir in enumerate(input_dirs):
-        try:
-            output_path = app_output_path / f"{dir.name}_output"
-            # Note: the work_dir should be mapped to the host drive when used in
-            #       a container for better performance.
-            work_dir = f".unetr_app_workdir{idx}"
-            work_dirs.extend(work_dir)
-
-            logging.info(f"Start processing input in: {dir} with results in: {output_path}")
-
-            # Run app with specific input and output path.
-            # Passing in the input and output do have the side effect of changing
-            # app context. This side effect will likely be eliminated in later releases.
-            app.run(input=dir, output=output_path, workdir=work_dir)
-
-            logging.info(f"Completed processing input in: {dir} with results in: {output_path}")
-        except Exception as ex:
-            logging.error(f"Failed processing input in {dir}, due to: {ex}\n")
-            traceback.print_exc()
-        finally:
-            # Remove the workdir; alternatively do this later, if storage space is not a concern.
-            shutil.rmtree(work_dir, ignore_errors=True)
-
-    # Alternative. Explicitly remove the working dirs at the end of main.
-    # [shutil.rmtree(work_dir, ignore_errors=True) for work_dir in work_dirs]
+
+    load_env_log_level()
+    logging.info(f"Begin {__name__}")
+    AIUnetrSegApp().run()
+    logging.info(f"End {__name__}")
diff --git a/examples/apps/ai_unetr_seg_app/unetr_seg_operator.py b/examples/apps/ai_unetr_seg_app/unetr_seg_operator.py
