Ehtesham123/OCR_AEB_Serial_Number

1# Scene Text Recognition Model Hub
2# Copyright 2022 Darwin Bautista
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     https://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import re
17from abc import ABC, abstractmethod
18from itertools import groupby
19from typing import Optional
20
21import torch
22from torch import Tensor
23from torch.nn.utils.rnn import pad_sequence
24
25
26class CharsetAdapter:
27    """Transforms labels according to the target charset."""
28
29    def __init__(self, target_charset) -> None:
30        super().__init__()
31        self.lowercase_only = target_charset == target_charset.lower()
32        self.uppercase_only = target_charset == target_charset.upper()
33        self.unsupported = re.compile(f'[^{re.escape(target_charset)}]')
34
35    def __call__(self, label):
36        if self.lowercase_only:
37            label = label.lower()
38        elif self.uppercase_only:
39            label = label.upper()
40        # Remove unsupported characters
41        label = self.unsupported.sub('', label)
42        return label
43
44
45class BaseTokenizer(ABC):
46
47    def __init__(self, charset: str, specials_first: tuple = (), specials_last: tuple = ()) -> None:
48        self._itos = specials_first + tuple(charset) + specials_last
49        self._stoi = {s: i for i, s in enumerate(self._itos)}
50
51    def __len__(self):
52        return len(self._itos)
53
54    def _tok2ids(self, tokens: str) -> list[int]:
55        return [self._stoi[s] for s in tokens]
56
57    def _ids2tok(self, token_ids: list[int], join: bool = True) -> str:
58        tokens = [self._itos[i] for i in token_ids]
59        return ''.join(tokens) if join else tokens
60
61    @abstractmethod
62    def encode(self, labels: list[str], device: Optional[torch.device] = None) -> Tensor:
63        """Encode a batch of labels to a representation suitable for the model.
64
65        Args:
66            labels: List of labels. Each can be of arbitrary length.
67            device: Create tensor on this device.
68
69        Returns:
70            Batched tensor representation padded to the max label length. Shape: N, L
71        """
72        raise NotImplementedError
73
74    @abstractmethod
75    def _filter(self, probs: Tensor, ids: Tensor) -> tuple[Tensor, list[int]]:
76        """Internal method which performs the necessary filtering prior to decoding."""
77        raise NotImplementedError
78
79    def decode(self, token_dists: Tensor, raw: bool = False) -> tuple[list[str], list[Tensor]]:
80        """Decode a batch of token distributions.
81
82        Args:
83            token_dists: softmax probabilities over the token distribution. Shape: N, L, C
84            raw: return unprocessed labels (will return list of list of strings)
85
86        Returns:
87            list of string labels (arbitrary length) and
88            their corresponding sequence probabilities as a list of Tensors
89        """
90        batch_tokens = []
91        batch_probs = []
92        for dist in token_dists:
93            probs, ids = dist.max(-1)  # greedy selection
94            if not raw:
95                probs, ids = self._filter(probs, ids)
96            tokens = self._ids2tok(ids, not raw)
97            batch_tokens.append(tokens)
98            batch_probs.append(probs)
99        return batch_tokens, batch_probs
100
101
102class Tokenizer(BaseTokenizer):
103    BOS = '[B]'
104    EOS = '[E]'
105    PAD = '[P]'
106
107    def __init__(self, charset: str) -> None:
108        specials_first = (self.EOS,)
109        specials_last = (self.BOS, self.PAD)
110        super().__init__(charset, specials_first, specials_last)
111        self.eos_id, self.bos_id, self.pad_id = [self._stoi[s] for s in specials_first + specials_last]
112
113    def encode(self, labels: list[str], device: Optional[torch.device] = None) -> Tensor:
114        batch = [
115            torch.as_tensor([self.bos_id] + self._tok2ids(y) + [self.eos_id], dtype=torch.long, device=device)
116            for y in labels
117        ]
118        return pad_sequence(batch, batch_first=True, padding_value=self.pad_id)
119
120    def _filter(self, probs: Tensor, ids: Tensor) -> tuple[Tensor, list[int]]:
121        ids = ids.tolist()
122        try:
123            eos_idx = ids.index(self.eos_id)
124        except ValueError:
125            eos_idx = len(ids)  # Nothing to truncate.
126        # Truncate after EOS
127        ids = ids[:eos_idx]
128        probs = probs[: eos_idx + 1]  # but include prob. for EOS (if it exists)
129        return probs, ids
130
131
132class CTCTokenizer(BaseTokenizer):
133    BLANK = '[B]'
134
135    def __init__(self, charset: str) -> None:
136        # BLANK uses index == 0 by default
137        super().__init__(charset, specials_first=(self.BLANK,))
138        self.blank_id = self._stoi[self.BLANK]
139
140    def encode(self, labels: list[str], device: Optional[torch.device] = None) -> Tensor:
141        # We use a padded representation since we don't want to use CUDNN's CTC implementation
142        batch = [torch.as_tensor(self._tok2ids(y), dtype=torch.long, device=device) for y in labels]
143        return pad_sequence(batch, batch_first=True, padding_value=self.blank_id)
144
145    def _filter(self, probs: Tensor, ids: Tensor) -> tuple[Tensor, list[int]]:
146        # Best path decoding:
147        ids = list(zip(*groupby(ids.tolist())))[0]  # Remove duplicate tokens
148        ids = [x for x in ids if x != self.blank_id]  # Remove BLANKs
149        # `probs` is just pass-through since all positions are considered part of the path
150        return probs, ids
151