| |
| import functools |
| import operator |
|
|
| import cv2 |
| import numpy as np |
| import pyclipper |
| from numpy.fft import ifft |
| from numpy.linalg import norm |
| from shapely.geometry import Polygon |
|
|
| from mmocr.core.evaluation.utils import boundary_iou |
|
|
|
|
| def filter_instance(area, confidence, min_area, min_confidence): |
| return bool(area < min_area or confidence < min_confidence) |
|
|
|
|
| def box_score_fast(bitmap, _box): |
| h, w = bitmap.shape[:2] |
| box = _box.copy() |
| xmin = np.clip(np.floor(box[:, 0].min()).astype(np.int32), 0, w - 1) |
| xmax = np.clip(np.ceil(box[:, 0].max()).astype(np.int32), 0, w - 1) |
| ymin = np.clip(np.floor(box[:, 1].min()).astype(np.int32), 0, h - 1) |
| ymax = np.clip(np.ceil(box[:, 1].max()).astype(np.int32), 0, h - 1) |
|
|
| mask = np.zeros((ymax - ymin + 1, xmax - xmin + 1), dtype=np.uint8) |
| box[:, 0] = box[:, 0] - xmin |
| box[:, 1] = box[:, 1] - ymin |
| cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1) |
| return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0] |
|
|
|
|
| def unclip(box, unclip_ratio=1.5): |
| poly = Polygon(box) |
| distance = poly.area * unclip_ratio / poly.length |
| offset = pyclipper.PyclipperOffset() |
| offset.AddPath(box, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON) |
| expanded = np.array(offset.Execute(distance)) |
| return expanded |
|
|
|
|
| def fill_hole(input_mask): |
| h, w = input_mask.shape |
| canvas = np.zeros((h + 2, w + 2), np.uint8) |
| canvas[1:h + 1, 1:w + 1] = input_mask.copy() |
|
|
| mask = np.zeros((h + 4, w + 4), np.uint8) |
|
|
| cv2.floodFill(canvas, mask, (0, 0), 1) |
| canvas = canvas[1:h + 1, 1:w + 1].astype(np.bool) |
|
|
| return ~canvas | input_mask |
|
|
|
|
| def centralize(points_yx, |
| normal_sin, |
| normal_cos, |
| radius, |
| contour_mask, |
| step_ratio=0.03): |
|
|
| h, w = contour_mask.shape |
| top_yx = bot_yx = points_yx |
| step_flags = np.ones((len(points_yx), 1), dtype=np.bool) |
| step = step_ratio * radius * np.hstack([normal_sin, normal_cos]) |
| while np.any(step_flags): |
| next_yx = np.array(top_yx + step, dtype=np.int32) |
| next_y, next_x = next_yx[:, 0], next_yx[:, 1] |
| step_flags = (next_y >= 0) & (next_y < h) & (next_x > 0) & ( |
| next_x < w) & contour_mask[np.clip(next_y, 0, h - 1), |
| np.clip(next_x, 0, w - 1)] |
| top_yx = top_yx + step_flags.reshape((-1, 1)) * step |
| step_flags = np.ones((len(points_yx), 1), dtype=np.bool) |
| while np.any(step_flags): |
| next_yx = np.array(bot_yx - step, dtype=np.int32) |
| next_y, next_x = next_yx[:, 0], next_yx[:, 1] |
| step_flags = (next_y >= 0) & (next_y < h) & (next_x > 0) & ( |
| next_x < w) & contour_mask[np.clip(next_y, 0, h - 1), |
| np.clip(next_x, 0, w - 1)] |
| bot_yx = bot_yx - step_flags.reshape((-1, 1)) * step |
| centers = np.array((top_yx + bot_yx) * 0.5, dtype=np.int32) |
| return centers |
|
|
|
|
| def merge_disks(disks, disk_overlap_thr): |
| xy = disks[:, 0:2] |
| radius = disks[:, 2] |
| scores = disks[:, 3] |
| order = scores.argsort()[::-1] |
|
|
| merged_disks = [] |
| while order.size > 0: |
| if order.size == 1: |
| merged_disks.append(disks[order]) |
| break |
| i = order[0] |
| d = norm(xy[i] - xy[order[1:]], axis=1) |
| ri = radius[i] |
| r = radius[order[1:]] |
| d_thr = (ri + r) * disk_overlap_thr |
|
|
| merge_inds = np.where(d <= d_thr)[0] + 1 |
| if merge_inds.size > 0: |
| merge_order = np.hstack([i, order[merge_inds]]) |
| merged_disks.append(np.mean(disks[merge_order], axis=0)) |
| else: |
| merged_disks.append(disks[i]) |
|
|
| inds = np.where(d > d_thr)[0] + 1 |
| order = order[inds] |
| merged_disks = np.vstack(merged_disks) |
|
|
| return merged_disks |
|
|
|
|
| def poly_nms(polygons, threshold): |
| assert isinstance(polygons, list) |
|
|
| polygons = np.array(sorted(polygons, key=lambda x: x[-1])) |
|
|
| keep_poly = [] |
| index = [i for i in range(polygons.shape[0])] |
|
|
| while len(index) > 0: |
| keep_poly.append(polygons[index[-1]].tolist()) |
| A = polygons[index[-1]][:-1] |
| index = np.delete(index, -1) |
|
|
| iou_list = np.zeros((len(index), )) |
| for i in range(len(index)): |
| B = polygons[index[i]][:-1] |
|
|
| iou_list[i] = boundary_iou(A, B, 1) |
| remove_index = np.where(iou_list > threshold) |
| index = np.delete(index, remove_index) |
|
|
| return keep_poly |
|
|
|
|
| def fourier2poly(fourier_coeff, num_reconstr_points=50): |
| """ Inverse Fourier transform |
| Args: |
| fourier_coeff (ndarray): Fourier coefficients shaped (n, 2k+1), |
| with n and k being candidates number and Fourier degree |
| respectively. |
| num_reconstr_points (int): Number of reconstructed polygon points. |
| Returns: |
| Polygons (ndarray): The reconstructed polygons shaped (n, n') |
| """ |
|
|
| a = np.zeros((len(fourier_coeff), num_reconstr_points), dtype='complex') |
| k = (len(fourier_coeff[0]) - 1) // 2 |
|
|
| a[:, 0:k + 1] = fourier_coeff[:, k:] |
| a[:, -k:] = fourier_coeff[:, :k] |
|
|
| poly_complex = ifft(a) * num_reconstr_points |
| polygon = np.zeros((len(fourier_coeff), num_reconstr_points, 2)) |
| polygon[:, :, 0] = poly_complex.real |
| polygon[:, :, 1] = poly_complex.imag |
| return polygon.astype('int32').reshape((len(fourier_coeff), -1)) |
|
|
|
|
| class Node: |
|
|
| def __init__(self, ind): |
| self.__ind = ind |
| self.__links = set() |
|
|
| @property |
| def ind(self): |
| return self.__ind |
|
|
| @property |
| def links(self): |
| return set(self.__links) |
|
|
| def add_link(self, link_node): |
| self.__links.add(link_node) |
| link_node.__links.add(self) |
|
|
|
|
| def graph_propagation(edges, scores, text_comps, edge_len_thr=50.): |
| """Propagate edge score information and construct graph. This code was |
| partially adapted from https://github.com/GXYM/DRRG licensed under the MIT |
| license. |
| |
| Args: |
| edges (ndarray): The edge array of shape N * 2, each row is a node |
| index pair that makes up an edge in graph. |
| scores (ndarray): The edge score array. |
| text_comps (ndarray): The text components. |
| edge_len_thr (float): The edge length threshold. |
| |
| Returns: |
| vertices (list[Node]): The Nodes in graph. |
| score_dict (dict): The edge score dict. |
| """ |
| assert edges.ndim == 2 |
| assert edges.shape[1] == 2 |
| assert edges.shape[0] == scores.shape[0] |
| assert text_comps.ndim == 2 |
| assert isinstance(edge_len_thr, float) |
|
|
| edges = np.sort(edges, axis=1) |
| score_dict = {} |
| for i, edge in enumerate(edges): |
| if text_comps is not None: |
| box1 = text_comps[edge[0], :8].reshape(4, 2) |
| box2 = text_comps[edge[1], :8].reshape(4, 2) |
| center1 = np.mean(box1, axis=0) |
| center2 = np.mean(box2, axis=0) |
| distance = norm(center1 - center2) |
| if distance > edge_len_thr: |
| scores[i] = 0 |
| if (edge[0], edge[1]) in score_dict: |
| score_dict[edge[0], edge[1]] = 0.5 * ( |
| score_dict[edge[0], edge[1]] + scores[i]) |
| else: |
| score_dict[edge[0], edge[1]] = scores[i] |
|
|
| nodes = np.sort(np.unique(edges.flatten())) |
| mapping = -1 * np.ones((np.max(nodes) + 1), dtype=np.int) |
| mapping[nodes] = np.arange(nodes.shape[0]) |
| order_inds = mapping[edges] |
| vertices = [Node(node) for node in nodes] |
| for ind in order_inds: |
| vertices[ind[0]].add_link(vertices[ind[1]]) |
|
|
| return vertices, score_dict |
|
|
|
|
| def connected_components(nodes, score_dict, link_thr): |
| """Conventional connected components searching. This code was partially |
| adapted from https://github.com/GXYM/DRRG licensed under the MIT license. |
| |
| Args: |
| nodes (list[Node]): The list of Node objects. |
| score_dict (dict): The edge score dict. |
| link_thr (float): The link threshold. |
| |
| Returns: |
| clusters (List[list[Node]]): The clustered Node objects. |
| """ |
| assert isinstance(nodes, list) |
| assert all([isinstance(node, Node) for node in nodes]) |
| assert isinstance(score_dict, dict) |
| assert isinstance(link_thr, float) |
|
|
| clusters = [] |
| nodes = set(nodes) |
| while nodes: |
| node = nodes.pop() |
| cluster = {node} |
| node_queue = [node] |
| while node_queue: |
| node = node_queue.pop(0) |
| neighbors = set([ |
| neighbor for neighbor in node.links if |
| score_dict[tuple(sorted([node.ind, neighbor.ind]))] >= link_thr |
| ]) |
| neighbors.difference_update(cluster) |
| nodes.difference_update(neighbors) |
| cluster.update(neighbors) |
| node_queue.extend(neighbors) |
| clusters.append(list(cluster)) |
| return clusters |
|
|
|
|
| def clusters2labels(clusters, num_nodes): |
| """Convert clusters of Node to text component labels. This code was |
| partially adapted from https://github.com/GXYM/DRRG licensed under the MIT |
| license. |
| |
| Args: |
| clusters (List[list[Node]]): The clusters of Node objects. |
| num_nodes (int): The total node number of graphs in an image. |
| |
| Returns: |
| node_labels (ndarray): The node label array. |
| """ |
| assert isinstance(clusters, list) |
| assert all([isinstance(cluster, list) for cluster in clusters]) |
| assert all( |
| [isinstance(node, Node) for cluster in clusters for node in cluster]) |
| assert isinstance(num_nodes, int) |
|
|
| node_labels = np.zeros(num_nodes) |
| for cluster_ind, cluster in enumerate(clusters): |
| for node in cluster: |
| node_labels[node.ind] = cluster_ind |
| return node_labels |
|
|
|
|
| def remove_single(text_comps, comp_pred_labels): |
| """Remove isolated text components. This code was partially adapted from |
| https://github.com/GXYM/DRRG licensed under the MIT license. |
| |
| Args: |
| text_comps (ndarray): The text components. |
| comp_pred_labels (ndarray): The clustering labels of text components. |
| |
| Returns: |
| filtered_text_comps (ndarray): The text components with isolated ones |
| removed. |
| comp_pred_labels (ndarray): The clustering labels with labels of |
| isolated text components removed. |
| """ |
| assert text_comps.ndim == 2 |
| assert text_comps.shape[0] == comp_pred_labels.shape[0] |
|
|
| single_flags = np.zeros_like(comp_pred_labels) |
| pred_labels = np.unique(comp_pred_labels) |
| for label in pred_labels: |
| current_label_flag = (comp_pred_labels == label) |
| if np.sum(current_label_flag) == 1: |
| single_flags[np.where(current_label_flag)[0][0]] = 1 |
| keep_ind = [i for i in range(len(comp_pred_labels)) if not single_flags[i]] |
| filtered_text_comps = text_comps[keep_ind, :] |
| filtered_labels = comp_pred_labels[keep_ind] |
|
|
| return filtered_text_comps, filtered_labels |
|
|
|
|
| def norm2(point1, point2): |
| return ((point1[0] - point2[0])**2 + (point1[1] - point2[1])**2)**0.5 |
|
|
|
|
| def min_connect_path(points): |
| """Find the shortest path to traverse all points. This code was partially |
| adapted from https://github.com/GXYM/DRRG licensed under the MIT license. |
| |
| Args: |
| points(List[list[int]]): The point sequence [[x0, y0], [x1, y1], ...]. |
| |
| Returns: |
| shortest_path(List[list[int]]): The shortest index path. |
| """ |
| assert isinstance(points, list) |
| assert all([isinstance(point, list) for point in points]) |
| assert all([isinstance(coord, int) for point in points for coord in point]) |
|
|
| points_queue = points.copy() |
| shortest_path = [] |
| current_edge = [[], []] |
|
|
| edge_dict0 = {} |
| edge_dict1 = {} |
| current_edge[0] = points_queue[0] |
| current_edge[1] = points_queue[0] |
| points_queue.remove(points_queue[0]) |
| while points_queue: |
| for point in points_queue: |
| length0 = norm2(point, current_edge[0]) |
| edge_dict0[length0] = [point, current_edge[0]] |
| length1 = norm2(current_edge[1], point) |
| edge_dict1[length1] = [current_edge[1], point] |
| key0 = min(edge_dict0.keys()) |
| key1 = min(edge_dict1.keys()) |
|
|
| if key0 <= key1: |
| start = edge_dict0[key0][0] |
| end = edge_dict0[key0][1] |
| shortest_path.insert(0, [points.index(start), points.index(end)]) |
| points_queue.remove(start) |
| current_edge[0] = start |
| else: |
| start = edge_dict1[key1][0] |
| end = edge_dict1[key1][1] |
| shortest_path.append([points.index(start), points.index(end)]) |
| points_queue.remove(end) |
| current_edge[1] = end |
|
|
| edge_dict0 = {} |
| edge_dict1 = {} |
|
|
| shortest_path = functools.reduce(operator.concat, shortest_path) |
| shortest_path = sorted(set(shortest_path), key=shortest_path.index) |
|
|
| return shortest_path |
|
|
|
|
| def in_contour(cont, point): |
| x, y = point |
| is_inner = cv2.pointPolygonTest(cont, (int(x), int(y)), False) > 0.5 |
| return is_inner |
|
|
|
|
| def fix_corner(top_line, bot_line, start_box, end_box): |
| """Add corner points to predicted side lines. This code was partially |
| adapted from https://github.com/GXYM/DRRG licensed under the MIT license. |
| |
| Args: |
| top_line (List[list[int]]): The predicted top sidelines of text |
| instance. |
| bot_line (List[list[int]]): The predicted bottom sidelines of text |
| instance. |
| start_box (ndarray): The first text component box. |
| end_box (ndarray): The last text component box. |
| |
| Returns: |
| top_line (List[list[int]]): The top sidelines with corner point added. |
| bot_line (List[list[int]]): The bottom sidelines with corner point |
| added. |
| """ |
| assert isinstance(top_line, list) |
| assert all(isinstance(point, list) for point in top_line) |
| assert isinstance(bot_line, list) |
| assert all(isinstance(point, list) for point in bot_line) |
| assert start_box.shape == end_box.shape == (4, 2) |
|
|
| contour = np.array(top_line + bot_line[::-1]) |
| start_left_mid = (start_box[0] + start_box[3]) / 2 |
| start_right_mid = (start_box[1] + start_box[2]) / 2 |
| end_left_mid = (end_box[0] + end_box[3]) / 2 |
| end_right_mid = (end_box[1] + end_box[2]) / 2 |
| if not in_contour(contour, start_left_mid): |
| top_line.insert(0, start_box[0].tolist()) |
| bot_line.insert(0, start_box[3].tolist()) |
| elif not in_contour(contour, start_right_mid): |
| top_line.insert(0, start_box[1].tolist()) |
| bot_line.insert(0, start_box[2].tolist()) |
| if not in_contour(contour, end_left_mid): |
| top_line.append(end_box[0].tolist()) |
| bot_line.append(end_box[3].tolist()) |
| elif not in_contour(contour, end_right_mid): |
| top_line.append(end_box[1].tolist()) |
| bot_line.append(end_box[2].tolist()) |
| return top_line, bot_line |
|
|
|
|
| def comps2boundaries(text_comps, comp_pred_labels): |
| """Construct text instance boundaries from clustered text components. This |
| code was partially adapted from https://github.com/GXYM/DRRG licensed under |
| the MIT license. |
| |
| Args: |
| text_comps (ndarray): The text components. |
| comp_pred_labels (ndarray): The clustering labels of text components. |
| |
| Returns: |
| boundaries (List[list[float]]): The predicted boundaries of text |
| instances. |
| """ |
| assert text_comps.ndim == 2 |
| assert len(text_comps) == len(comp_pred_labels) |
| boundaries = [] |
| if len(text_comps) < 1: |
| return boundaries |
| for cluster_ind in range(0, int(np.max(comp_pred_labels)) + 1): |
| cluster_comp_inds = np.where(comp_pred_labels == cluster_ind) |
| text_comp_boxes = text_comps[cluster_comp_inds, :8].reshape( |
| (-1, 4, 2)).astype(np.int32) |
| score = np.mean(text_comps[cluster_comp_inds, -1]) |
|
|
| if text_comp_boxes.shape[0] < 1: |
| continue |
|
|
| elif text_comp_boxes.shape[0] > 1: |
| centers = np.mean( |
| text_comp_boxes, axis=1).astype(np.int32).tolist() |
| shortest_path = min_connect_path(centers) |
| text_comp_boxes = text_comp_boxes[shortest_path] |
| top_line = np.mean( |
| text_comp_boxes[:, 0:2, :], axis=1).astype(np.int32).tolist() |
| bot_line = np.mean( |
| text_comp_boxes[:, 2:4, :], axis=1).astype(np.int32).tolist() |
| top_line, bot_line = fix_corner(top_line, bot_line, |
| text_comp_boxes[0], |
| text_comp_boxes[-1]) |
| boundary_points = top_line + bot_line[::-1] |
|
|
| else: |
| top_line = text_comp_boxes[0, 0:2, :].astype(np.int32).tolist() |
| bot_line = text_comp_boxes[0, 2:4:-1, :].astype(np.int32).tolist() |
| boundary_points = top_line + bot_line |
|
|
| boundary = [p for coord in boundary_points for p in coord] + [score] |
| boundaries.append(boundary) |
|
|
| return boundaries |
|
|
