启用ocr中的cls功能

photo_review/photo_review.py

@@ -18,7 +18,7 @@ from sqlalchemy import update
 from config.keys import PATIENT_NAME, ADMISSION_DATE, DISCHARGE_DATE, MEDICAL_EXPENSES, PERSONAL_CASH_PAYMENT, \
     PERSONAL_ACCOUNT_PAYMENT, PERSONAL_FUNDED_AMOUNT, MEDICAL_INSURANCE_TYPE, HOSPITAL, DEPARTMENT, DOCTOR
 from config.mysql import MysqlSession
-from config.photo_review import PHHD_BATCH_SIZE, SLEEP_MINUTES, SETTLEMENT_IE, DISCHARGE_IE, COST_IE
+from config.photo_review import PHHD_BATCH_SIZE, SLEEP_MINUTES, SETTLEMENT_IE, DISCHARGE_IE, COST_IE, OCR
 from photo_review.entity.bd_yljg import BdYljg
 from photo_review.entity.bd_ylks import BdYlks
 from photo_review.entity.zx_ie_cost import ZxIeCost
@@ -95,6 +95,36 @@ def merge_result(result1, result2):
     return result1
 
 
+# 获取图片OCR，并将其box转为两点矩形框
+def get_ocr_layout(img_path):
+    def _get_box(box):
+        box = [
+            min(box[0][0], box[3][0]),  # x1
+            min(box[0][1], box[1][1]),  # y1
+            max(box[1][0], box[2][0]),  # x2
+            max(box[2][1], box[3][1]),  # y2
+        ]
+        return box
+
+    def _normal_box(box):
+        # Ensure the height and width of bbox are greater than zero
+        if box[3] - box[1] < 0 or box[2] - box[0] < 0:
+            return False
+        return True
+
+    layout = []
+    ocr_result = OCR.ocr(img_path)
+    ocr_result = ocr_result[0]
+    for segment in ocr_result:
+        box = segment[0]
+        box = _get_box(box)
+        if not _normal_box(box):
+            continue
+        text = segment[1][0]
+        layout.append((box, text))
+    return layout
+
+
 # 关键信息提取
 def information_extraction(ie, phrecs):
     result = {}
@@ -107,7 +137,9 @@ def information_extraction(ie, phrecs):
             for img in split_result:
                 with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
                     cv2.imwrite(temp_file.name, img["img"])
-                    docs.append({"doc": temp_file.name})
+                    # 为使用ocr中的cls，单独调用ocr
+                    layout = get_ocr_layout(temp_file.name)
+                    docs.append({"doc": temp_file.name, "layout": layout})
                     doc_phrecs.append({"phrec": phrec, "x_offset": img["x_offset"], "y_offset": img["y_offset"]})
     if not docs:
         return result