Spaces:
Running
Running
| import cv2 | |
| import numpy as np | |
| class Text: | |
| def __init__(self, id, content, location): | |
| self.id = id | |
| self.content = content | |
| self.location = location | |
| self.width = self.location['right'] - self.location['left'] | |
| self.height = self.location['bottom'] - self.location['top'] | |
| self.area = self.width * self.height | |
| self.word_width = self.width / len(self.content) | |
| ''' | |
| ******************************** | |
| *** Relation with Other text *** | |
| ******************************** | |
| ''' | |
| def is_justified(self, ele_b, direction='h', max_bias_justify=4): | |
| ''' | |
| Check if the element is justified | |
| :param max_bias_justify: maximum bias if two elements to be justified | |
| :param direction: | |
| - 'v': vertical up-down connection | |
| - 'h': horizontal left-right connection | |
| ''' | |
| l_a = self.location | |
| l_b = ele_b.location | |
| # connected vertically - up and below | |
| if direction == 'v': | |
| # left and right should be justified | |
| if abs(l_a['left'] - l_b['left']) < max_bias_justify and abs(l_a['right'] - l_b['right']) < max_bias_justify: | |
| return True | |
| return False | |
| elif direction == 'h': | |
| # top and bottom should be justified | |
| if abs(l_a['top'] - l_b['top']) < max_bias_justify and abs(l_a['bottom'] - l_b['bottom']) < max_bias_justify: | |
| return True | |
| return False | |
| def is_on_same_line(self, text_b, direction='h', bias_gap=4, bias_justify=4): | |
| ''' | |
| Check if the element is on the same row(direction='h') or column(direction='v') with ele_b | |
| :param direction: | |
| - 'v': vertical up-down connection | |
| - 'h': horizontal left-right connection | |
| :return: | |
| ''' | |
| l_a = self.location | |
| l_b = text_b.location | |
| # connected vertically - up and below | |
| if direction == 'v': | |
| # left and right should be justified | |
| if self.is_justified(text_b, direction='v', max_bias_justify=bias_justify): | |
| # top and bottom should be connected (small gap) | |
| if abs(l_a['bottom'] - l_b['top']) < bias_gap or abs(l_a['top'] - l_b['bottom']) < bias_gap: | |
| return True | |
| return False | |
| elif direction == 'h': | |
| # top and bottom should be justified | |
| if self.is_justified(text_b, direction='h', max_bias_justify=bias_justify): | |
| # top and bottom should be connected (small gap) | |
| if abs(l_a['right'] - l_b['left']) < bias_gap or abs(l_a['left'] - l_b['right']) < bias_gap: | |
| return True | |
| return False | |
| def is_intersected(self, text_b, bias): | |
| l_a = self.location | |
| l_b = text_b.location | |
| left_in = max(l_a['left'], l_b['left']) + bias | |
| top_in = max(l_a['top'], l_b['top']) + bias | |
| right_in = min(l_a['right'], l_b['right']) | |
| bottom_in = min(l_a['bottom'], l_b['bottom']) | |
| w_in = max(0, right_in - left_in) | |
| h_in = max(0, bottom_in - top_in) | |
| area_in = w_in * h_in | |
| if area_in > 0: | |
| return True | |
| ''' | |
| *********************** | |
| *** Revise the Text *** | |
| *********************** | |
| ''' | |
| def merge_text(self, text_b): | |
| text_a = self | |
| top = min(text_a.location['top'], text_b.location['top']) | |
| left = min(text_a.location['left'], text_b.location['left']) | |
| right = max(text_a.location['right'], text_b.location['right']) | |
| bottom = max(text_a.location['bottom'], text_b.location['bottom']) | |
| self.location = {'left': left, 'top': top, 'right': right, 'bottom': bottom} | |
| self.width = self.location['right'] - self.location['left'] | |
| self.height = self.location['bottom'] - self.location['top'] | |
| self.area = self.width * self.height | |
| left_element = text_a | |
| right_element = text_b | |
| if text_a.location['left'] > text_b.location['left']: | |
| left_element = text_b | |
| right_element = text_a | |
| self.content = left_element.content + ' ' + right_element.content | |
| self.word_width = self.width / len(self.content) | |
| def shrink_bound(self, binary_map): | |
| bin_clip = binary_map[self.location['top']:self.location['bottom'], self.location['left']:self.location['right']] | |
| height, width = np.shape(bin_clip) | |
| shrink_top = 0 | |
| shrink_bottom = 0 | |
| for i in range(height): | |
| # top | |
| if shrink_top == 0: | |
| if sum(bin_clip[i]) == 0: | |
| shrink_top = 1 | |
| else: | |
| shrink_top = -1 | |
| elif shrink_top == 1: | |
| if sum(bin_clip[i]) != 0: | |
| self.location['top'] += i | |
| shrink_top = -1 | |
| # bottom | |
| if shrink_bottom == 0: | |
| if sum(bin_clip[height-i-1]) == 0: | |
| shrink_bottom = 1 | |
| else: | |
| shrink_bottom = -1 | |
| elif shrink_bottom == 1: | |
| if sum(bin_clip[height-i-1]) != 0: | |
| self.location['bottom'] -= i | |
| shrink_bottom = -1 | |
| if shrink_top == -1 and shrink_bottom == -1: | |
| break | |
| shrink_left = 0 | |
| shrink_right = 0 | |
| for j in range(width): | |
| # left | |
| if shrink_left == 0: | |
| if sum(bin_clip[:, j]) == 0: | |
| shrink_left = 1 | |
| else: | |
| shrink_left = -1 | |
| elif shrink_left == 1: | |
| if sum(bin_clip[:, j]) != 0: | |
| self.location['left'] += j | |
| shrink_left = -1 | |
| # right | |
| if shrink_right == 0: | |
| if sum(bin_clip[:, width-j-1]) == 0: | |
| shrink_right = 1 | |
| else: | |
| shrink_right = -1 | |
| elif shrink_right == 1: | |
| if sum(bin_clip[:, width-j-1]) != 0: | |
| self.location['right'] -= j | |
| shrink_right = -1 | |
| if shrink_left == -1 and shrink_right == -1: | |
| break | |
| self.width = self.location['right'] - self.location['left'] | |
| self.height = self.location['bottom'] - self.location['top'] | |
| self.area = self.width * self.height | |
| self.word_width = self.width / len(self.content) | |
| ''' | |
| ********************* | |
| *** Visualization *** | |
| ********************* | |
| ''' | |
| def visualize_element(self, img, color=(0, 0, 255), line=1, show=False): | |
| loc = self.location | |
| cv2.rectangle(img, (loc['left'], loc['top']), (loc['right'], loc['bottom']), color, line) | |
| if show: | |
| print(self.content) | |
| cv2.imshow('text', img) | |
| cv2.waitKey() | |
| cv2.destroyWindow('text') | |