MONAI Pipeline Generator

.github/workflows/pr.yml

@@ -55,3 +55,22 @@ jobs:
         with:
           fail_ci_if_error: false
           files: ./coverage.xml
+
+  test-pipeline-generator:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.10"
+      - name: Install uv
+        uses: astral-sh/setup-uv@v6
+      - name: Install dependencies
+        working-directory: tools/pipeline-generator
+        run: |
+          uv sync
+      - name: Run tests
+        working-directory: tools/pipeline-generator
+        run: |
+          uv run pytest

monai/deploy/operators/monai_bundle_inference_operator.py

@@ -1,4 +1,4 @@
-# Copyright 2002 MONAI Consortium
+# Copyright 2022-2025 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
@@ -13,6 +13,7 @@
 import logging
 import os
 import pickle
+import sys
 import tempfile
 import time
 import zipfile
@@ -101,7 +102,11 @@ def _read_from_archive(archive, root_name: str, config_name: str, do_search=True
         return content_text
 
     def _extract_from_archive(
-        archive, root_name: str, config_names: List[str], dest_folder: Union[str, Path], do_search=True
+        archive,
+        root_name: str,
+        config_names: List[str],
+        dest_folder: Union[str, Path],
+        do_search=True,
     ):
         """A helper function for extract files of configs from the archive to the destination folder
 
@@ -151,6 +156,45 @@ def _extract_from_archive(
     if isinstance(config_names, str):
         config_names = [config_names]
 
+    # Check if bundle_path is a directory (for directory-based bundles)
+    bundle_path_obj = Path(bundle_path)
+    if bundle_path_obj.is_dir():
+        # Handle directory-based bundles
+        parser = ConfigParser()
+
+        # Read metadata from configs/metadata.json
+        metadata_path = bundle_path_obj / "configs" / "metadata.json"
+        if not metadata_path.exists():
+            raise IOError(f"Cannot find metadata.json at {metadata_path}")
+
+        with open(metadata_path, "r") as f:
+            metadata_content = f.read()
+            parser.read_meta(f=json.loads(metadata_content))
+
+        # Read other config files
+        config_files = []
+        for config_name in config_names:
+            config_name_base = config_name.split(".")[0]  # Remove extension if present
+            # Validate config name to prevent path traversal
+            if ".." in config_name_base or "/" in config_name_base or "\\" in config_name_base:
+                raise ValueError(f"Invalid config name: {config_name_base}")
+            found = False
+            for suffix in bundle_suffixes:
+                config_path = bundle_path_obj / "configs" / f"{config_name_base}{suffix}"
+                if config_path.exists():
+                    ...
+                    config_files.append(config_path)
+                    found = True
+                    break
+            if not found:
+                raise IOError(f"Cannot find config file for {config_name} in {bundle_path_obj / 'configs'}")
+
+        parser.read_config(config_files)
+        parser.parse()
+
+        return parser
+
+    # Original ZIP file handling code
     name, _ = os.path.splitext(os.path.basename(bundle_path))  # bundle file name same archive folder name
     parser = ConfigParser()
 
@@ -363,6 +407,10 @@ def __init__(
             if self._bundle_path and self._bundle_path.is_file():
                 self._init_config(self._bundle_config_names.config_names)
                 self._init_completed = True
+            elif self._bundle_path and self._bundle_path.is_dir():
+                # For directory-based bundles, delay initialization to compute method
+                logging.debug(f"Bundle path {self._bundle_path} is a directory. Will initialize during execution.")
+                # Keep the bundle_path for directory-based bundles
             else:
                 logging.debug(
                     f"Bundle, at path {self._bundle_path}, not available. Will get it in the execution context."
@@ -420,6 +468,11 @@ def _init_config(self, config_names):
             config_names ([str]): Names of the config (files) in the bundle
         """
 
+        # Ensure bundle root is on sys.path so 'scripts.*' can be imported
+        bundle_root = str(self._bundle_path)
+        if bundle_root not in sys.path:
+            sys.path.insert(0, bundle_root)
+
         parser = get_bundle_config(str(self._bundle_path), config_names)
         self._parser = parser
 
@@ -562,7 +615,79 @@ def compute(self, op_input, op_output, context):
             # When run as a MAP docker, the bundle file is expected to be in the context, even if the model
             # network is loaded on a remote inference server (when the feature is introduced).
             logging.debug(f"Model network not loaded. Trying to load from model path: {self._bundle_path}")
-            self._model_network = torch.jit.load(self.bundle_path, map_location=self._device).eval()
+
+            # Check if bundle_path is a directory
+            if self._bundle_path.is_dir():
+                # For directory-based bundles, look for model in models/ subdirectory
+                model_path = self._bundle_path / "models" / "model.ts"
+                if not model_path.exists():
+                    # Try model.pt as fallback
+                    model_path = self._bundle_path / "models" / "model.pt"
+                if not model_path.exists():
+                    raise IOError(f"Cannot find model.ts or model.pt in {self._bundle_path / 'models'}")
+
+                # Ensure device is set
+                if not hasattr(self, "_device"):
+                    self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+                # Initialize config for directory bundles if not already done
+                if not self._init_completed:
+                    logging.info(f"Initializing config from directory bundle: {self._bundle_path}")
+                    self._init_config(self._bundle_config_names.config_names)
+                    self._init_completed = True
+
+                # Load model based on file type
+                if model_path.suffix == ".ts":
+                    # TorchScript bundle
+                    self._model_network = torch.jit.load(str(model_path), map_location=self._device).eval()
+                else:
+                    # .pt checkpoint: instantiate network from config and load state dict
+                    try:
+                        # Some .pt files may still be TorchScript; try jit first
+                        self._model_network = torch.jit.load(str(model_path), map_location=self._device).eval()
+                    except Exception:
+                        # Fallback to eager model with loaded weights
+                        if self._parser is None:
+                            # Ensure parser/config are initialized
+                            self._init_config(self._bundle_config_names.config_names)
+                        # Instantiate network from config
+                        # Ensure bundle root is on sys.path so 'scripts.*' can be imported
+                        bundle_root = str(self._bundle_path)
+                        if bundle_root not in sys.path:
+                            sys.path.insert(0, bundle_root)
+                        network = (
+                            self._parser.get_parsed_content("network")
+                            if self._parser.get("network") is not None
+                            else None
+                        )
+                        if network is None:
+                            # Backward compatibility: some bundles use "network_def" then to(device)
+                            network = (
+                                self._parser.get_parsed_content("network_def")
+                                if self._parser.get("network_def") is not None
+                                else None
+                            )
+                            if network is not None:
+                                network = network.to(self._device)
+                        if network is None:
+                            raise RuntimeError("Unable to instantiate network from bundle configs.") from None
+
+                        checkpoint = torch.load(str(model_path), map_location=self._device)
+                        # Determine the state dict layout
+                        state_dict = None
+                        if isinstance(checkpoint, dict):
+                            if "state_dict" in checkpoint and isinstance(checkpoint["state_dict"], dict):
+                                state_dict = checkpoint["state_dict"]
+                            elif "model" in checkpoint and isinstance(checkpoint["model"], dict):
+                                state_dict = checkpoint["model"]
+                        if state_dict is None:
+                            # Assume raw state dict
+                            state_dict = checkpoint
+                        network.load_state_dict(state_dict, strict=True)
+                        self._model_network = network.eval()
+            else:
+                # Original ZIP bundle handling
+                self._model_network = torch.jit.load(self._bundle_path, map_location=self._device).eval()
         else:
             raise IOError("Model network is not load and model file not found.")
 
@@ -701,7 +826,50 @@ def _receive_input(self, name: str, op_input, context):
             logging.debug(f"Shape of the converted input image: {value.shape}")
             logging.debug(f"Metadata of the converted input image: {metadata}")
         elif isinstance(value, np.ndarray):
+            # Keep numpy array as-is when possible and set metadata so downstream transforms handle channels.
+            # Use bundle metadata to infer expected number of channels and adjust conservatively.
+            ndims = value.ndim
+            expected_channels = None
+            try:
+                in_meta = self._inputs.get(name, {})
+                if isinstance(in_meta, dict):
+                    expected_channels = in_meta.get("num_channels")
+            except Exception:
+                expected_channels = None
+
+            if ndims == 3:
+                # No channel present (W, H, D)
+                if expected_channels is not None and expected_channels > 1:
+                    raise ValueError(
+                        f"Input for '{name!r}' has no channel dimension but bundle expects {expected_channels} channels. "
+                        "Provide multi-channel input or add a transform to stack channels before inference."
+                    )
+                # else expected 1 or unknown -> proceed without channel
+            elif ndims == 4:
+                # Channel-last assumed (W, H, D, C)
+                actual_channels = value.shape[-1]
+                if expected_channels is not None and expected_channels != actual_channels:
+                    if expected_channels == 1 and actual_channels > 1:
+                        logging.warning(
+                            "Input for '%s' has %d channels but bundle expects 1; selecting channel 0.",
+                            name,
+                            actual_channels,
+                        )
+                        value = value[..., 0]
+                        ndims = 3
+                    else:
+                        raise ValueError(
+                            f"Input for '{name!r}' has {actual_channels} channels but bundle expects {expected_channels}."
+                        )
+                # else exact match or unknown -> keep as-is
+            else:
+                # Unsupported rank for medical image input
+                raise ValueError(f"Unsupported input rank {ndims} for '{name!r}'. Expected 3D (W,H,D) or 4D (W,H,D,C).")
             value = torch.from_numpy(value).to(self._device)
+            if metadata is None:
+                metadata = {}
+            # Indicate whether there was a channel for EnsureChannelFirstd
+            metadata["original_channel_dim"] = "no_channel" if ndims == 3 else -1
 
         # else value is some other object from memory
 
@@ -732,7 +900,28 @@ def _send_output(self, value: Any, name: str, metadata: Dict, op_output, context
                 raise TypeError("arg 1 must be of type torch.Tensor or ndarray.")
 
             logging.debug(f"Output {name} numpy image shape: {value.shape}")
-            result: Any = Image(np.swapaxes(np.squeeze(value, 0), 0, 2).astype(np.uint8), metadata=metadata)
+
+            # Handle 2D masks and generic 2D tensors gracefully
+            if value.ndim == 2:
+                # Already HxW image; binarize/scale left to downstream operators
+                out_img = value.astype(np.uint8)
+                result: Any = Image(out_img, metadata=metadata)
+            elif value.ndim == 3:
+                # Could be (C, H, W) with C==1 or (H, W, C)
+                if value.shape[0] == 1:  # (1, H, W) -> (H, W)
+                    out_img = value[0].astype(np.uint8)
+                    result = Image(out_img, metadata=metadata)
+                elif value.shape[-1] == 1:  # (H, W, 1) -> (H, W)
+                    out_img = value[..., 0].astype(np.uint8)
+                    result = Image(out_img, metadata=metadata)
+                else:
+                    # Fallback to original behavior for 3D volumetric layout assumptions
+                    out_img = np.swapaxes(np.squeeze(value, 0), 0, 2).astype(np.uint8)
+                    result = Image(out_img, metadata=metadata)
+            else:
+                # Keep existing behavior for higher-dimensional data (e.g., 3D volumes)
+                out_img = np.swapaxes(np.squeeze(value, 0), 0, 2).astype(np.uint8)
+                result = Image(out_img, metadata=metadata)
             logging.debug(f"Converted Image shape: {result.asnumpy().shape}")
         elif otype == np.ndarray:
             result = np.asarray(value)

monai/deploy/operators/nii_data_loader_operator.py

@@ -80,7 +80,24 @@ def convert_and_save(self, nii_path):
         image_reader = SimpleITK.ImageFileReader()
         image_reader.SetFileName(str(nii_path))
         image = image_reader.Execute()
-        image_np = np.transpose(SimpleITK.GetArrayFromImage(image), [2, 1, 0])
+        image_np = SimpleITK.GetArrayFromImage(image)
+
+        # Handle different dimensionalities properly
+        if image_np.ndim == 3:
+            # Standard 3D volume: transpose from (z, y, x) to (x, y, z)
+            image_np = np.transpose(image_np, [2, 1, 0])
+        elif image_np.ndim == 4:
+            # 4D volume with channels: (c, z, y, x) to (c, x, y, z)
+            image_np = np.transpose(image_np, [0, 3, 2, 1])
+        elif image_np.ndim == 2:
+            # 2D slice: transpose from (y, x) to (x, y)
+            image_np = np.transpose(image_np, [1, 0])
+        else:
+            # For other dimensions, log a warning and return as-is
+            self._logger.warning(
+                f"Unexpected {image_np.ndim}D NIfTI file shape {image_np.shape} from {nii_path}, returning without transpose"
+            )
+
         return image_np
 
 

tools/pipeline-generator/.gitignore

@@ -0,0 +1,2 @@
+results*/
+test_*/