Blame for app/app.py

"""Gradio demo: pick or upload a tuber image and get morphometrics plus predicted plot yield.

Pipeline per image:
  1. YOLOv11-seg instance segmentation.
  2. Per-tuber polygon morphometrics.
  3. Pixel-to-inch calibration via the 30" conveyor belt.
  4. USDA size class assignment.
  5. Plot-level features (total and class areas, mean_area) feed an M3
     yield prediction using saved coefficients.

Two input modes:
  - Dataset image: pick from the bundled set under $POTATOPHENO_DATASET.
  - Custom upload: any top-down RGB image of tubers on a 30" conveyor.

Environment variables:
  POTATOPHENO_WEIGHTS  path to bestModel.pt (default: app/bestModel.pt)
  POTATOPHENO_COEFS    path to m3_coefficients.json (default: app/m3_coefficients.json)
  POTATOPHENO_DATASET  root with 2025<TRIAL>/ subdirs (default: app/dataset)
  POTATOPHENO_BELT_IN  conveyor belt width in inches (default: 30.0)
"""

from __future__ import annotations

import json
import os
from pathlib import Path

import gradio as gr
import numpy as np
import pandas as pd
from PIL import Image, ImageDraw

from potatopheno.calibration import BELT_WIDTH_IN_DEFAULT, apply_real_scale
from potatopheno.classification import assign_usda_class
from potatopheno.io import IMG_EXTS
from potatopheno.morphometrics import polygon_shape_metrics
from potatopheno.segmentation import load_model, predict_image

APP_DIR = Path(__file__).parent

WEIGHTS_PATH = Path(os.environ.get("POTATOPHENO_WEIGHTS", APP_DIR / "bestModel.pt"))
COEFS_PATH = Path(os.environ.get("POTATOPHENO_COEFS", APP_DIR / "m3_coefficients.json"))
DATASET_ROOT = Path(os.environ.get("POTATOPHENO_DATASET", APP_DIR / "dataset"))
BELT_WIDTH_IN = float(os.environ.get("POTATOPHENO_BELT_IN", BELT_WIDTH_IN_DEFAULT))

# USDA-class colour palette (B -> OV; small -> large). Used both for the overlay
# and the legend rendered next to the headline.
USDA_COLORS: dict[str, tuple[int, int, int]] = {
    "B":  (66, 135, 245),
    "A1": (52, 168, 83),
    "A2": (251, 188, 5),
    "OV": (234, 67, 53),
}
UNKNOWN_COLOR = (180, 180, 180)

import importlib.metadata as _md

# Mask geometry depends on these versions; log them so a Space build that
# drifts from the pinned environment is visible in the runtime logs.
print("[app] versions:     " + ", ".join(
    f"{_p}={_md.version(_p)}" for _p in ("ultralytics", "torch", "numpy", "opencv-python")
))
print(f"[app] weights:      {WEIGHTS_PATH}")
print(f"[app] coefficients: {COEFS_PATH}")
print(f"[app] dataset:      {DATASET_ROOT}")
_model = load_model(WEIGHTS_PATH)
_coefs = json.loads(Path(COEFS_PATH).read_text())["M3"]


def _scan_dataset(root: Path) -> dict[str, str]:
    """Return {label: absolute_path} for every image under root/2025<TRIAL>/."""
    out: dict[str, str] = {}
    if not root.exists():
        print(f"[app] WARNING: dataset root {root} does not exist; dataset tab will be empty")
        return out
    for trial_dir in sorted(root.glob("2025*")):
        trial = trial_dir.name.removeprefix("2025")
        for img in sorted(trial_dir.rglob("*")):
            if img.suffix.lower() in IMG_EXTS:
                out[f"{trial} / {img.stem}"] = str(img)
    return out


_dataset = _scan_dataset(DATASET_ROOT)
_dataset_labels = list(_dataset.keys())
print(f"[app] dataset images indexed: {len(_dataset_labels)}")


