diff --git a/photo_mask.py b/photo_mask.py index 2454814..6275a64 100644 --- a/photo_mask.py +++ b/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]: