176 lines
7.0 KiB
Python
176 lines
7.0 KiB
Python
import os
|
|
import csv
|
|
import cv2
|
|
import glob
|
|
import numpy as np
|
|
|
|
from table import Table
|
|
from pdf import get_pdf_info
|
|
from morph_transform import morph_transform
|
|
from utils import (translate, scale, merge_close_values, get_row_idx,
|
|
get_column_idx, reduce_index, outline, fill, remove_empty)
|
|
|
|
|
|
def spreadsheet(pdf_dir, filename, fill, s, jtol, mtol, invert, debug):
|
|
if debug:
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.patches as patches
|
|
print "working on", filename
|
|
imagename = os.path.join(pdf_dir, filename.split('.')[0] + '.png')
|
|
img = cv2.imread(imagename)
|
|
img_x, img_y = img.shape[1], img.shape[0]
|
|
text, pdf_x, pdf_y = get_pdf_info(
|
|
os.path.join(pdf_dir, filename), 'spreadsheet')
|
|
scaling_factor_x = pdf_x / float(img_x)
|
|
scaling_factor_y = pdf_y / float(img_y)
|
|
tables, v_segments, h_segments = morph_transform(imagename, s, invert)
|
|
|
|
if debug == ["contours"]:
|
|
for t in tables.keys():
|
|
cv2.rectangle(img, (t[0], t[1]), (t[2], t[3]), (255, 0, 0), 3)
|
|
plt.imshow(img)
|
|
if debug == ["joints"]:
|
|
x_coord = []
|
|
y_coord = []
|
|
for k in tables.keys():
|
|
for coord in tables[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)
|
|
|
|
# detect if vertical
|
|
num_v = [t for t in text if (not t.upright) and t.get_text().strip()]
|
|
num_h = [t for t in text if t.upright and t.get_text().strip()]
|
|
vger = len(num_v) / float(len(num_v) + len(num_h))
|
|
rotated = ''
|
|
if vger > 0.8:
|
|
clockwise = sum(t.matrix[1] < 0 and t.matrix[2] > 0 for t in text)
|
|
anticlockwise = sum(t.matrix[1] > 0 and t.matrix[2] < 0 for t in text)
|
|
rotated = 'left' if clockwise < anticlockwise else 'right'
|
|
|
|
tables_new = {}
|
|
for k in tables.keys():
|
|
x1, y1, x2, y2 = k
|
|
x1 = scale(x1, scaling_factor_x)
|
|
y1 = scale(abs(translate(-img_y, y1)), scaling_factor_y)
|
|
x2 = scale(x2, scaling_factor_x)
|
|
y2 = scale(abs(translate(-img_y, y2)), scaling_factor_y)
|
|
j_x, j_y = zip(*tables[k])
|
|
j_x = [scale(j, scaling_factor_x) for j in j_x]
|
|
j_y = [scale(abs(translate(-img_y, j)), scaling_factor_y) for j in j_y]
|
|
joints = zip(j_x, j_y)
|
|
tables_new[(x1, y1, x2, y2)] = joints
|
|
|
|
v_segments_new = []
|
|
for v in v_segments:
|
|
x1, x2 = scale(v[0], scaling_factor_x), scale(v[2], scaling_factor_x)
|
|
y1, y2 = scale(abs(translate(-img_y, v[1])), scaling_factor_y), scale(
|
|
abs(translate(-img_y, v[3])), scaling_factor_y)
|
|
v_segments_new.append((x1, y1, x2, y2))
|
|
|
|
h_segments_new = []
|
|
for h in h_segments:
|
|
x1, x2 = scale(h[0], scaling_factor_x), scale(h[2], scaling_factor_x)
|
|
y1, y2 = scale(abs(translate(-img_y, h[1])), scaling_factor_y), scale(
|
|
abs(translate(-img_y, h[3])), scaling_factor_y)
|
|
h_segments_new.append((x1, y1, x2, y2))
|
|
|
|
num_tables = 0
|
|
# sort tables based on y-coord
|
|
for k in sorted(tables_new.keys(), key=lambda x: x[1], reverse=True):
|
|
# find rows and columns that lie in table
|
|
lb = (k[0], k[1])
|
|
rt = (k[2], k[3])
|
|
v_s = [v for v in v_segments_new if v[1] > lb[1] - 2 and v[3]
|
|
< rt[1] + 2 and lb[0] - 2 <= v[0] <= rt[0] + 2]
|
|
h_s = [h for h in h_segments_new if h[0] > lb[0] - 2 and h[2]
|
|
< rt[0] + 2 and lb[1] - 2 <= h[1] <= rt[1] + 2]
|
|
|
|
if debug == ["lines"]:
|
|
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]])
|
|
|
|
columns, rows = zip(*tables_new[k])
|
|
columns, rows = list(columns), list(rows)
|
|
columns.extend([lb[0], rt[0]])
|
|
rows.extend([lb[1], rt[1]])
|
|
# sort horizontal and vertical segments
|
|
columns = merge_close_values(sorted(columns), mtol)
|
|
rows = merge_close_values(sorted(rows, reverse=True), mtol)
|
|
# make grid using x and y coord of shortlisted rows and columns
|
|
columns = [(columns[i], columns[i + 1])
|
|
for i in range(0, len(columns) - 1)]
|
|
rows = [(rows[i], rows[i + 1]) for i in range(0, len(rows) - 1)]
|
|
|
|
table = Table(columns, rows)
|
|
# light up cell edges
|
|
table = table.set_edges(v_s, h_s, jtol)
|
|
# table set span method
|
|
table = table.set_spanning()
|
|
# TODO
|
|
table = outline(table)
|
|
|
|
if debug == ["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]])
|
|
if debug:
|
|
plt.show()
|
|
|
|
# fill text after sorting it
|
|
if not rotated:
|
|
text.sort(key=lambda x: (-x.y0, x.x0))
|
|
elif rotated == 'left':
|
|
text.sort(key=lambda x: (x.x0, x.y0))
|
|
elif rotated == 'right':
|
|
text.sort(key=lambda x: (-x.x0, -x.y0))
|
|
|
|
for t in text:
|
|
r_idx = get_row_idx(t, rows)
|
|
c_idx = get_column_idx(t, columns)
|
|
if None in [r_idx, c_idx]:
|
|
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 fill:
|
|
table = fill(table, fill)
|
|
|
|
data = []
|
|
for i in range(len(table.cells)):
|
|
data.append([table.cells[i][j].get_text().strip().encode('utf-8')
|
|
for j in range(len(table.cells[i]))])
|
|
if rotated == 'left':
|
|
data = zip(*data[::-1])
|
|
elif rotated == 'right':
|
|
data = zip(*data[::1])
|
|
data.reverse()
|
|
data = remove_empty(data)
|
|
csvname = filename.split(
|
|
'.')[0] + ('_table_%d' % (num_tables + 1)) + '.csv'
|
|
csvpath = os.path.join(pdf_dir, csvname)
|
|
with open(csvpath, 'w') as outfile:
|
|
writer = csv.writer(outfile, quoting=csv.QUOTE_ALL)
|
|
for d in data:
|
|
writer.writerow(d)
|
|
print "saved as", csvname
|
|
print
|
|
num_tables += 1
|