71 lines
2.5 KiB
Python
71 lines
2.5 KiB
Python
# 可视化的模型对比测试
|
|
import re
|
|
import time
|
|
from pprint import pprint
|
|
|
|
from paddlenlp import Taskflow
|
|
from paddlenlp.utils.doc_parser import DocParser
|
|
|
|
from photo_review.util.ucloud import get_private_url
|
|
|
|
|
|
def visual_model_test(task_path, test_img, schema):
|
|
img = re.split(r'[\\/]', test_img)[-1]
|
|
img = img.split("?")[0]
|
|
img_name = ""
|
|
img_type = "jpg"
|
|
last_dot_index = img.rfind(".")
|
|
if last_dot_index != -1:
|
|
img_name = img[:last_dot_index]
|
|
img_type = img[last_dot_index + 1:]
|
|
|
|
# 默认模型
|
|
ie = Taskflow("information_extraction", schema=schema, model="uie-x-base")
|
|
results = ie({"doc": test_img})
|
|
pprint(results[0])
|
|
DocParser.write_image_with_results(
|
|
test_img,
|
|
result=results[0],
|
|
save_path="./img_result/" + img_name + "_default." + img_type)
|
|
# 自己训练的模型
|
|
my_ie = Taskflow("information_extraction", schema=schema, model="uie-x-base", task_path=task_path)
|
|
my_results = my_ie({"doc": test_img})
|
|
pprint(my_results[0])
|
|
DocParser.write_image_with_results(
|
|
test_img,
|
|
result=my_results[0],
|
|
save_path="./img_result/" + img_name + "_my." + img_type)
|
|
|
|
|
|
def main(model_type, pic_name=None):
|
|
# 开始时间
|
|
start_time = time.time()
|
|
|
|
# 结算清单
|
|
if model_type == "settlement":
|
|
task_path = "../../config/model/settlement_list_model"
|
|
test_img_path = 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"
|
|
test_img_path = 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"
|
|
test_img_path = get_private_url(pic_name) if pic_name else "img/PH20240511000648_4_094542_2.jpg"
|
|
schema = ["姓名", "入院日期", "出院日期", "费用总额"]
|
|
else:
|
|
print("请输入正确的类型！")
|
|
return
|
|
visual_model_test(task_path, test_img_path, schema)
|
|
|
|
# 结束时间
|
|
end_time = time.time()
|
|
pprint(f"处理时长：{end_time - start_time}秒")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main("settlement")
|
|
# main("discharge")
|
|
# main("cost")
|