Enhance pipeline generator functionality and improve bundle organization

- Updated the CLI output to reflect "Verified" models instead of "Tested" for better clarity.
- Added new model configurations for pancreas segmentation and spleen segmentation in the config file.
- Implemented a method to organize downloaded bundle structures into the standard MONAI format, improving file management.
- Enhanced dependency handling in the AppGenerator to resolve conflicts between configuration and metadata.
- Added unit tests to verify the new bundle organization functionality and ensure correct behavior under various scenarios.

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

Author: mocsharp

tools/pipeline-generator/pipeline_generator/cli/main.py

@@ -118,7 +118,7 @@ def list(ctx: click.Context, format: str, bundles_only: bool, tested_only: bool)
     bundle_count = sum(1 for m in models if m.is_monai_bundle)
     tested_count = sum(1 for m in models if m.model_id in tested_models)
     console.print(
-        f"\n[green]Total models: {len(models)} (MONAI Bundles: {bundle_count}, Tested: {tested_count})[/green]"
+        f"\n[green]Total models: {len(models)} (MONAI Bundles: {bundle_count}, Verified: {tested_count})[/green]"
     )
 
 

tools/pipeline-generator/pipeline_generator/config/config.yaml

@@ -40,6 +40,9 @@ endpoints:
       - model_id: "MONAI/swin_unetr_btcv_segmentation"
         input_type: "nifti"
         output_type: "nifti"
+      - model_id: "MONAI/pancreas_ct_dints_segmentation"
+        input_type: "nifti"
+        output_type: "nifti"
       - model_id: "MONAI/Llama3-VILA-M3-3B"
         input_type: "custom"
         output_type: "custom"
@@ -64,6 +67,14 @@ endpoints:
         - torch>=2.0.0
         - Pillow>=8.0.0
         - PyYAML>=6.0
+      - model_id: "MONAI/example_spleen_segmentation"
+        input_type: "nifti"
+        output_type: "nifti"
+        dependencies:
+        - torch>=1.11.0,<3.0.0
+        - numpy>=1.21.2,<2.0.0
+        - monai>=1.3.0
+        - nibabel>=3.0.0
 
 additional_models:
   - model_id: "LGAI-EXAONE/EXAONEPath"

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

