From b2e37a2d1b3fa6bb9d2c368e57a1fe54dca5ebeb Mon Sep 17 00:00:00 2001 From: myhloli Date: Thu, 21 Nov 2024 17:54:56 +0800 Subject: [PATCH] feat(ocr): improve text detection and OCR accuracy - Update OCR utils to handle different box formats and improve angle calculation - Modify PDF extraction kit to support OCR option and optimize processing flow - Enhance PPOCR model to sort and filter detection boxes, improving text splitting accuracy --- magic_pdf/model/pdf_extract_kit.py | 44 ++++++++++--------- .../sub_modules/ocr/paddleocr/ocr_utils.py | 15 ++++--- .../ocr/paddleocr/ppocr_273_mod.py | 11 ++++- 3 files changed, 44 insertions(+), 26 deletions(-) diff --git a/magic_pdf/model/pdf_extract_kit.py b/magic_pdf/model/pdf_extract_kit.py index 0c0e664b..f2adfadb 100644 --- a/magic_pdf/model/pdf_extract_kit.py +++ b/magic_pdf/model/pdf_extract_kit.py @@ -116,13 +116,12 @@ class CustomPEKModel: device=self.device ) # 初始化ocr - if self.apply_ocr: - self.ocr_model = atom_model_manager.get_atom_model( - atom_model_name=AtomicModel.OCR, - ocr_show_log=show_log, - det_db_box_thresh=0.3, - lang=self.lang - ) + self.ocr_model = atom_model_manager.get_atom_model( + atom_model_name=AtomicModel.OCR, + ocr_show_log=show_log, + det_db_box_thresh=0.3, + lang=self.lang + ) # init table model if self.apply_table: table_model_dir = self.configs["weights"][self.table_model_name] @@ -174,24 +173,29 @@ class CustomPEKModel: ocr_res_list, table_res_list, single_page_mfdetrec_res = get_res_list_from_layout_res(layout_res) # ocr识别 - if self.apply_ocr: - ocr_start = time.time() - # Process each area that requires OCR processing - for res in ocr_res_list: - new_image, useful_list = crop_img(res, pil_img, crop_paste_x=50, crop_paste_y=50) - adjusted_mfdetrec_res = get_adjusted_mfdetrec_res(single_page_mfdetrec_res, useful_list) + ocr_start = time.time() + # Process each area that requires OCR processing + for res in ocr_res_list: + new_image, useful_list = crop_img(res, pil_img, crop_paste_x=50, crop_paste_y=50) + adjusted_mfdetrec_res = get_adjusted_mfdetrec_res(single_page_mfdetrec_res, useful_list) - # OCR recognition - new_image = cv2.cvtColor(np.asarray(new_image), cv2.COLOR_RGB2BGR) + # OCR recognition + new_image = cv2.cvtColor(np.asarray(new_image), cv2.COLOR_RGB2BGR) + if self.apply_ocr: ocr_res = self.ocr_model.ocr(new_image, mfd_res=adjusted_mfdetrec_res)[0] + else: + ocr_res = self.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) - layout_res.extend(ocr_result_list) + # Integration results + if ocr_res: + ocr_result_list = get_ocr_result_list(ocr_res, useful_list) + layout_res.extend(ocr_result_list) - ocr_cost = round(time.time() - ocr_start, 2) + ocr_cost = round(time.time() - ocr_start, 2) + if self.apply_ocr: logger.info(f"ocr time: {ocr_cost}") + else: + logger.info(f"det time: {ocr_cost}") # 表格识别 table recognition if self.apply_table: diff --git a/magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py b/magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py index 73e68a2b..b84e39fb 100644 --- a/magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py +++ b/magic_pdf/model/sub_modules/ocr/paddleocr/ocr_utils.py @@ -211,16 +211,21 @@ def get_ocr_result_list(ocr_res, useful_list): ocr_result_list = [] for box_ocr_res in ocr_res: - p1, p2, p3, p4 = box_ocr_res[0] - text, score = box_ocr_res[1] + if len(box_ocr_res) == 2: + p1, p2, p3, p4 = box_ocr_res[0] + text, score = box_ocr_res[1] + else: + p1, p2, p3, p4 = box_ocr_res + text, score = "", 1 # average_angle_degrees = calculate_angle_degrees(box_ocr_res[0]) # if average_angle_degrees > 0.5: - if calculate_is_angle(box_ocr_res[0]): + poly = [p1, p2, p3, p4] + if calculate_is_angle(poly): # logger.info(f"average_angle_degrees: {average_angle_degrees}, text: {text}") # 与x轴的夹角超过0.5度，对边界做一下矫正 # 计算几何中心 - x_center = sum(point[0] for point in box_ocr_res[0]) / 4 - y_center = sum(point[1] for point in box_ocr_res[0]) / 4 + x_center = sum(point[0] for point in poly) / 4 + y_center = sum(point[1] for point in poly) / 4 new_height = ((p4[1] - p1[1]) + (p3[1] - p2[1])) / 2 new_width = p3[0] - p1[0] p1 = [x_center - new_width / 2, y_center - new_height / 2] diff --git a/magic_pdf/model/sub_modules/ocr/paddleocr/ppocr_273_mod.py b/magic_pdf/model/sub_modules/ocr/paddleocr/ppocr_273_mod.py index 1260d20b..e7f7331f 100644 --- a/magic_pdf/model/sub_modules/ocr/paddleocr/ppocr_273_mod.py +++ b/magic_pdf/model/sub_modules/ocr/paddleocr/ppocr_273_mod.py @@ -78,9 +78,18 @@ class ModifiedPaddleOCR(PaddleOCR): for idx, img in enumerate(imgs): img = preprocess_image(img) dt_boxes, elapse = self.text_detector(img) - if not dt_boxes: + if dt_boxes is None: ocr_res.append(None) continue + dt_boxes = sorted_boxes(dt_boxes) + # merge_det_boxes 和 update_det_boxes 都会把poly转成bbox再转回poly，因此需要过滤所有倾斜程度较大的文本框 + dt_boxes = merge_det_boxes(dt_boxes) + if mfd_res: + bef = time.time() + dt_boxes = update_det_boxes(dt_boxes, mfd_res) + aft = time.time() + logger.debug("split text box by formula, new dt_boxes num : {}, elapsed : {}".format( + len(dt_boxes), aft - bef)) tmp_res = [box.tolist() for box in dt_boxes] ocr_res.append(tmp_res) return ocr_res