Update app.py

app.py

@@ -1,40 +1,30 @@
 import gradio as gr
-import os
 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")
 
-# 定义图像OCR识别函数
-def ocr_images(images):
-    results = {}
-    for image in images:
-        # 确保图片是RGB格式
-        image = Image.open(image).convert("RGB")
-        pixel_values = processor(images=image, return_tensors="pt").pixel_values
-        output_ids = model.generate(pixel_values)
-        transcription = processor.batch_decode(output_ids, skip_special_tokens=True)[0]
-        results[image.filename] = transcription
-    return results
+# 定义图像文字识别函数
+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界面
-def ocr_interface(images):
-    results = ocr_images(images)
-    result_text = "\n\n".join([f"{filename}:\n{transcription}" for filename, transcription in results.items()])
-    return result_text
+# Gradio界面
+iface = gr.Interface(
+    fn=recognize_text_from_image,
+    inputs=gr.Image(type="file", label="上传文件图片"),
+    outputs="text",
+    title="桌面文件扫描器",
+    description="上传文件图片，将自动识别其中的文字内容"
+)
 
-# 创建Gradio应用
-with gr.Blocks() as demo:
-    gr.Markdown("## 多图片OCR识别")
-    with gr.Row():
-        image_input = gr.File(label="选择多张图片", file_count="multiple", type="file")
-    output_text = gr.Textbox(label="OCR 识别结果")
-    
-    # 添加按钮和功能绑定
-    submit_button = gr.Button("开始识别")
-    submit_button.click(ocr_interface, inputs=image_input, outputs=output_text)
-
-# 启动应用
-demo.launch()
+# 启动 Gradio
+iface.launch()