Spaces:
Runtime error
Runtime error
import gradio as gr | |
from transformers import TrOCRProcessor, VisionEncoderDecoderModel | |
import requests | |
from PIL import Image | |
from io import BytesIO | |
# Load TrOCR model | |
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten") | |
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten") | |
def process_image(image): | |
"""Processes an image with TrOCR and returns the extracted text.""" | |
# Prepare image | |
pixel_values = processor(image, return_tensors="pt").pixel_values | |
# Generate text | |
generated_ids = model.generate(pixel_values, max_new_tokens=50) | |
# Decode text | |
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
return generated_text | |
# Gradio interface | |
title = "Handwritten Text OCR with TrOCR" | |
description = "Upload a handwritten text image and get the extracted text using Microsoft's TrOCR model." | |
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2109.10282'>TrOCR: Transformer-based Optical Character Recognition</a> | <a href='https://github.com/microsoft/unilm/tree/master/trocr'>GitHub Repo</a></p>" | |
examples = [ | |
["https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg"], | |
["https://upload.wikimedia.org/wikipedia/commons/8/8d/Handwriting-sample.jpg"] | |
] | |
iface = gr.Interface(fn=process_image, | |
inputs=gr.Image(type="pil"), | |
outputs=gr.Textbox(), | |
title=title, | |
description=description, | |
article=article, | |
examples=examples) | |
iface.launch() | |