diff --git a/layout_detection/layout_detection.py b/layout_detection/layout_detection.py new file mode 100644 index 0000000..e0a0e62 --- /dev/null +++ b/layout_detection/layout_detection.py @@ -0,0 +1,54 @@ +from itertools import count + +import cv2 +import numpy as np +import pdf2image +from matplotlib import pyplot as plt +import imutils + + +def find_layout_boxes(image: np.array): + + gray_scale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + blurred = cv2.GaussianBlur(gray_scale, (5, 5), 1) + thresh = cv2.threshold(blurred, 253, 255, cv2.THRESH_BINARY)[1] + img_bin = ~thresh + + line_min_width = 10 + kernel_h = np.ones((10, line_min_width), np.uint8) + kernel_v = np.ones((line_min_width, 10), np.uint8) + + img_bin_h = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_h) + img_bin_v = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_v) + + img_bin_final = img_bin_h | img_bin_v + + contours = cv2.findContours(img_bin_final, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + contours = imutils.grab_contours(contours) + for c in contours: + peri = cv2.arcLength(c, True) + approx = cv2.approxPolyDP(c, 0.04 * peri, True) + yield cv2.boundingRect(approx) + + +def annotate_layout_boxes(image, rects): + for rect in rects: + (x, y, w, h) = rect + cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2) + + return image + + +def annotate_layout_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) + + layout_boxes = find_layout_boxes(page) + page = annotate_layout_boxes(page, layout_boxes) + + fig, ax = plt.subplots(1, 1) + fig.set_size_inches(20, 20) + ax.imshow(page) + plt.show() + diff --git a/scripts/annotate.py b/scripts/annotate.py index 74fa570..68e4dd2 100644 --- a/scripts/annotate.py +++ b/scripts/annotate.py @@ -2,13 +2,14 @@ import argparse from table_parsing.table_parsig import annotate_tables_in_pdf from box_detection.box_detection import annotate_boxes_in_pdf +from layout_detection.layout_detection import annotate_layout_in_pdf def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("pdf_path") parser.add_argument("page_index", type=int) - parser.add_argument("--object", choices=["table", "box"], default="table") + parser.add_argument("--object", choices=["table", "box", "layout"], default="table") args = parser.parse_args() @@ -21,3 +22,6 @@ if __name__ == "__main__": annotate_tables_in_pdf(args.pdf_path, page_index=args.page_index) elif args.object == "box": annotate_boxes_in_pdf(args.pdf_path, page_index=args.page_index) + elif args.object == "layout": + annotate_layout_in_pdf(args.pdf_path, page_index=args.page_index) +