import cv2 import numpy as np from funcy import lfilter, lmap # type: ignore from cv_analysis.layout_parsing import parse_layout from cv_analysis.utils.postprocessing import ( remove_isolated, # xywh_to_vecs, xywh_to_vec_rect, adjacent1d ) from cv_analysis.utils.structures import Rectangle from cv_analysis.utils.visual_logging import vizlogger def add_external_contours(image: np.ndarray, image_h_w_lines_only): contours, _ = cv2.findContours(image_h_w_lines_only, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) for cnt in contours: x, y, w, h = cv2.boundingRect(cnt) cv2.rectangle(image, (x, y), (x + w, y + h), 255, 1) return image def apply_motion_blur(image: np.ndarray, angle, size=80): """Solidifies and slightly extends detected lines. Args: image (np.array): page image as array angle: direction in which to apply blur, 0 or 90 size (int): kernel size; 80 found empirically to work well Returns: np.array """ kernel = np.zeros((size, size), dtype=np.float32) vizlogger.debug(kernel, "tables08_blur_kernel1.png") kernel[(size - 1) // 2, :] = np.ones(size, dtype=np.float32) vizlogger.debug(kernel, "tables09_blur_kernel2.png") new_kernel: np.ndarray = cv2.warpAffine( kernel, cv2.getRotationMatrix2D((size / 2 - 0.5, size / 2 - 0.5), angle, 1.0), (size, size), ) vizlogger.debug(new_kernel, "tables10_blur_kernel3.png") new_kernel = new_kernel * (1.0 / np.sum(new_kernel)) vizlogger.debug(new_kernel, "tables11_blur_kernel4.png") blurred = cv2.filter2D(image, -1, new_kernel) return blurred def isolate_vertical_and_horizontal_components(img_bin): """Identifies and reinforces horizontal and vertical lines in a binary image. Args: img_bin (np.array): array corresponding to single binarized page image bounding_rects (list): list of layout boxes of the form (x, y, w, h), potentially containing tables Returns: np.ndarray """ 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) img_lines_raw = img_bin_v | img_bin_h 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) img_bin_h = apply_motion_blur(img_bin_h, 0) img_bin_v = apply_motion_blur(img_bin_v, 90) img_bin_extended = img_bin_h | img_bin_v _, img_bin_extended = cv2.threshold(img_bin_extended, 120, 255, cv2.THRESH_BINARY) img_bin_final = cv2.dilate(img_bin_extended, np.ones((1, 1), np.uint8), iterations=1) # add contours before lines are extended by blurring img_bin_final = add_external_contours(img_bin_final, img_lines_raw) return img_bin_final def find_table_layout_boxes(image: np.ndarray): def is_large_enough(box): (_, _, w, h) = box if w * h >= 100000: return Rectangle.from_xywh(box) layout_boxes = parse_layout(image) a = lmap(is_large_enough, layout_boxes) return lmap(is_large_enough, layout_boxes) def preprocess(image: np.ndarray): image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) > 2 else image _, image = cv2.threshold(image, 195, 255, cv2.THRESH_BINARY) return ~image def turn_connected_components_into_rects(image: np.ndarray): def is_large_enough(stat): x1, y1, w, h, area = stat return area > 2000 and w > 35 and h > 25 _, _, stats_list, _ = cv2.connectedComponentsWithStats(~image, connectivity=8, ltype=cv2.CV_32S) rects_list: list[np.ndarray] = lfilter(is_large_enough, stats_list) if rects_list: stats: np.ndarray = np.vstack(rects_list) return stats[:, :-1][2:] return [] def parse_tables(image: np.ndarray, show=False): """Runs the full table parsing process. Args: image (np.array): single PDF page, converted to a numpy array Returns: list: list of rectangles corresponding to table cells """ image = preprocess(image) image = isolate_vertical_and_horizontal_components(image) rects = turn_connected_components_into_rects(image) rects = list(map(Rectangle.from_xywh, rects)) rects = remove_isolated(rects) return rects # def make_lines(image: np.array, horizontal=True, kernel_length=40) def detect_horizontal_lines(image_bin: np.ndarray, kernel_length=40): line_min_width = 48 kernel_h = np.ones((1, line_min_width), np.uint8) img_bin_h = cv2.morphologyEx(image_bin, cv2.MORPH_OPEN, kernel_h) kernel_h = np.ones((1, 30), np.uint8) img_bin_h = cv2.dilate(img_bin_h, kernel_h, iterations=2) img_bin_h = apply_motion_blur(img_bin_h, 0) _, img_bin_h = cv2.threshold(img_bin_h, 120, 255, cv2.THRESH_BINARY) # img_bin_h = cv2.dilate(img_bin_h, np.ones((1, 1), np.uint8), iterations=1) return img_bin_h def detect_vertical_lines(image_bin: np.ndarray, kernel_length=40): line_min_width = 48 kernel_v = np.ones((line_min_width, 1), np.uint8) img_bin_v = cv2.morphologyEx(image_bin, cv2.MORPH_OPEN, kernel_v) kernel_v = np.ones((30, 1), np.uint8) img_bin_v = cv2.dilate(img_bin_v, kernel_v, iterations=2) img_bin_v = apply_motion_blur(img_bin_v, 90) _, img_bin_v = cv2.threshold(img_bin_v, 120, 255, cv2.THRESH_BINARY) # img_bin_v = cv2.dilate(img_bin_v, np.ones((1, 1), np.uint8), iterations=1) return img_bin_v def detect_endpoints(image: np.ndarray, is_horizontal: bool) -> list[tuple[int, int, int, int]]: def are_collinear(quad1: tuple[int, int, int, int], quad2: tuple[int, int, int, int], index: int) -> bool: dist_a = abs(quad1[index] - quad2[index]) dist_b = abs(quad1[index + 2] - quad2[index + 2]) overlap = True if index else (quad1[1] >= quad2[3] or quad1[3] >= quad2[1]) return (dist_a < 15) and (dist_b < 15) and overlap points = cv2.HoughLinesP( image, # Input edge image 1, # Distance resolution in pixels np.pi / 180, # Angle resolution in radians threshold=100, # Min number of votes for valid line minLineLength=200, # Min allowed length of line maxLineGap=10, # Max allowed gap between line for joining them ) points = points if points is not None else [] lines = list(map(lambda x: tuple(x[0]), points)) if not lines: return lines index = int(is_horizontal) lines.sort(key=lambda q: q[index]) corrected = [lines[0]] for quad in lines[1:]: if are_collinear(corrected[-1], quad, bool(is_horizontal)): prev = corrected.pop(-1) corrected.append( ( min(prev[0], quad[0]), min(prev[1], quad[1]), max(prev[2], quad[2]), min(prev[3], quad[3]), ) if is_horizontal else ( min(prev[0], quad[0]), max(prev[1], quad[1]), min(prev[2], quad[2]), min(prev[3], quad[3]), ) ) else: corrected.append(quad) return corrected def parse_lines(image: np.ndarray, show=False) -> list[dict[str, float]]: image = preprocess(image) # kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2)) # image = cv2.dilate(image, kernel, iterations=4) horizontal_line_img = detect_horizontal_lines(image) vertical_line_img = detect_vertical_lines(image) horizontal_endpoints = detect_endpoints(horizontal_line_img, is_horizontal=True) vertical_endpoints = detect_endpoints(vertical_line_img, is_horizontal=False) def format_quad(quad: tuple[int, int, int, int], max_x: int, max_y: int) -> dict[str, float]: x1, y1, x2, y2 = quad if x1 > (x2 + 5): x1, y1, x2, y2 = x2, y2, x1, y1 elif y1 > (y2 + 5): x1, y1, x2, y2 = x2, y2, x1, y1 return {"x1": x1 / max_x, "y1": y1 / max_y, "x2": x2 / max_x, "y2": y2 / max_y} ymax, xmax = image.shape return list( map( lambda quad: format_quad(quad, xmax, ymax), horizontal_endpoints + vertical_endpoints, ) )