cv-analysis-service/vidocp/table_parsing.py

import cv2
import numpy as np
from pdf2image import pdf2image

from vidocp.utils import show_cv2, draw_stats


def add_external_contours(image, img):

    contours, hierarchy = cv2.findContours(img, 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 isolate_vertical_and_horizontal_components(img_bin):

    line_min_width = 30
    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_bin_final = img_bin_h | img_bin_v

    return img_bin_final


def parse_table(image: np.array):

    gray_scale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    th1, img_bin = cv2.threshold(gray_scale, 150, 255, cv2.THRESH_BINARY)
    img_bin = ~img_bin

    img_bin = isolate_vertical_and_horizontal_components(img_bin)
    img_bin_final = add_external_contours(img_bin, img_bin)

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

    return stats


def annotate_tables_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)

    stats = parse_table(page)
    page = draw_stats(page, stats)

    show_cv2(page)