import gradio as gr
import numpy as np

from ocr import Recognizer

with gr.Blocks() as app:
    recog = Recognizer("model/model.xml", "model/char_list.txt")

    with gr.Row():
        brush = gr.Brush(default_color="#000", default_size=15)
        sketch = gr.Sketchpad("bg.png", type="numpy", brush=brush, layers=False)
        result = gr.TextArea()
    check = gr.Button()

    def parse_item(item):
        prob = item["prob"]
        char = item["char"]
        return f"{char}: {prob}"

    def do_recog(img: dict[str, np.ndarray]):
        _, nbest = recog(img["composite"])
        return "\n".join(parse_item(i) for item in nbest for i in item)

    def clear():
        return "bg.png"

    check.click(do_recog, sketch, result)
    sketch.clear(clear, None, sketch)

app.launch()