Move Pdf class to core as FileHandler
parent
0c9e21d881
commit
5d29f0c21c
222
camelot/core.py
222
camelot/core.py
|
|
@ -1,69 +1,14 @@
|
||||||
|
import os
|
||||||
|
import tempfile
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import pandas as pd
|
import pandas as pd
|
||||||
|
from PyPDF2 import PdfFileReader, PdfFileWriter
|
||||||
|
|
||||||
|
from .utils import get_page_layout, get_text_objects, get_rotation
|
||||||
|
|
||||||
|
|
||||||
class Cell(object):
|
class Cell(object):
|
||||||
"""Cell.
|
|
||||||
Defines a cell object with coordinates relative to a left-bottom
|
|
||||||
origin, which is also PDFMiner's coordinate space.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
x1 : float
|
|
||||||
x-coordinate of left-bottom point.
|
|
||||||
|
|
||||||
y1 : float
|
|
||||||
y-coordinate of left-bottom point.
|
|
||||||
|
|
||||||
x2 : float
|
|
||||||
x-coordinate of right-top point.
|
|
||||||
|
|
||||||
y2 : float
|
|
||||||
y-coordinate of right-top point.
|
|
||||||
|
|
||||||
Attributes
|
|
||||||
----------
|
|
||||||
lb : tuple
|
|
||||||
Tuple representing left-bottom coordinates.
|
|
||||||
|
|
||||||
lt : tuple
|
|
||||||
Tuple representing left-top coordinates.
|
|
||||||
|
|
||||||
rb : tuple
|
|
||||||
Tuple representing right-bottom coordinates.
|
|
||||||
|
|
||||||
rt : tuple
|
|
||||||
Tuple representing right-top coordinates.
|
|
||||||
|
|
||||||
bbox : tuple
|
|
||||||
Tuple representing the cell's bounding box using the
|
|
||||||
lower-bottom and right-top coordinates.
|
|
||||||
|
|
||||||
left : bool
|
|
||||||
Whether or not cell is bounded on the left.
|
|
||||||
|
|
||||||
right : bool
|
|
||||||
Whether or not cell is bounded on the right.
|
|
||||||
|
|
||||||
top : bool
|
|
||||||
Whether or not cell is bounded on the top.
|
|
||||||
|
|
||||||
bottom : bool
|
|
||||||
Whether or not cell is bounded on the bottom.
|
|
||||||
|
|
||||||
text_objects : list
|
|
||||||
List of text objects assigned to cell.
|
|
||||||
|
|
||||||
text : string
|
|
||||||
Text assigned to cell.
|
|
||||||
|
|
||||||
spanning_h : bool
|
|
||||||
Whether or not cell spans/extends horizontally.
|
|
||||||
|
|
||||||
spanning_v : bool
|
|
||||||
Whether or not cell spans/extends vertically.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, x1, y1, x2, y2):
|
def __init__(self, x1, y1, x2, y2):
|
||||||
|
|
||||||
self.x1 = x1
|
self.x1 = x1
|
||||||
|
|
@ -86,76 +31,23 @@ class Cell(object):
|
||||||
self.image = None
|
self.image = None
|
||||||
|
|
||||||
def add_text(self, text):
|
def add_text(self, text):
|
||||||
"""Adds text to cell.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
text : string
|
|
||||||
"""
|
|
||||||
self.text = ''.join([self.text, text])
|
self.text = ''.join([self.text, text])
|
||||||
|
|
||||||
def get_text(self):
|
def get_text(self):
|
||||||
"""Returns text assigned to cell.
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
text : string
|
|
||||||
"""
|
|
||||||
return self.text
|
return self.text
|
||||||
|
|
||||||
def add_object(self, t_object):
|
def add_object(self, t_object):
|
||||||
"""Adds PDFMiner text object to cell.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
t_object : object
|
|
||||||
"""
|
|
||||||
self.text_objects.append(t_object)
|
self.text_objects.append(t_object)
|
||||||
|
|
||||||
def get_objects(self):
|
def get_objects(self):
|
||||||
"""Returns list of text objects assigned to cell.
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
text_objects : list
|
|
||||||
"""
|
|
||||||
return self.text_objects
|
return self.text_objects
|
||||||
|
|
||||||
def get_bounded_edges(self):
|
def get_bounded_edges(self):
|
||||||
"""Returns the number of edges by which a cell is bounded.
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
bounded_edges : int
|
|
||||||
"""
|
|
||||||
self.bounded_edges = self.top + self.bottom + self.left + self.right
|
self.bounded_edges = self.top + self.bottom + self.left + self.right
|
||||||
return self.bounded_edges
|
return self.bounded_edges
|
||||||
|
|
||||||
|
|
||||||
class Table(object):
|
class Table(object):
|
||||||
"""Table.
|
|
||||||
Defines a table object with coordinates relative to a left-bottom
|
|
||||||
origin, which is also PDFMiner's coordinate space.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
cols : list
|
|
||||||
List of tuples representing column x-coordinates in increasing
|
|
||||||
order.
|
|
||||||
|
|
||||||
rows : list
|
|
||||||
List of tuples representing row y-coordinates in decreasing
|
|
||||||
order.
|
|
||||||
|
|
||||||
Attributes
|
|
||||||
----------
|
|
||||||
cells : list
|
|
||||||
List of cell objects with row-major ordering.
|
|
||||||
|
|
||||||
nocont_ : int
|
|
||||||
Number of lines that did not contribute to setting cell edges.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, cols, rows):
|
def __init__(self, cols, rows):
|
||||||
|
|
||||||
self.cols = cols
|
self.cols = cols
|
||||||
|
|
@ -166,8 +58,6 @@ class Table(object):
|
||||||
self.image = None
|
self.image = None
|
||||||
|
|
||||||
def set_all_edges(self):
|
def set_all_edges(self):
|
||||||
"""Sets all table edges to True.
|
|
||||||
"""
|
|
||||||
for r in range(len(self.rows)):
|
for r in range(len(self.rows)):
|
||||||
for c in range(len(self.cols)):
|
for c in range(len(self.cols)):
|
||||||
self.cells[r][c].left = True
|
self.cells[r][c].left = True
|
||||||
|
|
@ -177,8 +67,6 @@ class Table(object):
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def set_border_edges(self):
|
def set_border_edges(self):
|
||||||
"""Sets table border edges to True.
|
|
||||||
"""
|
|
||||||
for r in range(len(self.rows)):
|
for r in range(len(self.rows)):
|
||||||
self.cells[r][0].left = True
|
self.cells[r][0].left = True
|
||||||
self.cells[r][len(self.cols) - 1].right = True
|
self.cells[r][len(self.cols) - 1].right = True
|
||||||
|
|
@ -188,19 +76,6 @@ class Table(object):
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def set_edges(self, vertical, horizontal, jtol=2):
|
def set_edges(self, vertical, horizontal, jtol=2):
|
||||||
"""Sets a cell's edges to True depending on whether they
|
|
||||||
overlap with lines found by imgproc.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
vertical : list
|
|
||||||
List of vertical lines detected by imgproc. Coordinates
|
|
||||||
scaled and translated to the PDFMiner's coordinate space.
|
|
||||||
|
|
||||||
horizontal : list
|
|
||||||
List of horizontal lines detected by imgproc. Coordinates
|
|
||||||
scaled and translated to the PDFMiner's coordinate space.
|
|
||||||
"""
|
|
||||||
for v in vertical:
|
for v in vertical:
|
||||||
# find closest x coord
|
# find closest x coord
|
||||||
# iterate over y coords and find closest points
|
# iterate over y coords and find closest points
|
||||||
|
|
@ -308,10 +183,6 @@ class Table(object):
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def set_spanning(self):
|
def set_spanning(self):
|
||||||
"""Sets a cell's spanning_h or spanning_v attribute to True
|
|
||||||
depending on whether the cell spans/extends horizontally or
|
|
||||||
vertically.
|
|
||||||
"""
|
|
||||||
for r in range(len(self.rows)):
|
for r in range(len(self.rows)):
|
||||||
for c in range(len(self.cols)):
|
for c in range(len(self.cols)):
|
||||||
bound = self.cells[r][c].get_bounded_edges()
|
bound = self.cells[r][c].get_bounded_edges()
|
||||||
|
|
@ -351,13 +222,6 @@ class Table(object):
|
||||||
return self
|
return self
|
||||||
|
|
||||||
def get_list(self):
|
def get_list(self):
|
||||||
"""Returns a two-dimensional list of text assigned to each
|
|
||||||
cell.
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
ar : list
|
|
||||||
"""
|
|
||||||
ar = []
|
ar = []
|
||||||
for r in range(len(self.rows)):
|
for r in range(len(self.rows)):
|
||||||
ar.append([self.cells[r][c].get_text().strip()
|
ar.append([self.cells[r][c].get_text().strip()
|
||||||
|
|
@ -366,4 +230,76 @@ class Table(object):
|
||||||
|
|
||||||
|
|
||||||
class TableSet(object):
|
class TableSet(object):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
class FileHandler(object):
|
||||||
|
def __init__(self, filename, pages='1'):
|
||||||
|
self.filename = filename
|
||||||
|
if not self.filename.endswith('.pdf'):
|
||||||
|
raise TypeError("File format not supported.")
|
||||||
|
self.pages = __get_pages(pages)
|
||||||
|
self.temp = tempfile.mkdtemp()
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def __get_pages(filename, pages):
|
||||||
|
p = {}
|
||||||
|
if pages == '1':
|
||||||
|
p.append({'start': 1, 'end': 1})
|
||||||
|
else:
|
||||||
|
infile = PdfFileReader(open(filename, 'rb'), strict=False)
|
||||||
|
if pages == 'all':
|
||||||
|
p.append({'start': 1, 'end': infile.getNumPages()})
|
||||||
|
else:
|
||||||
|
for r in pages.split(','):
|
||||||
|
if '-' in r:
|
||||||
|
a, b = r.split('-')
|
||||||
|
if b == 'end':
|
||||||
|
b = infile.getNumPages()
|
||||||
|
p.append({'start': int(a), 'end': int(b)})
|
||||||
|
else:
|
||||||
|
p.append({'start': int(r), 'end': int(r)})
|
||||||
|
return p
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def __save_page(filename, page, temp):
|
||||||
|
with open(filename, 'rb') as fileobj:
|
||||||
|
infile = PdfFileReader(fileobj, strict=False)
|
||||||
|
fpath = os.path.join(temp, 'page-{0}.pdf'.format(page))
|
||||||
|
fname, fext = os.path.splitext(fpath)
|
||||||
|
p = infile.getPage(page - 1)
|
||||||
|
outfile = PdfFileWriter()
|
||||||
|
outfile.addPage(p)
|
||||||
|
with open(fpath, 'wb') as f:
|
||||||
|
outfile.write(f)
|
||||||
|
layout, dim = get_page_layout(fpath)
|
||||||
|
# fix rotated pdf
|
||||||
|
lttextlh = get_text_objects(layout, ltype="lh")
|
||||||
|
lttextlv = get_text_objects(layout, ltype="lv")
|
||||||
|
ltchar = get_text_objects(layout, ltype="char")
|
||||||
|
rotation = get_rotation(lttextlh, lttextlv, ltchar)
|
||||||
|
if rotation != '':
|
||||||
|
fpath_new = ''.join([fname.replace('page', 'p'), '_rotated', fext])
|
||||||
|
os.rename(fpath, fpath_new)
|
||||||
|
infile = PdfFileReader(open(fpath_new, 'rb'), strict=False)
|
||||||
|
outfile = PdfFileWriter()
|
||||||
|
p = infile.getPage(0)
|
||||||
|
if rotation == 'left':
|
||||||
|
p.rotateClockwise(90)
|
||||||
|
elif rotation == 'right':
|
||||||
|
p.rotateCounterClockwise(90)
|
||||||
|
outfile.addPage(p)
|
||||||
|
with open(fpath, 'wb') as f:
|
||||||
|
outfile.write(f)
|
||||||
|
|
||||||
|
def parse(self):
|
||||||
|
for p in self.pages:
|
||||||
|
__save_page(self.filename, p, self.temp)
|
||||||
|
pages = [os.path.join(self.temp, 'page-{0}.pdf'.format(p))
|
||||||
|
for p in self.pagenos]
|
||||||
|
tables = {}
|
||||||
|
for p in pages:
|
||||||
|
table = self.parser.get_tables(p)
|
||||||
|
if table is not None:
|
||||||
|
tables.update(table)
|
||||||
|
return tables
|
||||||
268
camelot/pdf.py
268
camelot/pdf.py
|
|
@ -1,268 +0,0 @@
|
||||||
import os
|
|
||||||
import shutil
|
|
||||||
import tempfile
|
|
||||||
import itertools
|
|
||||||
import multiprocessing as mp
|
|
||||||
from functools import partial
|
|
||||||
|
|
||||||
import cv2
|
|
||||||
from PyPDF2 import PdfFileReader, PdfFileWriter
|
|
||||||
|
|
||||||
from .utils import get_page_layout, get_text_objects, get_rotation
|
|
||||||
|
|
||||||
|
|
||||||
__all__ = ['Pdf']
|
|
||||||
|
|
||||||
|
|
||||||
def _parse_page_numbers(pagenos):
|
|
||||||
"""Converts list of dicts to list of ints.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
pagenos : list
|
|
||||||
List of dicts representing page ranges. A dict must have only
|
|
||||||
two keys named 'start' and 'end' having int as their value.
|
|
||||||
|
|
||||||
Returns
|
|
||||||
-------
|
|
||||||
page_numbers : list
|
|
||||||
List of int page numbers.
|
|
||||||
"""
|
|
||||||
page_numbers = []
|
|
||||||
for p in pagenos:
|
|
||||||
page_numbers.extend(range(p['start'], p['end'] + 1))
|
|
||||||
page_numbers = sorted(set(page_numbers))
|
|
||||||
return page_numbers
|
|
||||||
|
|
||||||
|
|
||||||
def _save_page(temp, pdfname, pageno):
|
|
||||||
with open(pdfname, 'rb') as pdffile:
|
|
||||||
infile = PdfFileReader(pdffile, strict=False)
|
|
||||||
sp_path = os.path.join(temp, 'page-{0}.pdf'.format(pageno))
|
|
||||||
sp_name, sp_ext = os.path.splitext(sp_path)
|
|
||||||
page = infile.getPage(pageno - 1)
|
|
||||||
outfile = PdfFileWriter()
|
|
||||||
outfile.addPage(page)
|
|
||||||
with open(sp_path, 'wb') as f:
|
|
||||||
outfile.write(f)
|
|
||||||
layout, dim = get_page_layout(sp_path)
|
|
||||||
lttextlh = get_text_objects(layout, ltype="lh")
|
|
||||||
lttextlv = get_text_objects(layout, ltype="lv")
|
|
||||||
ltchar = get_text_objects(layout, ltype="char")
|
|
||||||
rotation = get_rotation(lttextlh, lttextlv, ltchar)
|
|
||||||
if rotation != '':
|
|
||||||
sp_new_path = ''.join([sp_name.replace('page', 'p'), '_rotated', sp_ext])
|
|
||||||
os.rename(sp_path, sp_new_path)
|
|
||||||
sp_in = PdfFileReader(open(sp_new_path, 'rb'),
|
|
||||||
strict=False)
|
|
||||||
sp_out = PdfFileWriter()
|
|
||||||
sp_page = sp_in.getPage(0)
|
|
||||||
if rotation == 'left':
|
|
||||||
sp_page.rotateClockwise(90)
|
|
||||||
elif rotation == 'right':
|
|
||||||
sp_page.rotateCounterClockwise(90)
|
|
||||||
sp_out.addPage(sp_page)
|
|
||||||
with open(sp_path, 'wb') as pdf_out:
|
|
||||||
sp_out.write(pdf_out)
|
|
||||||
|
|
||||||
|
|
||||||
class Pdf:
|
|
||||||
"""Pdf manager.
|
|
||||||
Handles all operations like temp directory creation, splitting file
|
|
||||||
into single page pdfs, running extraction using multiple processes
|
|
||||||
and removing the temp directory.
|
|
||||||
|
|
||||||
Parameters
|
|
||||||
----------
|
|
||||||
extractor : object
|
|
||||||
camelot.stream.Stream or camelot.lattice.Lattice extractor
|
|
||||||
object.
|
|
||||||
|
|
||||||
pdfname : string
|
|
||||||
Path to pdf file.
|
|
||||||
|
|
||||||
pagenos : list
|
|
||||||
List of dicts representing page ranges. A dict must have only
|
|
||||||
two keys named 'start' and 'end' having int as their value.
|
|
||||||
(optional, default: [{'start': 1, 'end': 1}])
|
|
||||||
|
|
||||||
parallel : bool
|
|
||||||
Whether or not to run using multiple processes.
|
|
||||||
(optional, default: False)
|
|
||||||
|
|
||||||
clean : bool
|
|
||||||
Whether or not to remove the temp directory.
|
|
||||||
(optional, default: 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("File format not supported.")
|
|
||||||
self.pagenos = _parse_page_numbers(pagenos)
|
|
||||||
self.parallel = parallel
|
|
||||||
if self.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 file into single page pdfs.
|
|
||||||
"""
|
|
||||||
if self.parallel:
|
|
||||||
pfunc = partial(_save_page, self.temp, self.pdfname)
|
|
||||||
self.pool.map(pfunc, self.pagenos)
|
|
||||||
else:
|
|
||||||
for p in self.pagenos:
|
|
||||||
_save_page(self.temp, self.pdfname, p)
|
|
||||||
|
|
||||||
|
|
||||||
def extract(self):
|
|
||||||
"""Runs table extraction by calling extractor.get_tables
|
|
||||||
on all single page pdfs.
|
|
||||||
"""
|
|
||||||
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 in ['lattice', 'ocrl']:
|
|
||||||
self.debug = self.extractor.debug
|
|
||||||
self.debug_images = []
|
|
||||||
self.debug_segments = []
|
|
||||||
self.debug_tables = []
|
|
||||||
elif self.extractor.method == 'ocrs':
|
|
||||||
self.debug = self.extractor.debug
|
|
||||||
self.debug_images = []
|
|
||||||
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 in ['lattice', 'ocr']:
|
|
||||||
self.debug_images.append(self.extractor.debug_images)
|
|
||||||
self.debug_segments.append(self.extractor.debug_segments)
|
|
||||||
self.debug_tables.append(self.extractor.debug_tables)
|
|
||||||
elif self.extractor.method == 'ocrs':
|
|
||||||
self.debug_images.append(self.extractor.debug_images)
|
|
||||||
if self.clean:
|
|
||||||
self.remove_tempdir()
|
|
||||||
return tables
|
|
||||||
|
|
||||||
def remove_tempdir(self):
|
|
||||||
"""Removes temporary directory that was created to save single
|
|
||||||
page pdfs and their images.
|
|
||||||
"""
|
|
||||||
shutil.rmtree(self.temp)
|
|
||||||
|
|
||||||
def debug_plot(self):
|
|
||||||
"""Generates a matplotlib plot based on the selected extractor
|
|
||||||
debug option.
|
|
||||||
"""
|
|
||||||
import matplotlib.pyplot as plt
|
|
||||||
import matplotlib.patches as patches
|
|
||||||
|
|
||||||
if self.debug is True:
|
|
||||||
if hasattr(self, 'debug_text'):
|
|
||||||
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()
|
|
||||||
elif hasattr(self, 'debug_images'):
|
|
||||||
for img in self.debug_images:
|
|
||||||
plt.imshow(img)
|
|
||||||
plt.show()
|
|
||||||
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 can 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 can 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 can only be used with Lattice.")
|
|
||||||
elif self.debug == 'table':
|
|
||||||
try:
|
|
||||||
for tables in self.debug_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.")
|
|
||||||
Loading…
Reference in New Issue