kavg's picture
Update app.py
cbcf7a3 verified
raw
history blame
No virus
903 Bytes
from transformers import VisionEncoderDecoderModel
from transformers import TrOCRProcessor, AutoTokenizer, ViTImageProcessor
import gradio as gr
from PIL import Image
def OCR(image):
model = VisionEncoderDecoderModel.from_pretrained("kavg/TrOCR-SIN-DeiT-Handwritten")
tokenizer = AutoTokenizer.from_pretrained('NLPC-UOM/SinBERT-large')
feature_extractor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224")
processor = TrOCRProcessor(image_processor=feature_extractor, tokenizer=tokenizer)
pixel_values = processor(image, return_tensors="pt").pixel_values
generated_ids = model.generate(pixel_values)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
return generated_text
demo = gr.Interface(fn=OCR, inputs=gr.Image(show_label=False, type="pil"),outputs=gr.Textbox())
demo.launch(share=True, debug=True)