import gradio as gr
import torch
from PIL import Image
from transformers import TrOCRProcessor, VisionEncoderDecoderModel

MODEL_ID = "kkatiz/thai-trocr-thaigov-v2"

_processor = None
_model = None


def load_model():
    global _processor, _model
    if _processor is None or _model is None:
        _processor = TrOCRProcessor.from_pretrained(MODEL_ID)
        _model = VisionEncoderDecoderModel.from_pretrained(MODEL_ID)
        _model.eval()
    return _processor, _model


def ocr(image: Image.Image) -> str:
    if image is None:
        return ""
    processor, model = load_model()
    pixel_values = processor(images=image, return_tensors="pt").pixel_values
    with torch.no_grad():
        generated_ids = model.generate(pixel_values)
    text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return text.strip()


with gr.Blocks(title="Thai TrOCR Demo") as demo:
    gr.Markdown(
        "# Thai TrOCR Demo\n"
        "Upload an image and the model will extract Thai text."
    )

    with gr.Row():
        image_input = gr.Image(type="pil", label="Input image")
        text_output = gr.Textbox(label="Predicted text")

    run_btn = gr.Button("Run OCR")
    run_btn.click(fn=ocr, inputs=image_input, outputs=text_output)

    gr.Examples(
        examples=[
            "https://i.ibb.co/QXZFSNx/test7.png",
        ],
        inputs=image_input,
        cache_examples=False,
    )


demo.launch()