| import torch |
| import torch.nn as nn |
| from util.util import PosCNN, PositionalEncoding |
| from .Attention import Block |
| from params import * |
| import torch.nn.functional as F |
|
|
| class LayerNorm(nn.Module): |
| def forward(self, x): |
| return F.layer_norm(x, x.size()[1:], weight=None, bias=None, eps=1e-05) |
| |
| class Writer(nn.Module): |
| def __init__( |
| self, |
| num_classes= NUM_WRITERS, |
| embed_dim=256, |
| num_heads=4, |
| mlp_ratio=4.0, |
| qkv_bias=True, |
| qk_scale=None, |
| drop_rate=0.0, |
| attn_drop_rate=0.0, |
| drop_path_rate=0.0, |
| norm_layer=nn.LayerNorm, |
| max_num_patch=1000, |
| ): |
| super(Writer, self).__init__() |
| self.embed_dim = embed_dim |
| depth = 3 |
| norm_layer = nn.LayerNorm |
| self.layer_norm = LayerNorm() |
| patch_size = 4 |
| self.patch = nn.Conv2d( |
| 1, |
| self.embed_dim, |
| kernel_size=patch_size * 2, |
| stride=patch_size, |
| padding=patch_size // 2, |
| ) |
| self.pos_block = PosCNN(self.embed_dim, self.embed_dim) |
| self.pos_enc = PositionalEncoding(embed_dim, drop_rate, max_num_patch) |
| self.norm = nn.LayerNorm(self.embed_dim) |
| self.downsample_blocks = nn.ModuleList( |
| [ |
| Block( |
| dim=embed_dim, |
| num_heads=num_heads, |
| mlp_ratio=mlp_ratio, |
| qkv_bias=qkv_bias, |
| norm_layer=norm_layer, |
| ) |
| for i in range(depth) |
| ] |
| ) |
| self.avgpool = nn.AdaptiveAvgPool1d(1) |
| self.head = ( |
| nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity() |
| ) |
|
|
| self.cross_entropy = nn.CrossEntropyLoss() |
| self.initialize_weights() |
|
|
| def initialize_weights(self): |
| self.apply(self._init_weights) |
|
|
| def _init_weights(self, m): |
| if isinstance(m, nn.Linear): |
| if isinstance(m, nn.Linear) and m.bias is not None: |
| nn.init.constant_(m.bias, 0) |
| elif isinstance(m, nn.LayerNorm): |
| nn.init.constant_(m.bias, 0) |
| nn.init.constant_(m.weight, 1.0) |
| elif isinstance(m, nn.BatchNorm1d): |
| nn.init.constant_(m.bias, 0) |
| nn.init.constant_(m.weight, 1.0) |
| elif isinstance(m, nn.InstanceNorm1d): |
| nn.init.constant_(m.bias, 0) |
| nn.init.constant_(m.weight, 1.0) |
|
|
| def forward(self, x, y = None,training=True): |
|
|
| x = self.layer_norm(x) |
| x = self.patch(x) |
|
|
| """ block 1""" |
| b, c, h, w = x.shape |
| x = x.view(b, c, -1).permute(0, 2, 1) |
| for j, blk in enumerate(self.downsample_blocks): |
| x = blk(x) |
| if j == 0: |
| x = self.pos_block(x, h, w) |
|
|
| """head""" |
| x = self.norm(x) |
| feature = x |
|
|
| x = self.avgpool(x.transpose(1, 2)) |
| x = torch.flatten(x, 1) |
| output = self.head(x) |
| if training: |
| output = self.cross_entropy(output, y.long()) |
| return feature, output |
| else: |
| return feature |
| |
|
|
|
|
| class strLabelConverter(object): |
| """Convert between str and label. |
| NOTE: |
| Insert `blank` to the alphabet for CTC. |
| Args: |
| alphabet (str): set of the possible characters. |
| ignore_case (bool, default=True): whether or not to ignore all of the case. |
| """ |
|
|
| def __init__(self, alphabet, ignore_case=False): |
| self._ignore_case = ignore_case |
| if self._ignore_case: |
| alphabet = alphabet.lower() |
| self.alphabet = alphabet + '-' |
|
|
| self.dict = {} |
| for i, char in enumerate(alphabet): |
| |
| self.dict[char] = i + 1 |
|
|
| def encode(self, text): |
| """Support batch or single str. |
| Args: |
| text (str or list of str): texts to convert. |
| Returns: |
| torch.IntTensor [length_0 + length_1 + ... length_{n - 1}]: encoded texts. |
| torch.IntTensor [n]: length of each text. |
| """ |
| ''' |
| if isinstance(text, str): |
| text = [ |
| self.dict[char.lower() if self._ignore_case else char] |
| for char in text |
| ] |
| length = [len(text)] |
| elif isinstance(text, collections.Iterable): |
| length = [len(s) for s in text] |
| text = ''.join(text) |
| text, _ = self.encode(text) |
| return (torch.IntTensor(text), torch.IntTensor(length)) |
| ''' |
| length = [] |
| result = [] |
| results = [] |
| for item in text: |
| item = item.decode('utf-8', 'strict') |
| length.append(len(item)) |
| for char in item: |
| index = self.dict[char] |
| result.append(index) |
| results.append(result) |
| result = [] |
|
|
| return (torch.nn.utils.rnn.pad_sequence([torch.LongTensor(text) for text in results], batch_first=True), torch.IntTensor(length)) |
|
|
| def decode(self, t, length, raw=False): |
| """Decode encoded texts back into strs. |
| Args: |
| torch.IntTensor [length_0 + length_1 + ... length_{n - 1}]: encoded texts. |
| torch.IntTensor [n]: length of each text. |
| Raises: |
| AssertionError: when the texts and its length does not match. |
| Returns: |
| text (str or list of str): texts to convert. |
| """ |
| if length.numel() == 1: |
| length = length[0] |
| assert t.numel() == length, "text with length: {} does not match declared length: {}".format(t.numel(), |
| length) |
| if raw: |
| return ''.join([self.alphabet[i - 1] for i in t]) |
| else: |
| char_list = [] |
| for i in range(length): |
| if t[i] != 0 and (not (i > 0 and t[i - 1] == t[i])): |
| char_list.append(self.alphabet[t[i] - 1]) |
| return ''.join(char_list) |
| else: |
| |
| assert t.numel() == length.sum(), "texts with length: {} does not match declared length: {}".format( |
| t.numel(), length.sum()) |
| texts = [] |
| index = 0 |
| for i in range(length.numel()): |
| l = length[i] |
| texts.append( |
| self.decode( |
| t[index:index + l], torch.IntTensor([l]), raw=raw)) |
| index += l |
| return texts |
