342 lines
12 KiB
Python
342 lines
12 KiB
Python
from __future__ import division
|
|
import os
|
|
import types
|
|
import copy_reg
|
|
import logging
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from wand.image import Image
|
|
|
|
from .table import Table
|
|
from .utils import (transform, segments_bbox, text_bbox, detect_vertical, merge_close_values,
|
|
get_row_index, get_column_index, get_score, reduce_index,
|
|
outline, fill_spanning, count_empty, encode_list, pdf_to_text)
|
|
|
|
|
|
__all__ = ['Lattice']
|
|
|
|
|
|
def _reduce_method(m):
|
|
if m.im_self is None:
|
|
return getattr, (m.im_class, m.im_func.func_name)
|
|
else:
|
|
return getattr, (m.im_self, m.im_func.func_name)
|
|
copy_reg.pickle(types.MethodType, _reduce_method)
|
|
|
|
|
|
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)
|
|
try:
|
|
__, contours, __ = cv2.findContours(
|
|
mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
except ValueError:
|
|
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]
|
|
try:
|
|
__, jc, __ = cv2.findContours(
|
|
roi, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE)
|
|
except ValueError:
|
|
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 = [], []
|
|
try:
|
|
_, vcontours, _ = cv2.findContours(
|
|
vertical, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
except ValueError:
|
|
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))
|
|
|
|
try:
|
|
_, hcontours, _ = cv2.findContours(
|
|
horizontal, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
except ValueError:
|
|
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 : string
|
|
Fill data in horizontal and/or vertical spanning
|
|
cells. (optional, default: None) {None, 'h', 'v', 'hv'}
|
|
|
|
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 : string
|
|
Debug by visualizing pdf geometry.
|
|
(optional, default: None) {'contour', 'line', 'joint', 'table'}
|
|
|
|
Attributes
|
|
----------
|
|
tables : dict
|
|
Dictionary with page number as key and list of tables on that
|
|
page as value.
|
|
"""
|
|
|
|
def __init__(self, fill=None, scale=15, jtol=2, mtol=2,
|
|
invert=False, pdf_margin=(2.0, 0.5, 0.1), debug=None):
|
|
|
|
self.method = 'lattice'
|
|
self.fill = fill
|
|
self.scale = scale
|
|
self.jtol = jtol
|
|
self.mtol = mtol
|
|
self.invert = invert
|
|
self.char_margin, self.line_margin, self.word_margin = pdf_margin
|
|
self.debug = debug
|
|
|
|
def get_tables(self, pdfname):
|
|
"""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.
|
|
"""
|
|
text, __, width, height = pdf_to_text(pdfname, self.char_margin,
|
|
self.line_margin, self.word_margin)
|
|
bname, __ = os.path.splitext(pdfname)
|
|
if not text:
|
|
logging.warning("{0}: PDF has no text. It may be an image.".format(
|
|
os.path.basename(bname)))
|
|
return None
|
|
imagename = ''.join([bname, '.png'])
|
|
with Image(filename=pdfname, depth=8, resolution=300) as png:
|
|
png.save(filename=imagename)
|
|
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:
|
|
self.debug_images = (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:
|
|
self.debug_segments = (v_segments, h_segments)
|
|
self.debug_tables = []
|
|
|
|
pdf_page = {}
|
|
page_tables = {}
|
|
table_no = 1
|
|
# 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
|
|
table_info = {}
|
|
v_s, h_s = segments_bbox(k, v_segments, h_segments)
|
|
t_bbox = text_bbox(k, text)
|
|
table_info['text_p'] = 100 * (1 - (len(t_bbox) / len(text)))
|
|
table_rotation = detect_vertical(t_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)
|
|
nouse = table.nocont_ / (len(v_s) + len(h_s))
|
|
table_info['line_p'] = 100 * (1 - nouse)
|
|
# set spanning cells to True
|
|
table = table.set_spanning()
|
|
# set table border edges to True
|
|
table = outline(table)
|
|
|
|
if self.debug:
|
|
self.debug_tables.append(table)
|
|
|
|
rerror = []
|
|
cerror = []
|
|
for t in text:
|
|
try:
|
|
r_idx, rass_error = get_row_index(t, rows)
|
|
except TypeError:
|
|
# couldn't assign LTChar to any cell
|
|
continue
|
|
try:
|
|
c_idx, cass_error = get_column_index(t, cols)
|
|
except TypeError:
|
|
# couldn't assign LTChar to any cell
|
|
continue
|
|
rerror.append(rass_error)
|
|
cerror.append(cass_error)
|
|
r_idx, c_idx = reduce_index(table, table_rotation, r_idx, c_idx)
|
|
table.cells[r_idx][c_idx].add_object(t)
|
|
|
|
for i in range(len(table.cells)):
|
|
for j in range(len(table.cells[i])):
|
|
t_bbox = table.cells[i][j].get_objects()
|
|
try:
|
|
cell_rotation = detect_vertical(t_bbox)
|
|
except ZeroDivisionError:
|
|
cell_rotation = ''
|
|
pass
|
|
# fill text after sorting it
|
|
if cell_rotation == '':
|
|
t_bbox.sort(key=lambda x: (-x.y0, x.x0))
|
|
elif cell_rotation == 'left':
|
|
t_bbox.sort(key=lambda x: (x.x0, x.y0))
|
|
elif cell_rotation == 'right':
|
|
t_bbox.sort(key=lambda x: (-x.x0, -x.y0))
|
|
table.cells[i][j].add_text(''.join([t.get_text()
|
|
for t in t_bbox]))
|
|
|
|
score = get_score([[50, rerror], [50, cerror]])
|
|
table_info['score'] = score
|
|
|
|
if self.fill is not None:
|
|
table = fill_spanning(table, fill=self.fill)
|
|
ar = table.get_list()
|
|
if table_rotation == 'left':
|
|
ar = zip(*ar[::-1])
|
|
elif table_rotation == 'right':
|
|
ar = zip(*ar[::1])
|
|
ar.reverse()
|
|
ar = encode_list(ar)
|
|
table_info['data'] = ar
|
|
empty_p, r_nempty_cells, c_nempty_cells = count_empty(ar)
|
|
table_info['empty_p'] = empty_p
|
|
table_info['r_nempty_cells'] = r_nempty_cells
|
|
table_info['c_nempty_cells'] = c_nempty_cells
|
|
table_info['nrows'] = len(ar)
|
|
table_info['ncols'] = len(ar[0])
|
|
page_tables['table_{0}'.format(table_no)] = table_info
|
|
table_no += 1
|
|
pdf_page[os.path.basename(bname)] = page_tables
|
|
|
|
if self.debug:
|
|
return None
|
|
|
|
return pdf_page |