优化涂抹的匹配与方向判断

photo_mask/photo_mask.py

@@ -67,12 +67,13 @@ def find_boxes(content, layout, offset=0, length=None, improve=False, image_path
         except Exception as e:
             # 如果出现其他错误，抛出
             raise e
-        if extra_content:
-            extra_content.append(content)
-        else:
-            extra_content = [content]
+
         for layout in layouts:
-            for c in extra_content:
+            if extra_content:
+                matches = re.findall(extra_content, layout[1])
+            else:
+                matches = [content]
+            for c in matches:
                 if c in layout[1]:
                     temp_box = find_boxes(c, layout)[0]
                     if temp_box:
@@ -110,36 +111,31 @@ def get_mask_layout(image, name, id_card_num):
             # 无识别结果
             return result
         else:
-            # 处理形近字
-            name_list = [name]
-            # 移除名字中重复的字
-            unique_name = "".join(dict.fromkeys(name))
-            for char in unique_name:
+            re_list = []
+            for char in name:
+                char_re = f"{char}"
                 if char in SIMILAR_CHAR:
-                    for sc in SIMILAR_CHAR[char]:
-                        name_list.append(name.replace(char, sc))
+                    char_re += "|" + "|".join(SIMILAR_CHAR[char])
+                re_list.append(char_re)
 
             name_len = len(name)
             name_offset = int(name_len / 2)
-            r = f"{name[:-1]}.?|.?{name[1:]}" + (
-                f"|{name[:name_offset]}.?{name[-name_offset:]}" if (name_len & 1) == 1 else "")
+            if name_len > 2:
+                r = f"[{']['.join(re_list[:-1])}].?|.?[{']['.join(re_list[1:])}]"
+            else:
+                r = f"[{']['.join(re_list)}]"
+            if (name_len & 1) == 1:
+                r += f"|[{']['.join(re_list[:name_offset])}].?[{']['.join(re_list[-name_offset:])}]"
 
             for layout in layouts:
                 find_name_by_key = True
                 find_id_card_num_by_key = True
-                for name in name_list:
-                    if name in layout[1]:
-                        result += find_boxes(name, layout, improve=True, image_path=temp_file.name)
-                        find_name_by_key = False
-                        break
-                if find_name_by_key and name_len > 2:
-                    matches = re.findall(r, layout[1])
-                    for match in matches:
-                        if match in layout[1]:
-                            result += find_boxes(match, layout, improve=True, image_path=temp_file.name,
-                                                 extra_content=name_list)
-                            find_name_by_key = False
-                            break
+                matches = re.findall(r, layout[1])
+                for match in matches:
+                    result += find_boxes(match, layout, improve=True, image_path=temp_file.name, extra_content=r)
+                    find_name_by_key = False
+                    break
+
                 if id_card_num in layout[1]:
                     result += find_boxes(id_card_num, layout, improve=True, image_path=temp_file.name)
                     find_id_card_num_by_key = False
@@ -176,27 +172,41 @@ def handle_image_for_mask(split_result):
 
 
 def mask_photo(img_url, name, id_card_num, color=(255, 255, 255)):
-    # 是否涂抹了
-    is_masked = False
+    def _mask(i, n, icn, c):
+        do_mask = False
+        split_results = image_util.split(i)
+        for split_result in split_results:
+            to_mask_img, x_offset, y_offset = handle_image_for_mask(split_result)
+            results = get_mask_layout(to_mask_img, n, icn)
+
+            if results:
+                do_mask = True
+
+            for result in results:
+                result = (
+                    result[0] + x_offset,
+                    result[1] + y_offset,
+                    result[2] + x_offset,
+                    result[3] + y_offset,
+                )
+                cv2.rectangle(i, (int(result[0]), int(result[1])), (int(result[2]), int(result[3])), c, -1, 0)
+        return do_mask, i
+
     # 打开图片
     image = image_util.read(img_url)
-    # image = image_util.rotate(image, 180)
-    split_results = image_util.split(image)
-    for split_result in split_results:
-        to_mask_img, x_offset, y_offset = handle_image_for_mask(split_result)
-        results = get_mask_layout(to_mask_img, name, id_card_num)
+    original_image = image
+    is_masked, image = _mask(image, name, id_card_num, color)
+    if not is_masked:
+        # 如果没有涂抹，可能是图片方向不对
+        angles = image_util.parse_rotation_angles(image)
+        angle = angles[0]
+        if angle != "0":
+            image = image_util.rotate(image, int(angle))
+            is_masked, image = _mask(image, name, id_card_num, color)
+        if not is_masked:
+            # 如果旋转后也没有涂抹，恢复原来的方向
+            image = original_image
 
-        if results:
-            is_masked = True
-
-        for result in results:
-            result = (
-                result[0] + x_offset,
-                result[1] + y_offset,
-                result[2] + x_offset,
-                result[3] + y_offset,
-            )
-            cv2.rectangle(image, (int(result[0]), int(result[1])), (int(result[2]), int(result[3])), color, -1, 0)
     return is_masked, image