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

# 加载模型
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")

# 定义图像文字识别函数
def recognize_text_from_image(image):
    # 图像预处理
    image = Image.open(image).convert("RGB")
    # 使用OCR模型预测
    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

# Gradio界面
iface = gr.Interface(
    fn=recognize_text_from_image,
    inputs=gr.Image(type="file", label="上传文件图片"),
    outputs="text",
    title="桌面文件扫描器",
    description="上传文件图片，将自动识别其中的文字内容"
)

# 启动 Gradio
iface.launch()