Update app.py

app.py

@@ -1,90 +1,157 @@
 import os
+import base64
 import requests
+from io import BytesIO
+from typing import List, Union
+
+from PIL import Image
+import pypdfium2 as pdfium
 import gradio as gr
 
-# 从环境变量中读取你的 HF API Token
+# 从环境变量中读取 HF API Token
 HF_API_TOKEN = os.environ.get("HF_API_TOKEN")
 if HF_API_TOKEN is None:
     raise RuntimeError(
         "环境变量 HF_API_TOKEN 未设置，请在 Space 的 Settings -> Variables 中添加一个名为 HF_API_TOKEN 的 Secret。"
     )
 
-# 想使用的模型 ID，可以自行替换为其他支持 Inference API 的模型
-# 比如 "meta-llama/Llama-3.2-1B-Instruct"、"Qwen/Qwen2.5-1.5B-Instruct" 等
-MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"
-
+# 使用 OCR 模型
+MODEL_ID = "tencent/HunyuanOCR"
 API_URL = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
 HEADERS = {"Authorization": f"Bearer {HF_API_TOKEN}"}
 
 
-def query_hf_api(prompt: str, max_new_tokens: int = 256, temperature: float = 0.7) -> str:
+def image_to_base64(image: Image.Image) -> str:
+    """把 PIL Image 转成 base64 字符串"""
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_bytes = buffered.getvalue()
+    img_b64 = base64.b64encode(img_bytes).decode("utf-8")
+    return img_b64
+
+
+def call_ocr_model(image: Image.Image) -> str:
+    """对单张图片调用 HunyuanOCR"""
+    img_b64 = image_to_base64(image)
     payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": max_new_tokens,
-            "temperature": temperature,
-            "do_sample": True,
-        },
+        "inputs": {
+            "image": img_b64
+        }
     }
-    response = requests.post(API_URL, headers=HEADERS, json=payload, timeout=120)
-    response.raise_for_status()
-    data = response.json()
 
-    # text-generation 类模型常见返回格式是 [{"generated_text": "..."}]
+    try:
+        response = requests.post(API_URL, headers=HEADERS, json=payload, timeout=120)
+        response.raise_for_status()
+    except Exception as e:
+        return f"[调用模型出错] {type(e).__name__}: {e}"
+
+    try:
+        data = response.json()
+    except Exception as e:
+        return f"[解析返回结果出错] {type(e).__name__}: {e}\n原始返回：{response.text[:1000]}"
+
+    # 尝试多种常见返回结构
     if isinstance(data, list) and len(data) > 0:
-        return data[0].get("generated_text", "").strip()
-    # 兜底：直接把返回内容转成字符串方便调试
+        first = data[0]
+        if isinstance(first, dict):
+            for key in ["generated_text", "text", "output", "label"]:
+                if key in first and isinstance(first[key], str):
+                    return first[key].strip()
+        return str(first)
+
+    if isinstance(data, dict):
+        for key in ["generated_text", "text", "output", "label"]:
+            if key in data and isinstance(data[key], str):
+                return data[key].strip()
+        return str(data)
+
     return str(data)
 
 
-def chat_fn(history, message, max_new_tokens, temperature):
-    # 简单地把历史对话拼成一个长 prompt
-    dialog = ""
-    if history:
-        for user_msg, bot_msg in history:
-            dialog += f"用户: {user_msg}\n助手: {bot_msg}\n"
-    dialog += f"用户: {message}\n助手:"
+def pdf_to_images(pdf_bytes: bytes, dpi: int = 200) -> List[Image.Image]:
+    """把 PDF 的每一页渲染成 PIL Image 列表"""
+    pdf = pdfium.PdfDocument(pdf_bytes)
+    n_pages = len(pdf)
+    images: List[Image.Image] = []
 
-    try:
-        output = query_hf_api(dialog, max_new_tokens=int(max_new_tokens), temperature=float(temperature))
-    except Exception as e:
-        output = f"[调用模型出错] {type(e).__name__}: {e}"
+    for i in range(n_pages):
+        page = pdf[i]
+        pil_image = page.render(scale=dpi / 72).to_pil()  # 72 dpi 是 PDF 默认分辨率
+        images.append(pil_image)
+
+    return images
+
+
+def run_ocr(file: Union[bytes, None], image: Union[Image.Image, None]) -> str:
+    """
+    总入口：可以上传 PDF 或 图片。
+    - 如果上传了 PDF（file），对 PDF 每一页做 OCR
+    - 如果只上传图片，对图片做 OCR
+    - 如果两个都没传，提示用户
+    """
+    if file is None and image is None:
+        return "请上传 PDF 文件或图片。"
+
+    results = []
+
+    # 1. 如果上传了 PDF
+    if file is not None:
+        try:
+            pdf_bytes = file
+            pages = pdf_to_images(pdf_bytes)
+        except Exception as e:
+            return f"[解析 PDF 出错] {type(e).__name__}: {e}"
 
-    history = history + [(message, output)]
-    return history, ""
+        if not pages:
+            return "PDF 中未检测到页面。"
+
+        for idx, page_img in enumerate(pages, start=1):
+            text = call_ocr_model(page_img)
+            results.append(f"===== 第 {idx} 页 =====\n{text}\n")
+
+    # 2. 如果上传了图片
+    if image is not None:
+        text = call_ocr_model(image)
+        # 如果前面已经有 PDF 结果，就在后面追加，否则单独一段
+        if results:
+            results.append("===== 图片识别结果 =====\n" + text)
+        else:
+            results.append(text)
+
+    return "\n".join(results)
 
 
 with gr.Blocks() as demo:
-    gr.Markdown(f"# 云端模型聊天 Demo\n使用模型：`{MODEL_ID}`（通过 Hugging Face Inference API）")
+    gr.Markdown(
+        f"""# 文档 OCR Demo（HunyuanOCR）
+使用模型：`{MODEL_ID}`
+
+你可以：
+- 上传 **PDF 文件**（多页会逐页识别，并按页分隔）
+- 或上传 **单张图片**（截图、拍照等）
+"""
+    )
 
     with gr.Row():
-        with gr.Column(scale=3):
-            chatbot = gr.Chatbot(label="对话", height=500)
-            msg = gr.Textbox(
-                label="你的问题",
-                placeholder="输入你想问的问题，回车或点击发送",
-                lines=2,
+        with gr.Column():
+            pdf_input = gr.File(
+                label="上传 PDF 文件（可选）",
+                file_types=[".pdf"],
+                type="binary",
             )
-            send_btn = gr.Button("发送")
-            clear_btn = gr.Button("清空对话")
-
-        with gr.Column(scale=1):
-            gr.Markdown("### 参数设置")
-            max_new_tokens = gr.Slider(16, 512, value=256, step=16, label="max_new_tokens")
-            temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="temperature")
-
-    send_btn.click(
-        chat_fn,
-        inputs=[chatbot, msg, max_new_tokens, temperature],
-        outputs=[chatbot, msg],
-    )
-    msg.submit(
-        chat_fn,
-        inputs=[chatbot, msg, max_new_tokens, temperature],
-        outputs=[chatbot, msg],
+            image_input = gr.Image(
+                type="pil",
+                label="上传图片（可选）",
+            )
+            run_button = gr.Button("开始识别")
+        with gr.Column():
+            output_text = gr.Textbox(label="识别结果", lines=25)
+
+    run_button.click(
+        fn=run_ocr,
+        inputs=[pdf_input, image_input],
+        outputs=output_text,
     )
-    clear_btn.click(lambda: ([], ""), None, [chatbot, msg])
 
 if __name__ == "__main__":
-    # 不要给 launch() 传额外参数，HF 会自己管理 host/port
     demo.launch()