wrapper.py

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

from sorghum_pipeline.pipeline import SorghumPipeline
from sorghum_pipeline.config import Config, Paths


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)

    # Use input path directly (already in work_dir from app.py)
    input_path = Path(input_image_path)

    # Build in-memory config pointing input/output to the working directory
    cfg = Config()
    cfg.paths = Paths(
        input_folder=str(work),
        output_folder=str(work),
        boundingbox_dir=str(work)
    )
    pipeline = SorghumPipeline(config=cfg)

    # Run the pipeline (single image minimal demo)
    os.environ['MINIMAL_DEMO'] = '1'
    os.environ['FAST_OUTPUT'] = '1'
    results = pipeline.run(single_image_path=str(input_path))

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

    # Return only the requested 7 images with fixed keys
    wanted = [
        work / 'Vegetation_indices_images/ndvi.png',
        work / 'Vegetation_indices_images/ari.png',
        work / 'Vegetation_indices_images/gndvi.png',
        work / 'texture_output/lbp.png',
        work / 'texture_output/hog.png',
        work / 'texture_output/lacunarity.png',
        work / 'results/size.size_analysis.png',
    ]
    labels = [
        'NDVI', 'ARI', 'GNDVI', 'LBP', 'HOG', 'Lacunarity', 'SizeAnalysis'
    ]
    for label, path in zip(labels, wanted):
        if path.exists():
            outputs[label] = str(path)

    return outputs