ThomasZilliox/equation-solver

1import warnings2warnings.filterwarnings("ignore")3import logging4logging.getLogger("asyncio").setLevel(logging.CRITICAL)5 6import gradio as gr7import cv28import numpy as np9import matplotlib10matplotlib.use("Agg")11import matplotlib.pyplot as plt12from PIL import Image13import io14 15# ── Import pipeline from core ──────────────────────────────────────────────16from core import (17    extract_characters,18    load_models,19    predict_char,20    predict_chars,21    solve_equation,22)23import tensorflow as tf24from pathlib import Path25 26 27# ── Visualisation helpers ──────────────────────────────────────────────────28 29def draw_segmentation(original_img_gray, bboxes, labels, premium=False):30    """Draw coloured bounding boxes + labels on the original image."""31    display = cv2.cvtColor(original_img_gray, cv2.COLOR_GRAY2RGB)32 33    palette = [34        (220,  50,  50),35        ( 50, 140, 220),36        ( 50, 200,  80),37        (220, 160,  50),38        (160,  50, 220),39        ( 50, 200, 200),40    ]41    # Premium uses gold-tinted palette42    if premium:43        palette = [44            (212, 175,  55),45            (255, 215,   0),46            (184, 134,  11),47            (218, 165,  32),48            (255, 193,   7),49            (205, 152,  15),50        ]51 52    for i, ((x, y, w, h), label) in enumerate(zip(bboxes, labels)):53        color = palette[i % len(palette)]54        cv2.rectangle(display, (x, y), (x + w, y + h), color, 2)55        font_scale = max(0.6, min(1.2, h / 40))56        (tw, th), baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, 2)57        label_y = y - 6 if y - 6 - th >= 0 else y + h + th + 458        cv2.rectangle(display,59                      (x, label_y - th - baseline),60                      (x + tw + 4, label_y + baseline),61                      color, -1)62        cv2.putText(display, label, (x + 2, label_y),63                    cv2.FONT_HERSHEY_SIMPLEX, font_scale,64                    (255, 255, 255), 2, cv2.LINE_AA)65 66    max_w, max_h = 800, 40067    h_img, w_img = display.shape[:2]68    scale = min(max_w / w_img, max_h / h_img, 1.0)69    if scale < 1.0:70        display = cv2.resize(display,71                             (int(w_img * scale), int(h_img * scale)),72                             interpolation=cv2.INTER_AREA)73    return display74 75 76def make_char_strip(segmented_chars, labels, confidences=None, premium=False):77    """Row of character tiles with label + optional confidence."""78    if not segmented_chars:79        return None80 81    if confidences is None:82        confidences = [None] * len(segmented_chars)83 84    n = len(segmented_chars)85    std_colors    = ["#dc3232","#328cdc","#32c850","#dca032","#a032dc","#32c8c8"]86    premium_colors= ["#d4af37","#ffd700","#b8860b","#daa520","#ffc107","#cd981a"]87    colors = premium_colors if premium else std_colors88 89    bg = "#1a1a2e" if not premium else "#12100a"90    fig, axes = plt.subplots(1, n, figsize=(max(2, n * 1.2), 2.2))91    if n == 1:92        axes = [axes]93    fig.patch.set_facecolor(bg)94 95    for i, (char_img, label, conf) in enumerate(zip(segmented_chars, labels, confidences)):96        ax    = axes[i]97        color = colors[i % len(colors)]98        ax.imshow(char_img, cmap="gray", vmin=0, vmax=255)99        title = f"{label}\n{conf:.0%}" if conf is not None else label100        ax.set_title(title, color=color, fontsize=12,101                     fontweight="bold", pad=3, linespacing=1.3)102        ax.axis("off")103        for spine in ax.spines.values():104            spine.set_edgecolor(color)105            spine.set_linewidth(2)106            spine.set_visible(True)107 108    plt.tight_layout(pad=0.5)109    buf = io.BytesIO()110    plt.savefig(buf, format="png", dpi=120, bbox_inches="tight",111                facecolor=fig.get_facecolor())112    plt.close(fig)113    buf.seek(0)114    return Image.open(buf).copy()115 116 117# ── Model caches (kept separate so both can coexist) ──────────────────────118 119_cnn_cache = None   # {"mode": "cnn", "digit": …, "operator": …}120_vit_cache = None   # {"mode": "vit", "vit": …, "class_names": …}121 122 123HF_REPO    = "ThomasZilliox/equation-solver"124MODEL_DIR  = Path("./model")125 126def _hf_download(filename):127    """Download a file from HF Hub into ./model/ and return its local path."""128    from huggingface_hub import hf_hub_download129    print(f"Downloading {filename} from HuggingFace Hub...")130    path = hf_hub_download(repo_id=HF_REPO, filename=filename,131                           local_dir=str(MODEL_DIR))132    print(f"  -> {path}")133    return Path(path)134 135 136def _get_cnn_models():137    global _cnn_cache138    if _cnn_cache is None:139        digit_path    = MODEL_DIR / "digit_recognizer.h5"140        operator_path = MODEL_DIR / "operator_recognizer.h5"141        if not digit_path.exists():142            digit_path    = _hf_download("digit_recognizer.h5")143        if not operator_path.exists():144            operator_path = _hf_download("operator_recognizer.h5")145        digit_model    = tf.keras.models.load_model(str(digit_path))146        operator_model = tf.keras.models.load_model(str(operator_path))147        _cnn_cache = {"mode": "cnn", "digit": digit_model, "operator": operator_model}148    return _cnn_cache149 150 151def _get_vit_models():152    global _vit_cache153    if _vit_cache is None:154        import json155        import onnxruntime as ort156        onnx_path = MODEL_DIR / "vit_math_recognizer.onnx"157        cnames    = MODEL_DIR / "class_names.json"158        if not onnx_path.exists():159            onnx_path = _hf_download("vit_math_recognizer.onnx")160        # Download the external weights file if present in the repo161        data_path = MODEL_DIR / "vit_math_recognizer.onnx.data"162        if not data_path.exists():163            try:164                _hf_download("vit_math_recognizer.onnx.data")165            except Exception:166                pass  # not all exports have an external data file167        if not cnames.exists():168            cnames = _hf_download("class_names.json")169        session = ort.InferenceSession(170            str(onnx_path),171            providers=["CPUExecutionProvider"]172        )173        with open(cnames) as f:174            class_names = {int(k): v for k, v in json.load(f).items()}175        _vit_cache = {"mode": "vit", "session": session, "class_names": class_names}176    return _vit_cache177 178 179# ── Shared inference logic ─────────────────────────────────────────────────180 181def _run_pipeline(pil_image, models, use_craft=False):182    """183    Core pipeline used by both buttons.184    use_craft=True  → CRAFT segmentation (premium)185    use_craft=False → CV2 segmentation  (standard)186    """187    if pil_image is None:188        return None, None, "⚠️ Please upload an image."189 190    img_gray = np.array(pil_image.convert("L"))191 192    segmented_chars, bboxes = extract_characters(img_gray)193 194    if not segmented_chars:195        return np.array(pil_image.convert("RGB")), None, "⚠️ No characters found in the image."196 197    labels, confidences = [], []198    for char in segmented_chars:199        label, conf = predict_char(char, models, log=True)200        labels.append(label)201        confidences.append(conf)202 203    equation_string = "".join(labels)204    result_text     = solve_equation(equation_string)205    avg_conf        = sum(confidences) / len(confidences) if confidences else 0206 207    premium = models["mode"] == "vit"208    annotated  = draw_segmentation(img_gray, bboxes, labels, premium=premium)209    char_strip = make_char_strip(segmented_chars, labels, confidences, premium=premium)210 211    seg_label  = "Adaptive CV2"212    rec_label  = "ViT (fine-tuned)" if premium else "CNN"213    display_result = (214        f"**Segmentation:** `{seg_label}`  |  "215        f"**Recognizer:** `{rec_label}`  |  "216        f"**Avg confidence:** `{avg_conf:.1%}`\n\n"217        f"**Detected equation:** `{equation_string}`\n\n"218        f"**Result:** `{result_text}`"219    )220    return annotated, char_strip, display_result221 222 223# ── Button handlers ────────────────────────────────────────────────────────224 225def solve_standard(pil_image):226    """CV2 segmentation + CNN recognizer."""227    try:228        models = _get_cnn_models()229    except Exception as e:230        return None, None, f"⚠️ Could not load CNN models: {e}"231    return _run_pipeline(pil_image, models)232 233 234def solve_premium(pil_image):235    """CRAFT segmentation + ViT recognizer."""236    try:237        models = _get_vit_models()238    except FileNotFoundError as e:239        return None, None, f"⚠️ {e}"240    except Exception as e:241        return None, None, f"⚠️ Could not load ViT model: {e}"242    return _run_pipeline(pil_image, models)243 244 245# ── UI ─────────────────────────────────────────────────────────────────────246 247css = """248#title        { text-align: center; }249#result_box   { font-size: 1.2em; }250.gr-image     { border-radius: 8px; }251 252/* Premium button — gold gradient */253#premium_btn  { background: linear-gradient(135deg, #b8860b, #ffd700, #b8860b);254                color: #1a1000 !important;255                font-weight: 700;256                border: none; }257#premium_btn:hover { background: linear-gradient(135deg, #daa520, #ffe066, #daa520); }258"""259 260with gr.Blocks(title="Handwritten Equation Solver") as demo:261 262    gr.Markdown(263        "# 🔢 Handwritten Equation Solver\n"264        "Upload a photo of a **handwritten arithmetic equation** (digits `0–9` and operators `+ - * /`).",265        elem_id="title",266    )267 268    with gr.Row():269        # ── Left: input + buttons ──────────────────────────────────────────270        with gr.Column(scale=1):271            input_image = gr.Image(272                type="pil",273                label="📷 Upload equation image",274                image_mode="RGB",275            )276            with gr.Row():277                run_btn = gr.Button("Solve ➜", variant="primary")278                premium_btn = gr.Button("⭐ Premium Solve", elem_id="premium_btn")279 280            gr.Markdown(281                "<small>"282                "**Solve** — Adaptive segmentation · CNN recognizer<br>"283                "**Premium Solve** — Adaptive segmentation · fine-tuned ViT"284                "</small>"285            )286 287        # ── Right: outputs ─────────────────────────────────────────────────288        with gr.Column(scale=2):289            annotated_image = gr.Image(290                label="🔍 Segmentation",291                type="numpy",292                interactive=False,293            )294            char_strip_image = gr.Image(295                label="✂️ Extracted characters",296                type="pil",297                interactive=False,298            )299            result_md = gr.Markdown(300                value="*Upload an image and click **Solve** or **Premium Solve**.*",301                elem_id="result_box",302            )303 304    outputs = [annotated_image, char_strip_image, result_md]305 306    run_btn.click(fn=solve_standard, inputs=[input_image], outputs=outputs)307    premium_btn.click(fn=solve_premium, inputs=[input_image], outputs=outputs)308 309    gr.Markdown(310        "---\n"311        "**Tips:** Use a white background with dark ink. "312        "Supported operators: `+`  `−`  `×`  `÷`"313    )314 315if __name__ == "__main__":316    demo.launch(css=css)