From 203f1c4e4a6372aa64cc6623e9fa8cc195bdea3f Mon Sep 17 00:00:00 2001
From: liuyebo <1515783401@qq.com>
Date: Fri, 28 Jun 2024 20:33:21 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=97=8B=E8=BD=AC=E5=9B=BE?=
 =?UTF-8?q?=E7=89=87=E7=9A=84=E8=AF=86=E5=88=AB=E6=9D=A1=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 photo_review/photo_review.py | 76 ++++++++++++++++++++----------------
 1 file changed, 42 insertions(+), 34 deletions(-)

diff --git a/photo_review/photo_review.py b/photo_review/photo_review.py
index d479ca8..4534634 100644
--- a/photo_review/photo_review.py
+++ b/photo_review/photo_review.py
@@ -35,8 +35,6 @@ from photo_review.util.data_util import handle_date, handle_decimal, parse_depar
 from photo_review.util.util import get_default_datetime
 from ucloud import ucloud
 
-OCR = PaddleOCR(use_angle_cls=True, lang="ch", show_log=False)
-
 
 # 获取图片
 def open_image(img_path):
@@ -137,7 +135,7 @@ def rotate_image(img, angle):
 
 
 # 获取图片OCR，并将其box转为两点矩形框
-def get_ocr_layout(img_path):
+def get_ocr_layout(ocr, img_path):
     def _get_box(old_box):
         new_box = [
             min(old_box[0][0], old_box[3][0]),  # x1
@@ -154,7 +152,7 @@ def get_ocr_layout(img_path):
         return True
 
     layout = []
-    ocr_result = OCR.ocr(img_path)
+    ocr_result = ocr.ocr(img_path)
     ocr_result = ocr_result[0]
     if not ocr_result:
         return layout
@@ -168,11 +166,34 @@ def get_ocr_layout(img_path):
     return layout
 
 
+def ie_temp_image(ie, ocr, image):
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
+        cv2.imwrite(temp_file.name, image)
+
+    ie_result = []
+    try:
+        layout = get_ocr_layout(ocr, temp_file.name)
+        if not layout:
+            # 无识别结果
+            ie_result = []
+        else:
+            ie_result = ie({"doc": temp_file.name, "layout": layout})[0]
+    except Exception as e:
+        logging.error("信息抽取时出错", exc_info=e)
+    finally:
+        try:
+            os.remove(temp_file.name)
+        except Exception as e:
+            logging.info(f"删除临时文件 {temp_file.name} 时出错", exc_info=e)
+    return ie_result
+
+
 # 关键信息提取
 def information_extraction(ie, phrecs):
     result = {}
     # 同一批图的标识
     identity = int(time.time())
+    ocr = PaddleOCR(use_angle_cls=True, lang="ch", show_log=False)
     for phrec in phrecs:
         pic_path = ucloud.get_private_url(phrec.cfjaddress)
         if not pic_path:
@@ -180,39 +201,26 @@ def information_extraction(ie, phrecs):
         split_result = split_image(pic_path)
 
         for img in split_result:
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file1:
-                cv2.imwrite(temp_file1.name, img["img"])
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file2:
+            ie_result1 = ie_temp_image(ie, ocr, img["img"])
+            if not ie_result1 or len(ie_result1) < len(ie.kwargs.get("schema")):
                 rotated_img = rotate_image(img["img"], 90)
-                cv2.imwrite(temp_file2.name, rotated_img)
-
-            ie_result = []
-            try:
-                layout1 = get_ocr_layout(temp_file1.name)
-                layout2 = get_ocr_layout(temp_file2.name)
-                if not (layout1 or layout2):
-                    # 无识别结果
+                ie_result2 = ie_temp_image(ie, ocr, rotated_img)
+                if not (ie_result1 or ie_result2):
                     continue
-                ie_result = ie([{"doc": temp_file1.name, "layout": layout1},
-                                {"doc": temp_file2.name, "layout": layout2}])
-            except Exception as e:
-                logging.error("信息抽取时出错", exc_info=e)
-            finally:
-                try:
-                    os.remove(temp_file1.name)
-                except Exception as e:
-                    logging.info(f"删除临时文件 {temp_file1.name} 时出错", exc_info=e)
-                try:
-                    os.remove(temp_file2.name)
-                except Exception as e:
-                    logging.info(f"删除临时文件 {temp_file2.name} 时出错", exc_info=e)
-
-            if len(ie_result[1]) > len(ie_result[0]):
-                # 旋转90度后识别效果更好
-                ie_result = ie_result[1]
-                angle = 90
+                elif not ie_result1:
+                    ie_result = ie_result2
+                    angle = 90
+                elif not ie_result2:
+                    ie_result = ie_result1
+                    angle = 0
+                elif len(ie_result2) > len(ie_result1):
+                    ie_result = ie_result2
+                    angle = 90
+                else:
+                    ie_result = ie_result1
+                    angle = 0
             else:
-                ie_result = ie_result[0]
+                ie_result = ie_result1
                 angle = 0
 
             now = get_default_datetime()