优化可视化测试功能架构

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 re
import sys
import tempfile
import time
from pprint import pprint
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.utils.doc_parser import DocParser
from paddleocr import PaddleOCR
from ucloud import ucloud
from util import image_util, util
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_type = img[last_dot_index + 1:]
img_array = open_image(image)
img_array = image_util.read(image)
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:
cv2.imwrite(temp_file.name, img_array)
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):
if model_type == "ocr":
imgs = split_image(test_img)
imgs = image_util.split(test_img)
layout = []
temp_files_paths = []
doc_parser = DocParser(layout_analysis=False)
for img in imgs:
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
angle = get_image_rotation_angle(img["img"])
rotated_img = rotate_image(img["img"], angle)
# angle = image_util.parse_rotation_angles(img["img"])[0]
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)
img["x_offset"] -= offset_x
img["y_offset"] -= offset_y
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"]:
for p in parsed_doc["layout"]:
for p in parsed_doc:
box = p[0]
box[0] += img["x_offset"]
box[1] += img["y_offset"]
box[2] += img["x_offset"]
box[3] += img["y_offset"]
layout += parsed_doc["layout"]
layout += parsed_doc
write_visual_result(test_img, angle, layout=layout)
else:
docs = []
split_result = split_image(test_img)
split_result = image_util.split(test_img)
temp_files_paths = []
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)
angle = int(image_util.parse_rotation_angles(img["img"])[0])
rotated_img = image_util.rotate(img["img"], angle)
cv2.imwrite(temp_file.name, rotated_img)
temp_files_paths.append(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"
schema = None
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"
schema = ["患者姓名", "入院日期", "出院日期", "费用总额", "个人现金支付", "个人账户支付", "自费金额",
"医保类型", "住院号", "医保结算单号码"]
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"
schema = ["医院", "科室", "患者姓名", "入院日期", "出院日期", "主治医生", "住院号", "年龄"]
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"
schema = ["患者姓名", "入院日期", "出院日期", "费用总额"]
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"
schema = {"名称": ["类别", "规格", "单价", "数量", "金额"]}
else: