Update app/main.py

app/main.py

@@ -7,8 +7,9 @@ import gradio as gr
 import re
 import random
 from datetime import datetime
+from io import BytesIO
+from urllib.parse import quote_plus
 from jinja2 import Template
-import urllib.parse
 import openai
 
 # —— 环境变量读取 —— #
@@ -30,31 +31,42 @@ def get_access_token(api_key, secret_key):
     resp.raise_for_status()
     return resp.json().get("access_token")
 
-def ocr_image(image_bytes, token):
-    # Base64 encode and URL-encode the image
+def ocr_image(image_bytes: bytes, token: str):
+    # Base64 编码后 URL 编码
     img_b64 = base64.b64encode(image_bytes).decode()
-    img_encoded = urllib.parse.quote_plus(img_b64)
-    url = f"https://aip.baidubce.com/rest/2.0/ocr/v1/handwriting?access_token={token}"
-    headers = {"Content-Type": "application/x-www-form-urlencoded"}
-    data = f"image={img_encoded}&language_type=ENG"
-    resp = requests.post(url, headers=headers, data=data)
+    img_encoded = quote_plus(img_b64)
+    resp = requests.post(
+        f"https://aip.baidubce.com/rest/2.0/ocr/v1/handwriting?access_token={token}",
+        headers={"Content-Type": "application/x-www-form-urlencoded"},
+        data={
+            "image": img_encoded,
+            "language_type": "ENG"
+        }
+    )
     resp.raise_for_status()
-    return resp.json().get("words_result", [])
+    data = resp.json()
+    return data.get("words_result", [])
 
 # —— 文本高亮 —— #
-def highlight_brackets(text):
+def highlight_brackets(text: str) -> str:
+    # 中括号绿色；小括号红色
     text = re.sub(r'\[([^\[\]]+)\]', r'<span class="highlight-bracket-green">\1</span>', text)
     text = re.sub(r'\(([^\(\)]+)\)', r'<span class="highlight-bracket">\1</span>', text)
     return text.replace("\n", "<br>")
 
 # —— 主处理函数 —— #
-def process(image_bytes: bytes):
-    # 1. OCR
+def process(image_pil):
+    # 1. PIL 转 bytes
+    buf = BytesIO()
+    image_pil.save(buf, format="PNG")
+    image_bytes = buf.getvalue()
+
+    # 2. OCR
     token = get_access_token(BAIDU_API_KEY, BAIDU_SECRET_KEY)
     words = ocr_image(image_bytes, token)
     essay_text = "\n".join([w["words"] for w in words])
 
-    # 2. 格式化原文
+    # 3. 格式化原文
     fmt_prompt = (
         "请帮我整理下面的英语作文文本格式，只整理英文正文部分，"
         "保证原汁原味（明显错误空格换行、乱码、非常用字符改正），"
@@ -66,7 +78,7 @@ def process(image_bytes: bytes):
     )
     revised = fm.choices[0].message.content
 
-    # 3. 批改
+    # 4. 批改
     corr_prompt = (
         "请帮我把下面的英语作文的语法错误改正，输出改正后的文章，"
         "原文错误用()括起来，修改部分用[]括起来：\n\n" + revised
@@ -77,10 +89,9 @@ def process(image_bytes: bytes):
     )
     corrected = cm.choices[0].message.content
 
-    # 4. 批改意见
+    # 5. 批改意见
     review_prompt = (
-        "下面是一份已经批改过的英语作文，请根据批注给出逐条批改意见：\n\n"
-        + corrected
+        "下面是一份已经批改过的英语作文，请根据批注给出逐条批改意见：\n\n" + corrected
     )
     rm = openai.ChatCompletion.create(
         model="gpt-4o-mini",
@@ -88,10 +99,9 @@ def process(image_bytes: bytes):
     )
     review = rm.choices[0].message.content
 
-    # 5. 评分
+    # 6. 评分
     rate_prompt = (
-        "请按照 IELTS/CEFR 写作评价体系，从语言通顺度、连贯度、词汇与语法三维度打分并给出原因：\n\n"
-        + revised
+        "请按照 IELTS/CEFR 写作评价体系，从语言通顺度、连贯度、词汇与语法三维度打分并给出原因：\n\n" + revised
     )
     rr = openai.ChatCompletion.create(
         model="gpt-4o-mini",
@@ -99,7 +109,7 @@ def process(image_bytes: bytes):
     )
     rating = rr.choices[0].message.content
 
-    # 6. 优秀范文
+    # 7. 优秀范文
     rewrite_prompt = (
         "请使用优秀表达重写下面这篇作文，加粗可供学习的部分：\n\n" + revised
     )
@@ -109,10 +119,12 @@ def process(image_bytes: bytes):
     )
     perfect = wm.choices[0].message.content
 
-    # 7. 渲染 HTML
+    # 8. 渲染 HTML
     code = f"{random.randint(0,9999):04}-{datetime.now().strftime('%Y%m%d%H%M%S')}"
-    with open("app/templates/base.html", encoding="utf-8") as f:
+    tpl_path = os.path.join("app", "templates", "base.html")
+    with open(tpl_path, encoding="utf-8") as f:
         tpl = Template(f.read())
+
     html_content = (
         "<h2>原文格式化</h2>"
         + markdown.markdown(revised)
@@ -127,25 +139,34 @@ def process(image_bytes: bytes):
     )
     full_html = tpl.render(code=code, content=html_content)
 
-    # 写文件
-    os.makedirs("app/output", exist_ok=True)
-    html_path = f"app/output/{code}.html"
-    pdf_path  = f"app/output/{code}.pdf"
+    # 9. 写文件并返回
+    output_dir = os.path.join("app", "output")
+    os.makedirs(output_dir, exist_ok=True)
+    html_path = os.path.join(output_dir, f"{code}.html")
+    pdf_path  = os.path.join(output_dir, f"{code}.pdf")
+
     with open(html_path, "w", encoding="utf-8") as f:
         f.write(full_html)
-    pdfkit.from_string(full_html, pdf_path, options={"enable-local-file-access": ""})
+
+    pdfkit.from_string(full_html, pdf_path,
+                       options={"enable-local-file-access": ""})
 
     return full_html, html_path, pdf_path
 
 # —— Gradio 接口 —— #
 with gr.Blocks(title="英语作文批改") as demo:
     gr.Markdown("## 上传英语作文照片，等待批改完成后下载 HTML 或 PDF")
-    image_in = gr.Image(type="bytes", label="上传照片")
+    image_in  = gr.Image(type="pil", label="上传照片")
     output_html = gr.HTML()
     btn = gr.Button("开始批改")
     file_html = gr.File(label="下载 HTML")
     file_pdf  = gr.File(label="下载 PDF")
-    btn.click(fn=process, inputs=image_in, outputs=[output_html, file_html, file_pdf])
+
+    btn.click(
+        fn=process,
+        inputs=image_in,
+        outputs=[output_html, file_html, file_pdf]
+    )
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)