调整模型载入

photo_review/__init__.py

@@ -0,0 +1,13 @@
+from paddlenlp import Taskflow
+from paddleocr import PaddleOCR
+
+from config import keys
+
+SETTLEMENT_IE = Taskflow("information_extraction", schema=keys.SETTLEMENT_LIST_SCHEMA, model="uie-x-base",
+                         task_path="config/model/settlement_list_model", layout_analysis=False)
+DISCHARGE_IE = Taskflow("information_extraction", schema=keys.DISCHARGE_RECORD_SCHEMA, model="uie-x-base",
+                        task_path="config/model/discharge_record_model", layout_analysis=False)
+COST_IE = Taskflow("information_extraction", schema=keys.COST_LIST_SCHEMA, model="uie-x-base", device_id=1,
+                   task_path="config/model/cost_list_model", layout_analysis=False)
+
+OCR = PaddleOCR(use_angle_cls=False, lang="ch", show_log=False)

photo_review/photo_review.py

@@ -1,4 +1,5 @@
 import concurrent.futures
+import json
 import logging
 import math
 import os
@@ -12,6 +13,10 @@ import numpy as np
 import paddleclas
 import requests
 
+from photo_review import OCR, SETTLEMENT_IE, DISCHARGE_IE, COST_IE
+from photo_review.entity.zx_ocr import ZxOcr
+from ucloud import ucloud
+
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 from time import sleep
@@ -31,7 +36,6 @@ from photo_review.entity.zx_phrec import ZxPhrec
 from photo_review.util.data_util import handle_date, handle_decimal, parse_department, handle_name, \
     handle_insurance_type, handle_original_data, handle_hospital, handle_department, handle_id, handle_age
 from photo_review.util.util import get_default_datetime
-from photo_review.enumeration.task import TaskEnum
 
 
 # 获取图片
@@ -214,6 +218,43 @@ def request_ie_result(task_enum, phrecs):
         raise Exception(f"请求信息抽取结果失败，状态码：{response.status_code}")
 
 
+# 关键信息提取
+def information_extraction(ie, phrecs):
+    result = {}
+    # 同一批图的标识
+    identity = int(time.time())
+    for phrec in phrecs:
+        pic_path = ucloud.get_private_url(phrec.cfjaddress)
+        if not pic_path:
+            continue
+        split_result = split_image(pic_path)
+
+        for img in split_result:
+            angles = get_image_rotation_angles(img["img"])
+            rotated_img = rotate_image(img["img"], int(angles[0]))
+            ie_results = [{"result": ie_temp_image(ie, OCR, rotated_img), "angle": angles[0]}]
+            if not ie_results[0] or len(ie_results[0]) < len(ie.kwargs.get("schema")):
+                rotated_img = rotate_image(img["img"], int(angles[1]))
+                ie_results.append({"result": ie_temp_image(ie, OCR, rotated_img), "angle": angles[1]})
+
+            now = get_default_datetime()
+            for ie_result in ie_results:
+                result_json = json.dumps(ie_result["result"], ensure_ascii=False)
+                if len(result_json) > 5000:
+                    result_json = result_json[:5000]
+                session = MysqlSession()
+                zx_ocr = ZxOcr(pk_phhd=phrec.pk_phhd, pk_phrec=phrec.pk_phrec, id=identity, cfjaddress=phrec.cfjaddress,
+                               content=result_json, rotation_angle=ie_result["angle"], x_offset=img["x_offset"],
+                               y_offset=img["y_offset"], create_time=now, update_time=now)
+                session.add(zx_ocr)
+                session.commit()
+                session.close()
+
+                result = merge_result(result, ie_result["result"])
+
+    return result
+
+
 # 从keys中获取准确率最高的value
 def get_best_value_in_keys(source, keys):
     # 最终结果
@@ -267,7 +308,7 @@ def save_or_update_ie(table, pk_phhd, data):
 
 
 def settlement_task(pk_phhd, settlement_list):
-    settlement_list_ie_result = request_ie_result(TaskEnum.SETTLEMENT, settlement_list)
+    settlement_list_ie_result = information_extraction(SETTLEMENT_IE, settlement_list)
     settlement_data = {
         "pk_phhd": pk_phhd,
         "name": handle_name(get_best_value_in_keys(settlement_list_ie_result, PATIENT_NAME)),
@@ -297,7 +338,7 @@ def settlement_task(pk_phhd, settlement_list):
 
 
 def discharge_task(pk_phhd, discharge_record):
-    discharge_record_ie_result = request_ie_result(TaskEnum.DISCHARGE, discharge_record)
+    discharge_record_ie_result = request_ie_result(DISCHARGE_IE, discharge_record)
     discharge_data = {
         "pk_phhd": pk_phhd,
         "hospital": handle_hospital(get_best_value_in_keys(discharge_record_ie_result, HOSPITAL)),
@@ -338,7 +379,7 @@ def discharge_task(pk_phhd, discharge_record):
 
 
 def cost_task(pk_phhd, cost_list):
-    cost_list_ie_result = request_ie_result(TaskEnum.COST, cost_list)
+    cost_list_ie_result = request_ie_result(COST_IE, cost_list)
     cost_data = {
         "pk_phhd": pk_phhd,
         "name": handle_name(get_best_value_in_keys(cost_list_ie_result, PATIENT_NAME)),