MohsenDehghani's picture
Update app.py
5692df3 verified
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()