Spaces:
Runtime error
Runtime error
File size: 3,223 Bytes
2366e36 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 |
# 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)
|