From 6b68a7f7ae802a66c9a2a1bdcc1b6665a1935c42 Mon Sep 17 00:00:00 2001 From: liuyebo <1515783401@qq.com> Date: Fri, 28 Jun 2024 17:01:27 +0800 Subject: [PATCH] =?UTF-8?q?=E5=9B=BE=E7=89=87=E6=96=B9=E5=90=91=E8=AF=86?= =?UTF-8?q?=E5=88=AB=E5=99=A8=E5=87=86=E7=A1=AE=E7=8E=87=E6=9C=89=E9=99=90?= =?UTF-8?q?=EF=BC=8C=E5=B0=9D=E8=AF=95=E5=90=8C=E6=97=B6=E8=AF=86=E5=88=AB?= =?UTF-8?q?=E5=8E=9F=E5=9B=BE=E5=92=8C=E6=97=8B=E8=BD=AC90=E5=BA=A6?= =?UTF-8?q?=E4=B9=8B=E5=90=8E=E7=9A=84=E5=9B=BE=E3=80=82=E4=BF=AE=E6=AD=A3?= =?UTF-8?q?=E6=97=A5=E5=BF=97=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- photo_review/photo_review.py | 72 +++++++++++++++++++++++++++++++----- ucloud/ucloud.py | 2 +- 2 files changed, 64 insertions(+), 10 deletions(-) diff --git a/photo_review/photo_review.py b/photo_review/photo_review.py index 7ac9573..d479ca8 100644 --- a/photo_review/photo_review.py +++ b/photo_review/photo_review.py @@ -11,6 +11,7 @@ import cv2 import numpy as np import paddleclas from paddlenlp import Taskflow +from paddleocr import PaddleOCR sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) @@ -34,6 +35,8 @@ 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): @@ -107,7 +110,7 @@ def get_image_rotation_angle(img): try: angle = int(next(result)[0]["label_names"][0]) except Exception as e: - logging.error("获取图片旋转角度失败", e) + logging.error("获取图片旋转角度失败", exc_info=e) return angle @@ -133,6 +136,38 @@ def rotate_image(img, angle): return rotated +# 获取图片OCR，并将其box转为两点矩形框 +def get_ocr_layout(img_path): + def _get_box(old_box): + new_box = [ + min(old_box[0][0], old_box[3][0]), # x1 + min(old_box[0][1], old_box[1][1]), # y1 + max(old_box[1][0], old_box[2][0]), # x2 + max(old_box[2][1], old_box[3][1]), # y2 + ] + return new_box + + def _normal_box(box_data): + # Ensure the height and width of bbox are greater than zero + if box_data[3] - box_data[1] < 0 or box_data[2] - box_data[0] < 0: + return False + return True + + layout = [] + ocr_result = OCR.ocr(img_path) + ocr_result = ocr_result[0] + if not ocr_result: + return layout + for segment in ocr_result: + box = segment[0] + box = _get_box(box) + if not _normal_box(box): + continue + text = segment[1][0] + layout.append((box, text)) + return layout + + # 关键信息提取 def information_extraction(ie, phrecs): result = {} @@ -145,21 +180,40 @@ 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_file: - angle = get_image_rotation_angle(img["img"]) - rotated_img = rotate_image(img["img"], angle) - cv2.imwrite(temp_file.name, rotated_img) + 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: + rotated_img = rotate_image(img["img"], 90) + cv2.imwrite(temp_file2.name, rotated_img) ie_result = [] try: - ie_result = ie({"doc": temp_file.name})[0] + layout1 = get_ocr_layout(temp_file1.name) + layout2 = get_ocr_layout(temp_file2.name) + if not (layout1 or layout2): + # 无识别结果 + continue + ie_result = ie([{"doc": temp_file1.name, "layout": layout1}, + {"doc": temp_file2.name, "layout": layout2}]) except Exception as e: - logging.error("信息抽取时出错:", e) + logging.error("信息抽取时出错", exc_info=e) finally: try: - os.remove(temp_file.name) + os.remove(temp_file1.name) except Exception as e: - logging.info(f"删除临时文件 {temp_file.name} 时出错: {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 + else: + ie_result = ie_result[0] + angle = 0 now = get_default_datetime() result_json = json.dumps(ie_result, ensure_ascii=False) diff --git a/ucloud/ucloud.py b/ucloud/ucloud.py index 23428b2..d30dca0 100644 --- a/ucloud/ucloud.py +++ b/ucloud/ucloud.py @@ -12,7 +12,7 @@ def get_private_url(key): # 判断文件是否存在 _, resp = get_ufile_handler.head_file(BUCKET, key) if resp.status_code != 200: - logging.warning("uCloud中未找到(%s)! status: %d error: %s", key, resp.status_code, resp.error) + logging.warning(f"uCloud中未找到({key})! status: {resp.status_code} error: {resp.error}") return None # 获取公有空间下载url