From baaee6f5b76b691860b4c73a9640503541e5b53e Mon Sep 17 00:00:00 2001 From: Julius Unverfehrt Date: Fri, 4 Feb 2022 16:08:39 +0100 Subject: [PATCH] Pull request #3: Layout detection happy little accident Merge in RR/table_parsing from layout_detetciton_happy_little_accident to master Squashed commit of the following: commit eb4452c9a488df16085a16eba08b7a182274d331 Author: Julius Unverfehrt Date: Fri Feb 4 16:08:12 2022 +0100 Empty line added commit d2fedf9a2f982af2157a408077654d388ca6cc6d Author: Julius Unverfehrt Date: Fri Feb 4 16:07:20 2022 +0100 Empty line added commit 638d14a4b6c7b4d34222fd3b4cbb8ce79bb32ef0 Author: Julius Unverfehrt Date: Fri Feb 4 16:05:43 2022 +0100 Quickfix typo commit 0271d2ba2e51227aa53e128bf857394e5b5b2d48 Author: Julius Unverfehrt Date: Fri Feb 4 15:57:06 2022 +0100 black commit c95c4ad3f3d01857e7dd1dde0802ed7f2a5837c1 Author: Julius Unverfehrt Date: Fri Feb 4 15:53:42 2022 +0100 Refactored layout_detection prototype commit 766bd0b916b532885e44a13581f100ffaa39bb55 Author: Julius Unverfehrt Date: Fri Feb 4 13:16:36 2022 +0100 reset table_parsing to table parsing functionality, moved layout detection accident to layout_detection commit 7c8955f56dfae2aef814caf4cbc6e903406994ba Merge: 9a065a0 af5c6d0 Author: Julius Unverfehrt Date: Fri Feb 4 13:11:28 2022 +0100 Merge branch 'master' of ssh://git.iqser.com:2222/rr/table_parsing into layout_detetciton_happy_little_accident commit 9a065a0e7f62823a3b18e301d12c80b1a74f0b3e Author: Julius Unverfehrt Date: Thu Feb 3 16:45:09 2022 +0100 Made Bob proud --- layout_detection/layout_detection.py | 54 ++++++++++++++++++++++++++++ scripts/annotate.py | 6 +++- 2 files changed, 59 insertions(+), 1 deletion(-) create mode 100644 layout_detection/layout_detection.py 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) +