summaryrefslogtreecommitdiff
path: root/text_recognizer/data/mappings/emnist.py
blob: 606d20090e8cd070ebbf524b16eb5bbde79052e8 (plain)
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
"""Emnist mapping."""
import json
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Union, Tuple

import torch
from torch import Tensor

ESSENTIALS_FILENAME = Path(__file__).parents[0].resolve() / "emnist_essentials.json"


class EmnistMapping:
    """Mapping for EMNIST labels."""

    def __init__(
        self,
        input_size: List[int],
        mapping: List[str],
        inverse_mapping: Dict[str, int],
        extra_symbols: Optional[Sequence[str]] = None,
        lower: bool = True,
    ) -> None:
        self.input_size = input_size
        self.mapping = mapping
        self.inverse_mapping = inverse_mapping
        self.extra_symbols = set(extra_symbols) if extra_symbols is not None else None
        self.mapping, self.inverse_mapping, self.input_size = self._load_mapping()
        if lower:
            self._to_lower()

    def _load_mapping(self) -> Tuple[List, Dict[str, int], List[int]]:
        """Return the EMNIST mapping."""
        with ESSENTIALS_FILENAME.open() as f:
            essentials = json.load(f)
        mapping = list(essentials["characters"])
        if self.extra_symbols is not None:
            mapping += self.extra_symbols
        inverse_mapping = {v: k for k, v in enumerate(mapping)}
        input_shape = essentials["input_shape"]
        return mapping, inverse_mapping, input_shape

    def _to_lower(self) -> None:
        """Converts mapping to lowercase letters only."""

        def _filter(x: int) -> int:
            if 40 <= x:
                return x - 26
            return x

        self.inverse_mapping = {v: _filter(k) for k, v in enumerate(self.mapping)}
        self.mapping = [c for c in self.mapping if not c.isupper()]

    def get_token(self, index: Union[int, Tensor]) -> str:
        """Returns token for index value."""
        if (index := int(index)) <= len(self.mapping):
            return self.mapping[index]
        raise KeyError(f"Index ({index}) not in mapping.")

    def get_index(self, token: str) -> Tensor:
        """Returns index value of token."""
        if token in self.inverse_mapping:
            return torch.LongTensor([self.inverse_mapping[token]])
        raise KeyError(f"Token ({token}) not found in inverse mapping.")

    def get_text(self, indices: Union[List[int], Tensor]) -> str:
        """Returns the text from a list of indices."""
        if isinstance(indices, Tensor):
            indices = indices.tolist()
        return "".join([self.mapping[index] for index in indices])

    def get_indices(self, text: str) -> Tensor:
        """Returns tensor of indices for a string."""
        return Tensor([self.inverse_mapping[token] for token in text])

    def __getitem__(self, x: Union[int, Tensor]) -> str:
        """Returns text for a list of indices."""
        return self.get_token(x)