app.py

import json
import os
from functools import lru_cache
from pathlib import Path
from typing import Any, Dict, Tuple

os.environ["CUDA_VISIBLE_DEVICES"] = ""

import gradio as gr
from huggingface_hub import hf_hub_download
from ultralytics import YOLO


SPACE_ROOT = Path(__file__).resolve().parent
DEFAULT_MODEL_FILE = "best.pt"
DEFAULT_MODEL_PATH = SPACE_ROOT / DEFAULT_MODEL_FILE


def _resolve_model_path() -> Path:
    if DEFAULT_MODEL_PATH.exists():
        return DEFAULT_MODEL_PATH

    return Path(
        hf_hub_download(
            repo_id="acwz/TeaLeafDetection",
            filename="best.pt",
        )
    )


@lru_cache(maxsize=1)
def get_model() -> YOLO:
    """Load the YOLO model once and reuse it for all requests."""
    model_path = _resolve_model_path()
    return YOLO(str(model_path), task="detect")


def predict(image, confidence: float = 0.25) -> Tuple[Any, Dict[str, Any]]:
    """Run inference and return the plotted image plus structured predictions."""
    if image is None:
        raise gr.Error("Please upload an image before running detection.")

    model = get_model()
    results = model.predict(source=image, conf=confidence, save=False, verbose=False)

    if not results:
        return image, {"error": "No predictions were produced."}

    result = results[0]
    plotted = result.plot()
    detections_raw = result.tojson()

    try:
        detections = json.loads(detections_raw) if isinstance(detections_raw, str) else detections_raw
    except json.JSONDecodeError:
        detections = {"raw": detections_raw}

    return plotted, detections


CSS = """
:root {
    --bg-start: #f2f8f1;
    --bg-end: #dbe8db;
    --panel: rgba(255, 255, 255, 0.82);
    --panel-border: rgba(39, 89, 59, 0.12);
    --text: #12311f;
    --muted: #4f6658;
    --accent: #1f7a4d;
}

body, .gradio-container {
    background: linear-gradient(160deg, var(--bg-start), var(--bg-end));
    color: var(--text);
}

.wrap {
    max-width: 1120px !important;
}

.hero {
    border: 1px solid var(--panel-border);
    border-radius: 28px;
    background: linear-gradient(135deg, rgba(31, 122, 77, 0.94), rgba(18, 49, 31, 0.94));
    color: white;
    padding: 1.35rem 1.5rem;
    box-shadow: 0 18px 40px rgba(18, 49, 31, 0.16);
}

.hero h1 {
    margin: 0;
    font-size: clamp(1.7rem, 3.8vw, 2.8rem);
    line-height: 1.05;
}

.hero p {
    margin: 0.5rem 0 0;
    opacity: 0.92;
}

.panel {
    border-radius: 24px !important;
    border: 1px solid var(--panel-border) !important;
    background: var(--panel) !important;
    backdrop-filter: blur(10px);
    box-shadow: 0 12px 30px rgba(18, 49, 31, 0.08);
}

.panel h3, .panel h4, .panel label {
    color: var(--text) !important;
}

.panel .wrap {
    max-width: none !important;
}

.gr-button-primary {
    background: linear-gradient(135deg, #1f7a4d, #2ea66b) !important;
    border: 0 !important;
}

.gr-button-secondary {
    border: 1px solid rgba(31, 122, 77, 0.26) !important;
}
"""


with gr.Blocks(title="Tea Leaf Detection API", theme=gr.themes.Soft(), css=CSS) as demo:
    with gr.Column(elem_classes=["wrap"]):
        gr.Markdown(
            """
            <div class="hero">
              <h1>Tea Leaf Detection API</h1>
              <p>Upload a tea leaf image, run YOLO inference, and call the same endpoint from code via the Gradio API.</p>
            </div>
            """
        )

        with gr.Row():
            with gr.Column(scale=1, elem_classes=["panel"]):
                input_image = gr.Image(type="pil", label="Tea leaf image")
                confidence = gr.Slider(
                    minimum=0.1,
                    maximum=0.9,
                    value=float(os.getenv("CONFIDENCE_THRESHOLD", "0.25")),
                    step=0.05,
                    label="Confidence threshold",
                )
                run_button = gr.Button("Run detection", variant="primary")

            with gr.Column(scale=1, elem_classes=["panel"]):
                output_image = gr.Image(type="numpy", label="Detection result")
                output_json = gr.JSON(label="Predictions")

        gr.Markdown(
            "The prediction function is exposed as an API endpoint, so you can call it with `api_name=\"/predict\"`."
        )

        run_button.click(
            fn=predict,
            inputs=[input_image, confidence],
            outputs=[output_image, output_json],
            api_name="predict",
        )

demo.queue()


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