Briq/camelot-py
8
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Add various metrics to score the quality of a parse

Browse Source 
Add various metrics to score the quality of a parse
This commit is contained in:
Vinayak Mehta
2016-08-30 14:52:49 +05:30
committed by GitHub
parent 43a009dab4
commit 552f9cf422
11 changed files with 1027 additions and 472 deletions
Download Patch File Download Diff File Expand all files Collapse all files
camelot/lattice.py
+157 -178
View File
@@ -1,18 +1,31 @@
from __future__ import print_function
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, elements_bbox, detect_vertical, merge_close_values,
get_row_index, get_column_index, reduce_index, outline,
fill_spanning, remove_empty, encode_list)
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
@@ -65,8 +78,8 @@ def _morph_transform(imagename, scale=15, invert=False):
vertical = threshold
horizontal = threshold
verticalsize = vertical.shape[0] / scale
horizontalsize = horizontal.shape[1] / scale
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))
@@ -79,8 +92,12 @@ def _morph_transform(imagename, scale=15, invert=False):
mask = vertical + horizontal
joints = np.bitwise_and(vertical, horizontal)
__, contours, __ = cv2.findContours(
mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
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 = {}
@@ -88,8 +105,12 @@ def _morph_transform(imagename, scale=15, invert=False):
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)
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 = []
@@ -100,16 +121,24 @@ def _morph_transform(imagename, scale=15, invert=False):
tables[(x, y + h, x + w, y)] = joint_coords
v_segments, h_segments = [], []
_, vcontours, _ = cv2.findContours(
vertical, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
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))
_, hcontours, _ = cv2.findContours(
horizontal, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
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
@@ -160,24 +189,19 @@ class Lattice:
page as value.
"""
def __init__(self, pdfobject, fill=None, scale=15, jtol=2, mtol=2,
invert=False, debug=None, verbose=False):
def __init__(self, fill=None, scale=15, jtol=2, mtol=2,
invert=False, pdf_margin=(2.0, 0.5, 0.1), debug=None):
self.pdfobject = pdfobject
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
self.verbose = verbose
self.tables = {}
if self.debug is not None:
self.debug_images = {}
self.debug_segments = {}
self.debug_tables = {}
def get_tables(self):
def get_tables(self, pdfname):
"""Returns all tables found in given pdf.
Returns
@@ -186,169 +210,124 @@ class Lattice:
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)
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 is not None:
self.debug_images[pkey] = (img, table_bbox)
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)
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:
self.debug_segments = (v_segments, h_segments)
self.debug_tables = []
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)
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 = {}
text_bbox, v_s, h_s = elements_bbox(k, text, v_segments,
h_segments)
table_info['text_p'] = 100 * (1 - (len(text_bbox) / len(text)))
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)
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 is not None:
debug_page_tables.append(table)
if self.debug:
self.debug_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'))
# 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))
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
rerror = []
cerror = []
for t in text_bbox:
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, rotated, r_idx, c_idx)
table.cells[r_idx][c_idx].add_text(
t.get_text().strip('\n'))
score = get_score([[50, rerror], [50, cerror]])
table_info['score'] = score
if self.debug is not None:
self.debug_tables[pkey] = debug_page_tables
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 = 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.pdfobject.clean:
self.pdfobject.remove_tempdir()
if self.debug is not None:
if self.debug:
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()
return pdf_page
camelot/pdf.py
+141 -83
View File
@@ -1,18 +1,11 @@
import os
import shutil
import tempfile
import itertools
import multiprocessing as mp
import cv2
from PyPDF2 import PdfFileReader, PdfFileWriter
from pdfminer.pdfparser import PDFParser
from pdfminer.pdfdocument import PDFDocument
from pdfminer.pdfpage import PDFPage
from pdfminer.pdfpage import PDFTextExtractionNotAllowed
from pdfminer.pdfinterp import PDFResourceManager
from pdfminer.pdfinterp import PDFPageInterpreter
from pdfminer.pdfdevice import PDFDevice
from pdfminer.converter import PDFPageAggregator
from pdfminer.layout import LAParams, LTChar, LTTextLineHorizontal
from wand.image import Image
__all__ = ['Pdf']
@@ -38,38 +31,6 @@ def _parse_page_numbers(pagenos):
return page_numbers
def _extract_text_objects(layout, LTObject, t=None):
"""Recursively parses pdf layout to get a list of
text objects.
Parameters
----------
layout : object
Layout object.
LTObject : object
Text object, either LTChar or LTTextLineHorizontal.
t : list (optional, default: None)
Returns
-------
t : list
List of text objects.
"""
if t is None:
t = []
try:
for obj in layout._objs:
if isinstance(obj, LTObject):
t.append(obj)
else:
t += _extract_text_objects(obj, LTObject)
except AttributeError:
pass
return t
class Pdf:
"""Handles all pdf operations which include:
@@ -99,66 +60,163 @@ class Pdf:
is greater than word_margin. (optional, default: 0.1)
"""
def __init__(self, pdfname, pagenos=[{'start': 1, 'end': 1}],
char_margin=2.0, line_margin=0.5, word_margin=0.1,
clean=False):
def __init__(self, extractor, pdfname, pagenos=[{'start': 1, 'end': 1}],
parallel=False, clean=False):
self.extractor = extractor
self.pdfname = pdfname
if not self.pdfname.endswith('.pdf'):
raise TypeError("Only PDF format is supported right now.")
self.pagenos = _parse_page_numbers(pagenos)
self.char_margin = char_margin
self.line_margin = line_margin
self.word_margin = word_margin
self.parallel = parallel
self.cpu_count = mp.cpu_count()
self.pool = mp.Pool(processes=self.cpu_count)
self.clean = clean
self.temp = tempfile.mkdtemp()
def split(self):
"""Splits pdf into single page pdfs.
"""
if not self.pdfname.endswith('.pdf'):
raise TypeError("Only PDF format is supported.")
infile = PdfFileReader(open(self.pdfname, 'rb'), strict=False)
for p in self.pagenos:
page = infile.getPage(p - 1)
outfile = PdfFileWriter()
outfile.addPage(page)
with open(os.path.join(self.temp, 'pg-{0}.pdf'.format(p)), 'wb') as f:
with open(os.path.join(self.temp, 'page-{0}.pdf'.format(p)), 'wb') as f:
outfile.write(f)
def remove_tempdir(self):
shutil.rmtree(self.temp)
def extract(self):
"""Extracts text objects, width, height from a pdf.
"""
for p in self.pagenos:
pkey = 'pg-{0}'.format(p)
pname = os.path.join(self.temp, '{}.pdf'.format(pkey))
with open(pname, 'r') as f:
parser = PDFParser(f)
document = PDFDocument(parser)
if not document.is_extractable:
raise PDFTextExtractionNotAllowed
laparams = LAParams(char_margin=self.char_margin,
line_margin=self.line_margin,
word_margin=self.word_margin)
rsrcmgr = PDFResourceManager()
device = PDFPageAggregator(rsrcmgr, laparams=laparams)
interpreter = PDFPageInterpreter(rsrcmgr, device)
for page in PDFPage.create_pages(document):
interpreter.process_page(page)
layout = device.get_result()
lattice_objects = _extract_text_objects(layout, LTChar)
stream_objects = _extract_text_objects(
layout, LTTextLineHorizontal)
width = layout.bbox[2]
height = layout.bbox[3]
yield p, lattice_objects, stream_objects, width, height
self.split()
pages = [os.path.join(self.temp, 'page-{0}.pdf'.format(p))
for p in self.pagenos]
if self.parallel:
tables = self.pool.map(self.extractor.get_tables, pages)
tables = {k: v for d in tables if d is not None for k, v in d.items()}
else:
tables = {}
if self.extractor.debug:
if self.extractor.method == 'stream':
self.debug = self.extractor.debug
self.debug_text = []
elif self.extractor.method == 'lattice':
self.debug = self.extractor.debug
self.debug_images = []
self.debug_segments = []
self.debug_tables = []
for p in pages:
table = self.extractor.get_tables(p)
if table is not None:
tables.update(table)
if self.extractor.debug:
if self.extractor.method == 'stream':
self.debug_text.append(self.extractor.debug_text)
elif self.extractor.method == 'lattice':
self.debug_images.append(self.extractor.debug_images)
self.debug_segments.append(self.extractor.debug_segments)
self.debug_tables.append(self.extractor.debug_tables)
if self.clean:
self.remove_tempdir()
return tables
def convert(self):
"""Converts single page pdfs to images.
def debug_plot(self):
"""Plots all text objects and various pdf geometries so that
user can choose number of columns, columns x-coordinates for
Stream or tweak Lattice parameters (scale, jtol, mtol).
"""
for p in self.pagenos:
pdfname = os.path.join(self.temp, 'pg-{0}.pdf'.format(p))
imagename = os.path.join(self.temp, 'pg-{0}.png'.format(p))
with Image(filename=pdfname, depth=8, resolution=300) as png:
png.save(filename=imagename)
import matplotlib.pyplot as plt
import matplotlib.patches as patches
def remove_tempdir(self):
shutil.rmtree(self.temp)
if self.debug is True:
try:
for text in self.debug_text:
fig = plt.figure()
ax = fig.add_subplot(111, aspect='equal')
xs, ys = [], []
for t in text:
xs.extend([t[0], t[1]])
ys.extend([t[2], t[3]])
ax.add_patch(
patches.Rectangle(
(t[0], t[1]),
t[2] - t[0],
t[3] - t[1]
)
)
ax.set_xlim(min(xs) - 10, max(xs) + 10)
ax.set_ylim(min(ys) - 10, max(ys) + 10)
plt.show()
except AttributeError:
raise ValueError("This option only be used with Stream.")
elif self.debug == 'contour':
try:
for img, table_bbox in self.debug_images:
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()
except AttributeError:
raise ValueError("This option only be used with Lattice.")
elif self.debug == 'joint':
try:
for img, table_bbox in self.debug_images:
x_coord = []
y_coord = []
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()
except AttributeError:
raise ValueError("This option only be used with Lattice.")
elif self.debug == 'line':
try:
for v_s, h_s in self.debug_segments:
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()
except AttributeError:
raise ValueError("This option only be used with Lattice.")
elif self.debug == 'table':
try:
for tables in self.debug_tables:
for table in tables:
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()
except AttributeError:
raise ValueError("This option only be used with Lattice.")
else:
raise UserWarning("This method can only be called after"
" debug has been specified.")
camelot/stream.py
+186 -91
View File
@@ -1,14 +1,26 @@
from __future__ import print_function
from __future__ import division
import os
import types
import copy_reg
import logging
import numpy as np
from .utils import get_column_index, encode_list
from .table import Table
from .utils import get_row_index, get_score, count_empty, encode_list, pdf_to_text
__all__ = ['Stream']
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 _group_rows(text, ytol=2):
"""Groups text objects into rows using ytol.
@@ -35,14 +47,16 @@ def _group_rows(text, ytol=2):
# type(obj) is LTChar]):
if t.get_text().strip():
if not np.isclose(row_y, t.y0, atol=ytol):
row_y = t.y0
rows.append(temp)
rows.append(sorted(temp, key=lambda t: t.x0))
temp = []
row_y = t.y0
temp.append(t)
rows.append(sorted(temp, key=lambda t: t.x0))
__ = rows.pop(0) # hacky
return rows
def _merge_columns(l):
def _merge_columns(l, mtol=2):
"""Merges overlapping columns and returns list with updated
columns boundaries.
@@ -62,7 +76,8 @@ def _merge_columns(l):
merged.append(higher)
else:
lower = merged[-1]
if higher[0] <= lower[1]:
if (higher[0] <= lower[1] or
np.isclose(higher[0], lower[1], atol=mtol)):
upper_bound = max(lower[1], higher[1])
lower_bound = min(lower[0], higher[0])
merged[-1] = (lower_bound, upper_bound)
@@ -71,6 +86,62 @@ def _merge_columns(l):
return merged
def _get_column_index(t, columns):
"""Gets index of the column in which the given object falls by
comparing their co-ordinates.
Parameters
----------
t : object
columns : list
Returns
-------
c : int
"""
offset1, offset2 = 0, 0
lt_col_overlap = []
for c in columns:
if c[0] <= t.x1 and c[1] >= t.x0:
left = t.x0 if c[0] <= t.x0 else c[0]
right = t.x1 if c[1] >= t.x1 else c[1]
lt_col_overlap.append(abs(left - right) / abs(c[0] - c[1]))
else:
lt_col_overlap.append(-1)
if len(filter(lambda x: x != -1, lt_col_overlap)) == 0:
logging.warning("Text doesn't fit any column.")
c_idx = lt_col_overlap.index(max(lt_col_overlap))
if t.x0 < columns[c_idx][0]:
offset1 = abs(t.x0 - columns[c_idx][0])
if t.x1 > columns[c_idx][1]:
offset2 = abs(t.x1 - columns[c_idx][1])
Y = abs(t.y0 - t.y1)
charea = abs(t.x0 - t.x1) * abs(t.y0 - t.y1)
error = (Y * (offset1 + offset2)) / charea
return c_idx, error
def _add_columns(cols, text, ytolerance):
if text:
text = _group_rows(text, ytol=ytolerance)
elements = [len(r) for r in text]
new_cols = [(t.x0, t.x1)
for r in text if len(r) == max(elements) for t in r]
cols.extend(_merge_columns(sorted(new_cols)))
return cols
def _join_columns(cols, width):
cols = sorted(cols)
cols = [(cols[i][0] + cols[i - 1][1]) / 2 for i in range(1, len(cols))]
cols.insert(0, 0)
cols.append(width) # or some tolerance
cols = [(cols[i], cols[i + 1])
for i in range(0, len(cols) - 1)]
return cols
class Stream:
"""Stream algorithm
@@ -105,20 +176,18 @@ class Stream:
page as value.
"""
def __init__(self, pdfobject, ncolumns=0, columns=None, ytol=2,
debug=False, verbose=False):
def __init__(self, ncolumns=0, columns=None, ytol=2, mtol=2,
pdf_margin=(2.0, 0.5, 0.1), debug=False):
self.pdfobject = pdfobject
self.method = 'stream'
self.ncolumns = ncolumns
self.columns = columns
self.ytol = ytol
self.mtol = mtol
self.char_margin, self.line_margin, self.word_margin = pdf_margin
self.debug = debug
self.verbose = verbose
self.tables = {}
if self.debug:
self.debug_text = {}
def get_tables(self):
def get_tables(self, pdfname):
"""Returns all tables found in given pdf.
Returns
@@ -127,86 +196,112 @@ class Stream:
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()
for page in self.pdfobject.extract():
p, __, text, __, __ = page
pkey = 'pg-{0}'.format(p)
text.sort(key=lambda x: (-x.y0, x.x0))
if self.debug:
self.debug_text[pkey] = text
rows = _group_rows(text, ytol=self.ytol)
elements = [len(r) for r in rows]
# a table can't have just 1 column, can it?
elements = filter(lambda x: x != 1, elements)
guess = False
if self.columns:
cols = self.columns.split(',')
cols = [(float(cols[i]), float(cols[i + 1]))
for i in range(0, len(cols) - 1)]
else:
guess = True
ncols = self.ncolumns if self.ncolumns else max(
set(elements), key=elements.count)
if ncols == 0:
# no tables detected
continue
cols = [(t.x0, t.x1)
for r in rows for t in r if len(r) == ncols]
cols = _merge_columns(sorted(cols))
cols = [(c[0] + c[1]) / 2.0 for c in cols]
ar = [['' for c in cols] for r in rows]
for r_idx, r in enumerate(rows):
for t in r:
if guess:
cog = (t.x0 + t.x1) / 2.0
diff = [abs(cog - c) for c in cols]
c_idx = diff.index(min(diff))
else:
c_idx = get_column_index(t, cols)
if None in [r_idx, c_idx]: # couldn't assign LTTextLH to any cell
continue
if ar[r_idx][c_idx]:
ar[r_idx][c_idx] = ' '.join(
[ar[r_idx][c_idx], t.get_text().strip()])
else:
ar[r_idx][c_idx] = t.get_text().strip()
vprint(pkey)
self.tables[pkey] = [encode_list(ar)]
if self.pdfobject.clean:
self.pdfobject.remove_tempdir()
__, 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
text.sort(key=lambda x: (-x.y0, x.x0))
if self.debug:
self.debug_text = [(t.x0, t.y0, t.x1, t.y1) for t in text]
return None
return self.tables
rows_grouped = _group_rows(text, ytol=self.ytol)
elements = [len(r) for r in rows_grouped]
row_mids = [sum([(t.y0 + t.y1) / 2 for t in r]) / len(r)
if len(r) > 0 else 0 for r in rows_grouped]
rows = [(row_mids[i] + row_mids[i - 1]) / 2 for i in range(1, len(row_mids))]
rows.insert(0, height) # or some tolerance
rows.append(0)
rows = [(rows[i], rows[i + 1])
for i in range(0, len(rows) - 1)]
def plot_text(self):
"""Plots all text objects so user can choose number of columns
or columns x-coordinates using the matplotlib interface.
"""
import matplotlib.pyplot as plt
import matplotlib.patches as patches
guess = False
if self.columns:
# user has to input boundary columns too
# take (0, width) by default
# similar to else condition
# len can't be 1
cols = self.columns.split(',')
cols = [(float(cols[i]), float(cols[i + 1]))
for i in range(0, len(cols) - 1)]
else:
if self.ncolumns:
ncols = self.ncolumns
cols = [(t.x0, t.x1)
for r in rows_grouped if len(r) == ncols for t in r]
cols = _merge_columns(sorted(cols), mtol=self.mtol)
if len(cols) != self.ncolumns:
logging.warning("{}: The number of columns after merge"
" isn't the same as what you specified."
" Change the value of mtol.".format(
os.path.basename(bname)))
cols = _join_columns(cols, width)
else:
guess = True
ncols = max(set(elements), key=elements.count)
len_non_mode = len(filter(lambda x: x != ncols, elements))
if ncols == 1 and not self.debug:
# no tables detected
logging.warning("{}: Only one column was detected, the PDF"
" may have no tables. Specify ncols if"
" the PDF has tables.".format(
os.path.basename(bname)))
cols = [(t.x0, t.x1)
for r in rows_grouped if len(r) == ncols for t in r]
cols = _merge_columns(sorted(cols), mtol=self.mtol)
inner_text = []
for i in range(1, len(cols)):
left = cols[i - 1][1]
right = cols[i][0]
inner_text.extend([t for t in text if t.x0 > left and t.x1 < right])
outer_text = [t for t in text if t.x0 > cols[-1][1] or t.x1 < cols[0][0]]
inner_text.extend(outer_text)
cols = _add_columns(cols, inner_text, self.ytol)
cols = _join_columns(cols, width)
for pkey in sorted(self.debug_text.keys()):
fig = plt.figure()
ax = fig.add_subplot(111, aspect='equal')
xs, ys = [], []
for t in self.debug_text[pkey]:
xs.extend([t.x0, t.x1])
ys.extend([t.y0, t.y1])
ax.add_patch(
patches.Rectangle(
(t.x0, t.y0),
t.x1 - t.x0,
t.y1 - t.y0
)
)
ax.set_xlim(min(xs) - 10, max(xs) + 10)
ax.set_ylim(min(ys) - 10, max(ys) + 10)
plt.show()
pdf_page = {}
page_tables = {}
table_info = {}
table = Table(cols, rows)
rerror = []
cerror = []
for row in rows_grouped:
for t in row:
try:
r_idx, rass_error = get_row_index(t, rows)
except ValueError as e:
# couldn't assign LTTextLH to any cell
vprint(e.message)
continue
try:
c_idx, cass_error = _get_column_index(t, cols)
except ValueError as e:
# couldn't assign LTTextLH to any cell
vprint(e.message)
continue
rerror.append(rass_error)
cerror.append(cass_error)
table.cells[r_idx][c_idx].add_text(
t.get_text().strip('\n'))
if guess:
score = get_score([[33, rerror], [33, cerror], [34, [len_non_mode / len(elements)]]])
else:
score = get_score([[50, rerror], [50, cerror]])
table_info['score'] = score
ar = table.get_list()
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_1'] = table_info
pdf_page[os.path.basename(bname)] = page_tables
return pdf_page
camelot/table.py
+3
View File
@@ -26,6 +26,7 @@ class Table:
self.rows = rows
self.cells = [[Cell(c[0], r[1], c[1], r[0])
for c in cols] for r in rows]
self.nocont_ = 0
def set_edges(self, vertical, horizontal, jtol=2):
"""Sets cell edges to True if corresponding line segments
@@ -53,6 +54,7 @@ class Table:
k = [k for k, t in enumerate(self.rows)
if np.isclose(v[1], t[0], atol=jtol)]
if not j:
self.nocont_ += 1
continue
J = j[0]
if i == [0]: # only left edge
@@ -104,6 +106,7 @@ class Table:
k = [k for k, t in enumerate(self.cols)
if np.isclose(h[2], t[0], atol=jtol)]
if not j:
self.nocont_ += 1
continue
J = j[0]
if i == [0]: # only top edge
camelot/utils.py
+165 -3
View File
@@ -1,5 +1,18 @@
from __future__ import division
import os
import numpy as np
from pdfminer.pdfparser import PDFParser
from pdfminer.pdfdocument import PDFDocument
from pdfminer.pdfpage import PDFPage
from pdfminer.pdfpage import PDFTextExtractionNotAllowed
from pdfminer.pdfinterp import PDFResourceManager
from pdfminer.pdfinterp import PDFPageInterpreter
from pdfminer.pdfdevice import PDFDevice
from pdfminer.converter import PDFPageAggregator
from pdfminer.layout import LAParams, LTChar, LTTextLineHorizontal
def translate(x1, x2):
"""Translates x2 by x1.
@@ -243,15 +256,24 @@ def get_row_index(t, rows):
----------
t : object
rows : list
rows : list, sorted in decreasing order
Returns
-------
r : int
"""
offset1, offset2 = 0, 0
for r in range(len(rows)):
if (t.y0 + t.y1) / 2.0 < rows[r][0] and (t.y0 + t.y1) / 2.0 > rows[r][1]:
return r
if t.y0 > rows[r][0]:
offset1 = abs(t.y0 - rows[r][0])
if t.y1 < rows[r][1]:
offset2 = abs(t.y1 - rows[r][1])
X = 1.0 if abs(t.x0 - t.x1) == 0.0 else abs(t.x0 - t.x1)
Y = 1.0 if abs(t.y0 - t.y1) == 0.0 else abs(t.y0 - t.y1)
charea = X * Y
error = (X * (offset1 + offset2)) / charea
return r, error
def get_column_index(t, columns):
@@ -268,9 +290,45 @@ def get_column_index(t, columns):
-------
c : int
"""
offset1, offset2 = 0, 0
for c in range(len(columns)):
if (t.x0 + t.x1) / 2.0 > columns[c][0] and (t.x0 + t.x1) / 2.0 < columns[c][1]:
return c
if t.x0 < columns[c][0]:
offset1 = abs(t.x0 - columns[c][0])
if t.x1 > columns[c][1]:
offset2 = abs(t.x1 - columns[c][1])
X = 1.0 if abs(t.x0 - t.x1) == 0.0 else abs(t.x0 - t.x1)
Y = 1.0 if abs(t.y0 - t.y1) == 0.0 else abs(t.y0 - t.y1)
charea = X * Y
error = (Y * (offset1 + offset2)) / charea
return c, error
def get_score(error_weights):
"""Calculates score based on weights assigned to various parameters,
and their error percentages.
Parameters
----------
error_weights : dict
Dict with a tuple of error percentages as key and weightage
assigned to them as value. Sum of all values should be equal
to 100.
Returns
-------
score : float
"""
SCORE_VAL = 100
score = 0
if sum([ew[0] for ew in error_weights]) != SCORE_VAL:
raise ValueError("Please assign a valid weightage to each parameter"
" such that their sum is equal to 100")
for ew in error_weights:
weight = ew[0] / len(ew[1])
for error_percentage in ew[1]:
score += weight * (1 - error_percentage)
return score
def reduce_index(t, rotated, r_idx, c_idx):
@@ -394,6 +452,110 @@ def remove_empty(d):
return d
def count_empty(d):
"""Counts empty rows and columns from list of lists.
Parameters
----------
d : list
Returns
-------
n_empty_rows : number of empty rows
n_empty_cols : number of empty columns
empty_p : percentage of empty cells
"""
empty_p = 0
r_nempty_cells, c_nempty_cells = [], []
for i in d:
for j in i:
if j.strip() == '':
empty_p += 1
empty_p = 100 * (empty_p / float(len(d) * len(d[0])))
for row in d:
r_nempty_c = 0
for r in row:
if r.strip() != '':
r_nempty_c += 1
r_nempty_cells.append(r_nempty_c)
d = zip(*d)
d = [list(col) for col in d]
for col in d:
c_nempty_c = 0
for c in col:
if c.strip() != '':
c_nempty_c += 1
c_nempty_cells.append(c_nempty_c)
return empty_p, r_nempty_cells, c_nempty_cells
def encode_list(ar):
"""Encodes list of text.
Parameters
----------
ar : list
Returns
-------
ar : list
"""
ar = [[r.encode('utf-8') for r in row] for row in ar]
return ar
def extract_text_objects(layout, LTObject, t=None):
"""Recursively parses pdf layout to get a list of
text objects.
Parameters
----------
layout : object
Layout object.
LTObject : object
Text object, either LTChar or LTTextLineHorizontal.
t : list (optional, default: None)
Returns
-------
t : list
List of text objects.
"""
if t is None:
t = []
try:
for obj in layout._objs:
if isinstance(obj, LTObject):
t.append(obj)
else:
t += extract_text_objects(obj, LTObject)
except AttributeError:
pass
return t
def pdf_to_text(pname, char_margin, line_margin, word_margin):
# pkey = 'page-{0}'.format(p)
# pname = os.path.join(self.temp, '{}.pdf'.format(pkey))
with open(pname, 'r') as f:
parser = PDFParser(f)
document = PDFDocument(parser)
if not document.is_extractable:
raise PDFTextExtractionNotAllowed
laparams = LAParams(char_margin=char_margin,
line_margin=line_margin,
word_margin=word_margin)
rsrcmgr = PDFResourceManager()
device = PDFPageAggregator(rsrcmgr, laparams=laparams)
interpreter = PDFPageInterpreter(rsrcmgr, device)
for page in PDFPage.create_pages(document):
interpreter.process_page(page)
layout = device.get_result()
lattice_objects = extract_text_objects(layout, LTChar)
stream_objects = extract_text_objects(
layout, LTTextLineHorizontal)
width = layout.bbox[2]
height = layout.bbox[3]
return lattice_objects, stream_objects, width, height