diff options
Diffstat (limited to 'text_recognizer/data/mappings/base_mapping.py')
-rw-r--r-- | text_recognizer/data/mappings/base_mapping.py | 37 |
1 files changed, 37 insertions, 0 deletions
diff --git a/text_recognizer/data/mappings/base_mapping.py b/text_recognizer/data/mappings/base_mapping.py new file mode 100644 index 0000000..572ac95 --- /dev/null +++ b/text_recognizer/data/mappings/base_mapping.py @@ -0,0 +1,37 @@ +"""Mapping to and from word pieces.""" +from abc import ABC, abstractmethod +from typing import Dict, List + +from torch import Tensor + + +class AbstractMapping(ABC): + def __init__( + self, input_size: List[int], mapping: List[str], inverse_mapping: Dict[str, int] + ) -> None: + self.input_size = input_size + self.mapping = mapping + self.inverse_mapping = inverse_mapping + + def __len__(self) -> int: + return len(self.mapping) + + @property + def num_classes(self) -> int: + return self.__len__() + + @abstractmethod + def get_token(self, *args, **kwargs) -> str: + ... + + @abstractmethod + def get_index(self, *args, **kwargs) -> Tensor: + ... + + @abstractmethod + def get_text(self, *args, **kwargs) -> str: + ... + + @abstractmethod + def get_indices(self, *args, **kwargs) -> Tensor: + ... |