added function for detecting external edges

This commit is contained in:
llocarnini 2022-02-01 19:25:05 +01:00
parent a68f89af03
commit bb5083d419

View File

@ -6,6 +6,7 @@ import pdf2image
from matplotlib import pyplot as plt from matplotlib import pyplot as plt
from timeit import timeit from timeit import timeit
def parse(image: np.array): def parse(image: np.array):
gray_scale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) gray_scale = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
th1, img_bin = cv2.threshold(gray_scale, 200, 255, cv2.THRESH_BINARY) 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_h = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_h)
img_bin_v = cv2.morphologyEx(img_bin, cv2.MORPH_OPEN, kernel_v) 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 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 return labels, stats
# def filter_unconnected_cells(stats): # def filter_unconnected_cells(stats):
# filtered_cells = [] # filtered_cells = []
# for i, val in enumerate(stats[2:]): # for left, middle, right in zip(stats[0:], stats[1:], list(stats[2:])+[None]):
# x, y, w, h, area = stats[i][0], stats[i][1], stats[i][2], stats[i][3], stats[i][4] # x, y, w, h, area = middle
# if w > 35 and h > 13 and area > 500: # if w > 35 and h > 13 and area > 500:
# # print(stats[i]) # if y == left[1] or y == right[1]:
# if y == stats[i - 1][1] or y == stats[i + 1][1]: # filtered_cells.append(middle)
# filtered_cells.append(stats[i])
# return filtered_cells # 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): # def annotate_image(image, stats):
# stats = filter_unconnected_cells(stats) # stats = filter_unconnected_cells(stats)
# for i,val in enumerate(stats): # for i,val in enumerate(stats):
@ -50,16 +65,11 @@ def parse(image: np.array):
# cv2.putText(image, anno, (xann, yann), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 255), 2) # cv2.putText(image, anno, (xann, yann), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 255), 2)
# #
# return image # return image
def annotate_image(image, stats): def annotate_image(image, stats):
print(stats.shape) stats = filter_unconnected_cells(stats)
for i in range(2, len(stats)): for stat in stats:
x, y, w, h, area = stats[i][0], stats[i][1], stats[i][2], stats[i][3], stats[i][4] x, y, w, h, area = stat
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) 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])): for i, (s, v) in enumerate(zip(["x", "y", "w", "h"], [x, y, w, h])):
anno = f"{s} = {v}" anno = f"{s} = {v}"
xann = int(x + 5) xann = int(x + 5)
@ -69,11 +79,33 @@ def annotate_image(image, stats):
return image 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): def parse_tables_in_pdf(pages):
return zip(map(parse, pages), count()) return zip(map(parse, pages), count())
def annotate_tables_in_pdf(pdf_path, page_index=1): def annotate_tables_in_pdf(pdf_path, page_index=1):
timeit() timeit()
page = pdf2image.convert_from_path(pdf_path, first_page=page_index + 1, last_page=page_index + 1)[0] 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) fig.set_size_inches(20, 20)
ax.imshow(page) ax.imshow(page)
plt.show() plt.show()