cv-analysis-service/cv_analysis/figure_detection.py

58 lines
1.8 KiB
Python

import cv2
import numpy as np
from pdf2image import pdf2image
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
#from PIL import Image
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, 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
# pages = []
# for i in range(0,16):
# pdf_path = "/home/lillian/ocr_docs/Report on spectra.pdf"
# page_index = i
# page = detect_figures_in_pdf(pdf_path,page_index)
# pages.append(Image.fromarray(page))
# p1, p = pages[0], pages[1:]
#
# out_pdf_path = "/home/lillian/ocr_docs/out.pdf"
#
# p1.save(
# out_pdf_path, "PDF", resolution=150.0, save_all=True, append_images=p
# )