From bb5083d419aa6abee7648051f6dd1017627d0fa2 Mon Sep 17 00:00:00 2001 From: llocarnini Date: Tue, 1 Feb 2022 19:25:05 +0100 Subject: [PATCH] added function for detecting external edges --- table_parsing/table_parsig.py | 83 +++++++++++++++++++++++------------ 1 file changed, 56 insertions(+), 27 deletions(-) diff --git a/table_parsing/table_parsig.py b/table_parsing/table_parsig.py index 7b486cf..9d21c14 100644 --- a/table_parsing/table_parsig.py +++ b/table_parsing/table_parsig.py @@ -6,6 +6,7 @@ import pdf2image from matplotlib import pyplot as plt from timeit import timeit + def parse(image: np.array): gray_scale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) th1, img_bin = cv2.threshold(gray_scale, 200, 255, cv2.THRESH_BINARY) @@ -17,26 +18,40 @@ def parse(image: np.array): img_bin_h = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_h) img_bin_v = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_v) + # print([cv2.countNonZero(row) for row in img_bin_v]) img_bin_final = img_bin_h | img_bin_v - _, labels, stats, _ = cv2.connectedComponentsWithStats(~img_bin_final, connectivity=8, ltype=cv2.CV_32S) + find_and_close_edges(img_bin_final) + _, labels, stats, _ = cv2.connectedComponentsWithStats(~img_bin_final, connectivity=8, ltype=cv2.CV_32S) return labels, stats # def filter_unconnected_cells(stats): # filtered_cells = [] -# for i, val in enumerate(stats[2:]): -# x, y, w, h, area = stats[i][0], stats[i][1], stats[i][2], stats[i][3], stats[i][4] +# for left, middle, right in zip(stats[0:], stats[1:], list(stats[2:])+[None]): +# x, y, w, h, area = middle # if w > 35 and h > 13 and area > 500: -# # print(stats[i]) -# if y == stats[i - 1][1] or y == stats[i + 1][1]: -# filtered_cells.append(stats[i]) +# if y == left[1] or y == right[1]: +# filtered_cells.append(middle) # return filtered_cells -# -# -# + +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 annotate_image(image, stats): # stats = filter_unconnected_cells(stats) # for i,val in enumerate(stats): @@ -50,30 +65,47 @@ def parse(image: np.array): # cv2.putText(image, anno, (xann, yann), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 255), 2) # # return image - def annotate_image(image, stats): - print(stats.shape) - for i in range(2, len(stats)): - x, y, w, h, area = stats[i][0], stats[i][1], stats[i][2], stats[i][3], stats[i][4] - if w > 35 and h > 13 and area > 500: - # print(stats[i]) - if y == stats[i - 1][1] or y == stats[i + 1][1]: - 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) + 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 find_and_close_edges(img_bin_final): + contours, hierarchy = cv2.findContours(img_bin_final, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) + # contoured_img = cv2.drawContours(img_bin_final,contours, -1,(255,255,255),2) + 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] + 1, top[1]] + bottomright = [right[0] - 1, 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: + cv2.rectangle(img_bin_final, tuple(topleft), tuple(bottomright), (255,255,255) , 2) + + 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] @@ -86,6 +118,3 @@ def annotate_tables_in_pdf(pdf_path, page_index=1): fig.set_size_inches(20, 20) ax.imshow(page) plt.show() - - -