Fix plot_geometry

pull/2/head
Vinayak Mehta 2018-09-07 06:25:13 +05:30
parent b3f840bba9
commit e615580e55
3 changed files with 59 additions and 71 deletions

View File

@ -232,7 +232,11 @@ class Lattice(BaseParser):
table = self._generate_table(table_idx, cols, rows, v_s=v_s, h_s=h_s) table = self._generate_table(table_idx, cols, rows, v_s=v_s, h_s=h_s)
_tables.append(table) _tables.append(table)
if self.debug: if self.debug is not None:
text = []
text.extend([(t.x0, t.y0, t.x1, t.y1) for t in self.horizontal_text])
text.extend([(t.x0, t.y0, t.x1, t.y1) for t in self.vertical_text])
self.g.text = text
self.g.images = (self.image, self.table_bbox_unscaled) self.g.images = (self.image, self.table_bbox_unscaled)
self.g.segments = (self.vertical_segments, self.horizontal_segments) self.g.segments = (self.vertical_segments, self.horizontal_segments)
self.g.tables = _tables self.g.tables = _tables

View File

@ -8,7 +8,7 @@ import pandas as pd
from .base import BaseParser from .base import BaseParser
from ..core import Table from ..core import Table
from ..utils import (text_in_bbox, get_table_index, compute_accuracy, from ..utils import (text_in_bbox, get_table_index, compute_accuracy,
count_empty_strings, encode_) count_empty_strings, encode_, setup_logging)
logger = setup_logging(__name__) logger = setup_logging(__name__)
@ -20,7 +20,7 @@ class Stream(BaseParser):
""" """
def __init__(self, table_area=None, columns=None, ytol=2, mtol=0, def __init__(self, table_area=None, columns=None, ytol=2, mtol=0,
margins=(1.0, 0.5, 0.1), split_text=False, flag_size=True, margins=(1.0, 0.5, 0.1), split_text=False, flag_size=True,
debug=False): debug=None):
self.table_area = table_area self.table_area = table_area
self.columns = columns self.columns = columns
self._validate_columns() self._validate_columns()
@ -244,10 +244,11 @@ class Stream(BaseParser):
table = self._generate_table(table_idx, cols, rows) table = self._generate_table(table_idx, cols, rows)
_tables.append(table) _tables.append(table)
if self.debug: if self.debug is not None:
text = [] text = []
text.extend([(t.x0, t.y0, t.x1, t.y1) for t in self.horizontal_text]) text.extend([(t.x0, t.y0, t.x1, t.y1) for t in self.horizontal_text])
text.extend([(t.x0, t.y0, t.x1, t.y1) for t in self.vertical_text]) text.extend([(t.x0, t.y0, t.x1, t.y1) for t in self.vertical_text])
self.g.text = text self.g.text = text
self.g.tables = _tables
return _tables, self.g return _tables, self.g

View File

@ -39,71 +39,54 @@ def plot_geometry(filepath, pages='1', mesh=False, geometry_type='text', **kwarg
ax.set_xlim(min(xs) - 10, max(xs) + 10) ax.set_xlim(min(xs) - 10, max(xs) + 10)
ax.set_ylim(min(ys) - 10, max(ys) + 10) ax.set_ylim(min(ys) - 10, max(ys) + 10)
plt.show() plt.show()
elif geometry_type == 'contour':
try:
for img, table_bbox in geometry.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 can only be used with Lattice.")
elif geometry_type == 'joint':
try:
for img, table_bbox in geometry.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 can only be used with Lattice.")
elif geometry_type == 'line':
try:
for v_s, h_s in geometry.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 can only be used with Lattice.")
elif geometry_type == 'table': elif geometry_type == 'table':
try: for tables in geometry.tables:
for tables in geometry.tables: for table in tables:
for table in tables: for row in table.cells:
for r in range(len(table.rows)): for cell in row:
for c in range(len(table.cols)): if cell.left:
if table.cells[r][c].left: plt.plot([cell.lb[0], cell.lt[0]],
plt.plot([table.cells[r][c].lb[0], [cell.lb[1], cell.lt[1]])
table.cells[r][c].lt[0]], if cell.right:
[table.cells[r][c].lb[1], plt.plot([cell.rb[0], cell.rt[0]],
table.cells[r][c].lt[1]]) [cell.rb[1], cell.rt[1]])
if table.cells[r][c].right: if cell.top:
plt.plot([table.cells[r][c].rb[0], plt.plot([cell.lt[0], cell.rt[0]],
table.cells[r][c].rt[0]], [cell.lt[1], cell.rt[1]])
[table.cells[r][c].rb[1], if cell.bottom:
table.cells[r][c].rt[1]]) plt.plot([cell.lb[0], cell.rb[0]],
if table.cells[r][c].top: [cell.lb[1], cell.rb[1]])
plt.plot([table.cells[r][c].lt[0], plt.show()
table.cells[r][c].rt[0]], elif geometry_type == 'contour':
[table.cells[r][c].lt[1], if not mesh:
table.cells[r][c].rt[1]]) raise ValueError("Use mesh=True")
if table.cells[r][c].bottom: for img, table_bbox in geometry.images:
plt.plot([table.cells[r][c].lb[0], for t in table_bbox.keys():
table.cells[r][c].rb[0]], cv2.rectangle(img, (t[0], t[1]),
[table.cells[r][c].lb[1], (t[2], t[3]), (255, 0, 0), 3)
table.cells[r][c].rb[1]]) plt.imshow(img)
plt.show() plt.show()
except AttributeError: elif geometry_type == 'joint':
raise ValueError("This option can only be used with Lattice.") if not mesh:
else: raise ValueError("Use mesh=True")
raise UserWarning("This method can only be called after" for img, table_bbox in geometry.images:
" debug has been specified.") 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()
elif geometry_type == 'line':
if not mesh:
raise ValueError("Use mesh=True")
for v_s, h_s in geometry.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()