Spaces:
Running
Running
| import numpy as np | |
| import cv2 | |
| class Element: | |
| def __init__(self, id, corner, category, text_content=None): | |
| self.id = id | |
| self.category = category | |
| self.col_min, self.row_min, self.col_max, self.row_max = corner | |
| self.width = self.col_max - self.col_min | |
| self.height = self.row_max - self.row_min | |
| self.area = self.width * self.height | |
| self.text_content = text_content | |
| self.parent_id = None | |
| self.children = [] # list of elements | |
| def init_bound(self): | |
| self.width = self.col_max - self.col_min | |
| self.height = self.row_max - self.row_min | |
| self.area = self.width * self.height | |
| def put_bbox(self): | |
| return self.col_min, self.row_min, self.col_max, self.row_max | |
| def wrap_info(self): | |
| info = {'id':self.id, 'class': self.category, 'height': self.height, 'width': self.width, | |
| 'position': {'column_min': self.col_min, 'row_min': self.row_min, 'column_max': self.col_max, | |
| 'row_max': self.row_max}} | |
| if self.text_content is not None: | |
| info['text_content'] = self.text_content | |
| if len(self.children) > 0: | |
| info['children'] = [] | |
| for child in self.children: | |
| info['children'].append(child.id) | |
| if self.parent_id is not None: | |
| info['parent'] = self.parent_id | |
| return info | |
| def resize(self, resize_ratio): | |
| self.col_min = int(self.col_min * resize_ratio) | |
| self.row_min = int(self.row_min * resize_ratio) | |
| self.col_max = int(self.col_max * resize_ratio) | |
| self.row_max = int(self.row_max * resize_ratio) | |
| self.init_bound() | |
| def element_merge(self, element_b, new_element=False, new_category=None, new_id=None): | |
| col_min_a, row_min_a, col_max_a, row_max_a = self.put_bbox() | |
| col_min_b, row_min_b, col_max_b, row_max_b = element_b.put_bbox() | |
| new_corner = (min(col_min_a, col_min_b), min(row_min_a, row_min_b), max(col_max_a, col_max_b), max(row_max_a, row_max_b)) | |
| if element_b.text_content is not None: | |
| self.text_content = element_b.text_content if self.text_content is None else self.text_content + '\n' + element_b.text_content | |
| if new_element: | |
| return Element(new_id, new_corner, new_category) | |
| else: | |
| self.col_min, self.row_min, self.col_max, self.row_max = new_corner | |
| self.init_bound() | |
| def calc_intersection_area(self, element_b, bias=(0, 0)): | |
| a = self.put_bbox() | |
| b = element_b.put_bbox() | |
| col_min_s = max(a[0], b[0]) - bias[0] | |
| row_min_s = max(a[1], b[1]) - bias[1] | |
| col_max_s = min(a[2], b[2]) | |
| row_max_s = min(a[3], b[3]) | |
| w = np.maximum(0, col_max_s - col_min_s) | |
| h = np.maximum(0, row_max_s - row_min_s) | |
| inter = w * h | |
| iou = inter / (self.area + element_b.area - inter) | |
| ioa = inter / self.area | |
| iob = inter / element_b.area | |
| return inter, iou, ioa, iob | |
| def element_relation(self, element_b, bias=(0, 0)): | |
| """ | |
| @bias: (horizontal bias, vertical bias) | |
| :return: -1 : a in b | |
| 0 : a, b are not intersected | |
| 1 : b in a | |
| 2 : a, b are identical or intersected | |
| """ | |
| inter, iou, ioa, iob = self.calc_intersection_area(element_b, bias) | |
| # area of intersection is 0 | |
| if ioa == 0: | |
| return 0 | |
| # a in b | |
| if ioa >= 1: | |
| return -1 | |
| # b in a | |
| if iob >= 1: | |
| return 1 | |
| return 2 | |
| def visualize_element(self, img, color=(0, 255, 0), line=1, show=False): | |
| loc = self.put_bbox() | |
| cv2.rectangle(img, loc[:2], loc[2:], color, line) | |
| # for child in self.children: | |
| # child.visualize_element(img, color=(255, 0, 255), line=line) | |
| if show: | |
| cv2.imshow('element', img) | |
| cv2.waitKey(0) | |
| cv2.destroyWindow('element') | |