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)