TrOCR-SIN-DeiT / app.py
kavg's picture
Create app.py
3abacca verified
raw
history blame contribute delete
No virus
879 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")
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(debug=True)