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)

    # Ensure demo env vars are set before pipeline construction
    os.environ['MINIMAL_DEMO'] = '1'
    os.environ['FAST_OUTPUT'] = '1'

    # 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)
    results = pipeline.run(single_image_path=str(input_path))

    # Collect outputs
    outputs: Dict[str, str] = {}
    try:
        # Log immediate output directory contents for debugging
        for sub in ['results', 'Vegetation_indices_images', 'texture_output']:
            p = work / sub
            if p.exists():
                files = sorted([str(x.name) for x in p.iterdir() if x.is_file()])
                print(f"Artifacts in {sub}: {files}")
    except Exception:
        pass

    # Collect desired vegetation indices (replace ARI with SAVI)
    wanted = [
        work / 'Vegetation_indices_images/ndvi.png',
        work / 'Vegetation_indices_images/gndvi.png',
        work / 'Vegetation_indices_images/savi.png',
    ]
    labels = [
        'NDVI', 'GNDVI', 'SAVI',
    ]
    for label, path in zip(labels, wanted):
        if path.exists():
            outputs[label] = str(path)

    # Also include overlay and mask if present
    overlay_path = work / 'results/overlay.png'
    mask_path = work / 'results/mask.png'
    composite_path = work / 'results/composite.png'
    yolo_tips_path = work / 'results/yolo_tips.png'
    input_img_path = work / 'results/input_image.png'
    if overlay_path.exists():
        outputs['Overlay'] = str(overlay_path)
    if mask_path.exists():
        outputs['Mask'] = str(mask_path)
    if composite_path.exists():
        outputs['Composite'] = str(composite_path)
    if yolo_tips_path.exists():
        outputs['YOLOTips'] = str(yolo_tips_path)
    if input_img_path.exists():
        outputs['InputImage'] = str(input_img_path)

    # Extract simple stats for display if present in pipeline results
    try:
        plants = results.get('plants', {}) if isinstance(results, dict) else {}
        if plants:
            _, pdata = next(iter(plants.items()))
            veg = pdata.get('vegetation_indices', {})
            stats_lines = []
            for name in ['NDVI', 'GNDVI', 'SAVI']:
                entry = veg.get(name, {})
                st = entry.get('statistics', {}) if isinstance(entry, dict) else {}
                if st:
                    stats_lines.append(f"{name}: mean={st.get('mean', 0):.3f}, std={st.get('std', 0):.3f}")
            # Morphology stats (height, yolo tips)
            morph = pdata.get('morphology_features', {}) if isinstance(pdata, dict) else {}
            traits = morph.get('traits', {}) if isinstance(morph, dict) else {}
            height_cm = traits.get('plant_height_cm')
            if isinstance(height_cm, (int, float)):
                stats_lines.append(f"Plant height: {height_cm:.2f} cm")
            num_tips = traits.get('num_yolo_tips')
            if isinstance(num_tips, (int, float)):
                stats_lines.append(f"YOLO tips: {int(num_tips)}")
            if stats_lines:
                outputs['StatsText'] = "\n".join(stats_lines)
    except Exception:
        pass

    return outputs