Spaces:
Running
Running
File size: 10,125 Bytes
9bf4bd7 |
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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Dict, List, Sequence
import cv2
import numpy as np
import torch
from mmengine.structures import InstanceData
from numpy.fft import ifft
from mmocr.registry import MODELS
from mmocr.structures import TextDetDataSample
from mmocr.utils import fill_hole
from .base import BaseTextDetPostProcessor
@MODELS.register_module()
class FCEPostprocessor(BaseTextDetPostProcessor):
"""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.
rescale_fields (list[str]): The bbox/polygon field names to
be rescaled. If None, no rescaling will be performed. Defaults to
['polygons'].
scales (list[int]) : The down-sample scale of each layer. Defaults
to [8, 16, 32].
text_repr_type (str): Boundary encoding type 'poly' or 'quad'. Defaults
to 'poly'.
alpha (float): The parameter to calculate final scores
:math:`Score_{final} = (Score_{text region} ^ alpha)
* (Score_{text center_region}^ beta)`. Defaults to 1.0.
beta (float): The parameter to calculate final score. Defaults to 2.0.
score_thr (float): The threshold used to filter out the final
candidates.Defaults to 0.3.
nms_thr (float): The threshold of nms. Defaults to 0.1.
"""
def __init__(self,
fourier_degree: int,
num_reconstr_points: int,
rescale_fields: Sequence[str] = ['polygons'],
scales: Sequence[int] = [8, 16, 32],
text_repr_type: str = 'poly',
alpha: float = 1.0,
beta: float = 2.0,
score_thr: float = 0.3,
nms_thr: float = 0.1,
**kwargs) -> None:
super().__init__(
text_repr_type=text_repr_type,
rescale_fields=rescale_fields,
**kwargs)
self.fourier_degree = fourier_degree
self.num_reconstr_points = num_reconstr_points
self.scales = scales
self.alpha = alpha
self.beta = beta
self.score_thr = score_thr
self.nms_thr = nms_thr
def split_results(self, pred_results: List[Dict]) -> List[List[Dict]]:
"""Split batched elements in pred_results along the first dimension
into ``batch_num`` sub-elements and regather them into a list of dicts.
Args:
pred_results (list[dict]): A list of dict with keys of ``cls_res``,
``reg_res`` corresponding to the classification result and
regression result computed from the input tensor with the
same index. They have the shapes of :math:`(N, C_{cls,i},
H_i, W_i)` and :math:`(N, C_{out,i}, H_i, W_i)`.
Returns:
list[list[dict]]: N lists. Each list contains three dicts from
different feature level.
"""
assert isinstance(pred_results, list) and len(pred_results) == len(
self.scales)
fields = list(pred_results[0].keys())
batch_num = len(pred_results[0][fields[0]])
level_num = len(pred_results)
results = []
for i in range(batch_num):
batch_list = []
for level in range(level_num):
feat_dict = {}
for field in fields:
feat_dict[field] = pred_results[level][field][i]
batch_list.append(feat_dict)
results.append(batch_list)
return results
def get_text_instances(self, pred_results: Sequence[Dict],
data_sample: TextDetDataSample
) -> TextDetDataSample:
"""Get text instance predictions of one image.
Args:
pred_results (List[dict]): A list of dict with keys of ``cls_res``,
``reg_res`` corresponding to the classification result and
regression result computed from the input tensor with the
same index. They have the shapes of :math:`(N, C_{cls,i}, H_i,
W_i)` and :math:`(N, C_{out,i}, H_i, W_i)`.
data_sample (TextDetDataSample): Datasample of an image.
Returns:
TextDetDataSample: A new DataSample with predictions filled in.
Polygons and results are saved in
``TextDetDataSample.pred_instances.polygons``. The confidence
scores are saved in ``TextDetDataSample.pred_instances.scores``.
"""
assert len(pred_results) == len(self.scales)
data_sample.pred_instances = InstanceData()
data_sample.pred_instances.polygons = []
data_sample.pred_instances.scores = []
result_polys = []
result_scores = []
for idx, pred_result in enumerate(pred_results):
# TODO: Scale can be calculated given image shape and feature
# shape. This param can be removed in the future.
polygons, scores = self._get_text_instances_single(
pred_result, self.scales[idx])
result_polys += polygons
result_scores += scores
result_polys, result_scores = self.poly_nms(result_polys,
result_scores,
self.nms_thr)
for result_poly, result_score in zip(result_polys, result_scores):
result_poly = np.array(result_poly, dtype=np.float32)
data_sample.pred_instances.polygons.append(result_poly)
data_sample.pred_instances.scores.append(result_score)
data_sample.pred_instances.scores = torch.FloatTensor(
data_sample.pred_instances.scores)
return data_sample
def _get_text_instances_single(self, pred_result: Dict, scale: int):
"""Get text instance predictions from one feature level.
Args:
pred_result (dict): A dict with keys of ``cls_res``, ``reg_res``
corresponding to the classification result and regression
result computed from the input tensor with the same index.
They have the shapes of :math:`(1, C_{cls,i}, H_i, W_i)` and
:math:`(1, C_{out,i}, H_i, W_i)`.
scale (int): Scale of current feature map which equals to
img_size / feat_size.
Returns:
result_polys (list[ndarray]): A list of polygons after postprocess.
result_scores (list[ndarray]): A list of scores after postprocess.
"""
cls_pred = pred_result['cls_res']
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 = pred_result['reg_res'].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)
result_polys = []
result_scores = []
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 = self._fourier2poly(c, self.num_reconstr_points)
scores = score_map[score_mask].reshape(-1, 1).tolist()
polygons, scores = self.poly_nms(polygons, scores, self.nms_thr)
result_polys += polygons
result_scores += scores
result_polys, result_scores = self.poly_nms(result_polys,
result_scores,
self.nms_thr)
if self.text_repr_type == 'quad':
new_polys = []
for poly in result_polys:
poly = np.array(poly).reshape(-1, 2).astype(np.float32)
points = cv2.boxPoints(cv2.minAreaRect(poly))
points = np.int0(points)
new_polys.append(points.reshape(-1))
return new_polys, result_scores
return result_polys, result_scores
def _fourier2poly(self,
fourier_coeff: np.ndarray,
num_reconstr_points: int = 50):
""" Inverse Fourier transform
Args:
fourier_coeff (ndarray): Fourier coefficients shaped (n, 2k+1),
with n and k being candidates number and Fourier degree
respectively.
num_reconstr_points (int): Number of reconstructed polygon
points. Defaults to 50.
Returns:
List[ndarray]: The reconstructed polygons.
"""
a = np.zeros((len(fourier_coeff), num_reconstr_points),
dtype='complex')
k = (len(fourier_coeff[0]) - 1) // 2
a[:, 0:k + 1] = fourier_coeff[:, k:]
a[:, -k:] = fourier_coeff[:, :k]
poly_complex = ifft(a) * num_reconstr_points
polygon = np.zeros((len(fourier_coeff), num_reconstr_points, 2))
polygon[:, :, 0] = poly_complex.real
polygon[:, :, 1] = poly_complex.imag
return polygon.astype('int32').reshape(
(len(fourier_coeff), -1)).tolist()
|