# ----- core analysis (shared by both input modes)
def _render_overlay(
    image: Image.Image,
    detections: list[tuple[int, np.ndarray]],
    usda_classes: list[str | float],
) -> Image.Image:
    """Draw per-tuber polygons colored by predicted USDA size class."""
    base = image.convert("RGBA")
    layer = Image.new("RGBA", base.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(layer)
    w, h = base.size
    for (_class_id, poly_norm), usda in zip(detections, usda_classes):
        coords = [(float(x) * w, float(y) * h) for x, y in poly_norm]
        if len(coords) < 3:
            continue
        color = USDA_COLORS.get(str(usda), UNKNOWN_COLOR)
        draw.polygon(coords, outline=(*color, 235), fill=(*color, 95))
    return Image.alpha_composite(base, layer).convert("RGB")


def _per_tuber_dataframe(detections: list[tuple[int, np.ndarray]], img_w: int, img_h: int) -> pd.DataFrame:
    rows: list[dict] = []
    for class_id, poly_norm in detections:
        xs = poly_norm[:, 0] * img_w
        ys = poly_norm[:, 1] * img_h
        metrics = polygon_shape_metrics(xs, ys)
        rows.append({"class_id": class_id, "img_width_px": img_w, "img_height_px": img_h, **metrics})
    if not rows:
        return pd.DataFrame()
    df = pd.DataFrame(rows)
    df = apply_real_scale(df, belt_width_in=BELT_WIDTH_IN)
    df = assign_usda_class(df)
    return df


def _aggregate_features(df: pd.DataFrame) -> dict[str, float]:
    by_class = df.groupby("class", observed=True)["area_poly_in2"].sum().to_dict()
    n_tubers = int(len(df))
    total = float(df["area_poly_in2"].sum())
    return {
        "area_B_in2": float(by_class.get("B", 0.0)),
        "area_A1_in2": float(by_class.get("A1", 0.0)),
        "area_A2_in2": float(by_class.get("A2", 0.0)),
        "area_OV_in2": float(by_class.get("OV", 0.0)),
        "total_area_in2": total,
        "mean_area": total / n_tubers if n_tubers else 0.0,
        "n_tubers": n_tubers,
    }


def _predict_m3(features: dict[str, float]) -> float:
    yhat = _coefs["intercept"]
    for predictor, coef in _coefs["coefficients"].items():
        yhat += coef * features[predictor]
    return float(yhat)


def _class_summary_df(df: pd.DataFrame, feats: dict[str, float]) -> pd.DataFrame:
    counts = df["class"].value_counts().reindex(["B", "A1", "A2", "OV"], fill_value=0)
    n_total = int(counts.sum())
    rows = []
    for cls in ["B", "A1", "A2", "OV"]:
        n = int(counts[cls])
        area = feats[f"area_{cls}_in2"]
        share = (n / n_total * 100.0) if n_total else 0.0
        rows.append({"class": cls, "tubers": n, "share (%)": round(share, 1), "area (in²)": round(area, 2)})
    return pd.DataFrame(rows)


def _empty_analysis(message: str):
    return None, message, pd.DataFrame(), pd.DataFrame()


def _analyze_pil(image: Image.Image, source_label: str | None = None):
    """Shared analysis path for both dataset and upload tabs.

    Returns (overlay PIL image, headline markdown, class summary df, per-tuber df).
    """
    # No retina_masks: the batch pipeline that produced the paper's tables runs
    # with the Ultralytics defaults, and full-resolution masks shift polygon
    # boundaries enough to move ~14% of tubers across a USDA class threshold.
    detections = predict_image(_model, image)
    df = _per_tuber_dataframe(detections, image.width, image.height)

    if df.empty:
        return _empty_analysis("**No tubers detected.**")

    overlay = _render_overlay(image, detections, df["class"].astype(str).tolist())
    feats = _aggregate_features(df)
    yhat = _predict_m3(feats)
    class_summary = _class_summary_df(df, feats)

    source_line = f"_Source: {source_label}_\n\n" if source_label else ""
    headline = (
        f"## Predicted yield (M3): **{yhat:.2f} kg**\n\n"
        f"{source_line}"
        f"- {feats['n_tubers']} tubers detected\n"
        f"- Total projected area: {feats['total_area_in2']:.2f} in²\n"
        f"- Mean tuber area: {feats['mean_area']:.2f} in²\n\n"
        f"**Mask color legend**: "
        f"<span style='color:rgb(66,135,245)'>■ B</span> "
        f"<span style='color:rgb(52,168,83)'>■ A1</span> "
        f"<span style='color:rgb(251,188,5)'>■ A2</span> "
        f"<span style='color:rgb(234,67,53)'>■ OV</span>"
    )

    display_cols = [
        "class", "area_poly_in2", "width_in", "height_in",
        "inMinDim", "perimeter_px", "circularity", "eccentricity",
    ]
    return overlay, headline, class_summary, df[display_cols].round(3)


# ----- handlers
def _load_dataset_preview(label: str | None):
    if not label or label not in _dataset:
        return None
    return Image.open(_dataset[label])


def analyze_from_dataset(label: str | None):
    if not label or label not in _dataset:
        return _empty_analysis("Pick a dataset image from the dropdown.")
    image = Image.open(_dataset[label])
    return _analyze_pil(image, source_label=label)


def analyze_from_upload(image: Image.Image | None):
    if image is None:
        return _empty_analysis("Upload an image to analyze.")
    return _analyze_pil(image, source_label=None)


# ----- UI
def _results_section() -> tuple[gr.Markdown, gr.Dataframe, gr.Dataframe]:
    """Full-width Results block: headline | class breakdown side-by-side, then detail full-width."""
    gr.Markdown("## Results")
    with gr.Row(equal_height=False):
        with gr.Column(scale=1):
            headline = gr.Markdown()
        with gr.Column(scale=1):
            class_summary = gr.Dataframe(
                label="Per class info",
                headers=["class", "tubers", "share (%)", "area (in²)"],
                interactive=False,
                wrap=True,
            )
    detail = gr.Dataframe(label="Per-tuber morphometrics", interactive=False, wrap=True)
    return headline, class_summary, detail


with gr.Blocks(title="Potato tuber phenotyping demo") as demo:
    gr.Markdown(
        "# Potato tuber phenotyping demo\n"
        "Pipeline: YOLOv11-seg instance segmentation → per-tuber morphometrics → USDA size class"
        " → plot-yield prediction via the M3 model detailed in the paper."
        "<br />"
        "***Images are assumed to be top-down RGB shots of harvested tubers on a 30″ wide conveyor belt.***"        
    )

    with gr.Tabs():
        with gr.Tab("Dataset image"):
            gr.Markdown(
                f"Pick one of **{len(_dataset_labels)}** images from the paper dataset. "
                "Start typing to filter."
            )
            with gr.Row(equal_height=False):
                with gr.Column(scale=1):
                    dataset_selector = gr.Dropdown(
                        choices=_dataset_labels,
                        label="Dataset image",
                        value=_dataset_labels[0] if _dataset_labels else None,
                        filterable=True,
                    )
                    dataset_preview = gr.Image(type="pil", label="Selected", interactive=False, height=320)
                    dataset_btn = gr.Button("Analyze", variant="primary")
                with gr.Column(scale=2):
                    ds_overlay = gr.Image(label="Detected tubers (coloured by USDA size class)", height=480)

            ds_headline, ds_class, ds_detail = _results_section()

            dataset_selector.change(
                fn=_load_dataset_preview,
                inputs=[dataset_selector],
                outputs=[dataset_preview],
            )
            dataset_btn.click(
                fn=analyze_from_dataset,
                inputs=[dataset_selector],
                outputs=[ds_overlay, ds_headline, ds_class, ds_detail],
            )
            if _dataset_labels:
                demo.load(fn=_load_dataset_preview, inputs=[dataset_selector], outputs=[dataset_preview])

        with gr.Tab("Upload your own"):
            gr.Markdown(
                "Drop in a top-down RGB image of tubers on a 30″ conveyor. "
                "Calibration assumes the 30″ width maps to the image height in pixels."
            )
            with gr.Row(equal_height=False):
                with gr.Column(scale=1):
                    upload_image = gr.Image(type="pil", label="Image", height=320)
                    upload_btn = gr.Button("Analyze", variant="primary")
                with gr.Column(scale=2):
                    up_overlay = gr.Image(label="Detected tubers (coloured by USDA size class)", height=480)

            up_headline, up_class, up_detail = _results_section()

            upload_btn.click(
                fn=analyze_from_upload,
                inputs=[upload_image],
                outputs=[up_overlay, up_headline, up_class, up_detail],
            )


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)