import gradio as gr
from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image

# Khởi tạo processor và model
processor = TrOCRProcessor.from_pretrained('microsoft/trocr-large-printed')
model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-large-printed')

def ocr_image(image):
    # Chuyển đổi hình ảnh để xử lý
    pixel_values = processor(images=image, return_tensors="pt").pixel_values
    # Sinh mã từ hình ảnh
    generated_ids = model.generate(pixel_values)
    # Giải mã thành văn bản
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return generated_text

# Tạo giao diện người dùng với Gradio
iface = gr.Interface(fn=ocr_image, 
                     inputs=gr.Image(type="pil"),  # Đã sửa ở đây
                     outputs="text",
                     title="OCR with TrOCR",
                     description="Tải lên một hình ảnh có văn bản in để nhận văn bản được trích xuất.")

# Chạy ứng dụng
iface.launch()