fix: process pdf once

app.py

@@ -2,15 +2,12 @@ import base64
 import os
 import uuid
 
+import gradio as gr
 import torch
-from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import JSONResponse
 from transformers import AutoConfig, AutoModel, AutoTokenizer
 
 from got_ocr import got_ocr
 
-app = FastAPI()
-
 # 初始化模型和分词器
 model_name = "ucaslcl/GOT-OCR2_0"
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -27,12 +24,13 @@ UPLOAD_FOLDER = "./uploads"
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 
 
-@app.post("/ocr")
-async def perform_ocr(image: UploadFile = File(...)):
+def perform_ocr(image):
+    if image is None:
+        return "请上传图片"
+
     # 保存上传的图片
     image_path = os.path.join(UPLOAD_FOLDER, f"{uuid.uuid4()}.png")
-    with open(image_path, "wb") as buffer:
-        buffer.write(await image.read())
+    image.save(image_path)
 
     # 执行OCR
     result, html_content = got_ocr(model, tokenizer, image_path, got_mode="format texts OCR")
@@ -40,15 +38,27 @@ async def perform_ocr(image: UploadFile = File(...)):
     # 删除临时文件
     os.remove(image_path)
 
-    # 准备响应
-    response = {"result": result}
     if html_content:
-        response["html_content"] = base64.b64encode(html_content.encode("utf-8")).decode("utf-8")
-
-    return JSONResponse(content=response)
-
+        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)
+
+
+# 创建 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结果")
+
+    ocr_button.click(fn=perform_ocr, inputs=image_input, outputs=output)
 
 if __name__ == "__main__":
-    import uvicorn
-
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    demo.launch()