101 lines
3.3 KiB
Python
101 lines
3.3 KiB
Python
import cv2
|
|
import numpy as np
|
|
from pdf2image import pdf2image
|
|
import pandas as pd
|
|
from PIL import Image
|
|
import timeit
|
|
|
|
from cv_analysis.utils.detection import detect_large_coherent_structures
|
|
from cv_analysis.utils.display import show_mpl
|
|
from cv_analysis.utils.draw import draw_rectangles
|
|
from cv_analysis.utils.post_processing import remove_included
|
|
from cv_analysis.utils.filters import is_large_enough, has_acceptable_format
|
|
from cv_analysis.utils.text import remove_primary_text_regions
|
|
from cv_analysis.utils.visual_logging import vizlogger
|
|
|
|
|
|
def is_likely_figure(cont, min_area=5000, max_width_to_hight_ratio=6):
|
|
return is_large_enough(cont, min_area) and has_acceptable_format(cont, max_width_to_hight_ratio)
|
|
|
|
|
|
def detect_figures(image: np.array):
|
|
|
|
image = image.copy()
|
|
vizlogger.debug(image, "figures01_start.png")
|
|
|
|
image = remove_primary_text_regions(image)
|
|
vizlogger.debug(image, "figures02_remove_text.png")
|
|
cnts = detect_large_coherent_structures(image)
|
|
|
|
cnts = filter(is_likely_figure, cnts)
|
|
rects = map(cv2.boundingRect, cnts)
|
|
rects = remove_included(rects)
|
|
|
|
return list(rects)
|
|
|
|
|
|
def detect_figures_in_pdf(pdf_path, page_index=1, show=False):
|
|
|
|
page = pdf2image.convert_from_path(pdf_path, dpi=300, first_page=page_index + 1, last_page=page_index + 1)[0]
|
|
page = np.array(page)
|
|
|
|
redaction_contours = detect_figures(page)
|
|
page = draw_rectangles(page, redaction_contours)
|
|
vizlogger.debug(page, "figures03_final.png")
|
|
if show:
|
|
show_mpl(page)
|
|
return page
|
|
|
|
|
|
def detect_figures_in_test_files():
|
|
def save_as_pdf(pages):
|
|
p1, p = pages[0], pages[1:]
|
|
out_pdf_path = "/home/lillian/ocr_docs/output_files/fig_detection_pdf.pdf"
|
|
p1.save(
|
|
out_pdf_path, "PDF", resolution=150.0, save_all=True, append_images=p
|
|
)
|
|
path = "/home/lillian/ocr_docs/"
|
|
ex_pages = pd.read_csv(path+"/metadata/metadata2.csv")
|
|
pages_detected = []
|
|
|
|
t0 = timeit.default_timer()
|
|
for name, page_nr in zip(ex_pages.pdf_name, ex_pages.page):
|
|
page = pdf2image.convert_from_path(path + "/original/" + name, dpi=300, first_page=page_nr, last_page=page_nr)[0]
|
|
page = np.array(page)
|
|
redaction_contours = detect_figures(page)
|
|
page = draw_rectangles(page, redaction_contours)
|
|
pages_detected.append(Image.fromarray(page))
|
|
print(timeit.default_timer()-t0)
|
|
|
|
save_as_pdf(pages_detected)
|
|
|
|
|
|
def detect_figures_in_png(pdf_path, show=False):
|
|
|
|
page = Image.open(pdf_path)
|
|
page = np.array(page)
|
|
|
|
redaction_contours = detect_figures(page)
|
|
page = draw_rectangles(page, redaction_contours)
|
|
vizlogger.debug(page, "figures03_final.png")
|
|
if show:
|
|
show_mpl(page)
|
|
return page
|
|
|
|
|
|
def detect_figures_in_test_files_png():
|
|
file_name = pd.read_csv("/home/lillian/ocr_docs/metadata/metadata2.csv")
|
|
path = "/home/lillian/ocr_docs/png_example_pages/"
|
|
pages = []
|
|
page_index = 0
|
|
t0 = timeit.default_timer()
|
|
for name in file_name.image_name:
|
|
page = detect_figures_in_png(path+name+".png", page_index, show=False)
|
|
pages.append(Image.fromarray(page))
|
|
t1 = timeit.default_timer()
|
|
print(t1-t0)
|
|
p1, p = pages[0], pages[1:]
|
|
out_pdf_path = "/home/lillian/ocr_docs/output_files/fig_detection_png2.pdf"
|
|
p1.save(
|
|
out_pdf_path, "PDF", resolution=150.0, save_all=True, append_images=p
|
|
) |