改进 响应

This commit is contained in:
ivamp 2024-10-25 00:30:25 +08:00
parent 22c3795f9f
commit 4e5e80e6b4
2 changed files with 28 additions and 9 deletions

36
main.py
View File

@ -111,7 +111,8 @@ async def ocr_image(lang_1: Optional[str] = None, lang_2: Optional[str] = None,
success=False, success=False,
message="OCR 处理失败", message="OCR 处理失败",
data=None, data=None,
error=error error=error,
full_text=None,
) )
try: try:
@ -122,7 +123,8 @@ async def ocr_image(lang_1: Optional[str] = None, lang_2: Optional[str] = None,
success=False, success=False,
message="无法读取图片文件", message="无法读取图片文件",
data=None, data=None,
error=error error=error,
full_text=None,
) )
# 调用通用的 OCR 处理函数 # 调用通用的 OCR 处理函数
@ -138,14 +140,13 @@ async def ocr_image(lang_1: Optional[str] = None, lang_2: Optional[str] = None,
) )
ocr_data_list.append(r) ocr_data_list.append(r)
# 返回 OCR 处理结果
pprint.pprint(ocr_data)
return OCRResponse( return OCRResponse(
success=True, success=True,
data=ocr_data_list, data=ocr_data_list,
error=None, error=None,
message=None, message=None,
full_text=full_text,
) )
@ -155,11 +156,11 @@ async def ocr_image_from_url(lang_1: Optional[str] = None, lang_2: Optional[str]
# 使用 aiohttp 进行异步 HTTP 请求 # 使用 aiohttp 进行异步 HTTP 请求
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
try: try:
async with session.get(url) as response: async with session.get(url) as http_response:
if response.status != 200: if http_response.status != 200:
raise Exception(f"无法获取图片, HTTP 状态码: {response.status}") raise Exception(f"无法获取图片, HTTP 状态码: {http_response.status}")
# 读取图片数据 # 读取图片数据
image_data = await response.read() image_data = await http_response.read()
except Exception as e: except Exception as e:
return OCRResponse( return OCRResponse(
@ -167,6 +168,7 @@ async def ocr_image_from_url(lang_1: Optional[str] = None, lang_2: Optional[str]
message="无法从 URL 获取图片", message="无法从 URL 获取图片",
data=None, data=None,
error=str(e), error=str(e),
full_text=None,
) )
# 获取 OCR 读取器并捕获可能的错误 # 获取 OCR 读取器并捕获可能的错误
@ -178,16 +180,32 @@ async def ocr_image_from_url(lang_1: Optional[str] = None, lang_2: Optional[str]
message="OCR 处理失败", message="OCR 处理失败",
data=None, data=None,
error=error, error=error,
full_text=None,
) )
# 调用通用的 OCR 处理函数 # 调用通用的 OCR 处理函数
ocr_data, status_code = await process_ocr(reader, image_data) ocr_data, status_code = await process_ocr(reader, image_data)
full_text = ""
ocr_data_list = list()
for entry in ocr_data["data"]:
full_text += entry["text"]
r = response.OCRData(
text=entry["text"],
bbox=entry["bbox"],
confidence=entry["confidence"]
)
ocr_data_list.append(r)
return OCRResponse( return OCRResponse(
success=True, success=True,
data=ocr_data, data=ocr_data_list,
error=None,
message=None,
full_text=full_text,
) )
if __name__ == "__main__": if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000) uvicorn.run(app, host="0.0.0.0", port=8000)

View File

@ -13,3 +13,4 @@ class OCRResponse(BaseModel):
error: Optional[str] error: Optional[str]
message: Optional[str] message: Optional[str]
data: Optional[List[OCRData]] = [] data: Optional[List[OCRData]] = []
full_text: Optional[str]