首次提交

test/visual_model_test/visual_model_test.py

@@ -0,0 +1,67 @@
+# 可视化的模型对比测试
+import re
+import time
+from pprint import pprint
+
+from paddlenlp import Taskflow
+from paddlenlp.utils.doc_parser import DocParser
+
+
+def visual_model_test(task_path, test_img, schema):
+    img = re.split(r'[\\/]', test_img)[-1]
+    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):
+    # 开始时间
+    start_time = time.time()
+
+    # 结算清单
+    if model_type == "settlement":
+        task_path = "../../config/model/settlement_list_model"
+        test_img_path = "img/PH20240511000638_1_094306_1.jpg"
+        schema = ["姓名", "入院日期", "出院日期", "费用总额", "个人现金支付", "个人账户支付", "自费", "医保类型"]
+    elif model_type == "discharge":
+        task_path = "../../config/model/discharge_record_model"
+        test_img_path = "img/PH20240401000003_3_001938_2.jpg"
+        schema = ["医院", "科别", "姓名", "入院日期", "出院日期", "主治医生"]
+    elif model_type == "cost":
+        task_path = "../../config/model/cost_list_model"
+        test_img_path = "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")