57 lines
2.2 KiB
Python
57 lines
2.2 KiB
Python
from cv_analysis.layout_parsing import annotate_layout_in_pdf
|
|
from cv_analysis.figure_detection import detect_figures
|
|
from cv_analysis.table_parsing import tables_in_image, parse_table
|
|
from cv_analysis.utils.text import find_primary_text_regions, remove_primary_text_regions
|
|
from cv_analysis.utils.draw import draw_rectangles
|
|
from cv_analysis.utils.display import show_mpl
|
|
|
|
|
|
def cut_out_content_structures(layout_rects, page):
|
|
large_rects = []
|
|
small_rects = []
|
|
for x, y, w, h in layout_rects:
|
|
rect = (x, y, w, h)
|
|
if w * h >= 50000:
|
|
cropped_page = page[y:(y + h), x:(x + w)]
|
|
large_rects.append([rect, cropped_page])
|
|
else:
|
|
cropped_page = page[y:(y + h), x:(x + w)]
|
|
small_rects.append([rect, cropped_page])
|
|
return large_rects, small_rects
|
|
|
|
|
|
def parse_content_structures(page, large_rects, small_rects):
|
|
for coordinates, cropped_image in large_rects:
|
|
non_text_rects = detect_figures(cropped_image)
|
|
|
|
if len(non_text_rects) == 0:
|
|
page = draw_rectangles(page, [coordinates], color=(0, 255, 0), annotate=True)
|
|
|
|
elif tables_in_image(cropped_image)[0]:
|
|
page = draw_rectangles(page, [coordinates], color=(255, 0, 0), annotate=True)
|
|
stats = parse_table(page)
|
|
page = draw_rectangles(page, stats, annotate=True)
|
|
|
|
else:
|
|
page = draw_rectangles(page, [coordinates], color=(0, 0, 255), annotate=True)
|
|
|
|
# for coordinates, cropped_image in small_rects:
|
|
# non_text_rects = detect_figures(cropped_image)
|
|
# if len(non_text_rects) == 0 and len(list(find_primary_text_regions(cropped_image))) > 0:
|
|
# page = draw_rectangles(page, [coordinates], color=(0, 255, 0), annotate=True)
|
|
# else:
|
|
# page = draw_rectangles(page, [coordinates], color=(0, 255, 255), annotate=True)
|
|
return page
|
|
|
|
|
|
def detect_figures_with_layout_parsing(pdf_path, page_index=1, show=False):
|
|
layout_rects, page = annotate_layout_in_pdf(pdf_path, page_index, return_rects=True)
|
|
big_structures, small_structures = cut_out_content_structures(layout_rects, page)
|
|
page = parse_content_structures(page, big_structures, small_structures)
|
|
|
|
if show:
|
|
show_mpl(page)
|
|
else:
|
|
return page
|
|
|