From 5193839c9ec231a20977d38b1fbeeb77bc13c3b5 Mon Sep 17 00:00:00 2001
From: liuyebo <1515783401@qq.com>
Date: Fri, 19 Jul 2024 12:15:02 +0800
Subject: [PATCH] =?UTF-8?q?=E8=B0=83=E6=95=B4=E6=B6=82=E6=8A=B9=E7=9A=84oc?=
 =?UTF-8?q?r=E6=A8=A1=E5=9E=8B=E5=8F=82=E6=95=B0=EF=BC=9B=E4=BC=98?=
 =?UTF-8?q?=E5=8C=96=E5=85=B3=E9=94=AE=E8=AF=8D?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 photo_mask/__init__.py   | 10 +++++-----
 photo_mask/photo_mask.py | 13 ++++++++-----
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/photo_mask/__init__.py b/photo_mask/__init__.py
index d0fc862..b08371f 100644
--- a/photo_mask/__init__.py
+++ b/photo_mask/__init__.py
@@ -1,4 +1,4 @@
-from paddlenlp.utils.doc_parser import DocParser
+from paddleocr import PaddleOCR
 
 """
 项目配置
@@ -19,11 +19,11 @@ UPLOAD_TRY_TIMES = 3
 """
 NAME_KEYS = [
     {"key": "姓名", "length": 4},
-    {"key": "交款人", "length": 4},
-    {"key": "文款人", "length": 4},
+    {"key": "娃名", "length": 4},
+    {"key": "款人", "length": 4},
     {"key": "购买方名称", "length": 4},
 ]
 ID_CARD_NUM_KEYS = [{"key": "身份证号", "length": 19}, ]
 
-# OCR = PaddleOCR(use_angle_cls=False, show_log=False, gpu_id=1)
-OCR = DocParser(use_gpu=True, layout_analysis=False)
+# 如果不希望识别出空格，可以设置use_space_char=False。做此项设置一定要测试，2.7.3版本此项设置有bug，会导致识别失败
+OCR = PaddleOCR(use_angle_cls=False, show_log=False, det_db_box_thresh=0.3)
diff --git a/photo_mask/photo_mask.py b/photo_mask/photo_mask.py
index 94b0bbd..89900cb 100644
--- a/photo_mask/photo_mask.py
+++ b/photo_mask/photo_mask.py
@@ -60,7 +60,7 @@ def find_boxes(content, layout, offset=0, length=None, improve=False, image_path
             captured_image, offset_x, offset_y = image_util.expand_to_a4_size(captured_image, True)
             cv2.imwrite(temp_file.name, captured_image)
         try:
-            layouts = OCR.parse({"doc": temp_file.name})["layout"]
+            layouts = util.get_ocr_layout(OCR, temp_file.name)
         except TypeError:
             # 如果是类型错误，大概率是没识别到文字
             layouts = []
@@ -91,8 +91,8 @@ def get_mask_layout(image, name, id_card_num):
     result = []
     try:
         try:
-            # layouts = util.get_ocr_layout(OCR, temp_file.name)
-            layouts = OCR.parse({"doc": temp_file.name})["layout"]
+            layouts = util.get_ocr_layout(OCR, temp_file.name)
+            # layouts = OCR.parse({"doc": temp_file.name})["layout"]
         except TypeError:
             # 如果是类型错误，大概率是没识别到文字
             layouts = []
@@ -118,6 +118,10 @@ def get_mask_layout(image, name, id_card_num):
                 keys = []
                 if find_name_by_key:
                     keys += NAME_KEYS
+                    if layout[1].startswith("名："):
+                        result += find_boxes("名：", layout, offset=2, length=3)
+                    if layout[1].startswith("名:"):
+                        result += find_boxes("名:", layout, offset=2, length=3)
                 if find_id_card_num_by_key:
                     keys += ID_CARD_NUM_KEYS
                 for key in keys:
@@ -135,8 +139,7 @@ def handle_image_for_mask(split_result):
     expand_img, offset_x, offset_y = image_util.expand_to_a4_size(split_result["img"], True)
     split_result["x_offset"] -= offset_x
     split_result["y_offset"] -= offset_y
-    gray_image = cv2.cvtColor(expand_img, cv2.COLOR_BGR2GRAY)
-    return gray_image, split_result["x_offset"], split_result["y_offset"]
+    return expand_img, split_result["x_offset"], split_result["y_offset"]
 
 
 def photo_mask(pk_phhd, name, id_card_num):