app.py

import gradio as gr
import tempfile
from pathlib import Path
from wrapper import run_pipeline_on_image
from PIL import Image

# Precomputed demo outputs inside the repo (placed by you once)
PRECOMPUTED_DIR = (Path(__file__).resolve().parent / "precomputed_demo")

##add the process function
def process(file_path):
    if not file_path:
        return None, None, None, None, None, [], ""

    with tempfile.TemporaryDirectory() as tmpdir:
        src = Path(file_path)
        ext = src.suffix.lstrip('.') or 'tif'
        img_path = Path(tmpdir) / f"input.{ext}"

        try:
            # Copy raw uploaded bytes
            img_bytes = src.read_bytes()
            img_path.write_bytes(img_bytes)
        except Exception:
            # Fallback: save via PIL if direct copy fails
            Image.open(src).save(img_path)

        # Run the full sorghum pipeline
        outputs = run_pipeline_on_image(str(img_path), tmpdir, save_artifacts=True)

        def load_pil(path_str):
            try:
                if not path_str:
                    return None
                im = Image.open(path_str)
                copied = im.copy()
                im.close()
                return copied
            except Exception:
                return None

        composite = load_pil(outputs.get('Composite'))
        overlay = load_pil(outputs.get('Overlay'))
        mask = load_pil(outputs.get('Mask'))
        input_img = load_pil(outputs.get('InputImage'))
        size_img = load_pil(str(Path(tmpdir) / 'results/size.size_analysis.png'))
        yolo_img = load_pil(str(Path(tmpdir) / 'results/yolo_tips.png'))

        # Texture images (green band)
        lbp_path = Path(tmpdir) / 'texture_output/lbp_green.png'
        hog_path = Path(tmpdir) / 'texture_output/hog_green.png'
        lac1_path = Path(tmpdir) / 'texture_output/lac1_green.png'
        texture_img = load_pil(str(lbp_path)) if lbp_path.exists() else None
        hog_img = load_pil(str(hog_path)) if hog_path.exists() else None
        lac1_img = load_pil(str(lac1_path)) if lac1_path.exists() else None

        # Vegetation indices
        order = ['NDVI', 'GNDVI', 'SAVI']
        gallery_items = [load_pil(outputs[k]) for k in order if k in outputs]

        stats_text = outputs.get('StatsText', '')

        # Output order matches UI components defined below
        # Row 1: Input image (slightly larger)
        # Row 2: Composite, Mask, Overlay
        # Row 3: Texture images (LBP, HOG, Lac1)
        # Row 4: Vegetation indices (gallery)
        # Row 5: Morphology Size and YOLO Tips
        # Final: Stats table
        return (
            input_img,
            composite,
            mask,
            overlay,
            texture_img,
            hog_img,
            lac1_img,
            gallery_items,
            size_img,
            yolo_img,
            stats_text,
        )

def _load_pil(path: Path):
    try:
        if not path or not path.exists():
            return None
        im = Image.open(path)
        out = im.copy()
        im.close()
        return out
    except Exception:
        return None

def _first_existing(paths):
    for p in paths:
        if p and p.exists():
            return p
    return None

def load_precomputed():
    base = PRECOMPUTED_DIR
    if not base.exists():
        return None, None, None, None, None, None, None, [], None, None, ""

    # Common subdirs
    results = base / "results"
    veg = base / "Vegetation_indices_images"
    tex = base / "texture_output"

    input_img = _load_pil(_first_existing([
        base / "input_image.png",
        base / "input_image.tif",
        base / "input.png",
    ]))
    composite = _load_pil(_first_existing([
        results / "composite.png",
    ]))
    mask = _load_pil(results / "mask.png")
    overlay = _load_pil(results / "overlay.png")

    texture_img = _load_pil(_first_existing([
        tex / "lbp_green.png",
        tex / "lbp.png",
    ]))
    hog_img = _load_pil(_first_existing([
        tex / "hog_green.png",
        tex / "hog.png",
    ]))
    lac1_img = _load_pil(_first_existing([
        tex / "lac1_green.png",
        tex / "lacunarity.png",
    ]))

    # Vegetation gallery in order
    veg_order = ["ndvi.png", "gndvi.png", "savi.png"]
    gallery_items = []
    for fname in veg_order:
        p = veg / fname
        img = _load_pil(p)
        if img is not None:
            gallery_items.append(img)

    size_img = _load_pil(results / "size.size_analysis.png")
    yolo_img = _load_pil(results / "yolo_tips.png")

    stats_txt_path = _first_existing([
        base / "stats.txt",
        results / "stats.txt",
    ])
    stats_text = ""
    if stats_txt_path and stats_txt_path.exists():
        try:
            stats_text = stats_txt_path.read_text()
        except Exception:
            stats_text = ""

    return (
        input_img,
        composite,
        mask,
        overlay,
        texture_img,
        hog_img,
        lac1_img,
        gallery_items,
        size_img,
        yolo_img,
        stats_text,
    )


with gr.Blocks() as demo:
    gr.Markdown("# 🌿 Automated Plant Analysis Demo")
    gr.Markdown("Upload a sorghum plant image (TIFF preferred) to compute and visualize composite, mask, overlay, texture (LBP), vegetation indices, and statistics.")

    with gr.Row():
        with gr.Column():
            # Use File input to preserve raw TIFFs
            inp = gr.File(
                type="filepath",
                file_types=[".tif", ".tiff", ".png", ".jpg"],
                label="Upload Image"
            )
            run = gr.Button("Run Pipeline", variant="primary")

    # Row 1: input image, slightly larger
    with gr.Row():
        input_img = gr.Image(type="pil", label="Input Image", interactive=False, height=380)
    # Row 2: composite, mask, overlay
    with gr.Row():
        composite_img = gr.Image(type="pil", label="Composite (Segmentation Input)", interactive=False)
        mask_img = gr.Image(type="pil", label="Mask", interactive=False)
        overlay_img = gr.Image(type="pil", label="Segmentation Overlay", interactive=False)

    # Row 3: textures
    with gr.Row():
        texture_img = gr.Image(type="pil", label="Texture LBP (Green Band)", interactive=False)
        hog_img = gr.Image(type="pil", label="Texture HOG (Green Band)", interactive=False)
        lac1_img = gr.Image(type="pil", label="Texture Lac1 (Green Band)", interactive=False)
    
    # Row 4: vegetation indices

    gallery = gr.Gallery(label="Vegetation Indices", columns=3, height="auto")

    # Row 5: morphology and YOLO tips
    with gr.Row():
        size_img = gr.Image(type="pil", label="Morphology Size", interactive=False)
        yolo_img = gr.Image(type="pil", label="YOLO Tips", interactive=False)

    # Final: statistics table
    stats = gr.Textbox(label="Statistics", lines=4)

    run.click(
        process,
        inputs=inp,
        outputs=[
            input_img,
            composite_img,
            mask_img,
            overlay_img,
            texture_img,
            hog_img,
            lac1_img,
            gallery,
            size_img,
            yolo_img,
            stats,
        ]
    )

    # No preloading: start with empty UI until the user runs the pipeline

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