Browse Source

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.
myhloli 1 year ago
parent
commit
36220d69d4
1 changed files with 4 additions and 1 deletions
  1. 4 1
      magic_pdf/libs/draw_bbox.py

+ 4 - 1
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']