cacode commited on
Commit
3af6c22
·
verified ·
1 Parent(s): 78fe6b5

Update app/main.py

Browse files
Files changed (1) hide show
  1. app/main.py +22 -19
app/main.py CHANGED
@@ -8,6 +8,7 @@ import re
8
  import random
9
  from datetime import datetime
10
  from jinja2 import Template
 
11
  import openai
12
 
13
  # —— 环境变量读取 —— #
@@ -21,34 +22,36 @@ def get_access_token(api_key, secret_key):
21
  resp = requests.post(
22
  "https://aip.baidubce.com/oauth/2.0/token",
23
  params={
24
- "grant_type":"client_credentials",
25
  "client_id": api_key,
26
  "client_secret": secret_key
27
  }
28
  )
 
29
  return resp.json().get("access_token")
30
 
31
  def ocr_image(image_bytes, token):
 
32
  img_b64 = base64.b64encode(image_bytes).decode()
33
- resp = requests.post(
34
- f"https://aip.baidubce.com/rest/2.0/ocr/v1/handwriting?access_token={token}",
35
- headers={"Content-Type":"application/x-www-form-urlencoded"},
36
- data={"image": img_b64, "language_type":"ENG"}
37
- )
 
38
  return resp.json().get("words_result", [])
39
 
40
  # —— 文本高亮 —— #
41
  def highlight_brackets(text):
42
- # 中括号绿色;小括号红色
43
  text = re.sub(r'\[([^\[\]]+)\]', r'<span class="highlight-bracket-green">\1</span>', text)
44
  text = re.sub(r'\(([^\(\)]+)\)', r'<span class="highlight-bracket">\1</span>', text)
45
  return text.replace("\n", "<br>")
46
 
47
  # —— 主处理函数 —— #
48
- def process(image):
49
  # 1. OCR
50
  token = get_access_token(BAIDU_API_KEY, BAIDU_SECRET_KEY)
51
- words = ocr_image(image.read(), token)
52
  essay_text = "\n".join([w["words"] for w in words])
53
 
54
  # 2. 格式化原文
@@ -59,7 +62,7 @@ def process(image):
59
  )
60
  fm = openai.ChatCompletion.create(
61
  model="gpt-4o-mini",
62
- messages=[{"role":"user","content":fmt_prompt}]
63
  )
64
  revised = fm.choices[0].message.content
65
 
@@ -70,7 +73,7 @@ def process(image):
70
  )
71
  cm = openai.ChatCompletion.create(
72
  model="gpt-4o-mini",
73
- messages=[{"role":"user","content":corr_prompt}]
74
  )
75
  corrected = cm.choices[0].message.content
76
 
@@ -81,7 +84,7 @@ def process(image):
81
  )
82
  rm = openai.ChatCompletion.create(
83
  model="gpt-4o-mini",
84
- messages=[{"role":"user","content":review_prompt}]
85
  )
86
  review = rm.choices[0].message.content
87
 
@@ -92,7 +95,7 @@ def process(image):
92
  )
93
  rr = openai.ChatCompletion.create(
94
  model="gpt-4o-mini",
95
- messages=[{"role":"user","content":rate_prompt}]
96
  )
97
  rating = rr.choices[0].message.content
98
 
@@ -102,7 +105,7 @@ def process(image):
102
  )
103
  wm = openai.ChatCompletion.create(
104
  model="gpt-4o-mini",
105
- messages=[{"role":"user","content":rewrite_prompt}]
106
  )
107
  perfect = wm.choices[0].message.content
108
 
@@ -125,19 +128,19 @@ def process(image):
125
  full_html = tpl.render(code=code, content=html_content)
126
 
127
  # 写文件
128
- html_path = f"/app/output/{code}.html"
129
- pdf_path = f"/app/output/{code}.pdf"
130
- os.makedirs("/app/output", exist_ok=True)
131
  with open(html_path, "w", encoding="utf-8") as f:
132
  f.write(full_html)
133
- pdfkit.from_string(full_html, pdf_path, options={"enable-local-file-access":""})
134
 
135
  return full_html, html_path, pdf_path
136
 
137
  # —— Gradio 接口 —— #
138
  with gr.Blocks(title="英语作文批改") as demo:
139
  gr.Markdown("## 上传英语作文照片,等待批改完成后下载 HTML 或 PDF")
