cv-analysis-service/vidocp/table_parsig.py

from itertools import count

import cv2
import imutils
import numpy as np
import pdf2image
from matplotlib import pyplot as plt


def parse(image: np.array):
    gray_scale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    #plt.imshow(gray_scale)
    blurred = cv2.GaussianBlur(gray_scale, (7, 7), 2)  #5 5 1
    thresh = cv2.threshold(blurred, 251, 255, cv2.THRESH_BINARY)[1]
    #plt.imshow(thresh)
    img_bin = ~thresh

    line_min_width = 7
    kernel_h = np.ones((10, line_min_width), np.uint8)
    kernel_v = np.ones((line_min_width, 10), 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)
    #plt.imshow(img_bin_h)
    #plt.imshow(img_bin_v)
    img_bin_final = img_bin_h | img_bin_v
    plt.imshow(img_bin_final)
    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)
        yield cv2.boundingRect(approx)

def parse_tables(image: np.array, rects: list):
    parsed_tables = []
    for rect in rects:
        (x,y,w,h) = rect
        region_of_interest = image[x:x+w, y:y+h]
        gray = cv2.cvtColor(region_of_interest, cv2.COLOR_BGR2GRAY)
        thresh = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY)[1]
        img_bin = ~thresh

        line_min_width = 5
        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)
    # find_and_close_internal_gaps(img_bin_v)
        img_bin_final = img_bin_h | img_bin_v
    #plt.imshow(img_bin_final)
    # find_and_close_internal_gaps(img_bin_final)
    # find_and_close_edges(img_bin_final)

        _, labels, stats, _ = cv2.connectedComponentsWithStats(~img_bin_final, connectivity=8, ltype=cv2.CV_32S)
        parsed_tables.append([(x,y,w,h), stats])
    return parsed_tables
        #yield (x,y,w,h), stats, region_of_interest
    # return stats

def annotate_table(image, parsed_tables):
    for table in parsed_tables:
        original_coordinates, stats = table
        stats = filter_unconnected_cells(stats)
        for stat in stats:
            x, y, w, h, area = stat
            cv2.rectangle(image, (x, y), (x + w, y + h), (255, 0, 255), 2)
            for i, (s, v) in enumerate(zip(["x", "y", "w", "h"], [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, (255, 0, 255), 2)

    return image



def filter_unconnected_cells(stats):
    filtered_cells = []
    # print(stats)
    for left, middle, right in zip(stats[0:], stats[1:],
                                   list(stats[2:]) + [np.array([None, None, None, None, None])]):
        x, y, w, h, area = middle
        if w > 35 and h > 13 and area > 500:
            if right[1] is None:
                if y == left[1] or x == left[0]:
                    filtered_cells.append(middle)
            else:
                if y == left[1] or y == right[1] or x == left[0] or x == right[0]:
                    filtered_cells.append(middle)
    return filtered_cells

def find_and_close_edges(img_bin_final):
    contours, hierarchy = cv2.findContours(img_bin_final, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    for cnt in contours:
        missing_external_edges = True
        left = tuple(cnt[cnt[:, :, 0].argmin()][0])
        right = tuple(cnt[cnt[:, :, 0].argmax()][0])
        top = tuple(cnt[cnt[:, :, 1].argmin()][0])
        bottom = tuple(cnt[cnt[:, :, 1].argmax()][0])
        topleft = [left[0], top[1]]
        bottomright = [right[0], bottom[1]]
        for arr in cnt:
            if np.array_equal(arr, np.array([bottomright])) or np.array_equal(arr, np.array([topleft])):
                missing_external_edges = False
                break

        if missing_external_edges and (bottomright[0] - topleft[0]) * (bottomright[1] - topleft[1]) >= 50000:
            cv2.rectangle(img_bin_final, tuple(topleft), tuple(bottomright), (255, 255, 255), 2)
            # print("missing cell detectet rectangle drawn")

    return img_bin_final



def parse_tables_in_pdf(pages):
    return zip(map(parse, pages), count())

# def annotate_tables_in_pdf(pdf_path, page_index=1):
#     # timeit()
#     page = pdf2image.convert_from_path(pdf_path, first_page=page_index + 1, last_page=page_index + 1)[0]
#     page = np.array(page)
#
#     _, stats = parse(page)
#     page = annotate_image(page, stats)
#     # print(timeit())
#     fig, ax = plt.subplots(1, 1)
#     fig.set_size_inches(20, 20)
#     ax.imshow(page)
#     plt.show()


def annotate_boxes(image, rects):
    print(type(rects))
    for rect in rects:
        (x, y, w, h) = rect
        cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 2)

    return image

def filter_tables_or_images(rects):
    filtered = []
    for rect in rects:
        (x,y,w,h) = rect
        print(w*h)
        if w * h > 10**6:
            filtered.append(rect)
    print(filtered)
    return filtered




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)

    layout_boxes = parse(page)
    page = annotate_boxes(page, layout_boxes)
    parsed_tables = parse_tables(page, filter_tables_or_images(layout_boxes))
    page = annotate_table(page, parsed_tables)



    fig, ax = plt.subplots(1, 1)
    fig.set_size_inches(20, 20)
    ax.imshow(page)
    plt.show()