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')