140
- image_in = gr.File(file_count="single", label="上传照片")
141
  output_html = gr.HTML()
142
  btn = gr.Button("开始批改")
143
  file_html = gr.File(label="下载 HTML")
 
8
  import random
9
  from datetime import datetime
10
  from jinja2 import Template
11
+ import urllib.parse
12
  import openai
13
 
14
  # —— 环境变量读取 —— #
 
22
  resp = requests.post(
23
  "https://aip.baidubce.com/oauth/2.0/token",
24
  params={
25
+ "grant_type": "client_credentials",
26
  "client_id": api_key,
27
  "client_secret": secret_key
28
  }
29
  )
30
+ resp.raise_for_status()
31
  return resp.json().get("access_token")
32
 
33
  def ocr_image(image_bytes, token):
34
+ # Base64 encode and URL-encode the image
35
  img_b64 = base64.b64encode(image_bytes).decode()
36
+ img_encoded = urllib.parse.quote_plus(img_b64)
37
+ url = f"https://aip.baidubce.com/rest/2.0/ocr/v1/handwriting?access_token={token}"
38
+ headers = {"Content-Type": "application/x-www-form-urlencoded"}
39
+ data = f"image={img_encoded}&language_type=ENG"
40
+ resp = requests.post(url, headers=headers, data=data)
41
+ resp.raise_for_status()
42
  return resp.json().get("words_result", [])
43
 
44
  # —— 文本高亮 —— #
45
  def highlight_brackets(text):
 
46
  text = re.sub(r'\[([^\[\]]+)\]', r'<span class="highlight-bracket-green">\1</span>', text)
47
  text = re.sub(r'\(([^\(\)]+)\)', r'<span class="highlight-bracket">\1</span>', text)
48
  return text.replace("\n", "<br>")
49
 
50
  # —— 主处理函数 —— #
51
+ def process(image_bytes: bytes):
52
  # 1. OCR
53
  token = get_access_token(BAIDU_API_KEY, BAIDU_SECRET_KEY)
54
+ words = ocr_image(image_bytes, token)
55
  essay_text = "\n".join([w["words"] for w in words])
56
 
57
  # 2. 格式化原文
 
62
  )
63
  fm = openai.ChatCompletion.create(
64
  model="gpt-4o-mini",
65
+ messages=[{"role": "user", "content": fmt_prompt}]
66
  )
67
  revised = fm.choices[0].message.content
68
 
 
73
  )
74
  cm = openai.ChatCompletion.create(
75
  model="gpt-4o-mini",
76
+ messages=[{"role": "user", "content": corr_prompt}]
77
  )
78
  corrected = cm.choices[0].message.content
79
 
 
84
  )
85
  rm = openai.ChatCompletion.create(
86
  model="gpt-4o-mini",
87
+ messages=[{"role": "user", "content": review_prompt}]
88
  )
89
  review = rm.choices[0].message.content
90
 
 
95
  )
96
  rr = openai.ChatCompletion.create(
97
  model="gpt-4o-mini",
98
+ messages=[{"role": "user", "content": rate_prompt}]
99
  )
100
  rating = rr.choices[0].message.content
101
 
 
105
  )
106
  wm = openai.ChatCompletion.create(
107
  model="gpt-4o-mini",
108
+ messages=[{"role": "user", "content": rewrite_prompt}]
109
  )
110
  perfect = wm.choices[0].message.content
111
 
 
128
  full_html = tpl.render(code=code, content=html_content)
129
 
130
  # 写文件
131
+ os.makedirs("app/output", exist_ok=True)
132
+ html_path = f"app/output/{code}.html"
133
+ pdf_path = f"app/output/{code}.pdf"
134
  with open(html_path, "w", encoding="utf-8") as f:
135
  f.write(full_html)
136
+ pdfkit.from_string(full_html, pdf_path, options={"enable-local-file-access": ""})
137
 
138
  return full_html, html_path, pdf_path
139
 
140
  # —— Gradio 接口 —— #
141
  with gr.Blocks(title="英语作文批改") as demo:
142
  gr.Markdown("## 上传英语作文照片,等待批改完成后下载 HTML 或 PDF")
143
+ image_in = gr.Image(type="bytes", label="上传照片")
144
  output_html = gr.HTML()
145
  btn = gr.Button("开始批改")
146
  file_html = gr.File(label="下载 HTML")