19fe6965fb
changes to fig_detection_with_layout.py so tables are getting parsed as well reusage of adding external contour in table_parsing.py
71 lines
2.9 KiB
Python
71 lines
2.9 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
|
|
from cv_analysis.utils.visual_logging import vizlogger
|
|
from PIL import Image
|
|
|
|
|
|
|
|
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 >= 75000:
|
|
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:
|
|
figure_rects = detect_figures(cropped_image)
|
|
if len(figure_rects) == 0: # text
|
|
page = draw_rectangles(page, [coordinates], color=(0, 255, 0), annotate=True)
|
|
elif tables_in_image(cropped_image)[0]: # table
|
|
stats = parse_table(page)
|
|
page = draw_rectangles(page, stats, color=(255, 0, 0), annotate=True)
|
|
else: # figure
|
|
page = draw_rectangles(page, [coordinates], color=(0, 0, 255), annotate=True)
|
|
|
|
# for coordinates, cropped_image in small_rects:
|
|
# figure_rects = detect_figures(cropped_image)
|
|
# if len(figure_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)
|
|
vizlogger.debug(page, "figures03_final.png")
|
|
if show:
|
|
show_mpl(page)
|
|
else:
|
|
return page
|
|
|
|
# pages = []
|
|
# for i in range(0,16):
|
|
# pdf_path = "/home/lillian/ocr_docs/Report on spectra.pdf"
|
|
# page_index = i
|
|
# 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)
|
|
# pages.append(Image.fromarray(page))
|
|
# p1, p = pages[0], pages[1:]
|
|
#
|
|
# out_pdf_path = "/home/lillian/ocr_docs/out1.pdf"
|
|
#
|
|
# p1.save(
|
|
# out_pdf_path, "PDF", resolution=150.0, save_all=True, append_images=p
|
|
# )
|