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 pseudo-RGB composite preview."""
    if image is None:
        return None
    
    try:
        # Process to pseudo-RGB composite (same as pipeline does)
        d = image.size[0] // 2
        boxes = [(j, i, j + d, i + d) for i, j in product(range(0, image.height, d), range(0, image.width, d))]
        stack = np.stack([np.array(image.crop(box), dtype=float) for box in boxes], axis=-1)
        green, red, red_edge, nir = np.split(stack, 4, axis=-1)
        
        # Pseudo-RGB: (green, red_edge, red)
        composite = np.concatenate([green, red_edge, red], axis=-1)
        
        # Normalize to uint8
        composite = np.nan_to_num(composite, nan=0.0, posinf=0.0, neginf=0.0)
        ptp = np.ptp(composite)
        if ptp > 0:
            normalized = (composite - composite.min()) / (ptp + 1e-6) * 255
        else:
            normalized = np.zeros_like(composite)
        composite_uint8 = np.clip(normalized, 0, 255).astype(np.uint8)
        
        # Convert to PIL for display
        return Image.fromarray(composite_uint8)
    except Exception as e:
        # Fallback: return original if processing fails
        return image

def process(image):
    if image is None:
        return []
    with tempfile.TemporaryDirectory() as tmpdir:
        # Save PIL image preserving original format
        # Determine extension from image 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)
        # Keep order consistent: return exactly the 7 images
        order = [
            'NDVI', 'ARI', 'GNDVI', 'LBP', 'HOG', 'Lacunarity', 'SizeAnalysis'
        ]
        return [outputs[k] for k in order if k in outputs]

with gr.Blocks() as demo:
    gr.Markdown("# 🌿 Sorghum Plant Analysis Demo")
    gr.Markdown("Upload a sorghum plant image to analyze vegetation indices, texture features, and morphology.")
    
    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)
    
    gallery = gr.Gallery(label="Analysis Results", columns=3, height="auto")
    
    # Update preview when image is uploaded
    inp.change(fn=show_preview, inputs=inp, outputs=preview)
    run.click(process, inputs=inp, outputs=gallery)

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