app.py

import gradio as gr
import tempfile
from pathlib import Path
from wrapper import run_pipeline_on_image
import numpy as np
from PIL import Image
from itertools import product

def show_preview(image):
    """Show uploaded image as-is."""
    return image

def process(image):
    if image is None:
        return None, None, [], ""
    with tempfile.TemporaryDirectory() as tmpdir:
        # Save PIL image preserving original format
        ext = image.format.lower() if image.format else 'png'
        img_path = Path(tmpdir) / f"input.{ext}"
        image.save(img_path)
        outputs = run_pipeline_on_image(str(img_path), tmpdir, save_artifacts=True)

        # Assemble displays
        overlay = outputs.get('Overlay')
        mask = outputs.get('Mask')
        order = ['NDVI', 'ARI', 'GNDVI']
        gallery_items = [outputs[k] for k in order if k in outputs]
        stats_text = outputs.get('StatsText', '')
        return overlay, mask, gallery_items, stats_text

with gr.Blocks() as demo:
    gr.Markdown("# 🌿 Sorghum Plant Analysis Demo")
    gr.Markdown("Upload a sorghum plant image to analyze vegetation indices, segmentation overlay, and stats.")

    with gr.Row():
        with gr.Column():
            inp = gr.Image(type="pil", label="Upload Image")
            run = gr.Button("Run Pipeline", variant="primary")
        with gr.Column():
            preview = gr.Image(type="pil", label="Uploaded Image Preview", interactive=False)

    with gr.Row():
        overlay_img = gr.Image(type="filepath", label="Segmentation Overlay", interactive=False)
        mask_img = gr.Image(type="filepath", label="Mask", interactive=False)

    gallery = gr.Gallery(label="Vegetation Indices", columns=3, height="auto")
    stats = gr.Textbox(label="Statistics", lines=4)

    # Update preview when image is uploaded
    inp.change(fn=show_preview, inputs=inp, outputs=preview)
    run.click(process, inputs=inp, outputs=[overlay_img, mask_img, gallery, stats])

if __name__ == "__main__":
    demo.launch()