diff --git a/box_detection/__init__.py b/box_detection/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/box_detection/box_detection.py b/box_detection/box_detection.py new file mode 100644 index 0000000..9bbe9dd --- /dev/null +++ b/box_detection/box_detection.py @@ -0,0 +1,42 @@ +from itertools import count + +import cv2 +import imutils +import numpy as np +import pdf2image +from matplotlib import pyplot as plt + + +def parse(image: np.array): + gray = ~cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + blurred = cv2.GaussianBlur(gray, (5, 5), 1) + thresh = cv2.threshold(blurred, 253, 255, cv2.THRESH_BINARY)[1] + cnts = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + cnts = imutils.grab_contours(cnts) + + for c in cnts: + peri = cv2.arcLength(c, True) + approx = cv2.approxPolyDP(c, 0.04 * peri, True) + yield cv2.boundingRect(approx) + + +def annotate_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_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 = parse(page) + page = annotate_boxes(page, asd) + + fig, ax = plt.subplots(1, 1) + fig.set_size_inches(20, 20) + ax.imshow(page) + plt.show() diff --git a/requirements.txt b/requirements.txt index c607322..2ca75a6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,3 +2,4 @@ opencv-python~=4.5.5.62 numpy~=1.22.1 pdf2image~=1.16.0 matplotlib~=3.5.1 +imutils==0.5.4 diff --git a/scripts/annotate.py b/scripts/annotate.py new file mode 100644 index 0000000..c369db0 --- /dev/null +++ b/scripts/annotate.py @@ -0,0 +1,24 @@ +import argparse + +from table_parsing.table_parsig import annotate_tables_in_pdf +from box_detection.box_detection import annotate_boxes_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") + + args = parser.parse_args() + + return args + + +if __name__ == "__main__": + args = parse_args() + if args.object == "table": + 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) + diff --git a/scripts/annotate_table.py b/scripts/annotate_table.py deleted file mode 100644 index 44e0985..0000000 --- a/scripts/annotate_table.py +++ /dev/null @@ -1,18 +0,0 @@ -import argparse - -from table_parsing.table_parsig import annotate_tables_in_pdf - - -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument("pdf_path") - parser.add_argument("page_index", type=int) - - args = parser.parse_args() - - return args - - -if __name__ == "__main__": - args = parse_args() - annotate_tables_in_pdf(args.pdf_path, page_index=args.page_index)