|
|
|
|
|
import cv2 |
|
|
import numpy as np |
|
|
|
|
|
from mmocr.core import points2boundary |
|
|
from mmocr.models.builder import POSTPROCESSOR |
|
|
from .base_postprocessor import BasePostprocessor |
|
|
from .utils import box_score_fast, unclip |
|
|
|
|
|
|
|
|
@POSTPROCESSOR.register_module() |
|
|
class DBPostprocessor(BasePostprocessor): |
|
|
"""Decoding predictions of DbNet to instances. This is partially adapted |
|
|
from https://github.com/MhLiao/DB. |
|
|
|
|
|
Args: |
|
|
text_repr_type (str): The boundary encoding type 'poly' or 'quad'. |
|
|
mask_thr (float): The mask threshold value for binarization. |
|
|
min_text_score (float): The threshold value for converting binary map |
|
|
to shrink text regions. |
|
|
min_text_width (int): The minimum width of boundary polygon/box |
|
|
predicted. |
|
|
unclip_ratio (float): The unclip ratio for text regions dilation. |
|
|
max_candidates (int): The maximum candidate number. |
|
|
""" |
|
|
|
|
|
def __init__(self, |
|
|
text_repr_type='poly', |
|
|
mask_thr=0.3, |
|
|
min_text_score=0.3, |
|
|
min_text_width=5, |
|
|
unclip_ratio=1.5, |
|
|
max_candidates=3000, |
|
|
**kwargs): |
|
|
super().__init__(text_repr_type) |
|
|
self.mask_thr = mask_thr |
|
|
self.min_text_score = min_text_score |
|
|
self.min_text_width = min_text_width |
|
|
self.unclip_ratio = unclip_ratio |
|
|
self.max_candidates = max_candidates |
|
|
|
|
|
def __call__(self, preds): |
|
|
""" |
|
|
Args: |
|
|
preds (Tensor): Prediction map with shape :math:`(C, H, W)`. |
|
|
|
|
|
Returns: |
|
|
list[list[float]]: The predicted text boundaries. |
|
|
""" |
|
|
assert preds.dim() == 3 |
|
|
|
|
|
prob_map = preds[0, :, :] |
|
|
text_mask = prob_map > self.mask_thr |
|
|
|
|
|
score_map = prob_map.data.cpu().numpy().astype(np.float32) |
|
|
text_mask = text_mask.data.cpu().numpy().astype(np.uint8) |
|
|
|
|
|
contours, _ = cv2.findContours((text_mask * 255).astype(np.uint8), |
|
|
cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE) |
|
|
|
|
|
boundaries = [] |
|
|
for i, poly in enumerate(contours): |
|
|
if i > self.max_candidates: |
|
|
break |
|
|
epsilon = 0.01 * cv2.arcLength(poly, True) |
|
|
approx = cv2.approxPolyDP(poly, epsilon, True) |
|
|
points = approx.reshape((-1, 2)) |
|
|
if points.shape[0] < 4: |
|
|
continue |
|
|
score = box_score_fast(score_map, points) |
|
|
if score < self.min_text_score: |
|
|
continue |
|
|
poly = unclip(points, unclip_ratio=self.unclip_ratio) |
|
|
if len(poly) == 0 or isinstance(poly[0], list): |
|
|
continue |
|
|
poly = poly.reshape(-1, 2) |
|
|
|
|
|
if self.text_repr_type == 'quad': |
|
|
poly = points2boundary(poly, self.text_repr_type, score, |
|
|
self.min_text_width) |
|
|
elif self.text_repr_type == 'poly': |
|
|
poly = poly.flatten().tolist() |
|
|
if score is not None: |
|
|
poly = poly + [score] |
|
|
if len(poly) < 8: |
|
|
poly = None |
|
|
|
|
|
if poly is not None: |
|
|
boundaries.append(poly) |
|
|
|
|
|
return boundaries |
|
|
|