@@ -112,6 +112,9 @@ def generate_app(
         # Download the bundle
         logger.info(f"Downloading bundle: {model_id}")
         bundle_path = self.downloader.download_bundle(model_id, output_dir)
+        
+        # Organize bundle into proper structure if needed
+        self.downloader.organize_bundle_structure(bundle_path)
 
         # Read bundle metadata and config
         metadata = self.downloader.get_bundle_metadata(bundle_path)
@@ -275,10 +278,35 @@ def _prepare_context(
         # 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 []
-        if metadata and "numpy_version" in metadata:
+        
+        # Handle dependency conflicts between config and metadata
+        config_deps = []
+        if extra_dependencies:
+            # Extract dependency names from config overrides
+            config_deps = [dep.split(">=")[0].split("==")[0].split("<")[0] for dep in extra_dependencies]
+            
+        # Add metadata dependencies only if not overridden by config
+        if metadata and "numpy_version" in metadata and "numpy" not in config_deps:
             extra_dependencies.append(f"numpy=={metadata['numpy_version']}")
-        if metadata and "pytorch_version" in metadata:
+        if metadata and "pytorch_version" in metadata and "torch" not in config_deps:
             extra_dependencies.append(f"torch=={metadata['pytorch_version']}")
+            
+        # Handle MONAI version - move logic from template to Python for better maintainability
+        has_monai_config = any(dep.startswith("monai") for dep in extra_dependencies)
+        if has_monai_config and metadata:
+            # Remove monai_version from metadata since we have config override
+            metadata = dict(metadata)  # Make a copy
+            metadata.pop("monai_version", None)
+        elif not has_monai_config:
+            # No config MONAI dependency - add one based on metadata or fallback
+            if metadata and "monai_version" in metadata:
+                extra_dependencies.append(f"monai=={metadata['monai_version']}")
+                # Remove from metadata since it's now in extra_dependencies
+                metadata = dict(metadata) if metadata else {}
+                metadata.pop("monai_version", None)
+            else:
+                # No metadata version, use fallback
+                extra_dependencies.append("monai>=1.5.0")
 
         return {
             "model_id": model_id,

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

@@ -144,3 +144,57 @@ def detect_model_file(self, bundle_path: Path) -> Optional[Path]:
 
         logger.warning(f"No model file found in bundle: {bundle_path}")
         return None
+
+    def organize_bundle_structure(self, bundle_path: Path) -> None:
+        """Organize bundle files into the expected MONAI Bundle structure.
+        
+        Creates the standard structure if files are in the root directory:
+        bundle_root/
+          configs/
+            metadata.json
+            inference.json
+          models/
+            model.pt
+            model.ts
+            
+        Args:
+            bundle_path: Path to the downloaded bundle
+        """
+        configs_dir = bundle_path / "configs"
+        models_dir = bundle_path / "models"
+        
+        # Check if structure already exists
+        has_configs_structure = (
+            configs_dir.exists() and 
+            (configs_dir / "metadata.json").exists()
+        )
+        has_models_structure = (
+            models_dir.exists() and 
+            any(models_dir.glob("model.*"))
+        )
+        
+        if has_configs_structure and has_models_structure:
+            logger.debug("Bundle already has proper structure")
+            return
+            
+        logger.info("Organizing bundle into standard structure")
+        
+        # Create directories
+        configs_dir.mkdir(exist_ok=True)
+        models_dir.mkdir(exist_ok=True)
+        
+        # Move config files to configs/
+        config_files = ["metadata.json", "inference.json"]
+        for config_file in config_files:
+            src_path = bundle_path / config_file
+            if src_path.exists() and not (configs_dir / config_file).exists():
+                src_path.rename(configs_dir / config_file)
+                logger.debug(f"Moved {config_file} to configs/")
+        
+        # Move model files to models/
+        model_extensions = [".pt", ".ts", ".onnx"]
+        for ext in model_extensions:
+            for model_file in bundle_path.glob(f"*{ext}"):
+                if model_file.is_file() and not (models_dir / model_file.name).exists():
+                    model_file.rename(models_dir / model_file.name)
+                    logger.debug(f"Moved {model_file.name} to models/")

tools/pipeline-generator/pipeline_generator/templates/requirements.txt.j2

@@ -3,11 +3,6 @@
 
 # MONAI Deploy App SDK and dependencies
 monai-deploy-app-sdk>=3.0.0
-{% if metadata.monai_version is defined %}
-monai=={{ metadata.monai_version }}
-{% else %}
-monai>=1.5.0
-{% endif %}
 
 
 # Required by MONAI Deploy SDK (always needed)

tools/pipeline-generator/tests/test_bundle_downloader.py

@@ -220,6 +220,78 @@ def test_detect_model_file_not_found(self, tmp_path):
 
         assert result is None
 
+    def test_organize_bundle_structure_flat_to_structured(self, tmp_path):
+        """Test organizing flat bundle structure into standard format."""
+        bundle_path = tmp_path / "bundle"
+        bundle_path.mkdir()
+
+        # Create files in flat structure
+        metadata_file = bundle_path / "metadata.json"
+        inference_file = bundle_path / "inference.json"
+        model_pt_file = bundle_path / "model.pt"
+        model_ts_file = bundle_path / "model.ts"
+        
+        metadata_file.write_text('{"name": "Test"}')
+        inference_file.write_text('{"config": "test"}')
+        model_pt_file.touch()
+        model_ts_file.touch()
+
+        # Organize structure
+        self.downloader.organize_bundle_structure(bundle_path)
+
+        # Check that files were moved to proper locations
+        assert (bundle_path / "configs" / "metadata.json").exists()
+        assert (bundle_path / "configs" / "inference.json").exists()
+        assert (bundle_path / "models" / "model.pt").exists()
+        assert (bundle_path / "models" / "model.ts").exists()
+        
+        # Check that original files were moved (not copied)
+        assert not metadata_file.exists()
+        assert not inference_file.exists()
+        assert not model_pt_file.exists()
+        assert not model_ts_file.exists()
+
+    def test_organize_bundle_structure_already_structured(self, tmp_path):
+        """Test organizing bundle that already has proper structure."""
+        bundle_path = tmp_path / "bundle"
+        configs_dir = bundle_path / "configs"
+        models_dir = bundle_path / "models"
+        configs_dir.mkdir(parents=True)
+        models_dir.mkdir(parents=True)
+
+        # Create files in proper structure
+        metadata_file = configs_dir / "metadata.json"
+        model_file = models_dir / "model.pt"
+        metadata_file.write_text('{"name": "Test"}')
+        model_file.touch()
+
+        # Should not change anything
+        self.downloader.organize_bundle_structure(bundle_path)
+
+        # Files should remain in place
+        assert metadata_file.exists()
+        assert model_file.exists()
+
+    def test_organize_bundle_structure_partial_structure(self, tmp_path):
+        """Test organizing bundle with partial structure."""
+        bundle_path = tmp_path / "bundle"
+        configs_dir = bundle_path / "configs"
+        configs_dir.mkdir(parents=True)
+
+        # Create metadata in configs but model in root
+        metadata_file = configs_dir / "metadata.json"
+        model_file = bundle_path / "model.pt"
+        metadata_file.write_text('{"name": "Test"}')
+        model_file.touch()
+
+        # Organize structure
+        self.downloader.organize_bundle_structure(bundle_path)
+
+        # Metadata should stay, model should move
+        assert metadata_file.exists()
+        assert (bundle_path / "models" / "model.pt").exists()
+        assert not model_file.exists()
+
     def test_detect_model_file_multiple_models(self, tmp_path):
         """Test detecting model file with multiple model files (returns first found)."""
         bundle_path = tmp_path / "bundle"