import gradio as gr
from transformers import TrOCRProcessor, VisionEncoderDecoderModel

name = "chanelcolgate/trocr-base-printed_captcha_ocr"
model = VisionEncoderDecoderModel.from_pretrained(name)
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-printed")


def process_image(image):
    # prepare image
    pixel_values = processor(image, return_tensors="pt").pixel_values

    # generate (no beam search)
    generated_ids = model.generate(pixel_values)

    # decode
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return generated_text


title = "Interactive demo: Captcha OCR"
description = "Demo tracuumasothue captcha"

interface = gr.Interface(fn=process_image,
                         inputs="image",
                         examples=[f"examples/captcha-{i}.png" for i in range(10)],
                         outputs="text",
                         title=title,
                         description=description).launch()