diff --git a/.gitignore b/.gitignore index eaf78ea..4e20f5c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ __pycache__/ *.egg-info/ deskew_model/ - +/pdfs/ +/results/ diff --git a/vidocp/layout_detection.py b/vidocp/layout_detection.py index d559df0..1d49684 100644 --- a/vidocp/layout_detection.py +++ b/vidocp/layout_detection.py @@ -23,6 +23,7 @@ def find_layout_boxes(image: np.array): 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) diff --git a/vidocp/table_parsing.py b/vidocp/table_parsing.py index be77a9c..27bb40c 100644 --- a/vidocp/table_parsing.py +++ b/vidocp/table_parsing.py @@ -1,15 +1,22 @@ +from functools import partial +from itertools import chain, starmap +from operator import attrgetter + import cv2 import numpy as np from pdf2image import pdf2image from vidocp.utils.display import show_mpl -from vidocp.utils.draw import draw_stats +from vidocp.utils.draw import draw_rectangles +from vidocp.utils.post_processing import xywh_to_vecs, xywh_to_vec_rect, adjacent1d, remove_isolated from vidocp.utils.deskew import deskew_histbased +from vidocp.layout_parsing import parse_layout def add_external_contours(image, img): + contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) - contours, hierarchy = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) + # contours = filter(partial(is_large_enough, min_area=5000000), contours) for cnt in contours: x, y, w, h = cv2.boundingRect(cnt) @@ -18,42 +25,131 @@ def add_external_contours(image, img): return image -def isolate_vertical_and_horizontal_components(img_bin): - - line_min_width = 30 +def isolate_vertical_and_horizontal_components(img_bin, bounding_rects): + line_min_width = 48 kernel_h = np.ones((1, line_min_width), np.uint8) kernel_v = np.ones((line_min_width, 1), 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) + show_mpl(img_bin_h | img_bin_v) + + kernel_h = np.ones((1, 30), np.uint8) + kernel_v = np.ones((30, 1), np.uint8) + img_bin_h = cv2.dilate(img_bin_h, kernel_h, iterations=2) + img_bin_v = cv2.dilate(img_bin_v, kernel_v, iterations=2) + show_mpl(img_bin_h | img_bin_v) + + #reduced filtersize from 100 to 80 to minimize splitting narrow cells + img_bin_h = apply_motion_blur(img_bin_h, 80, 0) + img_bin_v = apply_motion_blur(img_bin_v, 80, 90) img_bin_final = img_bin_h | img_bin_v + show_mpl(img_bin_final) + #changed threshold from 110 to 120 to minimize cell splitting + th1, img_bin_final = cv2.threshold(img_bin_final, 120, 255, cv2.THRESH_BINARY) + img_bin_final = cv2.dilate(img_bin_final, np.ones((1, 1), np.uint8), iterations=1) + show_mpl(img_bin_final) + # problem if layout parser detects too big of a layout box as in VV-748542.pdf p.22 + img_bin_final = disconnect_non_existing_cells(img_bin_final, bounding_rects) + show_mpl(img_bin_final) return img_bin_final +def disconnect_non_existing_cells(img_bin, bounding_rects): + for rect in bounding_rects: + x, y, w, h = rect + img_bin = cv2.rectangle(img_bin, (x, y), (x + w, y + h), (0, 0, 0), 5) + return img_bin + + +# FIXME: does not work yet +def has_table_shape(rects): + assert isinstance(rects, list) + + points = list(chain(*map(xywh_to_vecs, rects))) + brect = xywh_to_vec_rect(cv2.boundingRect(np.vstack(points))) + + rects = list(map(xywh_to_vec_rect, rects)) + + # print(rects) + # print(brect) + + def matches_bounding_rect_corner(rect, x, y): + corresp_coords = list(zip(*map(attrgetter(x, y), [brect, rect]))) + ret = all(starmap(partial(adjacent1d, tolerance=30), corresp_coords)) + # print() + # print(x, y) + # print(brect) + # print(rect) + # print(corresp_coords) + # print(ret) + + return ret + + return all( + ( + any(matches_bounding_rect_corner(r, "xmin", "ymin") for r in rects), + any(matches_bounding_rect_corner(r, "xmin", "ymax") for r in rects), + any(matches_bounding_rect_corner(r, "xmax", "ymax") for r in rects), + any(matches_bounding_rect_corner(r, "xmax", "ymin") for r in rects), + ) + ) + + +def apply_motion_blur(image, size, angle): + k = np.zeros((size, size), dtype=np.float32) + k[(size - 1) // 2, :] = np.ones(size, dtype=np.float32) + k = cv2.warpAffine(k, cv2.getRotationMatrix2D((size / 2 - 0.5, size / 2 - 0.5), angle, 1.0), (size, size)) + k = k * (1.0 / np.sum(k)) + return cv2.filter2D(image, -1, k) + + +def find_table_layout_boxes(image: np.array): + layout_boxes = parse_layout(image) + table_boxes = [] + for box in layout_boxes: + (x, y, w, h) = box + if w * h >= 100000: + table_boxes.append(box) + return table_boxes + + def parse_table(image: np.array): + def is_large_enough(stat): + x1, y1, w, h, area = stat + return area > 2000 and w > 35 and h > 25 gray_scale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape)>2 else image - th1, img_bin = cv2.threshold(gray_scale, 150, 255, cv2.THRESH_BINARY) + th1, img_bin = cv2.threshold(gray_scale, 195, 255, cv2.THRESH_BINARY) img_bin = ~img_bin + show_mpl(img_bin) - img_bin = isolate_vertical_and_horizontal_components(img_bin) + table_layout_boxes = find_table_layout_boxes(image) + img_bin = isolate_vertical_and_horizontal_components(img_bin, table_layout_boxes) img_bin_final = add_external_contours(img_bin, img_bin) - _, labels, stats, _ = cv2.connectedComponentsWithStats(~img_bin_final, connectivity=8, ltype=cv2.CV_32S) + _, _, stats, _ = cv2.connectedComponentsWithStats(~img_bin_final, connectivity=8, ltype=cv2.CV_32S) - return stats + stats = np.vstack(list(filter(is_large_enough, stats))) + rects = stats[:, :-1][2:] + + # FIXME: produces false negatives for `data0/043d551b4c4c768b899eaece4466c836.pdf 1 --type table` + rects = list(remove_isolated(rects, input_sorted=True)) + + return rects def annotate_tables_in_pdf(pdf_path, page_index=0, deskew=False): - page = pdf2image.convert_from_path(pdf_path, first_page=page_index + 1, last_page=page_index + 1)[0] page = np.array(page) if deskew: page = deskew_histbased(page) stats = parse_table(page) - page = draw_stats(page, stats) + page = draw_rectangles(page, stats, annotate=True) + # if stats: + # page = draw_rectangles(page, stats, annotate=True) show_mpl(page) diff --git a/vidocp/utils/draw.py b/vidocp/utils/draw.py index 32c66f6..2f7ef06 100644 --- a/vidocp/utils/draw.py +++ b/vidocp/utils/draw.py @@ -13,7 +13,9 @@ def draw_contours(image, contours): return image -def draw_rectangles(image, rectangles, color=None): +def draw_rectangles(image, rectangles, color=None, annotate=False): + def annotate_rect(x, y, w, h): + cv2.putText(image, "+", (x + (w // 2) - 12, y + (h // 2) + 9), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2) image = copy_and_normalize_channels(image) @@ -24,33 +26,7 @@ def draw_rectangles(image, rectangles, color=None): x, y, w, h = rect cv2.rectangle(image, (x, y), (x + w, y + h), color, 2) - return image - - -def draw_stats(image, stats, annotate=False): - - image = copy_and_normalize_channels(image) - - keys = ["x", "y", "w", "h"] - - def annotate_stat(x, y, w, h): - - for i, (s, v) in enumerate(zip(keys, [x, y, w, h])): - anno = f"{s} = {v}" - xann = int(x + 5) - yann = int(y + h - (20 * (i + 1))) - cv2.putText(image, anno, (xann, yann), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) - - def draw_stat(stat): - - x, y, w, h, area = stat - - cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2) - if annotate: - annotate_stat(x, y, w, h) - - for stat in stats[2:]: - draw_stat(stat) + annotate_rect(x, y, w, h) return image