wehub-resource-sync
469 lines
20 KiB
Python
469 lines
20 KiB
Python
"""
|
||
Data structures used by the evaluation process.
|
||
Yu Fang - March 2019
|
||
"""
|
||
|
||
from collections import Iterable
|
||
|
||
import numpy as np
|
||
from shapely.geometry import Polygon
|
||
|
||
|
||
# helper functions
|
||
def flatten(lis):
|
||
for item in lis:
|
||
if isinstance(item, Iterable) and not isinstance(item, str):
|
||
for x in flatten(item):
|
||
yield x
|
||
else:
|
||
yield item
|
||
|
||
# derived from https://blog.csdn.net/u012433049/article/details/82909484
|
||
def compute_poly_iou(list1, list2):
|
||
a1 = np.array(list1, dtype=int).reshape(-1, 2)
|
||
poly1 = Polygon(a1)
|
||
poly1_clean = poly1.buffer(0)
|
||
|
||
a2 = np.array(list2, dtype=int).reshape(-1, 2)
|
||
poly2 = Polygon(a2)
|
||
poly2_clean = poly2.buffer(0)
|
||
|
||
try:
|
||
# iou = poly1.intersection(poly2).area / poly1.union(poly2).area
|
||
iou = poly1_clean.intersection(poly2_clean).area / poly1_clean.union(poly2_clean).area
|
||
except ZeroDivisionError:
|
||
iou = 0
|
||
return iou
|
||
|
||
|
||
class Cell(object):
|
||
# @:param start_row : start row index of the Cell
|
||
# @:param start_col : start column index of the Cell
|
||
# @:param end-row : end row index of the Cell
|
||
# @:param end-col : end column index of the Cell
|
||
# @:param cell_box: bounding-box of the Cell (coordinates are saved as a string)
|
||
# @:param content_box: bounding-box of the text content within Cell (unused variable)
|
||
# @:param cell_id: unique id of the Cell
|
||
|
||
def __init__(self, table_id, start_row, start_col, cell_box, end_row, end_col, content_box=""):
|
||
self._start_row = int(start_row)
|
||
self._start_col = int(start_col)
|
||
self._cell_box = cell_box
|
||
self._content_box = content_box
|
||
self._table_id = table_id # the table_id this cell belongs to
|
||
# self._cell_name = cell_id # specify the cell using passed-in cell_id
|
||
self._cell_id = id(self)
|
||
# self._region = region
|
||
|
||
# check for end-row and end-col special case
|
||
if end_row == -1:
|
||
self._end_row = self.start_row
|
||
else:
|
||
self._end_row = int(end_row)
|
||
if end_col == -1:
|
||
self._end_col = self._start_col
|
||
else:
|
||
self._end_col = int(end_col)
|
||
|
||
@property
|
||
def start_row(self):
|
||
return self._start_row
|
||
|
||
@property
|
||
def start_col(self):
|
||
return self._start_col
|
||
|
||
@property
|
||
def end_row(self):
|
||
return self._end_row
|
||
|
||
@property
|
||
def end_col(self):
|
||
return self._end_col
|
||
|
||
@property
|
||
def cell_box(self):
|
||
return self._cell_box
|
||
|
||
@property
|
||
def content_box(self):
|
||
return self._content_box
|
||
|
||
@property
|
||
def cell_id(self):
|
||
return self._cell_id
|
||
|
||
@property
|
||
def table_id(self):
|
||
return self._table_id
|
||
|
||
def __str__(self):
|
||
return "CELL row=[%d, %d] col=[%d, %d] (coords=%s)" %(self.start_row, self.end_row
|
||
, self.start_col, self.end_col
|
||
, self.cell_box)
|
||
|
||
# return the IoU value of two cell blocks
|
||
def compute_cell_iou(self, another_cell):
|
||
cell_box_1_temp = []
|
||
for el in self.cell_box.split():
|
||
cell_box_1_temp.append((el.split(",")))
|
||
cell_box_1 = list(flatten(cell_box_1_temp))
|
||
cell_box_1 = [int(x) for x in cell_box_1]
|
||
|
||
cell_box_2_temp = []
|
||
for el in another_cell.cell_box.split():
|
||
cell_box_2_temp.append((el.split(",")))
|
||
cell_box_2 = list(flatten(cell_box_2_temp))
|
||
cell_box_2 = [int(x) for x in cell_box_2]
|
||
|
||
return compute_poly_iou(cell_box_1, cell_box_2)
|
||
|
||
# check if the two cell object denotes same cell area in table
|
||
def check_same(self, another_cell):
|
||
return self._start_row == another_cell.start_row and self._end_row == another_cell.end_row and \
|
||
self._start_col == another_cell.start_col and self._end_col == another_cell.end_col
|
||
|
||
|
||
# Note: currently save the relation with two cell object involved,
|
||
# can be replaced by cell_id in follow-up memory clean up
|
||
class AdjRelation:
|
||
|
||
DIR_HORIZ = 1
|
||
DIR_VERT = 2
|
||
|
||
def __init__(self, fromText, toText, direction):
|
||
# @param: fromText, toText are Cell objects （may be changed to cell-ID for further development）
|
||
self._fromText = fromText
|
||
self._toText = toText
|
||
self._direction = direction
|
||
|
||
@property
|
||
def fromText(self):
|
||
return self._fromText
|
||
|
||
@property
|
||
def toText(self):
|
||
return self._toText
|
||
|
||
@property
|
||
def direction(self):
|
||
return self._direction
|
||
|
||
def __str__(self):
|
||
if self.direction == self.DIR_VERT:
|
||
dir = "vertical"
|
||
else:
|
||
dir = "horizontal"
|
||
return 'ADJ_RELATION: ' + str(self._fromText) + ' ' + str(self._toText) + ' ' + dir
|
||
|
||
def isEqual(self, otherRelation):
|
||
return self.fromText.cell_id == otherRelation.fromText.cell_id and \
|
||
self.toText.cell_id == otherRelation.toText.cell_id and self.direction == otherRelation.direction
|
||
|
||
|
||
class Table:
|
||
|
||
def __init__(self, tableNode):
|
||
self._root = tableNode
|
||
self._id = id(self)
|
||
self._table_coords = ""
|
||
self._maxRow = 0 # PS: indexing from 0
|
||
self._maxCol = 0
|
||
self._cells = [] # save a table as list of <Cell>s
|
||
self.adj_relations = [] # save the adj_relations for the table
|
||
self.parsed = False
|
||
self.found = False # check if the find_adj_relations() has been called once
|
||
|
||
self.parse_table()
|
||
|
||
def __str__(self):
|
||
return "TABLE object - {} row x {} col".format(self._maxRow+1, self._maxCol+1)
|
||
|
||
@property
|
||
def id(self):
|
||
return self._id
|
||
|
||
@property
|
||
def table_coords(self):
|
||
return self._table_coords
|
||
|
||
@property
|
||
def table_cells(self):
|
||
return self._cells
|
||
|
||
# parse input xml to cell lists
|
||
def parse_table(self):
|
||
# get the table bbox
|
||
self._table_coords = str(self._root.getElementsByTagName("Coords")[0].getAttribute("points"))
|
||
|
||
# get info for each cell
|
||
cells = self._root.getElementsByTagName("cell")
|
||
max_row = max_col = 0
|
||
for cell in cells:
|
||
sr = cell.getAttribute("start-row")
|
||
sc = cell.getAttribute("start-col")
|
||
cell_id = cell.getAttribute("id")
|
||
b_points = str(cell.getElementsByTagName("Coords")[0].getAttribute("points"))
|
||
# try:
|
||
# try:
|
||
# text = cell.getElementsByTagName("content")[0].firstChild.nodeValue
|
||
# except AttributeError:
|
||
# text = ""
|
||
# except IndexError:
|
||
# text = "initialized cell as no content"
|
||
er = cell.getAttribute("end-row") if cell.hasAttribute("end-row") else -1
|
||
ec = cell.getAttribute("end-col") if cell.hasAttribute("end-col") else -1
|
||
new_cell = Cell(table_id=str(self.id), start_row=sr, start_col=sc, cell_box=b_points,
|
||
end_row=er, end_col=ec)
|
||
max_row = max(max_row, int(sr), int(er))
|
||
max_col = max(max_col, int(sc), int(ec))
|
||
self._cells.append(new_cell)
|
||
self._maxCol = max_col
|
||
self._maxRow = max_row
|
||
self.parsed = True
|
||
|
||
# generate a table-like structure for finding adj_relations
|
||
def convert_2d(self):
|
||
table = [[0 for x in range(self._maxCol+1)] for y in range(self._maxRow+1)] # init blank cell with int 0
|
||
for cell in self._cells:
|
||
cur_row = cell.start_row
|
||
while cur_row <= cell.end_row:
|
||
cur_col = cell.start_col
|
||
while cur_col <= cell.end_col:
|
||
temp = table[cur_row][cur_col]
|
||
if temp == 0:
|
||
table[cur_row][cur_col] = cell
|
||
elif type(temp) == list:
|
||
temp.append(cell)
|
||
table[cur_row][cur_col] = temp
|
||
else:
|
||
table[cur_row][cur_col] = [temp, cell]
|
||
cur_col += 1
|
||
cur_row += 1
|
||
|
||
return table
|
||
|
||
def find_adj_relations(self):
|
||
if self.found:
|
||
return self.adj_relations
|
||
else:
|
||
# if len(self._cells) == 0:
|
||
if self.parsed == False:
|
||
# fix: cases where there's no cell in table?
|
||
print("table is not parsed for further steps.")
|
||
self.parse_table()
|
||
self.find_adj_relations()
|
||
else:
|
||
retVal = []
|
||
tab = self.convert_2d()
|
||
|
||
# find horizontal relations
|
||
for r in range(self._maxRow+1):
|
||
for c_from in range(self._maxCol):
|
||
temp_pos = tab[r][c_from]
|
||
if temp_pos == 0:
|
||
continue
|
||
elif type(temp_pos) == list:
|
||
for cell in temp_pos:
|
||
c_to = c_from + 1
|
||
if tab[r][c_to] != 0:
|
||
# find relation between two adjacent cells
|
||
if type(tab[r][c_to]) == list:
|
||
for cell_to in tab[r][c_to]:
|
||
if cell != cell_to and (not cell.check_same(cell_to)):
|
||
adj_relation = AdjRelation(cell, cell_to, AdjRelation.DIR_HORIZ)
|
||
retVal.append(adj_relation)
|
||
else:
|
||
if cell != tab[r][c_to]:
|
||
adj_relation = AdjRelation(cell, tab[r][c_to], AdjRelation.DIR_HORIZ)
|
||
retVal.append(adj_relation)
|
||
else:
|
||
# find the next non-blank cell, if exists
|
||
for temp in range(c_from + 1, self._maxCol + 1):
|
||
if tab[r][temp] != 0:
|
||
if type(tab[r][temp]) == list:
|
||
for cell_to in tab[r][temp]:
|
||
adj_relation = AdjRelation(cell, cell_to,
|
||
AdjRelation.DIR_HORIZ)
|
||
retVal.append(adj_relation)
|
||
else:
|
||
adj_relation = AdjRelation(cell, tab[r][temp],
|
||
AdjRelation.DIR_HORIZ)
|
||
retVal.append(adj_relation)
|
||
break
|
||
else:
|
||
c_to = c_from + 1
|
||
if tab[r][c_to] != 0:
|
||
# find relation between two adjacent cells
|
||
if type(tab[r][c_to]) == list:
|
||
for cell_to in tab[r][c_to]:
|
||
if temp_pos != cell_to:
|
||
adj_relation = AdjRelation(temp_pos, cell_to, AdjRelation.DIR_HORIZ)
|
||
retVal.append(adj_relation)
|
||
else:
|
||
if temp_pos != tab[r][c_to]:
|
||
adj_relation = AdjRelation(temp_pos, tab[r][c_to], AdjRelation.DIR_HORIZ)
|
||
retVal.append(adj_relation)
|
||
else:
|
||
# find the next non-blank cell, if exists
|
||
for temp in range(c_from + 1, self._maxCol + 1):
|
||
if tab[r][temp] != 0:
|
||
if type(tab[r][temp]) == list:
|
||
for cell_to in tab[r][temp]:
|
||
adj_relation = AdjRelation(temp_pos, cell_to,
|
||
AdjRelation.DIR_HORIZ)
|
||
retVal.append(adj_relation)
|
||
else:
|
||
adj_relation = AdjRelation(temp_pos, tab[r][temp], AdjRelation.DIR_HORIZ)
|
||
retVal.append(adj_relation)
|
||
break
|
||
|
||
# find vertical relations
|
||
for c in range(self._maxCol+1):
|
||
for r_from in range(self._maxRow):
|
||
temp_pos = tab[r_from][c]
|
||
if temp_pos == 0:
|
||
continue
|
||
elif type(temp_pos) == list:
|
||
for cell in temp_pos:
|
||
r_to = r_from + 1
|
||
if tab[r_to][c] != 0:
|
||
# find relation between two adjacent cells
|
||
if type(tab[r_to][c]) == list:
|
||
for cell_to in tab[r_to][c]:
|
||
if cell != cell_to and (not cell.check_same(cell_to)):
|
||
adj_relation = AdjRelation(cell, cell_to, AdjRelation.DIR_VERT)
|
||
retVal.append(adj_relation)
|
||
else:
|
||
if cell != tab[r_to][c]:
|
||
adj_relation = AdjRelation(cell, tab[r_to][c], AdjRelation.DIR_VERT)
|
||
retVal.append(adj_relation)
|
||
else:
|
||
# find the next non-blank cell, if exists
|
||
for temp in range(r_from + 1, self._maxRow + 1):
|
||
if tab[temp][c] != 0:
|
||
if type(tab[temp][c]) == list:
|
||
for cell_to in tab[temp][c]:
|
||
adj_relation = AdjRelation(cell, cell_to,
|
||
AdjRelation.DIR_VERT)
|
||
retVal.append(adj_relation)
|
||
else:
|
||
adj_relation = AdjRelation(cell, tab[temp][c],
|
||
AdjRelation.DIR_VERT)
|
||
retVal.append(adj_relation)
|
||
break
|
||
else:
|
||
r_to = r_from + 1
|
||
if tab[r_to][c] != 0:
|
||
# find relation between two adjacent cells
|
||
if type(tab[r_to][c]) == list:
|
||
for cell_to in tab[r_to][c]:
|
||
if temp_pos != cell_to:
|
||
adj_relation = AdjRelation(temp_pos, cell_to, AdjRelation.DIR_VERT)
|
||
retVal.append(adj_relation)
|
||
else:
|
||
if temp_pos != tab[r_to][c]:
|
||
adj_relation = AdjRelation(temp_pos, tab[r_to][c], AdjRelation.DIR_VERT)
|
||
retVal.append(adj_relation)
|
||
else:
|
||
# find the next non-blank cell, if exists
|
||
for temp in range(r_from + 1, self._maxRow + 1):
|
||
if tab[temp][c] != 0:
|
||
if type(tab[temp][c]) == list:
|
||
for cell_to in tab[temp][c]:
|
||
adj_relation = AdjRelation(temp_pos, cell_to, AdjRelation.DIR_VERT)
|
||
retVal.append(adj_relation)
|
||
else:
|
||
adj_relation = AdjRelation(temp_pos, tab[temp][c], AdjRelation.DIR_VERT)
|
||
retVal.append(adj_relation)
|
||
break
|
||
|
||
# eliminate duplicates
|
||
repeat = True
|
||
while repeat:
|
||
repeat = False
|
||
duplicates = []
|
||
|
||
for ar1 in retVal:
|
||
for ar2 in retVal:
|
||
if ar1 != ar2:
|
||
if ar1.direction == ar2.direction and ar1.fromText == ar2.fromText and\
|
||
ar1.toText == ar2.toText:
|
||
duplicates.append(ar2)
|
||
break
|
||
else:
|
||
continue
|
||
break
|
||
|
||
if len(duplicates) > 0:
|
||
repeat = True
|
||
retVal.remove(duplicates[0])
|
||
|
||
self.found = True
|
||
self.adj_relations = retVal
|
||
return self.adj_relations
|
||
|
||
# compute the IOU of table, pass-in var is another Table object
|
||
def compute_table_iou(self, another_table):
|
||
table_box_1_temp = []
|
||
for el in self.table_coords.split():
|
||
table_box_1_temp.append((el.split(",")))
|
||
table_box_1 = list(flatten(table_box_1_temp))
|
||
table_box_1 = [int(x) for x in table_box_1]
|
||
|
||
table_box_2_temp = []
|
||
for el in another_table.table_coords.split():
|
||
table_box_2_temp.append((el.split(",")))
|
||
table_box_2 = list(flatten(table_box_2_temp))
|
||
table_box_2 = [int(x) for x in table_box_2]
|
||
|
||
return compute_poly_iou(table_box_1, table_box_2)
|
||
|
||
# find the cell mapping of tables as dictionary, pass-in var is another table and the desired IOU value
|
||
def find_cell_mapping(self, target_table, iou_value):
|
||
mapped_cell = [] # store the matches as tuples - (gt, result) mind the order of table when passing in
|
||
for cell_1 in self.table_cells:
|
||
for cell_2 in target_table.table_cells:
|
||
if cell_1.compute_cell_iou(cell_2) >= iou_value:
|
||
mapped_cell.append((cell_1, cell_2))
|
||
break
|
||
ret = dict(mapped_cell)
|
||
# print(ret)
|
||
return ret
|
||
|
||
# to print a table cell mapping
|
||
@classmethod
|
||
def printCellMapping(cls, dMappedCell):
|
||
print("-"*25)
|
||
for cell1, cell2 in dMappedCell.items():
|
||
print(" ", cell1, " --> ", cell2)
|
||
|
||
# to print a table set of adjacency relations
|
||
@classmethod
|
||
def printAdjacencyRelationList(cls, lAdjRel, title=""):
|
||
print("--- %s "%title + "-"*25)
|
||
for adj in lAdjRel:
|
||
print(adj)
|
||
|
||
|
||
class ResultStructure:
|
||
|
||
def __init__(self, truePos, gtTotal, resTotal):
|
||
self._truePos = truePos
|
||
self._gtTotal = gtTotal
|
||
self._resTotal = resTotal
|
||
|
||
@property
|
||
def truePos(self):
|
||
return self._truePos
|
||
|
||
@property
|
||
def gtTotal(self):
|
||
return self._gtTotal
|
||
|
||
@property
|
||
def resTotal(self):
|
||
return self._resTotal
|
||
|
||
def __str__(self):
|
||
return "true: {}, gt: {}, res: {}".format(self._truePos, self._gtTotal, self._resTotal) |