Forráskód Böngészése

fix(batch): refactor OCR detection integration and area ratio calculation

myhloli 5 hónapja
szülő
commit
ddf5a8781a
1 módosított fájl, 23 hozzáadás és 23 törlés
  1. 23 23
      magic_pdf/model/batch_analyze.py

+ 23 - 23
magic_pdf/model/batch_analyze.py

@@ -251,29 +251,29 @@ class BatchAnalyze:
                         ocr_res_list_dict['single_page_mfdetrec_res'], useful_list
                     )
 
-                # OCR-det
-                new_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)
-                ocr_res = ocr_model.ocr(
-                    new_image, mfd_res=adjusted_mfdetrec_res, rec=False
-                )[0]
-
-                # Integration results
-                if ocr_res:
-                    ocr_result_list = get_ocr_result_list(ocr_res, useful_list, ocr_res_list_dict['ocr_enable'], new_image, _lang)
-
-                    if res["category_id"] == 3:
-                        # ocr_result_list中所有bbox的面积之和
-                        ocr_res_area = sum(get_coords_and_area(ocr_res_item)[4] for ocr_res_item in ocr_result_list if 'poly' in ocr_res_item)
-                        # 求ocr_res_area和res的面积的比值
-                        res_area = get_coords_and_area(res)[4]
-                        if res_area > 0:
-                            ratio = ocr_res_area / res_area
-                            if ratio > 0.25:
-                                res["category_id"] = 1
-                            else:
-                                continue
-
-                    ocr_res_list_dict['layout_res'].extend(ocr_result_list)
+                    # OCR-det
+                    new_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)
+                    ocr_res = ocr_model.ocr(
+                        new_image, mfd_res=adjusted_mfdetrec_res, rec=False
+                    )[0]
+
+                    # Integration results
+                    if ocr_res:
+                        ocr_result_list = get_ocr_result_list(ocr_res, useful_list, ocr_res_list_dict['ocr_enable'], new_image, _lang)
+
+                        if res["category_id"] == 3:
+                            # ocr_result_list中所有bbox的面积之和
+                            ocr_res_area = sum(get_coords_and_area(ocr_res_item)[4] for ocr_res_item in ocr_result_list if 'poly' in ocr_res_item)
+                            # 求ocr_res_area和res的面积的比值
+                            res_area = get_coords_and_area(res)[4]
+                            if res_area > 0:
+                                ratio = ocr_res_area / res_area
+                                if ratio > 0.25:
+                                    res["category_id"] = 1
+                                else:
+                                    continue
+
+                        ocr_res_list_dict['layout_res'].extend(ocr_result_list)
 
             # det_count += len(ocr_res_list_dict['ocr_res_list'])
         # logger.info(f'ocr-det time: {round(time.time()-det_start, 2)}, image num: {det_count}')