优化可视化测试功能架构

This commit is contained in:
2024-07-19 14:25:06 +08:00
parent 893d4a7685
commit 60c61f318a

View File

@@ -1,19 +1,17 @@
# 可视化的模型对比测试 # 可视化的模型对比测试
import os import os
import re import re
import sys
import tempfile import tempfile
import time import time
from pprint import pprint from pprint import pprint
import cv2 import cv2
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from photo_review.photo_review import split_image, get_image_rotation_angle, rotate_image, open_image
from paddlenlp import Taskflow from paddlenlp import Taskflow
from paddlenlp.utils.doc_parser import DocParser from paddlenlp.utils.doc_parser import DocParser
from paddleocr import PaddleOCR
from ucloud import ucloud from ucloud import ucloud
from util import image_util, util
def write_visual_result(image, angle=0, layout=None, result=None): def write_visual_result(image, angle=0, layout=None, result=None):
@@ -26,9 +24,9 @@ def write_visual_result(image, angle=0, layout=None, result=None):
img_name = img[:last_dot_index] img_name = img[:last_dot_index]
img_type = img[last_dot_index + 1:] img_type = img[last_dot_index + 1:]
img_array = open_image(image) img_array = image_util.read(image)
if angle != 0: if angle != 0:
img_array = rotate_image(img_array, angle) img_array = image_util.rotate(img_array, angle)
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file: with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
cv2.imwrite(temp_file.name, img_array) cv2.imwrite(temp_file.name, img_array)
if layout: if layout:
@@ -49,35 +47,42 @@ def write_visual_result(image, angle=0, layout=None, result=None):
def visual_model_test(model_type, test_img, task_path, schema): def visual_model_test(model_type, test_img, task_path, schema):
if model_type == "ocr": if model_type == "ocr":
imgs = split_image(test_img) imgs = image_util.split(test_img)
layout = [] layout = []
temp_files_paths = [] temp_files_paths = []
doc_parser = DocParser(layout_analysis=False) doc_parser = DocParser(layout_analysis=False)
for img in imgs: for img in imgs:
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file: with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
angle = get_image_rotation_angle(img["img"]) # angle = image_util.parse_rotation_angles(img["img"])[0]
rotated_img = rotate_image(img["img"], angle) angle = 0
rotated_img = image_util.rotate(img["img"], angle)
rotated_img, offset_x, offset_y = image_util.expand_to_a4_size(rotated_img, True)
cv2.imwrite(temp_file.name, rotated_img) cv2.imwrite(temp_file.name, rotated_img)
img["x_offset"] -= offset_x
img["y_offset"] -= offset_y
temp_files_paths.append(temp_file.name) temp_files_paths.append(temp_file.name)
parsed_doc = doc_parser.parse({"doc": temp_file.name}) parsed_doc = util.get_ocr_layout(PaddleOCR(det_db_box_thresh=0.3), temp_file.name)
# parsed_doc = doc_parser.parse({"doc": temp_file.name})["layout"]
if img["x_offset"] or img["y_offset"]: if img["x_offset"] or img["y_offset"]:
for p in parsed_doc["layout"]: for p in parsed_doc:
box = p[0] box = p[0]
box[0] += img["x_offset"] box[0] += img["x_offset"]
box[1] += img["y_offset"] box[1] += img["y_offset"]
box[2] += img["x_offset"] box[2] += img["x_offset"]
box[3] += img["y_offset"] box[3] += img["y_offset"]
layout += parsed_doc["layout"] layout += parsed_doc
write_visual_result(test_img, angle, layout=layout) write_visual_result(test_img, angle, layout=layout)
else: else:
docs = [] docs = []
split_result = split_image(test_img) split_result = image_util.split(test_img)
temp_files_paths = [] temp_files_paths = []
for img in split_result: for img in split_result:
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file: with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
angle = get_image_rotation_angle(img["img"]) angle = int(image_util.parse_rotation_angles(img["img"])[0])
rotated_img = rotate_image(img["img"], angle) rotated_img = image_util.rotate(img["img"], angle)
cv2.imwrite(temp_file.name, rotated_img) cv2.imwrite(temp_file.name, rotated_img)
temp_files_paths.append(temp_file.name) temp_files_paths.append(temp_file.name)
docs.append({"doc": temp_file.name}) docs.append({"doc": temp_file.name})
@@ -116,20 +121,20 @@ def main(model_type, pic_name=None):
test_img_path = ucloud.get_private_url(pic_name) if pic_name else "img/PH20240428000832_1_093844_2.jpg" test_img_path = ucloud.get_private_url(pic_name) if pic_name else "img/PH20240428000832_1_093844_2.jpg"
schema = None schema = None
elif model_type == "settlement": elif model_type == "settlement":
task_path = "../config/model/settlement_list_model" task_path = "../model/settlement_list_model"
test_img_path = ucloud.get_private_url(pic_name) if pic_name else "img/PH20240511000638_1_094306_1.jpg" test_img_path = ucloud.get_private_url(pic_name) if pic_name else "img/PH20240511000638_1_094306_1.jpg"
schema = ["患者姓名", "入院日期", "出院日期", "费用总额", "个人现金支付", "个人账户支付", "自费金额", schema = ["患者姓名", "入院日期", "出院日期", "费用总额", "个人现金支付", "个人账户支付", "自费金额",
"医保类型", "住院号", "医保结算单号码"] "医保类型", "住院号", "医保结算单号码"]
elif model_type == "discharge": elif model_type == "discharge":
task_path = "../config/model/discharge_record_model" task_path = "../model/discharge_record_model"
test_img_path = ucloud.get_private_url(pic_name) if pic_name else "img/PH20240401000003_3_001938_2.jpg" test_img_path = ucloud.get_private_url(pic_name) if pic_name else "img/PH20240401000003_3_001938_2.jpg"
schema = ["医院", "科室", "患者姓名", "入院日期", "出院日期", "主治医生", "住院号", "年龄"] schema = ["医院", "科室", "患者姓名", "入院日期", "出院日期", "主治医生", "住院号", "年龄"]
elif model_type == "cost": elif model_type == "cost":
task_path = "../config/model/cost_list_model" task_path = "../model/cost_list_model"
test_img_path = ucloud.get_private_url(pic_name) if pic_name else "img/PH20240511000648_4_094542_2.jpg" test_img_path = ucloud.get_private_url(pic_name) if pic_name else "img/PH20240511000648_4_094542_2.jpg"
schema = ["患者姓名", "入院日期", "出院日期", "费用总额"] schema = ["患者姓名", "入院日期", "出院日期", "费用总额"]
elif model_type == "cost_detail": elif model_type == "cost_detail":
task_path = "../config/model/cost_list_detail_model" task_path = "../model/cost_list_detail_model"
test_img_path = ucloud.get_private_url(pic_name) if pic_name else "img/PH20240511000648_4_094542_2.jpg" test_img_path = ucloud.get_private_url(pic_name) if pic_name else "img/PH20240511000648_4_094542_2.jpg"
schema = {"名称": ["类别", "规格", "单价", "数量", "金额"]} schema = {"名称": ["类别", "规格", "单价", "数量", "金额"]}
else: else: