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