Echo/fcb_photo_review
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

优化旋转图片的识别条件

Browse Source
This commit is contained in:
liuyebo
2024-06-28 20:33:21 +08:00
parent 6b68a7f7ae
commit 203f1c4e4a
1 changed files with 42 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

76
photo_review/photo_review.py
View File

@@ -35,8 +35,6 @@ 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):
@@ -137,7 +135,7 @@ def rotate_image(img, angle):
# 获取图片OCR，并将其box转为两点矩形框
def get_ocr_layout(img_path):
def get_ocr_layout(ocr, img_path):
def _get_box(old_box):
new_box = [
min(old_box[0][0], old_box[3][0]), # x1
@@ -154,7 +152,7 @@ def get_ocr_layout(img_path):
return True
layout = []
ocr_result = OCR.ocr(img_path)
ocr_result = ocr.ocr(img_path)
ocr_result = ocr_result[0]
if not ocr_result:
return layout
@@ -168,11 +166,34 @@ def get_ocr_layout(img_path):
return layout
def ie_temp_image(ie, ocr, image):
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
cv2.imwrite(temp_file.name, image)
ie_result = []
try:
layout = get_ocr_layout(ocr, temp_file.name)
if not layout:
# 无识别结果
ie_result = []
else:
ie_result = ie({"doc": temp_file.name, "layout": layout})[0]
except Exception as e:
logging.error("信息抽取时出错", exc_info=e)
finally:
try:
os.remove(temp_file.name)
except Exception as e:
logging.info(f"删除临时文件 {temp_file.name} 时出错", exc_info=e)
return ie_result
# 关键信息提取
def information_extraction(ie, phrecs):
result = {}
# 同一批图的标识
identity = int(time.time())
ocr = PaddleOCR(use_angle_cls=True, lang="ch", show_log=False)
for phrec in phrecs:
pic_path = ucloud.get_private_url(phrec.cfjaddress)
if not pic_path:
@@ -180,39 +201,26 @@ 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_file1:
cv2.imwrite(temp_file1.name, img["img"])
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file2:
ie_result1 = ie_temp_image(ie, ocr, img["img"])
if not ie_result1 or len(ie_result1) < len(ie.kwargs.get("schema")):
rotated_img = rotate_image(img["img"], 90)
cv2.imwrite(temp_file2.name, rotated_img)
ie_result = []
try:
layout1 = get_ocr_layout(temp_file1.name)
layout2 = get_ocr_layout(temp_file2.name)
if not (layout1 or layout2):
# 无识别结果
ie_result2 = ie_temp_image(ie, ocr, rotated_img)
if not (ie_result1 or ie_result2):
continue
ie_result = ie([{"doc": temp_file1.name, "layout": layout1},
{"doc": temp_file2.name, "layout": layout2}])
except Exception as e:
logging.error("信息抽取时出错", exc_info=e)
finally:
try:
os.remove(temp_file1.name)
except Exception as 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
elif not ie_result1:
ie_result = ie_result2
angle = 90
elif not ie_result2:
ie_result = ie_result1
angle = 0
elif len(ie_result2) > len(ie_result1):
ie_result = ie_result2
angle = 90
else:
ie_result = ie_result1
angle = 0
else:
ie_result = ie_result[0]
ie_result = ie_result1
angle = 0
now = get_default_datetime()