Spaces:
Runtime error
Runtime error
| from __future__ import absolute_import, division, print_function | |
| import os | |
| import sys | |
| __dir__ = os.path.dirname(__file__) | |
| sys.path.append(__dir__) | |
| sys.path.append(os.path.join(__dir__, "..")) | |
| import time | |
| import cv2 | |
| import numpy as np | |
| import paddle | |
| from .locality_aware_nms import nms_locality | |
| class SASTPostProcess(object): | |
| """ | |
| The post process for SAST. | |
| """ | |
| def __init__( | |
| self, | |
| score_thresh=0.5, | |
| nms_thresh=0.2, | |
| sample_pts_num=2, | |
| shrink_ratio_of_width=0.3, | |
| expand_scale=1.0, | |
| tcl_map_thresh=0.5, | |
| **kwargs | |
| ): | |
| self.score_thresh = score_thresh | |
| self.nms_thresh = nms_thresh | |
| self.sample_pts_num = sample_pts_num | |
| self.shrink_ratio_of_width = shrink_ratio_of_width | |
| self.expand_scale = expand_scale | |
| self.tcl_map_thresh = tcl_map_thresh | |
| # c++ la-nms is faster, but only support python 3.5 | |
| self.is_python35 = False | |
| if sys.version_info.major == 3 and sys.version_info.minor == 5: | |
| self.is_python35 = True | |
| def point_pair2poly(self, point_pair_list): | |
| """ | |
| Transfer vertical point_pairs into poly point in clockwise. | |
| """ | |
| # constract poly | |
| point_num = len(point_pair_list) * 2 | |
| point_list = [0] * point_num | |
| for idx, point_pair in enumerate(point_pair_list): | |
| point_list[idx] = point_pair[0] | |
| point_list[point_num - 1 - idx] = point_pair[1] | |
| return np.array(point_list).reshape(-1, 2) | |
| def shrink_quad_along_width(self, quad, begin_width_ratio=0.0, end_width_ratio=1.0): | |
| """ | |
| Generate shrink_quad_along_width. | |
| """ | |
| ratio_pair = np.array( | |
| [[begin_width_ratio], [end_width_ratio]], dtype=np.float32 | |
| ) | |
| p0_1 = quad[0] + (quad[1] - quad[0]) * ratio_pair | |
| p3_2 = quad[3] + (quad[2] - quad[3]) * ratio_pair | |
| return np.array([p0_1[0], p0_1[1], p3_2[1], p3_2[0]]) | |
| def expand_poly_along_width(self, poly, shrink_ratio_of_width=0.3): | |
| """ | |
| expand poly along width. | |
| """ | |
| point_num = poly.shape[0] | |
| left_quad = np.array([poly[0], poly[1], poly[-2], poly[-1]], dtype=np.float32) | |
| left_ratio = ( | |
| -shrink_ratio_of_width | |
| * np.linalg.norm(left_quad[0] - left_quad[3]) | |
| / (np.linalg.norm(left_quad[0] - left_quad[1]) + 1e-6) | |
| ) | |
| left_quad_expand = self.shrink_quad_along_width(left_quad, left_ratio, 1.0) | |
| right_quad = np.array( | |
| [ | |
| poly[point_num // 2 - 2], | |
| poly[point_num // 2 - 1], | |
| poly[point_num // 2], | |
| poly[point_num // 2 + 1], | |
| ], | |
| dtype=np.float32, | |
| ) | |
| right_ratio = 1.0 + shrink_ratio_of_width * np.linalg.norm( | |
| right_quad[0] - right_quad[3] | |
| ) / (np.linalg.norm(right_quad[0] - right_quad[1]) + 1e-6) | |
| right_quad_expand = self.shrink_quad_along_width(right_quad, 0.0, right_ratio) | |
| poly[0] = left_quad_expand[0] | |
| poly[-1] = left_quad_expand[-1] | |
| poly[point_num // 2 - 1] = right_quad_expand[1] | |
| poly[point_num // 2] = right_quad_expand[2] | |
| return poly | |
| def restore_quad(self, tcl_map, tcl_map_thresh, tvo_map): | |
| """Restore quad.""" | |
| xy_text = np.argwhere(tcl_map[:, :, 0] > tcl_map_thresh) | |
| xy_text = xy_text[:, ::-1] # (n, 2) | |
| # Sort the text boxes via the y axis | |
| xy_text = xy_text[np.argsort(xy_text[:, 1])] | |
| scores = tcl_map[xy_text[:, 1], xy_text[:, 0], 0] | |
| scores = scores[:, np.newaxis] | |
| # Restore | |
| point_num = int(tvo_map.shape[-1] / 2) | |
| assert point_num == 4 | |
| tvo_map = tvo_map[xy_text[:, 1], xy_text[:, 0], :] | |
| xy_text_tile = np.tile(xy_text, (1, point_num)) # (n, point_num * 2) | |
| quads = xy_text_tile - tvo_map | |
| return scores, quads, xy_text | |
| def quad_area(self, quad): | |
| """ | |
| compute area of a quad. | |
| """ | |
| edge = [ | |
| (quad[1][0] - quad[0][0]) * (quad[1][1] + quad[0][1]), | |
| (quad[2][0] - quad[1][0]) * (quad[2][1] + quad[1][1]), | |
| (quad[3][0] - quad[2][0]) * (quad[3][1] + quad[2][1]), | |
| (quad[0][0] - quad[3][0]) * (quad[0][1] + quad[3][1]), | |
| ] | |
| return np.sum(edge) / 2.0 | |
| def nms(self, dets): | |
| if self.is_python35: | |
| import lanms | |
| dets = lanms.merge_quadrangle_n9(dets, self.nms_thresh) | |
| else: | |
| dets = nms_locality(dets, self.nms_thresh) | |
| return dets | |
| def cluster_by_quads_tco(self, tcl_map, tcl_map_thresh, quads, tco_map): | |
| """ | |
| Cluster pixels in tcl_map based on quads. | |
| """ | |
| instance_count = quads.shape[0] + 1 # contain background | |
| instance_label_map = np.zeros(tcl_map.shape[:2], dtype=np.int32) | |
| if instance_count == 1: | |
| return instance_count, instance_label_map | |
| # predict text center | |
| xy_text = np.argwhere(tcl_map[:, :, 0] > tcl_map_thresh) | |
| n = xy_text.shape[0] | |
| xy_text = xy_text[:, ::-1] # (n, 2) | |
| tco = tco_map[xy_text[:, 1], xy_text[:, 0], :] # (n, 2) | |
| pred_tc = xy_text - tco | |
| # get gt text center | |
| m = quads.shape[0] | |
| gt_tc = np.mean(quads, axis=1) # (m, 2) | |
| pred_tc_tile = np.tile(pred_tc[:, np.newaxis, :], (1, m, 1)) # (n, m, 2) | |
| gt_tc_tile = np.tile(gt_tc[np.newaxis, :, :], (n, 1, 1)) # (n, m, 2) | |
| dist_mat = np.linalg.norm(pred_tc_tile - gt_tc_tile, axis=2) # (n, m) | |
| xy_text_assign = np.argmin(dist_mat, axis=1) + 1 # (n,) | |
| instance_label_map[xy_text[:, 1], xy_text[:, 0]] = xy_text_assign | |
| return instance_count, instance_label_map | |
| def estimate_sample_pts_num(self, quad, xy_text): | |
| """ | |
| Estimate sample points number. | |
| """ | |
| eh = ( | |
| np.linalg.norm(quad[0] - quad[3]) + np.linalg.norm(quad[1] - quad[2]) | |
| ) / 2.0 | |
| ew = ( | |
| np.linalg.norm(quad[0] - quad[1]) + np.linalg.norm(quad[2] - quad[3]) | |
| ) / 2.0 | |
| dense_sample_pts_num = max(2, int(ew)) | |
| dense_xy_center_line = xy_text[ | |
| np.linspace( | |
| 0, | |
| xy_text.shape[0] - 1, | |
| dense_sample_pts_num, | |
| endpoint=True, | |
| dtype=np.float32, | |
| ).astype(np.int32) | |
| ] | |
| dense_xy_center_line_diff = dense_xy_center_line[1:] - dense_xy_center_line[:-1] | |
| estimate_arc_len = np.sum(np.linalg.norm(dense_xy_center_line_diff, axis=1)) | |
| sample_pts_num = max(2, int(estimate_arc_len / eh)) | |
| return sample_pts_num | |
| def detect_sast( | |
| self, | |
| tcl_map, | |
| tvo_map, | |
| tbo_map, | |
| tco_map, | |
| ratio_w, | |
| ratio_h, | |
| src_w, | |
| src_h, | |
| shrink_ratio_of_width=0.3, | |
| tcl_map_thresh=0.5, | |
| offset_expand=1.0, | |
| out_strid=4.0, | |
| ): | |
| """ | |
| first resize the tcl_map, tvo_map and tbo_map to the input_size, then restore the polys | |
| """ | |
| # restore quad | |
| scores, quads, xy_text = self.restore_quad(tcl_map, tcl_map_thresh, tvo_map) | |
| dets = np.hstack((quads, scores)).astype(np.float32, copy=False) | |
| dets = self.nms(dets) | |
| if dets.shape[0] == 0: | |
| return [] | |
| quads = dets[:, :-1].reshape(-1, 4, 2) | |
| # Compute quad area | |
| quad_areas = [] | |
| for quad in quads: | |
| quad_areas.append(-self.quad_area(quad)) | |
| # instance segmentation | |
| # instance_count, instance_label_map = cv2.connectedComponents(tcl_map.astype(np.uint8), connectivity=8) | |
| instance_count, instance_label_map = self.cluster_by_quads_tco( | |
| tcl_map, tcl_map_thresh, quads, tco_map | |
| ) | |
| # restore single poly with tcl instance. | |
| poly_list = [] | |
| for instance_idx in range(1, instance_count): | |
| xy_text = np.argwhere(instance_label_map == instance_idx)[:, ::-1] | |
| quad = quads[instance_idx - 1] | |
| q_area = quad_areas[instance_idx - 1] | |
| if q_area < 5: | |
| continue | |
| # | |
| len1 = float(np.linalg.norm(quad[0] - quad[1])) | |
| len2 = float(np.linalg.norm(quad[1] - quad[2])) | |
| min_len = min(len1, len2) | |
| if min_len < 3: | |
| continue | |
| # filter small CC | |
| if xy_text.shape[0] <= 0: | |
| continue | |
| # filter low confidence instance | |
| xy_text_scores = tcl_map[xy_text[:, 1], xy_text[:, 0], 0] | |
| if np.sum(xy_text_scores) / quad_areas[instance_idx - 1] < 0.1: | |
| # if np.sum(xy_text_scores) / quad_areas[instance_idx - 1] < 0.05: | |
| continue | |
| # sort xy_text | |
| left_center_pt = np.array( | |
| [[(quad[0, 0] + quad[-1, 0]) / 2.0, (quad[0, 1] + quad[-1, 1]) / 2.0]] | |
| ) # (1, 2) | |
| right_center_pt = np.array( | |
| [[(quad[1, 0] + quad[2, 0]) / 2.0, (quad[1, 1] + quad[2, 1]) / 2.0]] | |
| ) # (1, 2) | |
| proj_unit_vec = (right_center_pt - left_center_pt) / ( | |
| np.linalg.norm(right_center_pt - left_center_pt) + 1e-6 | |
| ) | |
| proj_value = np.sum(xy_text * proj_unit_vec, axis=1) | |
| xy_text = xy_text[np.argsort(proj_value)] | |
| # Sample pts in tcl map | |
| if self.sample_pts_num == 0: | |
| sample_pts_num = self.estimate_sample_pts_num(quad, xy_text) | |
| else: | |
| sample_pts_num = self.sample_pts_num | |
| xy_center_line = xy_text[ | |
| np.linspace( | |
| 0, | |
| xy_text.shape[0] - 1, | |
| sample_pts_num, | |
| endpoint=True, | |
| dtype=np.float32, | |
| ).astype(np.int32) | |
| ] | |
| point_pair_list = [] | |
| for x, y in xy_center_line: | |
| # get corresponding offset | |
| offset = tbo_map[y, x, :].reshape(2, 2) | |
| if offset_expand != 1.0: | |
| offset_length = np.linalg.norm(offset, axis=1, keepdims=True) | |
| expand_length = np.clip( | |
| offset_length * (offset_expand - 1), a_min=0.5, a_max=3.0 | |
| ) | |
| offset_detal = offset / offset_length * expand_length | |
| offset = offset + offset_detal | |
| # original point | |
| ori_yx = np.array([y, x], dtype=np.float32) | |
| point_pair = ( | |
| (ori_yx + offset)[:, ::-1] | |
| * out_strid | |
| / np.array([ratio_w, ratio_h]).reshape(-1, 2) | |
| ) | |
| point_pair_list.append(point_pair) | |
| # ndarry: (x, 2), expand poly along width | |
| detected_poly = self.point_pair2poly(point_pair_list) | |
| detected_poly = self.expand_poly_along_width( | |
| detected_poly, shrink_ratio_of_width | |
| ) | |
| detected_poly[:, 0] = np.clip(detected_poly[:, 0], a_min=0, a_max=src_w) | |
| detected_poly[:, 1] = np.clip(detected_poly[:, 1], a_min=0, a_max=src_h) | |
| poly_list.append(detected_poly) | |
| return poly_list | |
| def __call__(self, outs_dict, shape_list): | |
| score_list = outs_dict["f_score"] | |
| border_list = outs_dict["f_border"] | |
| tvo_list = outs_dict["f_tvo"] | |
| tco_list = outs_dict["f_tco"] | |
| if isinstance(score_list, paddle.Tensor): | |
| score_list = score_list.numpy() | |
| border_list = border_list.numpy() | |
| tvo_list = tvo_list.numpy() | |
| tco_list = tco_list.numpy() | |
| img_num = len(shape_list) | |
| poly_lists = [] | |
| for ino in range(img_num): | |
| p_score = score_list[ino].transpose((1, 2, 0)) | |
| p_border = border_list[ino].transpose((1, 2, 0)) | |
| p_tvo = tvo_list[ino].transpose((1, 2, 0)) | |
| p_tco = tco_list[ino].transpose((1, 2, 0)) | |
| src_h, src_w, ratio_h, ratio_w = shape_list[ino] | |
| poly_list = self.detect_sast( | |
| p_score, | |
| p_tvo, | |
| p_border, | |
| p_tco, | |
| ratio_w, | |
| ratio_h, | |
| src_w, | |
| src_h, | |
| shrink_ratio_of_width=self.shrink_ratio_of_width, | |
| tcl_map_thresh=self.tcl_map_thresh, | |
| offset_expand=self.expand_scale, | |
| ) | |
| poly_lists.append({"points": np.array(poly_list)}) | |
| return poly_lists | |