cv-analysis-service/cv_analysis/table_parsing.py

from functools import partial
from itertools import chain, starmap
from operator import attrgetter

import cv2
import numpy as np
from funcy import lmap, lfilter

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, 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.array, 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

    """
    k = np.zeros((size, size), dtype=np.float32)
    vizlogger.debug(k, "tables08_blur_kernel1.png")
    k[(size - 1) // 2, :] = np.ones(size, dtype=np.float32)
    vizlogger.debug(k, "tables09_blur_kernel2.png")
    k = cv2.warpAffine(
        k,
        cv2.getRotationMatrix2D((size / 2 - 0.5, size / 2 - 0.5), angle, 1.0),
        (size, size),
    )
    vizlogger.debug(k, "tables10_blur_kernel3.png")
    k = k * (1.0 / np.sum(k))
    vizlogger.debug(k, "tables11_blur_kernel4.png")
    blurred = cv2.filter2D(image, -1, k)
    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.array
    """
    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

    th1, 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.array):
    def is_large_enough(box):
        (x, y, 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.array):
    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.array):
    def is_large_enough(stat):
        x1, y1, w, h, area = stat
        return area > 2000 and w > 35 and h > 25

    _, _, stats, _ = cv2.connectedComponentsWithStats(~image, connectivity=8, ltype=cv2.CV_32S)

    stats = lfilter(is_large_enough, stats)
    if stats:
        stats = np.vstack(stats)
        return stats[:, :-1][2:]
    return []


def parse_tables(image: np.array, 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)
    #print(rects, "\n\n")
    rects = list(map(Rectangle.from_xywh, rects))
    #print(rects, "\n\n")
    rects = remove_isolated(rects)
    #print(rects, "\n\n")
    
    return rects