Spaces:
Running
Running
File size: 4,005 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 |
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Dict, Optional, Sequence, Union
import torch
import torch.nn as nn
from mmengine.model import Sequential
from mmocr.models.common.dictionary import Dictionary
from mmocr.models.textrecog.layers import BidirectionalLSTM
from mmocr.registry import MODELS
from mmocr.structures import TextRecogDataSample
from .base import BaseDecoder
@MODELS.register_module()
class CRNNDecoder(BaseDecoder):
"""Decoder for CRNN.
Args:
in_channels (int): Number of input channels.
dictionary (dict or :obj:`Dictionary`): The config for `Dictionary` or
the instance of `Dictionary`.
rnn_flag (bool): Use RNN or CNN as the decoder. Defaults to False.
module_loss (dict, optional): Config to build module_loss. Defaults
to None.
postprocessor (dict, optional): Config to build postprocessor.
Defaults to None.
init_cfg (dict or list[dict], optional): Initialization configs.
Defaults to None.
"""
def __init__(self,
in_channels: int,
dictionary: Union[Dictionary, Dict],
rnn_flag: bool = False,
module_loss: Dict = None,
postprocessor: Dict = None,
init_cfg=dict(type='Xavier', layer='Conv2d'),
**kwargs):
super().__init__(
init_cfg=init_cfg,
dictionary=dictionary,
module_loss=module_loss,
postprocessor=postprocessor)
self.rnn_flag = rnn_flag
if rnn_flag:
self.decoder = Sequential(
BidirectionalLSTM(in_channels, 256, 256),
BidirectionalLSTM(256, 256, self.dictionary.num_classes))
else:
self.decoder = nn.Conv2d(
in_channels,
self.dictionary.num_classes,
kernel_size=1,
stride=1)
self.softmax = nn.Softmax(dim=-1)
def forward_train(
self,
feat: torch.Tensor,
out_enc: Optional[torch.Tensor] = None,
data_samples: Optional[Sequence[TextRecogDataSample]] = None
) -> torch.Tensor:
"""
Args:
feat (Tensor): A Tensor of shape :math:`(N, C, 1, W)`.
out_enc (torch.Tensor, optional): Encoder output. Defaults to None.
data_samples (list[TextRecogDataSample], optional): Batch of
TextRecogDataSample, containing gt_text information. Defaults
to None.
Returns:
Tensor: The raw logit tensor. Shape :math:`(N, W, C)` where
:math:`C` is ``num_classes``.
"""
assert feat.size(2) == 1, 'feature height must be 1'
if self.rnn_flag:
x = feat.squeeze(2) # [N, C, W]
x = x.permute(2, 0, 1) # [W, N, C]
x = self.decoder(x) # [W, N, C]
outputs = x.permute(1, 0, 2).contiguous()
else:
x = self.decoder(feat)
x = x.permute(0, 3, 1, 2).contiguous()
n, w, c, h = x.size()
outputs = x.view(n, w, c * h)
return outputs
def forward_test(
self,
feat: Optional[torch.Tensor] = None,
out_enc: Optional[torch.Tensor] = None,
data_samples: Optional[Sequence[TextRecogDataSample]] = None
) -> torch.Tensor:
"""
Args:
feat (Tensor): A Tensor of shape :math:`(N, C, 1, W)`.
out_enc (torch.Tensor, optional): Encoder output. Defaults to None.
data_samples (list[TextRecogDataSample]): Batch of
TextRecogDataSample, containing ``gt_text`` information.
Defaults to None.
Returns:
Tensor: Character probabilities. of shape
:math:`(N, self.max_seq_len, C)` where :math:`C` is
``num_classes``.
"""
return self.softmax(self.forward_train(feat, out_enc, data_samples))
|