from mmdet.models.builder import DETECTORS
from mmdet.models.detectors import SingleStageDetector
[docs]@DETECTORS.register_module()
class SingleStageTextDetector(SingleStageDetector):
"""The class for implementing single stage text detector.
It is the parent class of PANet, PSENet, and DBNet.
"""
def __init__(self,
backbone,
neck,
bbox_head,
train_cfg=None,
test_cfg=None,
pretrained=None):
SingleStageDetector.__init__(self, backbone, neck, bbox_head,
train_cfg, test_cfg, pretrained)
[docs] def forward_train(self, img, img_metas, **kwargs):
"""
Args:
img (Tensor): Input images of shape (N, C, H, W).
Typically these should be mean centered and std scaled.
img_metas (list[dict]): A list of image info dict where each dict
has: 'img_shape', 'scale_factor', 'flip', and may also contain
'filename', 'ori_shape', 'pad_shape', and 'img_norm_cfg'.
For details on the values of these keys, see
:class:`mmdet.datasets.pipelines.Collect`.
Returns:
dict[str, Tensor]: A dictionary of loss components.
"""
x = self.extract_feat(img)
preds = self.bbox_head(x)
losses = self.bbox_head.loss(preds, **kwargs)
return losses
def simple_test(self, img, img_metas, rescale=False):
x = self.extract_feat(img)
outs = self.bbox_head(x)
boundaries = self.bbox_head.get_boundary(*outs, img_metas, rescale)
return [boundaries]