From 4e5e80e6b45f6fc8407a29de36ecf5a7163a23b4 Mon Sep 17 00:00:00 2001
From: ivamp <im@ivampiresp.com>
Date: Fri, 25 Oct 2024 00:30:25 +0800
Subject: [PATCH] =?UTF-8?q?=E6=94=B9=E8=BF=9B=20=E5=93=8D=E5=BA=94?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 main.py     | 36 +++++++++++++++++++++++++++---------
 response.py |  1 +
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/main.py b/main.py
index 4dc790f..87a331e 100644
--- a/main.py
+++ b/main.py
@@ -111,7 +111,8 @@ async def ocr_image(lang_1: Optional[str] = None, lang_2: Optional[str] = None,
             success=False,
             message="OCR 处理失败",
             data=None,
-            error=error
+            error=error,
+            full_text=None,
         )
 
     try:
@@ -122,7 +123,8 @@ async def ocr_image(lang_1: Optional[str] = None, lang_2: Optional[str] = None,
             success=False,
             message="无法读取图片文件",
             data=None,
-            error=error
+            error=error,
+            full_text=None,
         )
 
     # 调用通用的 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 处理结果
-    pprint.pprint(ocr_data)
 
     return OCRResponse(
         success=True,
         data=ocr_data_list,
         error=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 请求
     async with aiohttp.ClientSession() as session:
         try:
-            async with session.get(url) as response:
-                if response.status != 200:
-                    raise Exception(f"无法获取图片, HTTP 状态码: {response.status}")
+            async with session.get(url) as http_response:
+                if http_response.status != 200:
+                    raise Exception(f"无法获取图片, HTTP 状态码: {http_response.status}")
                 # 读取图片数据
-                image_data = await response.read()
+                image_data = await http_response.read()
 
         except Exception as e:
             return OCRResponse(
@@ -167,6 +168,7 @@ async def ocr_image_from_url(lang_1: Optional[str] = None, lang_2: Optional[str]
                 message="无法从 URL 获取图片",
                 data=None,
                 error=str(e),
+                full_text=None,
             )
 
     # 获取 OCR 读取器并捕获可能的错误
@@ -178,16 +180,32 @@ async def ocr_image_from_url(lang_1: Optional[str] = None, lang_2: Optional[str]
             message="OCR 处理失败",
             data=None,
             error=error,
+            full_text=None,
         )
 
     # 调用通用的 OCR 处理函数
     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(
         success=True,
-        data=ocr_data,
+        data=ocr_data_list,
+        error=None,
+        message=None,
+        full_text=full_text,
     )
 
 
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)
diff --git a/response.py b/response.py
index 9712028..c1c9272 100644
--- a/response.py
+++ b/response.py
@@ -13,3 +13,4 @@ class OCRResponse(BaseModel):
     error: Optional[str]
     message: Optional[str]
     data: Optional[List[OCRData]] = []
+    full_text: Optional[str]