Spaces:
Sleeping
Sleeping
| """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) | |