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