Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import cv2 | |
| import numpy as np | |
| import mmocr.utils as utils | |
| def points2boundary(points, text_repr_type, text_score=None, min_width=-1): | |
| """Convert a text mask represented by point coordinates sequence into a | |
| text boundary. | |
| Args: | |
| points (ndarray): Mask index of size (n, 2). | |
| text_repr_type (str): Text instance encoding type | |
| ('quad' for quadrangle or 'poly' for polygon). | |
| text_score (float): Text score. | |
| Returns: | |
| boundary (list[float]): The text boundary point coordinates (x, y) | |
| list. Return None if no text boundary found. | |
| """ | |
| assert isinstance(points, np.ndarray) | |
| assert points.shape[1] == 2 | |
| assert text_repr_type in ['quad', 'poly'] | |
| assert text_score is None or 0 <= text_score <= 1 | |
| if text_repr_type == 'quad': | |
| rect = cv2.minAreaRect(points) | |
| vertices = cv2.boxPoints(rect) | |
| boundary = [] | |
| if min(rect[1]) > min_width: | |
| boundary = [p for p in vertices.flatten().tolist()] | |
| elif text_repr_type == 'poly': | |
| height = np.max(points[:, 1]) + 10 | |
| width = np.max(points[:, 0]) + 10 | |
| mask = np.zeros((height, width), np.uint8) | |
| mask[points[:, 1], points[:, 0]] = 255 | |
| contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, | |
| cv2.CHAIN_APPROX_SIMPLE) | |
| boundary = list(contours[0].flatten().tolist()) | |
| if text_score is not None: | |
| boundary = boundary + [text_score] | |
| if len(boundary) < 8: | |
| return None | |
| return boundary | |
| def seg2boundary(seg, text_repr_type, text_score=None): | |
| """Convert a segmentation mask to a text boundary. | |
| Args: | |
| seg (ndarray): The segmentation mask. | |
| text_repr_type (str): Text instance encoding type | |
| ('quad' for quadrangle or 'poly' for polygon). | |
| text_score (float): The text score. | |
| Returns: | |
| boundary (list): The text boundary. Return None if no text found. | |
| """ | |
| assert isinstance(seg, np.ndarray) | |
| assert isinstance(text_repr_type, str) | |
| assert text_score is None or 0 <= text_score <= 1 | |
| points = np.where(seg) | |
| # x, y order | |
| points = np.concatenate([points[1], points[0]]).reshape(2, -1).transpose() | |
| boundary = None | |
| if len(points) != 0: | |
| boundary = points2boundary(points, text_repr_type, text_score) | |
| return boundary | |
| def extract_boundary(result): | |
| """Extract boundaries and their scores from result. | |
| Args: | |
| result (dict): The detection result with the key 'boundary_result' | |
| of one image. | |
| Returns: | |
| boundaries_with_scores (list[list[float]]): The boundary and score | |
| list. | |
| boundaries (list[list[float]]): The boundary list. | |
| scores (list[float]): The boundary score list. | |
| """ | |
| assert isinstance(result, dict) | |
| assert 'boundary_result' in result.keys() | |
| boundaries_with_scores = result['boundary_result'] | |
| assert utils.is_2dlist(boundaries_with_scores) | |
| boundaries = [b[:-1] for b in boundaries_with_scores] | |
| scores = [b[-1] for b in boundaries_with_scores] | |
| return (boundaries_with_scores, boundaries, scores) | |