Spaces:
Runtime error
Runtime error
File size: 4,087 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 104 105 106 107 108 109 110 111 |
# Copyright (c) OpenMMLab. All rights reserved.
import cv2
import numpy as np
from mmocr.models.builder import POSTPROCESSOR
from .base_postprocessor import BasePostprocessor
from .utils import fill_hole, fourier2poly, poly_nms
@POSTPROCESSOR.register_module()
class FCEPostprocessor(BasePostprocessor):
"""Decoding predictions of FCENet to instances.
Args:
fourier_degree (int): The maximum Fourier transform degree k.
num_reconstr_points (int): The points number of the polygon
reconstructed from predicted Fourier coefficients.
text_repr_type (str): Boundary encoding type 'poly' or 'quad'.
scale (int): The down-sample scale of the prediction.
alpha (float): The parameter to calculate final scores. Score_{final}
= (Score_{text region} ^ alpha)
* (Score_{text center region}^ beta)
beta (float): The parameter to calculate final score.
score_thr (float): The threshold used to filter out the final
candidates.
nms_thr (float): The threshold of nms.
"""
def __init__(self,
fourier_degree,
num_reconstr_points,
text_repr_type='poly',
alpha=1.0,
beta=2.0,
score_thr=0.3,
nms_thr=0.1,
**kwargs):
super().__init__(text_repr_type)
self.fourier_degree = fourier_degree
self.num_reconstr_points = num_reconstr_points
self.alpha = alpha
self.beta = beta
self.score_thr = score_thr
self.nms_thr = nms_thr
def __call__(self, preds, scale):
"""
Args:
preds (list[Tensor]): Classification prediction and regression
prediction.
scale (float): Scale of current layer.
Returns:
list[list[float]]: The instance boundary and confidence.
"""
assert isinstance(preds, list)
assert len(preds) == 2
cls_pred = preds[0][0]
tr_pred = cls_pred[0:2].softmax(dim=0).data.cpu().numpy()
tcl_pred = cls_pred[2:].softmax(dim=0).data.cpu().numpy()
reg_pred = preds[1][0].permute(1, 2, 0).data.cpu().numpy()
x_pred = reg_pred[:, :, :2 * self.fourier_degree + 1]
y_pred = reg_pred[:, :, 2 * self.fourier_degree + 1:]
score_pred = (tr_pred[1]**self.alpha) * (tcl_pred[1]**self.beta)
tr_pred_mask = (score_pred) > self.score_thr
tr_mask = fill_hole(tr_pred_mask)
tr_contours, _ = cv2.findContours(
tr_mask.astype(np.uint8), cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE) # opencv4
mask = np.zeros_like(tr_mask)
boundaries = []
for cont in tr_contours:
deal_map = mask.copy().astype(np.int8)
cv2.drawContours(deal_map, [cont], -1, 1, -1)
score_map = score_pred * deal_map
score_mask = score_map > 0
xy_text = np.argwhere(score_mask)
dxy = xy_text[:, 1] + xy_text[:, 0] * 1j
x, y = x_pred[score_mask], y_pred[score_mask]
c = x + y * 1j
c[:, self.fourier_degree] = c[:, self.fourier_degree] + dxy
c *= scale
polygons = fourier2poly(c, self.num_reconstr_points)
score = score_map[score_mask].reshape(-1, 1)
polygons = poly_nms(
np.hstack((polygons, score)).tolist(), self.nms_thr)
boundaries = boundaries + polygons
boundaries = poly_nms(boundaries, self.nms_thr)
if self.text_repr_type == 'quad':
new_boundaries = []
for boundary in boundaries:
poly = np.array(boundary[:-1]).reshape(-1,
2).astype(np.float32)
score = boundary[-1]
points = cv2.boxPoints(cv2.minAreaRect(poly))
points = np.int0(points)
new_boundaries.append(points.reshape(-1).tolist() + [score])
return boundaries
|