Extraction / app.py
sriusha's picture
Create app.py
4437fbf verified
Raw History Blame Contribute Delete
879 Bytes
import gradio as gr
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image
# Load processor + model
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
def ocr_image(image):
# Convert to PIL Image if uploaded
if not isinstance(image, Image.Image):
image = Image.fromarray(image)
pixel_values = processor(images=image, return_tensors="pt").pixel_values
generated_ids = model.generate(pixel_values)
text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
return text
demo = gr.Interface(
fn=ocr_image,
inputs=gr.Image(type="pil"), # Image upload box
outputs=gr.Textbox(label="Extracted Text") # Show extracted text
)
if __name__ == "__main__":
demo.launch()