Spaces:
Runtime error
Runtime error
File size: 2,948 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 |
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
from mmocr.models.builder import HEADS, build_loss, build_postprocessor
from mmocr.utils import check_argument
@HEADS.register_module()
class HeadMixin:
"""Base head class for text detection, including loss calcalation and
postprocess.
Args:
loss (dict): Config to build loss.
postprocessor (dict): Config to build postprocessor.
"""
def __init__(self, loss, postprocessor):
assert isinstance(loss, dict)
assert isinstance(postprocessor, dict)
self.loss_module = build_loss(loss)
self.postprocessor = build_postprocessor(postprocessor)
def resize_boundary(self, boundaries, scale_factor):
"""Rescale boundaries via scale_factor.
Args:
boundaries (list[list[float]]): The boundary list. Each boundary
has :math:`2k+1` elements with :math:`k>=4`.
scale_factor (ndarray): The scale factor of size :math:`(4,)`.
Returns:
list[list[float]]: The scaled boundaries.
"""
assert check_argument.is_2dlist(boundaries)
assert isinstance(scale_factor, np.ndarray)
assert scale_factor.shape[0] == 4
for b in boundaries:
sz = len(b)
check_argument.valid_boundary(b, True)
b[:sz -
1] = (np.array(b[:sz - 1]) *
(np.tile(scale_factor[:2], int(
(sz - 1) / 2)).reshape(1, sz - 1))).flatten().tolist()
return boundaries
def get_boundary(self, score_maps, img_metas, rescale):
"""Compute text boundaries via post processing.
Args:
score_maps (Tensor): The text score map.
img_metas (dict): The image meta info.
rescale (bool): Rescale boundaries to the original image resolution
if true, and keep the score_maps resolution if false.
Returns:
dict: A dict where boundary results are stored in
``boundary_result``.
"""
assert check_argument.is_type_list(img_metas, dict)
assert isinstance(rescale, bool)
score_maps = score_maps.squeeze()
boundaries = self.postprocessor(score_maps)
if rescale:
boundaries = self.resize_boundary(
boundaries,
1.0 / self.downsample_ratio / img_metas[0]['scale_factor'])
results = dict(
boundary_result=boundaries, filename=img_metas[0]['filename'])
return results
def loss(self, pred_maps, **kwargs):
"""Compute the loss for scene text detection.
Args:
pred_maps (Tensor): The input score maps of shape
:math:`(NxCxHxW)`.
Returns:
dict: The dict for losses.
"""
losses = self.loss_module(pred_maps, self.downsample_ratio, **kwargs)
return losses
|