Refactor ImageOverlayWriter documentation and enhance model_id validation

- Moved the documentation for ImageOverlayWriter into the class docstring for better organization and clarity.
- Improved the model_id validation logic in AppGenerator to prevent code injection and path traversal, ensuring stricter input checks.
- Updated the generated application template to reflect changes in the channel_first logic.
- Added unit tests to verify the correctness of the refactored channel_first logic.

Signed-off-by: Victor Chang <[email protected]>

Author: mocsharp

monai/deploy/operators/image_overlay_writer_operator.py

@@ -9,18 +9,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""
-Image Overlay Writer
-
-Blends a segmentation mask onto an RGB image and saves the result as a PNG.
-
-Named inputs:
-- image: original RGB frame as Image or ndarray (HWC, uint8/float)
-- pred: predicted mask as Image or ndarray (H x W or 1 x H x W). If multi-channel
-        probability tensor is provided, you may pre-argmax before this operator.
-- filename: base name (stem) for output file
-"""
-
 import logging
 from pathlib import Path
 from typing import Optional, Tuple
@@ -34,6 +22,17 @@
 
 
 class ImageOverlayWriter(Operator):
+    """
+    Image Overlay Writer
+
+    Blends a segmentation mask onto an RGB image and saves the result as a PNG.
+
+    Named inputs:
+    - image: original RGB frame as Image or ndarray (HWC, uint8/float)
+    - pred: predicted mask as Image or ndarray (H x W or 1 x H x W). If multi-channel
+            probability tensor is provided, you may pre-argmax before this operator.
+    - filename: base name (stem) for output file
+    """
     def __init__(
         self,
         fragment: Fragment,

tools/pipeline-generator/pipeline_generator/generator/app_generator.py

@@ -20,6 +20,8 @@
 from ..config.settings import Settings, load_config
 from .bundle_downloader import BundleDownloader
 
+import re
+
 logger = logging.getLogger(__name__)
 
 
@@ -92,9 +94,14 @@ def generate_app(
         Returns:
             Path to the generated application directory
         """
-        # Validate model_id to prevent code injection
-        if not model_id or not all(c.isalnum() or c in "/-_" for c in model_id):
-            raise ValueError(f"Invalid model_id: {model_id}. Only alphanumeric characters, /, -, and _ are allowed.")
+        # Validate model_id to prevent code injection and path traversal
+        # Only allow model IDs like "owner/model-name" or "model_name", no leading/trailing slash, no "..", no empty segments
+        model_id_pattern = r"^(?!.*\.\.)(?!/)(?!.*//)(?!.*\/$)[A-Za-z0-9_-]+(\/[A-Za-z0-9_-]+)*$"
+
+        if not model_id or not re.match(model_id_pattern, model_id):
+            raise ValueError(
+                f"Invalid model_id: {model_id}. Only alphanumeric characters, hyphens, underscores, and single slashes between segments are allowed. No leading/trailing slashes, consecutive slashes, or '..' allowed."
+            )
 
         # Create output directory
         output_dir.mkdir(parents=True, exist_ok=True)
@@ -250,6 +257,18 @@ def _prepare_context(
             elif isinstance(cfgs, dict):
                 resolved_channel_first = cfgs.get("channel_first", None)
 
+        # Determine final channel_first value
+        if resolved_channel_first is not None:
+            # Use explicit override from configuration
+            channel_first = resolved_channel_first
+        else:
+            # Apply default logic: False for image input classification, True otherwise
+            input_type_resolved = input_type or ("dicom" if use_dicom else ("image" if use_image else "nifti"))
+            if input_type_resolved == "image" and "classification" not in task.lower():
+                channel_first = False
+            else:
+                channel_first = True
+
         # Collect dependency hints from metadata.json
         required_packages_version = metadata.get("required_packages_version", {}) if metadata else {}
         extra_dependencies = getattr(model_config, "dependencies", []) if model_config else []
@@ -280,7 +299,7 @@ def _prepare_context(
             "authors": metadata.get("authors", "MONAI"),
             "output_postfix": output_postfix,
             "model_type": model_type,
-            "channel_first_override": resolved_channel_first,
+            "channel_first": channel_first,
             "required_packages_version": required_packages_version,
             "extra_dependencies": extra_dependencies,
         }

tools/pipeline-generator/pipeline_generator/templates/app.py.j2

@@ -146,7 +146,7 @@ class {{ app_name }}(Application):
         loader_op = ImageDirectoryLoader(
             self,
             input_folder=app_input_path,
-            channel_first={% if channel_first_override is not none %}{{ 'True' if channel_first_override else 'False' }}{% else %}{{ 'False' if input_type == 'image' and 'classification' not in task.lower() else 'True' }}{% endif %},
+            channel_first={{ channel_first }},
             name="image_loader"
         )
         {% elif input_type == "custom" %}

tools/pipeline-generator/tests/test_generator.py

@@ -306,9 +306,9 @@ def test_model_config_with_channel_first_override(self):
                                                 data_format="auto",
                                             )
 
-                                            # This covers lines 201-210
+                                            # Verify channel_first logic is computed correctly
                                             call_args = mock_app_py.call_args[0][1]
-                                            assert call_args["channel_first_override"] is False
+                                            assert call_args["channel_first"] is False
 
     def test_metadata_with_numpy_pytorch_versions(self):
         """Test metadata with numpy_version and pytorch_version."""
@@ -450,7 +450,59 @@ def test_model_config_with_dict_configs(self):
                                             )
 
                                             call_args = mock_app_py.call_args[0][1]
-                                            assert call_args["channel_first_override"] is True
+                                            assert call_args["channel_first"] is True
+
+    def test_channel_first_logic_refactoring(self):
+        """Test the refactored channel_first logic works correctly."""
+        generator = AppGenerator()
+
+        # Test case 1: image input, non-classification task -> should be False
+        context1 = generator._prepare_context(
+            model_id="test/model",
+            metadata={"task": "segmentation", "name": "Test Model"},
+            inference_config={},
+            model_file=None,
+            app_name="TestApp",
+            input_type="image",
+            output_type="nifti"
+        )
+        assert context1["channel_first"] is False
+
+        # Test case 2: image input, classification task -> should be True
+        context2 = generator._prepare_context(
+            model_id="test/model", 
+            metadata={"task": "classification", "name": "Test Model"},
+            inference_config={},
+            model_file=None,
+            app_name="TestApp",
+            input_type="image",
+            output_type="json"
+        )
+        assert context2["channel_first"] is True
+
+        # Test case 3: dicom input -> should be True
+        context3 = generator._prepare_context(
+            model_id="test/model",
+            metadata={"task": "segmentation", "name": "Test Model"},
+            inference_config={},
+            model_file=None,
+            app_name="TestApp",
+            input_type="dicom",
+            output_type="nifti"
+        )
+        assert context3["channel_first"] is True
+
+        # Test case 4: nifti input -> should be True
+        context4 = generator._prepare_context(
+            model_id="test/model",
+            metadata={"task": "segmentation", "name": "Test Model"},
+            inference_config={},
+            model_file=None,
+            app_name="TestApp",
+            input_type="nifti",
+            output_type="nifti"
+        )
+        assert context4["channel_first"] is True
 
     def test_get_default_metadata(self):
         """Test _get_default_metadata method directly."""