import numpy as np import cv2 from data.boxEnlarge import enlargebox class GaussianBuilder(object): def __init__(self, init_size, sigma, enlarge_region, enlarge_affinity): self.init_size = init_size self.sigma = sigma self.enlarge_region = enlarge_region self.enlarge_affinity = enlarge_affinity self.gaussian_map, self.gaussian_map_color = self.generate_gaussian_map() def generate_gaussian_map(self): circle_mask = self.generate_circle_mask() gaussian_map = np.zeros((self.init_size, self.init_size), np.float32) for i in range(self.init_size): for j in range(self.init_size): gaussian_map[i, j] = ( 1 / 2 / np.pi / (self.sigma ** 2) * np.exp( -1 / 2 * ( (i - self.init_size / 2) ** 2 / (self.sigma ** 2) + (j - self.init_size / 2) ** 2 / (self.sigma ** 2) ) ) ) gaussian_map = gaussian_map * circle_mask gaussian_map = (gaussian_map / np.max(gaussian_map)).astype(np.float32) gaussian_map_color = (gaussian_map * 255).astype(np.uint8) gaussian_map_color = cv2.applyColorMap(gaussian_map_color, cv2.COLORMAP_JET) return gaussian_map, gaussian_map_color def generate_circle_mask(self): zero_arr = np.zeros((self.init_size, self.init_size), np.float32) circle_mask = cv2.circle( img=zero_arr, center=(self.init_size // 2, self.init_size // 2), radius=self.init_size // 2, color=1, thickness=-1, ) return circle_mask def four_point_transform(self, bbox): """ Using the bbox, standard 2D gaussian map, returns Transformed 2d Gaussian map """ width, height = ( np.max(bbox[:, 0]).astype(np.int32), np.max(bbox[:, 1]).astype(np.int32), ) init_points = np.array( [ [0, 0], [self.init_size, 0], [self.init_size, self.init_size], [0, self.init_size], ], dtype="float32", ) M = cv2.getPerspectiveTransform(init_points, bbox) warped_gaussian_map = cv2.warpPerspective(self.gaussian_map, M, (width, height)) return warped_gaussian_map, width, height def add_gaussian_map_to_score_map( self, score_map, bbox, enlarge_size, horizontal_text_bool, map_type=None ): """ Mapping 2D Gaussian to the character box coordinates of the score_map. :param score_map: Target map to put 2D gaussian on character box :type score_map: np.float32 :param bbox: character boxes :type bbox: np.float32 :param enlarge_size: Enlarge size of gaussian map to fit character shape :type enlarge_size: list of enlarge size [x dim, y dim] :param horizontal_text_bool: Flag that bbox is horizontal text or not :type horizontal_text_bool: bool :param map_type: Whether map's type is "region" | "affinity" :type map_type: str :return score_map: score map that all 2D gaussian put on character box :rtype: np.float32 """ map_h, map_w = score_map.shape bbox = enlargebox(bbox, map_h, map_w, enlarge_size, horizontal_text_bool) # If any one point of character bbox is out of range, don't put in on map if np.any(bbox < 0) or np.any(bbox[:, 0] > map_w) or np.any(bbox[:, 1] > map_h): return score_map bbox_left, bbox_top = np.array([np.min(bbox[:, 0]), np.min(bbox[:, 1])]).astype( np.int32 ) bbox -= (bbox_left, bbox_top) warped_gaussian_map, width, height = self.four_point_transform( bbox.astype(np.float32) ) try: bbox_area_of_image = score_map[ bbox_top : bbox_top + height, bbox_left : bbox_left + width, ] high_value_score = np.where( warped_gaussian_map > bbox_area_of_image, warped_gaussian_map, bbox_area_of_image, ) score_map[ bbox_top : bbox_top + height, bbox_left : bbox_left + width, ] = high_value_score except Exception as e: print("Error : {}".format(e)) print( "On generating {} map, strange box came out. (width: {}, height: {})".format( map_type, width, height ) ) return score_map def calculate_affinity_box_points(self, bbox_1, bbox_2, vertical=False): center_1, center_2 = np.mean(bbox_1, axis=0), np.mean(bbox_2, axis=0) if vertical: tl = (bbox_1[0] + bbox_1[-1] + center_1) / 3 tr = (bbox_1[1:3].sum(0) + center_1) / 3 br = (bbox_2[1:3].sum(0) + center_2) / 3 bl = (bbox_2[0] + bbox_2[-1] + center_2) / 3 else: tl = (bbox_1[0:2].sum(0) + center_1) / 3 tr = (bbox_2[0:2].sum(0) + center_2) / 3 br = (bbox_2[2:4].sum(0) + center_2) / 3 bl = (bbox_1[2:4].sum(0) + center_1) / 3 affinity_box = np.array([tl, tr, br, bl]).astype(np.float32) return affinity_box def generate_region( self, img_h, img_w, word_level_char_bbox, horizontal_text_bools ): region_map = np.zeros([img_h, img_w], dtype=np.float32) for i in range( len(word_level_char_bbox) ): # shape : [word_num, [char_num_in_one_word, 4, 2]] for j in range(len(word_level_char_bbox[i])): region_map = self.add_gaussian_map_to_score_map( region_map, word_level_char_bbox[i][j].copy(), self.enlarge_region, horizontal_text_bools[i], map_type="region", ) return region_map def generate_affinity( self, img_h, img_w, word_level_char_bbox, horizontal_text_bools ): affinity_map = np.zeros([img_h, img_w], dtype=np.float32) all_affinity_bbox = [] for i in range(len(word_level_char_bbox)): for j in range(len(word_level_char_bbox[i]) - 1): affinity_bbox = self.calculate_affinity_box_points( word_level_char_bbox[i][j], word_level_char_bbox[i][j + 1] ) affinity_map = self.add_gaussian_map_to_score_map( affinity_map, affinity_bbox.copy(), self.enlarge_affinity, horizontal_text_bools[i], map_type="affinity", ) all_affinity_bbox.append(np.expand_dims(affinity_bbox, axis=0)) if len(all_affinity_bbox) > 0: all_affinity_bbox = np.concatenate(all_affinity_bbox, axis=0) return affinity_map, all_affinity_bbox