diff --git a/README.md b/README.md index 45de78b..dca1bc4 100644 --- a/README.md +++ b/README.md @@ -85,3 +85,15 @@ python scripts/annotate.py data/test_pdf.pdf 7 --type layout The below image shows the detected layout elements on a page. ![](data/layout_parsing.png) + + +#### Figure Detection + +The figure detection utility detects figures specifically, which can be missed by the generic layout parsing utility. +```bash +python scripts/annotate.py data/test_pdf.pdf 3 --type figure +``` + +The below image shows the detected figure on a page. + +![](data/figure_detection.png) diff --git a/data/figure_detection.png b/data/figure_detection.png new file mode 100644 index 0000000..7716ade Binary files /dev/null and b/data/figure_detection.png differ diff --git a/scripts/annotate.py b/scripts/annotate.py index d690988..10d40cc 100644 --- a/scripts/annotate.py +++ b/scripts/annotate.py @@ -3,13 +3,14 @@ import argparse from vidocp.table_parsing import annotate_tables_in_pdf from vidocp.redaction_detection import annotate_boxes_in_pdf from vidocp.layout_parsing import annotate_layout_in_pdf +from vidocp.figure_detection import remove_text_in_pdf def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("pdf_path") parser.add_argument("page_index", type=int) - parser.add_argument("--type", choices=["table", "redaction", "layout"], default="table") + parser.add_argument("--type", choices=["table", "redaction", "layout", "figure"]) args = parser.parse_args() @@ -24,3 +25,5 @@ if __name__ == "__main__": annotate_boxes_in_pdf(args.pdf_path, page_index=args.page_index) elif args.type == "layout": annotate_layout_in_pdf(args.pdf_path, page_index=args.page_index) + elif args.type == "figure": + remove_text_in_pdf(args.pdf_path, page_index=args.page_index) diff --git a/vidocp/figure_detection.py b/vidocp/figure_detection.py new file mode 100644 index 0000000..e852646 --- /dev/null +++ b/vidocp/figure_detection.py @@ -0,0 +1,73 @@ +import cv2 +import numpy as np +from pdf2image import pdf2image + +from vidocp.utils import draw_contours, show_mpl, draw_rectangles, remove_included, remove_overlapping, show_cv2 + + +def is_large_enough(cont, min_area=10000): + return cv2.contourArea(cont, False) > min_area + + +def has_acceptable_format(cont, max_width_to_hight_ratio=6): + _, _, w, h = cv2.boundingRect(cont) + return max_width_to_hight_ratio >= w / h >= (1 / max_width_to_hight_ratio) + + +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() + + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + thresh = cv2.threshold(gray, 253, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1] + + close_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (20, 3)) + close = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, close_kernel, iterations=1) + + dilate_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 3)) + dilate = cv2.dilate(close, dilate_kernel, iterations=1) + + cnts, _ = cv2.findContours(dilate, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) + + def filter_rects(): + for c in cnts: + area = cv2.contourArea(c) + if area > 800 and area < 15000: + yield cv2.boundingRect(c) + + for rect in filter_rects(): + x, y, w, h = rect + cv2.rectangle(image, (x, y), (x + w, y + h), (255, 255, 255), -1) + + gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + + thresh = cv2.threshold(gray, 253, 255, cv2.THRESH_BINARY)[1] + + dilate_kernel = cv2.getStructuringElement(cv2.MORPH_OPEN, (5, 5)) + dilate = cv2.dilate(~thresh, dilate_kernel, iterations=4) + + close_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (20, 20)) + close = cv2.morphologyEx(dilate, cv2.MORPH_CLOSE, close_kernel, iterations=1) + + cnts, _ = cv2.findContours(close, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) + + cnts = filter(is_likely_figure, cnts) + rects = [cv2.boundingRect(c) for c in cnts] + rects = remove_included(rects) + + return rects + + +def remove_text_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) + + redaction_contours = detect_figures(page) + page = draw_rectangles(page, redaction_contours) + + show_mpl(page) diff --git a/vidocp/layout_parsing.py b/vidocp/layout_parsing.py index 62f6f26..67cd89e 100644 --- a/vidocp/layout_parsing.py +++ b/vidocp/layout_parsing.py @@ -1,5 +1,3 @@ -from collections import namedtuple -from functools import partial from itertools import compress from itertools import starmap from operator import __and__ @@ -8,72 +6,13 @@ import cv2 import numpy as np from pdf2image import pdf2image -from vidocp.utils import draw_rectangles, show_mpl - -Rectangle = namedtuple("Rectangle", "xmin ymin xmax ymax") - - -def make_box(x1, y1, x2, y2): - keys = "x1", "y1", "x2", "y2" - return dict(zip(keys, [x1, y1, x2, y2])) - - -def compute_intersection(a, b): - - dx = min(a.xmax, b.xmax) - max(a.xmin, b.xmin) - dy = min(a.ymax, b.ymax) - max(a.ymin, b.ymin) - - return dx * dy if (dx >= 0) and (dy >= 0) else 0 +from vidocp.utils import draw_rectangles, show_mpl, remove_overlapping, remove_included, has_no_parent def is_likely_segment(rect, min_area=100): return cv2.contourArea(rect, False) > min_area -def has_no_parent(hierarchy): - return hierarchy[-1] <= 0 - - -def xywh_to_vec_rect(rect): - x1, y1, w, h = rect - x2 = x1 + w - y2 = y1 + h - return Rectangle(x1, y1, x2, y2) - - -def vec_rect_to_xywh(rect): - x, y, x2, y2 = rect - w = x2 - x - h = y2 - y - return x, y, w, h - - -def remove_overlapping(rectangles): - def overlap(a, b): - return compute_intersection(a, b) > 0 - - def does_not_overlap(rect, rectangles): - return not any(overlap(rect, r2) for r2 in rectangles if not rect == r2) - - rectangles = list(map(xywh_to_vec_rect, rectangles)) - rectangles = filter(partial(does_not_overlap, rectangles=rectangles), rectangles) - rectangles = map(vec_rect_to_xywh, rectangles) - return rectangles - - -def remove_included(rectangles): - def included(a, b): - return b.xmin >= a.xmin and b.ymin >= a.ymin and b.xmax <= a.xmax and b.ymax <= a.ymax - - def is_not_included(rect, rectangles): - return not any(included(r2, rect) for r2 in rectangles if not rect == r2) - - rectangles = list(map(xywh_to_vec_rect, rectangles)) - rectangles = filter(partial(is_not_included, rectangles=rectangles), rectangles) - rectangles = map(vec_rect_to_xywh, rectangles) - return rectangles - - def find_segments(image): contours, hierarchies = cv2.findContours(image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) diff --git a/vidocp/utils.py b/vidocp/utils.py index 370ba27..ee528b4 100644 --- a/vidocp/utils.py +++ b/vidocp/utils.py @@ -1,3 +1,6 @@ +from collections import namedtuple +from functools import partial + import cv2 from matplotlib import pyplot as plt @@ -78,3 +81,63 @@ def draw_stats(image, stats, annotate=False): draw_stat(stat) return image + + +def remove_overlapping(rectangles): + def overlap(a, b): + return compute_intersection(a, b) > 0 + + def does_not_overlap(rect, rectangles): + return not any(overlap(rect, r2) for r2 in rectangles if not rect == r2) + + rectangles = list(map(xywh_to_vec_rect, rectangles)) + rectangles = filter(partial(does_not_overlap, rectangles=rectangles), rectangles) + rectangles = map(vec_rect_to_xywh, rectangles) + return rectangles + + +def remove_included(rectangles): + def included(a, b): + return b.xmin >= a.xmin and b.ymin >= a.ymin and b.xmax <= a.xmax and b.ymax <= a.ymax + + def is_not_included(rect, rectangles): + return not any(included(r2, rect) for r2 in rectangles if not rect == r2) + + rectangles = list(map(xywh_to_vec_rect, rectangles)) + rectangles = filter(partial(is_not_included, rectangles=rectangles), rectangles) + rectangles = map(vec_rect_to_xywh, rectangles) + return rectangles + + +Rectangle = namedtuple("Rectangle", "xmin ymin xmax ymax") + + +def make_box(x1, y1, x2, y2): + keys = "x1", "y1", "x2", "y2" + return dict(zip(keys, [x1, y1, x2, y2])) + + +def compute_intersection(a, b): + + dx = min(a.xmax, b.xmax) - max(a.xmin, b.xmin) + dy = min(a.ymax, b.ymax) - max(a.ymin, b.ymin) + + return dx * dy if (dx >= 0) and (dy >= 0) else 0 + + +def has_no_parent(hierarchy): + return hierarchy[-1] <= 0 + + +def xywh_to_vec_rect(rect): + x1, y1, w, h = rect + x2 = x1 + w + y2 = y1 + h + return Rectangle(x1, y1, x2, y2) + + +def vec_rect_to_xywh(rect): + x, y, x2, y2 = rect + w = x2 - x + h = y2 - y + return x, y, w, h