2022-02-04 17:44:02 +01:00

64 lines
1.8 KiB
Python

from functools import partial
import cv2
import numpy as np
import pdf2image
from iteration_utilities import starfilter, first
from matplotlib import pyplot as plt
def is_filled(hierarchy):
# See https://stackoverflow.com/questions/60095520/how-to-distinguish-filled-circle-contour-and-unfilled-circle-contour-in-opencv
return hierarchy[3] <= 0 and hierarchy[2] == -1
def is_boxy(contour):
epsilon = 0.01 * cv2.arcLength(contour, True)
approx = cv2.approxPolyDP(contour, epsilon, True)
return len(approx) <= 10
def is_large_enough(contour, min_area):
return cv2.contourArea(contour, False) > min_area
def is_likely_redaction(contour, hierarchy, min_area):
return is_filled(hierarchy) and is_boxy(contour) and is_large_enough(contour, min_area)
def find_redactions(image: np.array, min_nomralized_area=200000):
min_nomralized_area /= 200 # Assumes 200 DPI PDF -> image conversion resolution
gray = ~cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
blurred = cv2.GaussianBlur(gray, (5, 5), 1)
thresh = cv2.threshold(blurred, 252, 255, cv2.THRESH_BINARY)[1]
contours, hierarchies = cv2.findContours(thresh.copy(), cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
contours = map(
first, starfilter(partial(is_likely_redaction, min_area=min_nomralized_area), zip(contours, hierarchies[0]))
)
return contours
def annotate_poly(image, conts):
for cont in conts:
cv2.drawContours(image, cont, -1, (0, 255, 0), 2)
return image
def annotate_boxes_in_pdf(pdf_path, page_index=1):
page = pdf2image.convert_from_path(pdf_path, first_page=page_index + 1, last_page=page_index + 1)[0]
page = np.array(page)
asd = find_redactions(page)
page = annotate_poly(page, asd)
fig, ax = plt.subplots(1, 1)
fig.set_size_inches(20, 20)
ax.imshow(page)
plt.show()