From 36220d69d465ba6f34cd41e5c9062bb45f5cf0f6 Mon Sep 17 00:00:00 2001 From: myhloli Date: Fri, 27 Sep 2024 16:06:15 +0800 Subject: [PATCH] refactor(draw_bbox): clear cuda cache and update bbox sorting - Added CUDA cache clearing after layoutreader prediction to free up GPU memory. - Modified the bbox sorting logic to sort text and title blocks separately. - Adjusted drawing colors for better distinction in debug visualizations. --- magic_pdf/libs/draw_bbox.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/magic_pdf/libs/draw_bbox.py b/magic_pdf/libs/draw_bbox.py index 6316c75c..972c948b 100644 --- a/magic_pdf/libs/draw_bbox.py +++ b/magic_pdf/libs/draw_bbox.py @@ -341,10 +341,13 @@ def draw_layout_sort_bbox(pdf_info, pdf_bytes, out_path, filename): for page in pdf_info: page_line_list = [] for block in page['preproc_blocks']: - if block['type'] == 'text' or block['type'] == 'title': + if block['type'] == 'text' or block['type'] == 'title' or block['type'] == 'interline_equation': for line in block['lines']: bbox = line['bbox'] page_line_list.append(bbox) + if block['type'] == 'table' or block['type'] == 'image': + bbox = block['bbox'] + page_line_list.append(bbox) # 使用layoutreader排序 page_size = page['page_size']