354 lines
13 KiB
Python
354 lines
13 KiB
Python
from __future__ import print_function
|
|
import os
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from .table import Table
|
|
from .utils import (transform, elements_bbox, detect_vertical, merge_close_values,
|
|
get_row_index, get_column_index, reduce_index, outline,
|
|
fill_spanning, remove_empty, encode_list)
|
|
|
|
|
|
__all__ = ['Lattice']
|
|
|
|
|
|
def _morph_transform(imagename, scale=15, invert=False):
|
|
"""Morphological Transformation
|
|
|
|
Applies a series of morphological operations on the image
|
|
to find table contours and line segments.
|
|
http://answers.opencv.org/question/63847/how-to-extract-tables-from-an-image/
|
|
|
|
Empirical result for adaptiveThreshold's blockSize=5 and C=-0.2
|
|
taken from http://pequan.lip6.fr/~bereziat/pima/2012/seuillage/sezgin04.pdf
|
|
|
|
Parameters
|
|
----------
|
|
imagename : Path to image.
|
|
|
|
scale : int
|
|
Scaling factor. Large scaling factor leads to smaller lines
|
|
being detected. (optional, default: 15)
|
|
|
|
invert : bool
|
|
Invert pdf image to make sure that lines are in foreground.
|
|
(optional, default: False)
|
|
|
|
Returns
|
|
-------
|
|
img : ndarray
|
|
|
|
tables : dict
|
|
Dictionary with table bounding box as key and list of
|
|
joints found in the table as value.
|
|
|
|
v_segments : list
|
|
List of vertical line segments found in the image.
|
|
|
|
h_segments : list
|
|
List of horizontal line segments found in the image.
|
|
"""
|
|
img = cv2.imread(imagename)
|
|
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
|
|
|
if invert:
|
|
threshold = cv2.adaptiveThreshold(
|
|
gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,
|
|
15, -0.2)
|
|
else:
|
|
threshold = cv2.adaptiveThreshold(
|
|
np.invert(gray), 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
|
|
cv2.THRESH_BINARY,
|
|
15, -0.2)
|
|
|
|
vertical = threshold
|
|
horizontal = threshold
|
|
|
|
verticalsize = vertical.shape[0] / scale
|
|
horizontalsize = horizontal.shape[1] / scale
|
|
|
|
ver = cv2.getStructuringElement(cv2.MORPH_RECT, (1, verticalsize))
|
|
hor = cv2.getStructuringElement(cv2.MORPH_RECT, (horizontalsize, 1))
|
|
|
|
vertical = cv2.erode(vertical, ver, (-1, -1))
|
|
vertical = cv2.dilate(vertical, ver, (-1, -1))
|
|
|
|
horizontal = cv2.erode(horizontal, hor, (-1, -1))
|
|
horizontal = cv2.dilate(horizontal, hor, (-1, -1))
|
|
|
|
mask = vertical + horizontal
|
|
joints = np.bitwise_and(vertical, horizontal)
|
|
__, contours, __ = cv2.findContours(
|
|
mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
contours = sorted(contours, key=cv2.contourArea, reverse=True)[:10]
|
|
|
|
tables = {}
|
|
for c in contours:
|
|
c_poly = cv2.approxPolyDP(c, 3, True)
|
|
x, y, w, h = cv2.boundingRect(c_poly)
|
|
roi = joints[y : y + h, x : x + w]
|
|
__, jc, __ = cv2.findContours(
|
|
roi, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
|
|
if len(jc) <= 4: # remove contours with less than <=4 joints
|
|
continue
|
|
joint_coords = []
|
|
for j in jc:
|
|
jx, jy, jw, jh = cv2.boundingRect(j)
|
|
c1, c2 = x + (2 * jx + jw) / 2, y + (2 * jy + jh) / 2
|
|
joint_coords.append((c1, c2))
|
|
tables[(x, y + h, x + w, y)] = joint_coords
|
|
|
|
v_segments, h_segments = [], []
|
|
_, vcontours, _ = cv2.findContours(
|
|
vertical, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
for vc in vcontours:
|
|
x, y, w, h = cv2.boundingRect(vc)
|
|
x1, x2 = x, x + w
|
|
y1, y2 = y, y + h
|
|
v_segments.append(((x1 + x2) / 2, y2, (x1 + x2) / 2, y1))
|
|
|
|
_, hcontours, _ = cv2.findContours(
|
|
horizontal, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
for hc in hcontours:
|
|
x, y, w, h = cv2.boundingRect(hc)
|
|
x1, x2 = x, x + w
|
|
y1, y2 = y, y + h
|
|
h_segments.append((x1, (y1 + y2) / 2, x2, (y1 + y2) / 2))
|
|
|
|
return img, tables, v_segments, h_segments
|
|
|
|
|
|
class Lattice:
|
|
"""Lattice algorithm
|
|
|
|
Makes use of pdf geometry by processing its image, to make a table
|
|
and fills text objects in table cells.
|
|
|
|
Parameters
|
|
----------
|
|
pdfobject : camelot.pdf.Pdf
|
|
|
|
fill : None, 'h', 'v', 'hv'
|
|
Fill data in horizontal and/or vertical spanning
|
|
cells. (optional, default: None)
|
|
|
|
scale : int
|
|
Scaling factor. Large scaling factor leads to smaller lines
|
|
being detected. (optional, default: 15)
|
|
|
|
jtol : int
|
|
Tolerance to account for when comparing joint and line
|
|
coordinates. (optional, default: 2)
|
|
|
|
mtol : int
|
|
Tolerance to account for when merging lines which are
|
|
very close. (optional, default: 2)
|
|
|
|
invert : bool
|
|
Invert pdf image to make sure that lines are in foreground.
|
|
(optional, default: False)
|
|
|
|
debug : 'contour', 'line', 'joint', 'table'
|
|
Debug by visualizing pdf geometry.
|
|
(optional, default: None)
|
|
|
|
Attributes
|
|
----------
|
|
tables : dict
|
|
Dictionary with page number as key and list of tables on that
|
|
page as value.
|
|
"""
|
|
|
|
def __init__(self, pdfobject, fill=None, scale=15, jtol=2, mtol=2,
|
|
invert=False, debug=None, verbose=False):
|
|
|
|
self.pdfobject = pdfobject
|
|
self.fill = fill
|
|
self.scale = scale
|
|
self.jtol = jtol
|
|
self.mtol = mtol
|
|
self.invert = invert
|
|
self.debug = debug
|
|
self.verbose = verbose
|
|
self.tables = {}
|
|
if self.debug is not None:
|
|
self.debug_images = {}
|
|
self.debug_segments = {}
|
|
self.debug_tables = {}
|
|
|
|
def get_tables(self):
|
|
"""Returns all tables found in given pdf.
|
|
|
|
Returns
|
|
-------
|
|
tables : dict
|
|
Dictionary with page number as key and list of tables on that
|
|
page as value.
|
|
"""
|
|
vprint = print if self.verbose else lambda *a, **k: None
|
|
self.pdfobject.split()
|
|
self.pdfobject.convert()
|
|
for page in self.pdfobject.extract():
|
|
p, text, __, width, height = page
|
|
pkey = 'pg-{0}'.format(p)
|
|
imagename = os.path.join(
|
|
self.pdfobject.temp, '{}.png'.format(pkey))
|
|
pdf_x = width
|
|
pdf_y = height
|
|
img, table_bbox, v_segments, h_segments = _morph_transform(
|
|
imagename, scale=self.scale, invert=self.invert)
|
|
img_x = img.shape[1]
|
|
img_y = img.shape[0]
|
|
scaling_factor_x = pdf_x / float(img_x)
|
|
scaling_factor_y = pdf_y / float(img_y)
|
|
|
|
if self.debug is not None:
|
|
self.debug_images[pkey] = (img, table_bbox)
|
|
|
|
factors = (scaling_factor_x, scaling_factor_y, img_y)
|
|
table_bbox, v_segments, h_segments = transform(table_bbox, v_segments,
|
|
h_segments, factors)
|
|
|
|
if self.debug is not None:
|
|
self.debug_segments[pkey] = (v_segments, h_segments)
|
|
|
|
if self.debug is not None:
|
|
debug_page_tables = []
|
|
page_tables = []
|
|
# sort tables based on y-coord
|
|
for k in sorted(table_bbox.keys(), key=lambda x: x[1], reverse=True):
|
|
# select edges which lie within table_bbox
|
|
text_bbox, v_s, h_s = elements_bbox(k, text, v_segments,
|
|
h_segments)
|
|
rotated = detect_vertical(text_bbox)
|
|
cols, rows = zip(*table_bbox[k])
|
|
cols, rows = list(cols), list(rows)
|
|
cols.extend([k[0], k[2]])
|
|
rows.extend([k[1], k[3]])
|
|
# sort horizontal and vertical segments
|
|
cols = merge_close_values(sorted(cols), mtol=self.mtol)
|
|
rows = merge_close_values(
|
|
sorted(rows, reverse=True), mtol=self.mtol)
|
|
# make grid using x and y coord of shortlisted rows and cols
|
|
cols = [(cols[i], cols[i + 1])
|
|
for i in range(0, len(cols) - 1)]
|
|
rows = [(rows[i], rows[i + 1])
|
|
for i in range(0, len(rows) - 1)]
|
|
table = Table(cols, rows)
|
|
# set table edges to True using ver+hor lines
|
|
table = table.set_edges(v_s, h_s, jtol=self.jtol)
|
|
# set spanning cells to True
|
|
table = table.set_spanning()
|
|
# set table border edges to True
|
|
table = outline(table)
|
|
|
|
if self.debug is not None:
|
|
debug_page_tables.append(table)
|
|
|
|
# fill text after sorting it
|
|
if rotated == '':
|
|
text_bbox.sort(key=lambda x: (-x.y0, x.x0))
|
|
elif rotated == 'left':
|
|
text_bbox.sort(key=lambda x: (x.x0, x.y0))
|
|
elif rotated == 'right':
|
|
text_bbox.sort(key=lambda x: (-x.x0, -x.y0))
|
|
for t in text_bbox:
|
|
r_idx = get_row_index(t, rows)
|
|
c_idx = get_column_index(t, cols)
|
|
if None in [r_idx, c_idx]:
|
|
# couldn't assign LTChar to any cell
|
|
pass
|
|
else:
|
|
r_idx, c_idx = reduce_index(
|
|
table, rotated, r_idx, c_idx)
|
|
table.cells[r_idx][c_idx].add_text(
|
|
t.get_text().strip('\n'))
|
|
|
|
if self.fill is not None:
|
|
table = fill_spanning(table, fill=self.fill)
|
|
ar = table.get_list()
|
|
if rotated == 'left':
|
|
ar = zip(*ar[::-1])
|
|
elif rotated == 'right':
|
|
ar = zip(*ar[::1])
|
|
ar.reverse()
|
|
ar = remove_empty(ar)
|
|
ar = [list(o) for o in ar]
|
|
page_tables.append(encode_list(ar))
|
|
vprint(pkey)
|
|
self.tables[pkey] = page_tables
|
|
|
|
if self.debug is not None:
|
|
self.debug_tables[pkey] = debug_page_tables
|
|
|
|
if self.pdfobject.clean:
|
|
self.pdfobject.remove_tempdir()
|
|
|
|
if self.debug is not None:
|
|
return None
|
|
|
|
return self.tables
|
|
|
|
def plot_geometry(self, geometry):
|
|
"""Plots various pdf geometries that are detected so user can choose
|
|
tweak scale, jtol, mtol parameters.
|
|
"""
|
|
import matplotlib.pyplot as plt
|
|
|
|
if geometry == 'contour':
|
|
for pkey in self.debug_images.keys():
|
|
img, table_bbox = self.debug_images[pkey]
|
|
for t in table_bbox.keys():
|
|
cv2.rectangle(img, (t[0], t[1]),
|
|
(t[2], t[3]), (255, 0, 0), 3)
|
|
plt.imshow(img)
|
|
plt.show()
|
|
elif geometry == 'joint':
|
|
x_coord = []
|
|
y_coord = []
|
|
for pkey in self.debug_images.keys():
|
|
img, table_bbox = self.debug_images[pkey]
|
|
for k in table_bbox.keys():
|
|
for coord in table_bbox[k]:
|
|
x_coord.append(coord[0])
|
|
y_coord.append(coord[1])
|
|
max_x, max_y = max(x_coord), max(y_coord)
|
|
plt.plot(x_coord, y_coord, 'ro')
|
|
plt.axis([0, max_x + 100, max_y + 100, 0])
|
|
plt.imshow(img)
|
|
plt.show()
|
|
elif geometry == 'line':
|
|
for pkey in self.debug_segments.keys():
|
|
v_s, h_s = self.debug_segments[pkey]
|
|
for v in v_s:
|
|
plt.plot([v[0], v[2]], [v[1], v[3]])
|
|
for h in h_s:
|
|
plt.plot([h[0], h[2]], [h[1], h[3]])
|
|
plt.show()
|
|
elif geometry == 'table':
|
|
for pkey in self.debug_tables.keys():
|
|
for table in self.debug_tables[pkey]:
|
|
for i in range(len(table.cells)):
|
|
for j in range(len(table.cells[i])):
|
|
if table.cells[i][j].left:
|
|
plt.plot([table.cells[i][j].lb[0],
|
|
table.cells[i][j].lt[0]],
|
|
[table.cells[i][j].lb[1],
|
|
table.cells[i][j].lt[1]])
|
|
if table.cells[i][j].right:
|
|
plt.plot([table.cells[i][j].rb[0],
|
|
table.cells[i][j].rt[0]],
|
|
[table.cells[i][j].rb[1],
|
|
table.cells[i][j].rt[1]])
|
|
if table.cells[i][j].top:
|
|
plt.plot([table.cells[i][j].lt[0],
|
|
table.cells[i][j].rt[0]],
|
|
[table.cells[i][j].lt[1],
|
|
table.cells[i][j].rt[1]])
|
|
if table.cells[i][j].bottom:
|
|
plt.plot([table.cells[i][j].lb[0],
|
|
table.cells[i][j].rb[0]],
|
|
[table.cells[i][j].lb[1],
|
|
table.cells[i][j].rb[1]])
|
|
plt.show() |