From e615580e5541005e17bb8008d1d6a70ffa50d789 Mon Sep 17 00:00:00 2001 From: Vinayak Mehta Date: Fri, 7 Sep 2018 06:25:13 +0530 Subject: [PATCH] Fix plot_geometry --- camelot/parsers/lattice.py | 6 +- camelot/parsers/stream.py | 7 ++- camelot/plotting.py | 117 ++++++++++++++++--------------------- 3 files changed, 59 insertions(+), 71 deletions(-) diff --git a/camelot/parsers/lattice.py b/camelot/parsers/lattice.py index 727e00b..bf4fca3 100644 --- a/camelot/parsers/lattice.py +++ b/camelot/parsers/lattice.py @@ -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 diff --git a/camelot/parsers/stream.py b/camelot/parsers/stream.py index 9d12384..1976505 100644 --- a/camelot/parsers/stream.py +++ b/camelot/parsers/stream.py @@ -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 \ No newline at end of file diff --git a/camelot/plotting.py b/camelot/plotting.py index 793764e..7d4e4d3 100644 --- a/camelot/plotting.py +++ b/camelot/plotting.py @@ -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.") \ No newline at end of file + 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() \ No newline at end of file