fix: process pdf once

app.py

@@ -1,99 +1,39 @@
-import base64
-import os
-import uuid
-
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoConfig, AutoModel, AutoTokenizer
+from transformers import AutoModel, AutoTokenizer
 
-# 初始化模型和分词器
-model_name = "stepfun-ai/GOT-OCR2_0"
+model_name = "ucaslcl/GOT-OCR2_0"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModel.from_pretrained(model_name, trust_remote_code=True, low_cpu_mem_usage=True, device_map="cuda", use_safetensors=True)
+model = AutoModel.from_pretrained(model_name, trust_remote_code=True, device_map=device)
 model = model.eval().to(device)
-model.config.pad_token_id = tokenizer.eos_token_id
-
-UPLOAD_FOLDER = "./uploads"
-
-# 确保上传文件夹存在
-os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 
 
 @spaces.GPU()
-def got_ocr(model, tokenizer, image_path, got_mode="format texts OCR", fine_grained_mode="", ocr_color="", ocr_box=""):
-    # 执行OCR
-    try:
-        if got_mode == "plain texts OCR":
-            res = model.chat(tokenizer, image_path, ocr_type="ocr")
-            return res, None
-        elif got_mode == "format texts OCR":
-            result_path = f"{os.path.splitext(image_path)[0]}_result.html"
-            res = model.chat(tokenizer, image_path, ocr_type="format", render=True, save_render_file=result_path)
-        elif got_mode == "plain multi-crop OCR":
-            res = model.chat_crop(tokenizer, image_path, ocr_type="ocr")
-            return res, None
-        elif got_mode == "format multi-crop OCR":
-            result_path = f"{os.path.splitext(image_path)[0]}_result.html"
-            res = model.chat_crop(tokenizer, image_path, ocr_type="format", render=True, save_render_file=result_path)
-        elif got_mode == "plain fine-grained OCR":
-            res = model.chat(tokenizer, image_path, ocr_type="ocr", ocr_box=ocr_box, ocr_color=ocr_color)
-            return res, None
-        elif got_mode == "format fine-grained OCR":
-            result_path = f"{os.path.splitext(image_path)[0]}_result.html"
-            res = model.chat(tokenizer, image_path, ocr_type="format", ocr_box=ocr_box, ocr_color=ocr_color, render=True, save_render_file=result_path)
-
-        # 处理格式化结果
-        if "format" in got_mode and os.path.exists(result_path):
-            with open(result_path, "r") as f:
-                html_content = f.read()
-            encoded_html = base64.b64encode(html_content.encode("utf-8")).decode("utf-8")
-            return res, encoded_html
-        else:
-            return res, None
-
-    except Exception as e:
-        return f"错误: {str(e)}", None
-
-
-def perform_ocr(image):
+def ocr_process(image):
     if image is None:
-        return "请上传图片"
-
-    # 保存上传的图片
-    image_path = os.path.join(UPLOAD_FOLDER, f"{uuid.uuid4()}.png")
-    image.save(image_path)
+        return "错误:未提供图片"
 
-    # 执行OCR
-    result, html_content = got_ocr(model, tokenizer, image_path, got_mode="format texts OCR")
-
-    # 删除临时文件
-    os.remove(image_path)
-
-    if html_content:
-        encoded_html = base64.b64encode(html_content.encode("utf-8")).decode("utf-8")
-        iframe_src = f"data:text/html;base64,{encoded_html}"
-        iframe = f'<iframe src="{iframe_src}" width="100%" height="600px"></iframe>'
-        download_link = f'<a href="data:text/html;base64,{encoded_html}" download="result.html">下载完整结果</a>'
-        return gr.HTML(f"{download_link}<br>{iframe}")
-    else:
-        return gr.Markdown(result)
+    try:
+        res = model.chat(tokenizer, image, ocr_type="ocr")
+        return res
+    except Exception as e:
+        return f"错误: {str(e)}"
 
 
-# 创建 Gradio 界面
 with gr.Blocks() as demo:
     gr.Markdown("# OCR 图像识别")
+
     with gr.Row():
-        image_input = gr.Image(type="pil", label="上传图片")
-    with gr.Row():
-        ocr_button = gr.Button("开始OCR识别")
-    with gr.Row():
-        output = gr.HTML(label="OCR结果")
+        image_input = gr.Image(type="filepath", label="上传图片")
+
+    submit_button = gr.Button("开始OCR识别")
+
+    output_text = gr.Textbox(label="识别结果")
 
-    ocr_button.click(fn=perform_ocr, inputs=image_input, outputs=output)
+    submit_button.click(ocr_process, inputs=[image_input], outputs=[output_text])
 
 if __name__ == "__main__":
     demo.launch()