wrapper.py

from pathlib import Path
from typing import Dict
import shutil
from PIL import Image
import glob
import tempfile

from sorghum_pipeline.pipeline import SorghumPipeline


def run_pipeline_on_image(input_image_path: str, work_dir: str, save_artifacts: bool = True) -> Dict[str, str]:
    """
    Run sorghum pipeline on a single image (no instance segmentation).
    Returns dict[label -> image_path] for gallery display.
    """

    work = Path(work_dir)
    work.mkdir(parents=True, exist_ok=True)

    # Copy input to work dir
    input_copy = work / Path(input_image_path).name
    shutil.copy(input_image_path, input_copy)

    # Initialize pipeline with config
    #  adjust this if you have a YAML config file (e.g., "configs/demo.yaml")
    pipeline = SorghumPipeline(
        config_path=str(Path("sorghum_pipeline/config.py")),
        enable_occlusion_handling=False,
        enable_instance_integration=False
    )

    # Run the pipeline (single image, no frames, no SAM2Long)
    results = pipeline.run(
        load_all_frames=False,
        segmentation_only=False,
        run_instance_segmentation=False,
        features_frame_only=None
    )

    # Collect outputs
    outputs: Dict[str, str] = {}

    # Save original for reference
    original = work / "original.png"
    Image.open(input_copy).convert("RGB").save(original)
    outputs["Original"] = str(original)

    # Gather all PNG files created by OutputManager
    for f in glob.glob(str(work / "**/*.png"), recursive=True):
        name = Path(f).stem
        if name.lower() not in outputs:  # avoid duplicate "Original"
            outputs[name] = f

    return outputs