Spaces:
Running
Running
File size: 8,406 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 |
# Copyright (c) OpenMMLab. All rights reserved.
from functools import partial
from typing import Dict, List, Optional, Sequence, Tuple, Union
import mmengine
import numpy as np
from torch import Tensor
from mmocr.structures import TextDetDataSample
from mmocr.utils import boundary_iou, rescale_polygons
class BaseTextDetPostProcessor:
"""Base postprocessor for text detection models.
Args:
text_repr_type (str): The boundary encoding type, 'poly' or 'quad'.
Defaults to 'poly'.
rescale_fields (list[str], optional): The bbox/polygon field names to
be rescaled. If None, no rescaling will be performed.
train_cfg (dict, optional): The parameters to be passed to
``self.get_text_instances`` in training. Defaults to None.
test_cfg (dict, optional): The parameters to be passed to
``self.get_text_instances`` in testing. Defaults to None.
"""
def __init__(self,
text_repr_type: str = 'poly',
rescale_fields: Optional[Sequence[str]] = None,
train_cfg: Optional[Dict] = None,
test_cfg: Optional[Dict] = None) -> None:
assert text_repr_type in ['poly', 'quad']
assert rescale_fields is None or isinstance(rescale_fields, list)
assert train_cfg is None or isinstance(train_cfg, dict)
assert test_cfg is None or isinstance(test_cfg, dict)
self.text_repr_type = text_repr_type
self.rescale_fields = rescale_fields
self.train_cfg = train_cfg
self.test_cfg = test_cfg
def __call__(self,
pred_results: Union[Tensor, List[Tensor]],
data_samples: Sequence[TextDetDataSample],
training: bool = False) -> Sequence[TextDetDataSample]:
"""Postprocess pred_results according to metainfos in data_samples.
Args:
pred_results (Union[Tensor, List[Tensor]]): The prediction results
stored in a tensor or a list of tensor. Usually each item to
be post-processed is expected to be a batched tensor.
data_samples (list[TextDetDataSample]): Batch of data_samples,
each corresponding to a prediction result.
training (bool): Whether the model is in training mode. Defaults to
False.
Returns:
list[TextDetDataSample]: Batch of post-processed datasamples.
"""
cfg = self.train_cfg if training else self.test_cfg
if cfg is None:
cfg = {}
pred_results = self.split_results(pred_results)
process_single = partial(self._process_single, **cfg)
results = list(map(process_single, pred_results, data_samples))
return results
def _process_single(self, pred_result: Union[Tensor, List[Tensor]],
data_sample: TextDetDataSample,
**kwargs) -> TextDetDataSample:
"""Process prediction results from one image.
Args:
pred_result (Union[Tensor, List[Tensor]]): Prediction results of an
image.
data_sample (TextDetDataSample): Datasample of an image.
"""
results = self.get_text_instances(pred_result, data_sample, **kwargs)
if self.rescale_fields and len(self.rescale_fields) > 0:
assert isinstance(self.rescale_fields, list)
assert set(self.rescale_fields).issubset(
set(results.pred_instances.keys()))
results = self.rescale(results, data_sample.scale_factor)
return results
def rescale(self, results: TextDetDataSample,
scale_factor: Sequence[int]) -> TextDetDataSample:
"""Rescale results in ``results.pred_instances`` according to
``scale_factor``, whose keys are defined in ``self.rescale_fields``.
Usually used to rescale bboxes and/or polygons.
Args:
results (TextDetDataSample): The post-processed prediction results.
scale_factor (tuple(int)): (w_scale, h_scale)
Returns:
TextDetDataSample: Prediction results with rescaled results.
"""
scale_factor = np.asarray(scale_factor)
for key in self.rescale_fields:
results.pred_instances[key] = rescale_polygons(
results.pred_instances[key], scale_factor, mode='div')
return results
def get_text_instances(self, pred_results: Union[Tensor, List[Tensor]],
data_sample: TextDetDataSample,
**kwargs) -> TextDetDataSample:
"""Get text instance predictions of one image.
Args:
pred_result (tuple(Tensor)): Prediction results of an image.
data_sample (TextDetDataSample): Datasample of an image.
**kwargs: Other parameters. Configurable via ``__init__.train_cfg``
and ``__init__.test_cfg``.
Returns:
TextDetDataSample: A new DataSample with predictions filled in.
The polygon/bbox results are usually saved in
``TextDetDataSample.pred_instances.polygons`` or
``TextDetDataSample.pred_instances.bboxes``. The confidence scores
are saved in ``TextDetDataSample.pred_instances.scores``.
"""
raise NotImplementedError
def split_results(
self, pred_results: Union[Tensor, List[Tensor]]
) -> Union[List[Tensor], List[List[Tensor]]]:
"""Split batched tensor(s) along the first dimension pack split tensors
into a list.
Args:
pred_results (tensor or list[tensor]): Raw result tensor(s) from
detection head. Each tensor usually has the shape of (N, ...)
Returns:
list[tensor] or list[list[tensor]]: N tensors if ``pred_results``
is a tensor, or a list of N lists of tensors if
``pred_results`` is a list of tensors.
"""
assert isinstance(pred_results, Tensor) or mmengine.is_seq_of(
pred_results, Tensor)
if mmengine.is_seq_of(pred_results, Tensor):
for i in range(1, len(pred_results)):
assert pred_results[0].shape[0] == pred_results[i].shape[0], \
'The first dimension of all tensors should be the same'
batch_num = len(pred_results) if isinstance(pred_results, Tensor) else\
len(pred_results[0])
results = []
for i in range(batch_num):
if isinstance(pred_results, Tensor):
results.append(pred_results[i])
else:
results.append([])
for tensor in pred_results:
results[i].append(tensor[i])
return results
def poly_nms(self, polygons: List[np.ndarray], scores: List[float],
threshold: float) -> Tuple[List[np.ndarray], List[float]]:
"""Non-maximum suppression for text detection.
Args:
polygons (list[ndarray]): List of polygons.
scores (list[float]): List of scores.
threshold (float): Threshold for NMS.
Returns:
tuple(keep_polys, keep_scores):
- keep_polys (list[ndarray]): List of preserved polygons after NMS.
- keep_scores (list[float]): List of preserved scores after NMS.
"""
assert isinstance(polygons, list)
assert isinstance(scores, list)
assert len(polygons) == len(scores)
polygons = [
np.hstack((polygon, score))
for polygon, score in zip(polygons, scores)
]
polygons = np.array(sorted(polygons, key=lambda x: x[-1]))
keep_polys = []
keep_scores = []
index = [i for i in range(len(polygons))]
while len(index) > 0:
keep_polys.append(polygons[index[-1]][:-1].tolist())
keep_scores.append(polygons[index[-1]][-1])
A = polygons[index[-1]][:-1]
index = np.delete(index, -1)
iou_list = np.zeros((len(index), ))
for i in range(len(index)):
B = polygons[index[i]][:-1]
iou_list[i] = boundary_iou(A, B, 1)
remove_index = np.where(iou_list > threshold)
index = np.delete(index, remove_index)
return keep_polys, keep_scores
|