Fix plot_geometry
Vinayak Mehta
parent
b3f840bba9
commit
e615580e55
|
|
@ -232,7 +232,11 @@ class Lattice(BaseParser):
|
|||
table = self._generate_table(table_idx, cols, rows, v_s=v_s, h_s=h_s)
|
||||
_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.segments = (self.vertical_segments, self.horizontal_segments)
|
||||
self.g.tables = _tables
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ import pandas as pd
|
|||
from .base import BaseParser
|
||||
from ..core import Table
|
||||
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__)
|
||||
|
|
@ -20,7 +20,7 @@ class Stream(BaseParser):
|
|||
"""
|
||||
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,
|
||||
debug=False):
|
||||
debug=None):
|
||||
self.table_area = table_area
|
||||
self.columns = columns
|
||||
self._validate_columns()
|
||||
|
|
@ -244,10 +244,11 @@ class Stream(BaseParser):
|
|||
table = self._generate_table(table_idx, cols, rows)
|
||||
_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.tables = _tables
|
||||
|
||||
return _tables, self.g
|
||||
|
|
@ -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_ylim(min(ys) - 10, max(ys) + 10)
|
||||
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':
|
||||
try:
|
||||
for tables in geometry.tables:
|
||||
for table in tables:
|
||||
for r in range(len(table.rows)):
|
||||
for c in range(len(table.cols)):
|
||||
if table.cells[r][c].left:
|
||||
plt.plot([table.cells[r][c].lb[0],
|
||||
table.cells[r][c].lt[0]],
|
||||
[table.cells[r][c].lb[1],
|
||||
table.cells[r][c].lt[1]])
|
||||
if table.cells[r][c].right:
|
||||
plt.plot([table.cells[r][c].rb[0],
|
||||
table.cells[r][c].rt[0]],
|
||||
[table.cells[r][c].rb[1],
|
||||
table.cells[r][c].rt[1]])
|
||||
if table.cells[r][c].top:
|
||||
plt.plot([table.cells[r][c].lt[0],
|
||||
table.cells[r][c].rt[0]],
|
||||
[table.cells[r][c].lt[1],
|
||||
table.cells[r][c].rt[1]])
|
||||
if table.cells[r][c].bottom:
|
||||
plt.plot([table.cells[r][c].lb[0],
|
||||
table.cells[r][c].rb[0]],
|
||||
[table.cells[r][c].lb[1],
|
||||
table.cells[r][c].rb[1]])
|
||||
plt.show()
|
||||
except AttributeError:
|
||||
raise ValueError("This option can only be used with Lattice.")
|
||||
else:
|
||||
raise UserWarning("This method can only be called after"
|
||||
" debug has been specified.")
|
||||
for tables in geometry.tables:
|
||||
for table in tables:
|
||||
for row in table.cells:
|
||||
for cell in row:
|
||||
if cell.left:
|
||||
plt.plot([cell.lb[0], cell.lt[0]],
|
||||
[cell.lb[1], cell.lt[1]])
|
||||
if cell.right:
|
||||
plt.plot([cell.rb[0], cell.rt[0]],
|
||||
[cell.rb[1], cell.rt[1]])
|
||||
if cell.top:
|
||||
plt.plot([cell.lt[0], cell.rt[0]],
|
||||
[cell.lt[1], cell.rt[1]])
|
||||
if cell.bottom:
|
||||
plt.plot([cell.lb[0], cell.rb[0]],
|
||||
[cell.lb[1], cell.rb[1]])
|
||||
plt.show()
|
||||
elif geometry_type == 'contour':
|
||||
if not mesh:
|
||||
raise ValueError("Use mesh=True")
|
||||
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()
|
||||
elif geometry_type == 'joint':
|
||||
if not mesh:
|
||||
raise ValueError("Use mesh=True")
|
||||
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()
|
||||
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()
|
||||
Loading…
Reference in New Issue