提高涂抹的准确性

photo_mask.py

@@ -76,6 +76,11 @@ def split_image(img, max_ratio=1.41, best_ration=1.41, overlap=0.05):
     return split_result
 
 
+def capture_image(img, layout):
+    x1, y1, x2, y2 = layout
+    return img[int(y1):int(y2), int(x1):int(x2)]
+
+
 # 获取图片旋转角度
 def get_image_rotation_angles(img):
     angles = ['0', '90']
@@ -182,7 +187,16 @@ def get_ocr_layout(ocr, img_path):
     return layout
 
 
-def find_box_of_content(content, layout):
+def zoom_box(box, ratio):
+    x1, y1, x2, y2 = box
+    x1 = round(x1 - x1 * ratio)
+    y1 = round(y1 - y1 * ratio)
+    x2 = round(x2 + x2 * ratio)
+    y2 = round(y2 + y2 * ratio)
+    return [x1, y1, x2, y2]
+
+
+def find_box_of_content(content, layout, img_path):
     full_box = layout[0]
     x_len = full_box[2] - full_box[0]
     y_len = full_box[3] - full_box[1]
@@ -201,20 +215,55 @@ def find_box_of_content(content, layout):
 
     if direction == "x":
         # 横向排布
-        return (
+        box = [
             full_box[0] + index * char_len,
             full_box[1],
             full_box[0] + (index + len(content) + 1) * char_len,
             full_box[3],
-        )
+        ]
+        is_abnormal = box[2] - box[0] < (box[3] - box[1]) * len(content) / 2
     else:
         # 纵向排布
-        return (
+        box = [
             full_box[0],
             full_box[1] + index * char_len,
             full_box[2],
             full_box[1] + (index + len(content) + 1) * char_len,
-        )
+        ]
+        is_abnormal = box[3] - box[1] < (box[2] - box[0]) * len(content) / 2
+
+    if is_abnormal:
+        # 比例异常，再次识别
+        image = cv2.imread(img_path)
+        # 截图时偏大一点
+        capture_box = zoom_box(box, 0.2)
+        captured_image = capture_image(image, capture_box)
+        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
+            cv2.imwrite(temp_file.name, captured_image)
+        try:
+            layouts = DOC_PARSER.parse({"doc": temp_file.name})["layout"]
+        except TypeError:
+            # 如果是类型错误，大概率是没识别到文字
+            layouts = []
+        except Exception as e:
+            # 如果出现其他错误，抛出
+            raise e
+        for layout in layouts:
+            if content in layout[1]:
+                temp_box = find_box_of_content(content, layout, temp_file.name)
+                if temp_box:
+                    box = [
+                        temp_box[0] + capture_box[0],
+                        temp_box[1] + capture_box[1],
+                        temp_box[2] + capture_box[0],
+                        temp_box[3] + capture_box[1],
+                    ]
+                    break
+        try:
+            os.remove(temp_file.name)
+        except Exception as e:
+            logging.info(f"删除临时文件 {temp_file.name} 时出错", exc_info=e)
+    return box
 
 
 def find_box_of_value(key, layout, length):
@@ -275,7 +324,7 @@ def get_mask_layout(image, contents):
             for layout in layouts:
                 for content in contents:
                     if content in layout[1]:
-                        result.append(find_box_of_content(content, layout))
+                        result.append(find_box_of_content(content, layout, temp_file.name))
                 if "姓名" in layout[1]:
                     result.append(find_box_of_value("姓名", layout, 4))
                 if "交款人" in layout[1]: