From 8291a87c64f9a5f18caec82201bea15579b49730 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Sun, 10 Oct 2021 18:04:50 +0200 Subject: Move data utils to submodules --- text_recognizer/data/base_mapping.py | 37 --- text_recognizer/data/build_transitions.py | 261 --------------------- text_recognizer/data/download_utils.py | 73 ------ text_recognizer/data/emnist_essentials.json | 1 - text_recognizer/data/emnist_mapping.py | 60 ----- text_recognizer/data/iam_preprocessor.py | 211 ----------------- text_recognizer/data/image_utils.py | 49 ---- text_recognizer/data/make_wordpieces.py | 112 --------- text_recognizer/data/mappings/base_mapping.py | 37 +++ .../data/mappings/emnist_essentials.json | 1 + text_recognizer/data/mappings/emnist_mapping.py | 60 +++++ .../data/mappings/word_piece_mapping.py | 98 ++++++++ text_recognizer/data/sentence_generator.py | 89 ------- text_recognizer/data/transforms.py | 49 ---- text_recognizer/data/transforms/word_piece.py | 48 ++++ text_recognizer/data/utils/build_transitions.py | 261 +++++++++++++++++++++ text_recognizer/data/utils/download_utils.py | 73 ++++++ text_recognizer/data/utils/iam_preprocessor.py | 209 +++++++++++++++++ text_recognizer/data/utils/image_utils.py | 49 ++++ text_recognizer/data/utils/make_wordpieces.py | 112 +++++++++ text_recognizer/data/utils/sentence_generator.py | 89 +++++++ text_recognizer/data/word_piece_mapping.py | 98 -------- 22 files changed, 1037 insertions(+), 1040 deletions(-) delete mode 100644 text_recognizer/data/base_mapping.py delete mode 100644 text_recognizer/data/build_transitions.py delete mode 100644 text_recognizer/data/download_utils.py delete mode 100644 text_recognizer/data/emnist_essentials.json delete mode 100644 text_recognizer/data/emnist_mapping.py delete mode 100644 text_recognizer/data/iam_preprocessor.py delete mode 100644 text_recognizer/data/image_utils.py delete mode 100644 text_recognizer/data/make_wordpieces.py create mode 100644 text_recognizer/data/mappings/base_mapping.py create mode 100644 text_recognizer/data/mappings/emnist_essentials.json create mode 100644 text_recognizer/data/mappings/emnist_mapping.py create mode 100644 text_recognizer/data/mappings/word_piece_mapping.py delete mode 100644 text_recognizer/data/sentence_generator.py delete mode 100644 text_recognizer/data/transforms.py create mode 100644 text_recognizer/data/transforms/word_piece.py create mode 100644 text_recognizer/data/utils/build_transitions.py create mode 100644 text_recognizer/data/utils/download_utils.py create mode 100644 text_recognizer/data/utils/iam_preprocessor.py create mode 100644 text_recognizer/data/utils/image_utils.py create mode 100644 text_recognizer/data/utils/make_wordpieces.py create mode 100644 text_recognizer/data/utils/sentence_generator.py delete mode 100644 text_recognizer/data/word_piece_mapping.py diff --git a/text_recognizer/data/base_mapping.py b/text_recognizer/data/base_mapping.py deleted file mode 100644 index 572ac95..0000000 --- a/text_recognizer/data/base_mapping.py +++ /dev/null @@ -1,37 +0,0 @@ -"""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: - ... diff --git a/text_recognizer/data/build_transitions.py b/text_recognizer/data/build_transitions.py deleted file mode 100644 index 0f987ca..0000000 --- a/text_recognizer/data/build_transitions.py +++ /dev/null @@ -1,261 +0,0 @@ -"""Builds transition graph. - -Most code stolen from here: - https://github.com/facebookresearch/gtn_applications/blob/master/scripts/build_transitions.py -""" - -import collections -import itertools -from pathlib import Path -from typing import Any, Dict, List, Optional - -import click -import gtn -from loguru import logger - - -START_IDX = -1 -END_IDX = -2 -WORDSEP = "▁" - - -def build_graph(ngrams: List, disable_backoff: bool = False) -> gtn.Graph: - """Returns a gtn Graph based on the ngrams.""" - graph = gtn.Graph(False) - ngram = len(ngrams) - state_to_node = {} - - def get_node(state: Optional[List]) -> Any: - node = state_to_node.get(state, None) - - if node is not None: - return node - - start = state == tuple([START_IDX]) if ngram > 1 else True - end = state == tuple([END_IDX]) if ngram > 1 else True - node = graph.add_node(start, end) - state_to_node[state] = node - - if not disable_backoff and not end: - # Add back off when adding node. - for n in range(1, len(state) + 1): - backoff_node = state_to_node.get(state[n:], None) - - # Epsilon transition to the back-off state. - if backoff_node is not None: - graph.add_arc(node, backoff_node, gtn.epsilon) - break - return node - - for grams in ngrams: - for gram in grams: - istate, ostate = gram[:-1], gram[len(gram) - ngram + 1 :] - inode = get_node(istate) - - if END_IDX not in gram[1:] and gram[1:] not in state_to_node: - raise ValueError( - "Ill formed counts: if (x, y_1, ..., y_{n-1}) is above" - "the n-gram threshold, then (y_1, ..., y_{n-1}) must be" - "above the (n-1)-gram threshold" - ) - - if END_IDX in ostate: - # Merge all state having into one as final graph generated - # will be similar. - ostate = tuple([END_IDX]) - - onode = get_node(ostate) - # p(gram[-1] | gram[:-1]) - graph.add_arc( - inode, onode, gtn.epsilon if gram[-1] == END_IDX else gram[-1] - ) - return graph - - -def count_ngrams(lines: List, ngram: List, tokens_to_index: Dict) -> List: - """Counts the number of ngrams.""" - counts = [collections.Counter() for _ in range(ngram)] - for line in lines: - # Prepend implicit start token. - token_line = [START_IDX] - for t in line: - token_line.append(tokens_to_index[t]) - token_line.append(END_IDX) - for n, counter in enumerate(counts): - start_offset = n == 0 - end_offset = ngram == 1 - for e in range(n + start_offset, len(token_line) - end_offset): - counter[tuple(token_line[e - n : e + 1])] += 1 - - return counts - - -def prune_ngrams(ngrams: List, prune: List) -> List: - """Prunes ngrams.""" - pruned_ngrams = [] - for n, grams in enumerate(ngrams): - grams = grams.most_common() - pruned_grams = [gram for gram, c in grams if c > prune[n]] - pruned_ngrams.append(pruned_grams) - return pruned_ngrams - - -def add_blank_grams(pruned_ngrams: List, num_tokens: int, blank: str) -> List: - """Adds blank token to grams.""" - all_grams = [gram for grams in pruned_ngrams for gram in grams] - maxorder = len(pruned_ngrams) - blank_grams = {} - if blank == "forced": - pruned_ngrams = [pruned_ngrams[0] if i == 0 else [] for i in range(maxorder)] - pruned_ngrams[0].append(tuple([num_tokens])) - blank_grams[tuple([num_tokens])] = True - - for gram in all_grams: - # Iterate over all possibilities by using a vector of 0s, 1s to - # denote whether a blank is being used at each position. - if blank == "optional": - # Given a gram ab.. if order n, we have n + 1 positions - # available whether to use blank or not. - onehot_vectors = itertools.product([0, 1], repeat=len(gram) + 1) - elif blank == "forced": - # Must include a blank token in between. - onehot_vectors = [[1] * (len(gram) + 1)] - else: - raise ValueError( - "Invalid value specificed for blank. Must be in |optional|forced|none|" - ) - - for j in onehot_vectors: - new_array = [] - for idx, oz in enumerate(j[:-1]): - if oz == 1 and gram[idx] != START_IDX: - new_array.append(num_tokens) - new_array.append(gram[idx]) - if j[-1] == 1 and gram[-1] != END_IDX: - new_array.append(num_tokens) - for n in range(maxorder): - for e in range(n, len(new_array)): - cur_gram = tuple(new_array[e - n : e + 1]) - if num_tokens in cur_gram and cur_gram not in blank_grams: - pruned_ngrams[n].append(cur_gram) - blank_grams[cur_gram] = True - - return pruned_ngrams - - -def add_self_loops(pruned_ngrams: List) -> List: - """Adds self loops to the ngrams.""" - maxorder = len(pruned_ngrams) - - # Use dict for fast search. - all_grams = set([gram for grams in pruned_ngrams for gram in grams]) - for o in range(1, maxorder): - for gram in pruned_ngrams[o - 1]: - # Repeat one of the tokens. - for pos in range(len(gram)): - if gram[pos] == START_IDX or gram[pos] == END_IDX: - continue - new_gram = gram[:pos] + (gram[pos],) + gram[pos:] - - if new_gram not in all_grams: - pruned_ngrams[o].append(new_gram) - all_grams.add(new_gram) - return pruned_ngrams - - -def parse_lines(lines: List, lexicon: Path) -> List: - """Parses lines with a lexicon.""" - with open(lexicon, "r") as f: - lex = (line.strip().split() for line in f) - lex = {line[0]: line[1:] for line in lex} - print(len(lex)) - return [[t for w in line.split(WORDSEP) for t in lex[w]] for line in lines] - - -@click.command() -@click.option("--data_dir", type=str, default=None, help="Path to dataset root.") -@click.option( - "--tokens", type=str, help="Path to token list (in order used with training)." -) -@click.option("--lexicon", type=str, default=None, help="Path to lexicon") -@click.option( - "--prune", - nargs=2, - type=int, - help="Threshold values for prune unigrams, bigrams, etc.", -) -@click.option( - "--blank", - default=click.Choice(["none", "optional", "forced"]), - help="Specifies the usage of blank token" - "'none' - do not use blank token " - "'optional' - allow an optional blank inbetween tokens" - "'forced' - force a blank inbetween tokens (also referred to as garbage token)", -) -@click.option("--self_loops", is_flag=True, help="Add self loops for tokens") -@click.option("--disable_backoff", is_flag=True, help="Disable backoff transitions") -@click.option("--save_path", default=None, help="Path to save transition graph.") -def cli( - data_dir: str, - tokens: str, - lexicon: str, - prune: List[int], - blank: str, - self_loops: bool, - disable_backoff: bool, - save_path: str, -) -> None: - """CLI for creating the transitions.""" - logger.info(f"Building {len(prune)}-gram transition models.") - - if data_dir is None: - data_dir = ( - Path(__file__).resolve().parents[2] / "data" / "processed" / "iam_lines" - ) - logger.debug(f"Using data dir: {data_dir}") - if not data_dir.exists(): - raise RuntimeError(f"Could not locate iamdb directory at {data_dir}") - else: - data_dir = Path(data_dir) - - # Build table of counts and the back-off if below threshold. - with open(data_dir / "train.txt", "r") as f: - lines = [line.strip() for line in f] - - with open(data_dir / tokens, "r") as f: - tokens = [line.strip() for line in f] - - if lexicon is not None: - lexicon = data_dir / lexicon - lines = parse_lines(lines, lexicon) - - tokens_to_idx = {t: e for e, t in enumerate(tokens)} - - ngram = len(prune) - - logger.info("Counting data...") - ngrams = count_ngrams(lines, ngram, tokens_to_idx) - - pruned_ngrams = prune_ngrams(ngrams, prune) - - for n in range(ngram): - logger.info(f"Kept {len(pruned_ngrams[n])} of {len(ngrams[n])} {n + 1}-grams") - - if blank == "none": - pruned_ngrams = add_blank_grams(pruned_ngrams, len(tokens_to_idx), blank) - - if self_loops: - pruned_ngrams = add_self_loops(pruned_ngrams) - - logger.info("Building graph from pruned ngrams...") - graph = build_graph(pruned_ngrams, disable_backoff) - logger.info(f"Graph has {graph.num_arcs()} arcs and {graph.num_nodes()} nodes.") - - save_path = str(data_dir / save_path) - - logger.info(f"Saving graph to {save_path}") - gtn.save(save_path, graph) - - -if __name__ == "__main__": - cli() diff --git a/text_recognizer/data/download_utils.py b/text_recognizer/data/download_utils.py deleted file mode 100644 index a5a5360..0000000 --- a/text_recognizer/data/download_utils.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Util functions for downloading datasets.""" -import hashlib -from pathlib import Path -from typing import Dict, Optional -from urllib.request import urlretrieve - -from loguru import logger as log -from tqdm import tqdm - - -def _compute_sha256(filename: Path) -> str: - """Returns the SHA256 checksum of a file.""" - with filename.open(mode="rb") as f: - return hashlib.sha256(f.read()).hexdigest() - - -class TqdmUpTo(tqdm): - """TQDM progress bar when downloading files. - - From https://github.com/tqdm/tqdm/blob/master/examples/tqdm_wget.py - - """ - - def update_to( - self, blocks: int = 1, block_size: int = 1, total_size: Optional[int] = None - ) -> None: - """Updates the progress bar. - - Args: - blocks (int): Number of blocks transferred so far. Defaults to 1. - block_size (int): Size of each block, in tqdm units. Defaults to 1. - total_size (Optional[int]): Total size in tqdm units. Defaults to None. - """ - if total_size is not None: - self.total = total_size - self.update(blocks * block_size - self.n) - - -def _download_url(url: str, filename: str) -> None: - """Downloads a file from url to filename, with a progress bar.""" - with TqdmUpTo(unit="B", unit_scale=True, unit_divisor=1024, miniters=1) as t: - urlretrieve(url, filename, reporthook=t.update_to, data=None) # nosec - - -def download_dataset(metadata: Dict, dl_dir: Path) -> Optional[Path]: - """Downloads dataset using a metadata file. - - Args: - metadata (Dict): A metadata file of the dataset. - dl_dir (Path): Download directory for the dataset. - - Returns: - Optional[Path]: Returns filename if dataset is downloaded, None if it already - exists. - - Raises: - ValueError: If the SHA-256 value is not the same between the dataset and - the metadata file. - - """ - dl_dir.mkdir(parents=True, exist_ok=True) - filename = dl_dir / metadata["filename"] - if filename.exists(): - return - log.info(f"Downloading raw dataset from {metadata['url']} to {filename}...") - _download_url(metadata["url"], filename) - log.info("Computing the SHA-256...") - sha256 = _compute_sha256(filename) - if sha256 != metadata["sha256"]: - raise ValueError( - "Downloaded data file SHA-256 does not match that listed in metadata document." - ) - return filename diff --git a/text_recognizer/data/emnist_essentials.json b/text_recognizer/data/emnist_essentials.json deleted file mode 100644 index c412425..0000000 --- a/text_recognizer/data/emnist_essentials.json +++ /dev/null @@ -1 +0,0 @@ -{"characters": ["", "", "", "

", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", " ", "!", "\"", "#", "&", "'", "(", ")", "*", "+", ",", "-", ".", "/", ":", ";", "?"], "input_shape": [28, 28]} \ No newline at end of file diff --git a/text_recognizer/data/emnist_mapping.py b/text_recognizer/data/emnist_mapping.py deleted file mode 100644 index b2165d2..0000000 --- a/text_recognizer/data/emnist_mapping.py +++ /dev/null @@ -1,60 +0,0 @@ -"""Emnist mapping.""" -from typing import List, Optional, Set, Union - -import torch -from torch import Tensor - -from text_recognizer.data.base_mapping import AbstractMapping -from text_recognizer.data.emnist import emnist_mapping - - -class EmnistMapping(AbstractMapping): - """Mapping for EMNIST labels.""" - - def __init__( - self, extra_symbols: Optional[Set[str]] = None, lower: bool = True - ) -> None: - self.extra_symbols = set(extra_symbols) if extra_symbols is not None else None - self.mapping, self.inverse_mapping, self.input_size = emnist_mapping( - self.extra_symbols - ) - if lower: - self._to_lower() - super().__init__(self.input_size, self.mapping, self.inverse_mapping) - - 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) diff --git a/text_recognizer/data/iam_preprocessor.py b/text_recognizer/data/iam_preprocessor.py deleted file mode 100644 index 700944e..0000000 --- a/text_recognizer/data/iam_preprocessor.py +++ /dev/null @@ -1,211 +0,0 @@ -"""Preprocessor for extracting word letters from the IAM dataset. - -The code is mostly stolen from: - https://github.com/facebookresearch/gtn_applications/blob/master/datasets/iamdb.py -""" -import collections -import itertools -from pathlib import Path -import re -from typing import List, Optional, Set, Union - -import click -from loguru import logger as log -import torch - - -def load_metadata( - data_dir: Path, wordsep: str, use_words: bool = False -) -> collections.defaultdict: - """Loads IAM metadata and returns it as a dictionary.""" - forms = collections.defaultdict(list) - filename = "words.txt" if use_words else "lines.txt" - - with open(data_dir / "ascii" / filename, "r") as f: - lines = (line.strip().split() for line in f if line[0] != "#") - for line in lines: - # Skip word segmentation errors. - if use_words and line[1] == "err": - continue - text = " ".join(line[8:]) - - # Remove garbage tokens: - text = text.replace("#", "") - - # Swap word sep form | to wordsep - text = re.sub(r"\|+|\s", wordsep, text).strip(wordsep) - form_key = "-".join(line[0].split("-")[:2]) - line_key = "-".join(line[0].split("-")[:3]) - box_idx = 4 - use_words - box = tuple(int(val) for val in line[box_idx : box_idx + 4]) - forms[form_key].append({"key": line_key, "box": box, "text": text}) - return forms - - -class Preprocessor: - """A preprocessor for the IAM dataset.""" - - # TODO: attrs - - def __init__( - self, - data_dir: Union[str, Path], - num_features: int, - tokens_path: Optional[Union[str, Path]] = None, - lexicon_path: Optional[Union[str, Path]] = None, - use_words: bool = False, - prepend_wordsep: bool = False, - special_tokens: Optional[Set[str]] = None, - ) -> None: - self.wordsep = "▁" - self._use_word = use_words - self._prepend_wordsep = prepend_wordsep - self.special_tokens = special_tokens if special_tokens is not None else None - self.data_dir = Path(data_dir) - self.forms = load_metadata(self.data_dir, self.wordsep, use_words=use_words) - - # Load the set of graphemes: - graphemes = set() - for _, form in self.forms.items(): - for line in form: - graphemes.update(line["text"].lower()) - self.graphemes = sorted(graphemes) - - # Build the token-to-index and index-to-token maps. - if tokens_path is not None: - with open(tokens_path, "r") as f: - self.tokens = [line.strip() for line in f] - else: - self.tokens = self.graphemes - - if lexicon_path is not None: - with open(lexicon_path, "r") as f: - lexicon = (line.strip().split() for line in f) - lexicon = {line[0]: line[1:] for line in lexicon} - self.lexicon = lexicon - else: - self.lexicon = None - - if self.special_tokens is not None: - special_tokens_ = (*self.special_tokens, "#", "*") - self.tokens += special_tokens_ - self.graphemes += special_tokens_ - - self.graphemes_to_index = {t: i for i, t in enumerate(self.graphemes)} - self.tokens_to_index = {t: i for i, t in enumerate(self.tokens)} - self.num_features = num_features - self.text = [] - - @property - def num_tokens(self) -> int: - """Returns the number or tokens.""" - return len(self.tokens) - - @property - def use_words(self) -> bool: - """If words are used.""" - return self._use_word - - def extract_train_text(self) -> None: - """Extracts training text.""" - keys = [] - with open(self.data_dir / "task" / "trainset.txt") as f: - keys.extend((line.strip() for line in f)) - - for _, examples in self.forms.items(): - for example in examples: - if example["key"] not in keys: - continue - self.text.append(example["text"].lower()) - - def _to_index(self, line: str) -> torch.LongTensor: - if self.special_tokens is not None and line in self.special_tokens: - return torch.LongTensor([self.tokens_to_index[line]]) - token_to_index = self.graphemes_to_index - if self.lexicon is not None: - if len(line) > 0: - # If the word is not found in the lexicon, fall back to letters. - tokens = [ - t - for w in line.split(self.wordsep) - for t in self.lexicon.get(w, self.wordsep + w) - ] - token_to_index = self.tokens_to_index - if self._prepend_wordsep: - tokens = itertools.chain([self.wordsep], tokens) - return torch.LongTensor([token_to_index[t] for t in tokens]) - - def to_index(self, line: str) -> torch.LongTensor: - """Converts text to a tensor of indices.""" - if self.special_tokens is not None: - pattern = f"({'|'.join(self.special_tokens)})" - lines = list(filter(None, re.split(pattern, line))) - return torch.cat([self._to_index(line) for line in lines]) - return self._to_index(line) - - def to_text(self, indices: List[int]) -> str: - """Converts indices to text.""" - # Roughly the inverse of `to_index` - encoding = self.graphemes - if self.lexicon is not None: - encoding = self.tokens - return self._post_process(encoding[i] for i in indices) - - def tokens_to_text(self, indices: List[int]) -> str: - """Converts tokens to text.""" - return self._post_process(self.tokens[i] for i in indices) - - def _post_process(self, indices: List[int]) -> str: - """A list join.""" - return "".join(indices).strip(self.wordsep) - - -@click.command() -@click.option("--data_dir", type=str, default=None, help="Path to iam dataset") -@click.option( - "--use_words", is_flag=True, help="Load word segmented dataset instead of lines" -) -@click.option( - "--save_text", type=str, default=None, help="Path to save parsed train text" -) -@click.option("--save_tokens", type=str, default=None, help="Path to save tokens") -def cli( - data_dir: Optional[str], - use_words: bool, - save_text: Optional[str], - save_tokens: Optional[str], -) -> None: - """CLI for extracting text data from the iam dataset.""" - if data_dir is None: - data_dir = ( - Path(__file__).resolve().parents[2] - / "data" - / "downloaded" - / "iam" - / "iamdb" - ) - log.debug(f"Using data dir: {data_dir}") - if not data_dir.exists(): - raise RuntimeError(f"Could not locate iamdb directory at {data_dir}") - else: - data_dir = Path(data_dir) - - preprocessor = Preprocessor(data_dir, 64, use_words=use_words) - preprocessor.extract_train_text() - - processed_dir = data_dir.parents[2] / "processed" / "iam_lines" - log.debug(f"Saving processed files at: {processed_dir}") - - if save_text is not None: - log.info("Saving training text") - with open(processed_dir / save_text, "w") as f: - f.write("\n".join(t for t in preprocessor.text)) - - if save_tokens is not None: - log.info("Saving tokens") - with open(processed_dir / save_tokens, "w") as f: - f.write("\n".join(preprocessor.tokens)) - - -if __name__ == "__main__": - cli() diff --git a/text_recognizer/data/image_utils.py b/text_recognizer/data/image_utils.py deleted file mode 100644 index c2b8915..0000000 --- a/text_recognizer/data/image_utils.py +++ /dev/null @@ -1,49 +0,0 @@ -"""Image util functions for loading and saving images.""" -from pathlib import Path -from typing import Union -from urllib.request import urlopen - -import cv2 -import numpy as np -from PIL import Image - - -def read_image_pil(image_uri: Union[Path, str], grayscale: bool = False) -> Image: - """Return PIL image.""" - image = Image.open(image_uri) - if grayscale: - image = image.convert("L") - return image - - -def read_image(image_uri: Union[Path, str], grayscale: bool = False) -> np.array: - """Read image_uri.""" - - if isinstance(image_uri, str): - image_uri = Path(image_uri) - - def read_image_from_filename(image_filename: Path, imread_flag: int) -> np.array: - return cv2.imread(str(image_filename), imread_flag) - - def read_image_from_url(image_url: Path, imread_flag: int) -> np.array: - url_response = urlopen(str(image_url)) # nosec - image_array = np.array(bytearray(url_response.read()), dtype=np.uint8) - return cv2.imdecode(image_array, imread_flag) - - imread_flag = cv2.IMREAD_GRAYSCALE if grayscale else cv2.IMREAD_COLOR - image = None - - if image_uri.exists(): - image = read_image_from_filename(image_uri, imread_flag) - else: - image = read_image_from_url(image_uri, imread_flag) - - if image is None: - raise ValueError(f"Could not load image at {image_uri}") - - return image - - -def write_image(image: np.ndarray, filename: Union[Path, str]) -> None: - """Write image to file.""" - cv2.imwrite(str(filename), image) diff --git a/text_recognizer/data/make_wordpieces.py b/text_recognizer/data/make_wordpieces.py deleted file mode 100644 index 8e53815..0000000 --- a/text_recognizer/data/make_wordpieces.py +++ /dev/null @@ -1,112 +0,0 @@ -"""Creates word pieces from a text file. - -Most code stolen from: - - https://github.com/facebookresearch/gtn_applications/blob/master/scripts/make_wordpieces.py - -""" -import io -from pathlib import Path -from typing import List, Optional, Union - -import click -from loguru import logger as log -import sentencepiece as spm - - -def iamdb_pieces( - data_dir: Path, text_file: str, num_pieces: int, output_prefix: str -) -> None: - """Creates word pieces from the iamdb train text.""" - # Load training text. - with open(data_dir / text_file, "r") as f: - text = [line.strip() for line in f] - - sp = train_spm_model( - iter(text), - num_pieces + 1, # To account for - user_symbols=["/"], # added so token is in the output set - ) - - vocab = sorted(set(w for t in text for w in t.split("▁") if w)) - if "move" not in vocab: - raise RuntimeError("`MOVE` not in vocab") - - save_pieces(sp, num_pieces, data_dir, output_prefix, vocab) - - -def train_spm_model( - sentences: iter, vocab_size: int, user_symbols: Union[str, List[str]] = "" -) -> spm.SentencePieceProcessor: - """Trains the sentence piece model.""" - model = io.BytesIO() - spm.SentencePieceTrainer.train( - sentence_iterator=sentences, - model_writer=model, - vocab_size=vocab_size, - bos_id=-1, - eos_id=-1, - character_coverage=1.0, - user_defined_symbols=user_symbols, - ) - sp = spm.SentencePieceProcessor(model_proto=model.getvalue()) - return sp - - -def save_pieces( - sp: spm.SentencePieceProcessor, - num_pieces: int, - data_dir: Path, - output_prefix: str, - vocab: set, -) -> None: - """Saves word pieces to disk.""" - log.info(f"Generating word piece list of size {num_pieces}.") - pieces = [sp.id_to_piece(i) for i in range(1, num_pieces + 1)] - log.info(f"Encoding vocabulary of size {len(vocab)}.") - encoded_vocab = [sp.encode_as_pieces(v) for v in vocab] - - # Save pieces to file. - with open(data_dir / f"{output_prefix}_tokens_{num_pieces}.txt", "w") as f: - f.write("\n".join(pieces)) - - # Save lexicon to a file. - with open(data_dir / f"{output_prefix}_lex_{num_pieces}.txt", "w") as f: - for v, p in zip(vocab, encoded_vocab): - f.write(f"{v} {' '.join(p)}\n") - - -@click.command() -@click.option("--data_dir", type=str, default=None, help="Path to processed iam dir.") -@click.option( - "--text_file", type=str, default=None, help="Name of sentence piece training text." -) -@click.option( - "--output_prefix", - type=str, - default="word_pieces", - help="Prefix name to store tokens and lexicon.", -) -@click.option("--num_pieces", type=int, default=1000, help="Number of word pieces.") -def cli( - data_dir: Optional[str], - text_file: Optional[str], - output_prefix: Optional[str], - num_pieces: Optional[int], -) -> None: - """CLI for training the sentence piece model.""" - if data_dir is None: - data_dir = ( - Path(__file__).resolve().parents[2] / "data" / "processed" / "iam_lines" - ) - log.debug(f"Using data dir: {data_dir}") - if not data_dir.exists(): - raise RuntimeError(f"Could not locate iamdb directory at {data_dir}") - else: - data_dir = Path(data_dir) - - iamdb_pieces(data_dir, text_file, num_pieces, output_prefix) - - -if __name__ == "__main__": - cli() 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: + ... diff --git a/text_recognizer/data/mappings/emnist_essentials.json b/text_recognizer/data/mappings/emnist_essentials.json new file mode 100644 index 0000000..c412425 --- /dev/null +++ b/text_recognizer/data/mappings/emnist_essentials.json @@ -0,0 +1 @@ +{"characters": ["", "", "", "

", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9", "A", "B", "C", "D", "E", "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", "S", "T", "U", "V", "W", "X", "Y", "Z", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", " ", "!", "\"", "#", "&", "'", "(", ")", "*", "+", ",", "-", ".", "/", ":", ";", "?"], "input_shape": [28, 28]} \ No newline at end of file diff --git a/text_recognizer/data/mappings/emnist_mapping.py b/text_recognizer/data/mappings/emnist_mapping.py new file mode 100644 index 0000000..3eed3d8 --- /dev/null +++ b/text_recognizer/data/mappings/emnist_mapping.py @@ -0,0 +1,60 @@ +"""Emnist mapping.""" +from typing import List, Optional, Set, Union + +import torch +from torch import Tensor + +from text_recognizer.data.mappings.base_mapping import AbstractMapping +from text_recognizer.data.emnist import emnist_mapping + + +class EmnistMapping(AbstractMapping): + """Mapping for EMNIST labels.""" + + def __init__( + self, extra_symbols: Optional[Set[str]] = None, lower: bool = True + ) -> None: + self.extra_symbols = set(extra_symbols) if extra_symbols is not None else None + self.mapping, self.inverse_mapping, self.input_size = emnist_mapping( + self.extra_symbols + ) + if lower: + self._to_lower() + super().__init__(self.input_size, self.mapping, self.inverse_mapping) + + 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) diff --git a/text_recognizer/data/mappings/word_piece_mapping.py b/text_recognizer/data/mappings/word_piece_mapping.py new file mode 100644 index 0000000..6f1790e --- /dev/null +++ b/text_recognizer/data/mappings/word_piece_mapping.py @@ -0,0 +1,98 @@ +"""Word piece mapping.""" +from pathlib import Path +from typing import List, Optional, Set, Union + +from loguru import logger as log +import torch +from torch import Tensor + +from text_recognizer.data.mappings.emnist_mapping import EmnistMapping +from text_recognizer.data.utils.iam_preprocessor import Preprocessor + + +class WordPieceMapping(EmnistMapping): + """Word piece mapping.""" + + def __init__( + self, + data_dir: Optional[Path] = None, + num_features: int = 1000, + tokens: str = "iamdb_1kwp_tokens_1000.txt", + lexicon: str = "iamdb_1kwp_lex_1000.txt", + use_words: bool = False, + prepend_wordsep: bool = False, + special_tokens: Set[str] = {"", "", "

"}, + extra_symbols: Set[str] = {"\n"}, + ) -> None: + super().__init__(extra_symbols=extra_symbols) + self.data_dir = ( + ( + Path(__file__).resolve().parents[3] + / "data" + / "downloaded" + / "iam" + / "iamdb" + ) + if data_dir is None + else Path(data_dir) + ) + log.debug(f"Using data dir: {self.data_dir}") + if not self.data_dir.exists(): + raise RuntimeError(f"Could not locate iamdb directory at {self.data_dir}") + + processed_path = ( + Path(__file__).resolve().parents[3] / "data" / "processed" / "iam_lines" + ) + + tokens_path = processed_path / tokens + lexicon_path = processed_path / lexicon + + special_tokens = set(special_tokens) + if self.extra_symbols is not None: + special_tokens = special_tokens | set(extra_symbols) + + self.wordpiece_processor = Preprocessor( + data_dir=self.data_dir, + num_features=num_features, + tokens_path=tokens_path, + lexicon_path=lexicon_path, + use_words=use_words, + prepend_wordsep=prepend_wordsep, + special_tokens=special_tokens, + ) + + def __len__(self) -> int: + """Return number of word pieces.""" + return len(self.wordpiece_processor.tokens) + + def get_token(self, index: Union[int, Tensor]) -> str: + """Returns token for index.""" + if (index := int(index)) <= self.wordpiece_processor.num_tokens: + return self.wordpiece_processor.tokens[index] + raise KeyError(f"Index ({index}) not in mapping.") + + def get_index(self, token: str) -> Tensor: + """Returns index of token.""" + if token in self.wordpiece_processor.tokens: + return torch.LongTensor([self.wordpiece_processor.tokens_to_index[token]]) + raise KeyError(f"Token ({token}) not found in inverse mapping.") + + def get_text(self, indices: Union[List[int], Tensor]) -> str: + """Returns text from indices.""" + if isinstance(indices, Tensor): + indices = indices.tolist() + return self.wordpiece_processor.to_text(indices) + + def get_indices(self, text: str) -> Tensor: + """Returns indices of text.""" + return self.wordpiece_processor.to_index(text) + + def emnist_to_wordpiece_indices(self, x: Tensor) -> Tensor: + """Returns word pieces indices from emnist indices.""" + text = "".join([self.mapping[i] for i in x]) + text = text.lower().replace(" ", "▁") + return torch.LongTensor(self.wordpiece_processor.to_index(text)) + + def __getitem__(self, x: Union[int, Tensor]) -> str: + """Returns token for word piece index.""" + return self.get_token(x) diff --git a/text_recognizer/data/sentence_generator.py b/text_recognizer/data/sentence_generator.py deleted file mode 100644 index 8567e6d..0000000 --- a/text_recognizer/data/sentence_generator.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Downloading the Brown corpus with NLTK for sentence generating.""" -import itertools -import re -import string -from typing import Optional - -import nltk -from nltk.corpus.reader.util import ConcatenatedCorpusView -import numpy as np - -from text_recognizer.data.base_data_module import BaseDataModule - -NLTK_DATA_DIRNAME = BaseDataModule.data_dirname() / "downloaded" / "nltk" - - -class SentenceGenerator: - """Generates text sentences using the Brown corpus.""" - - def __init__(self, max_length: Optional[int] = None) -> None: - """Loads the corpus and sets word start indices.""" - self.corpus = brown_corpus() - self.word_start_indices = [0] + [ - _.start(0) + 1 for _ in re.finditer(" ", self.corpus) - ] - self.max_length = max_length - - def generate(self, max_length: Optional[int] = None) -> str: - r"""Generates a word or sentences from the Brown corpus. - - Sample a string from the Brown corpus of length at least one word and at most - max_length, padding to max_length with the '_' characters if sentence is - shorter. - - Args: - max_length (Optional[int]): The maximum number of characters in the sentence. - Defaults to None. - - Returns: - str: A sentence from the Brown corpus. - - Raises: - ValueError: If max_length was not specified at initialization and not - given as an argument. - - RuntimeError: If a valid string was not generated. - - """ - if max_length is None: - max_length = self.max_length - if max_length is None: - raise ValueError( - "Must provide max_length to this method or when making this object." - ) - - for _ in range(10): - try: - index = np.random.randint(0, len(self.word_start_indices) - 1) - start_index = self.word_start_indices[index] - end_index_candidates = [] - for index in range(index + 1, len(self.word_start_indices)): - if self.word_start_indices[index] - start_index > max_length: - break - end_index_candidates.append(self.word_start_indices[index]) - end_index = np.random.choice(end_index_candidates) - sampled_text = self.corpus[start_index:end_index].strip() - return sampled_text - except Exception: - pass - raise RuntimeError("Was not able to generate a valid string") - - -def brown_corpus() -> str: - """Returns a single string with the Brown corpus with all punctuations stripped.""" - sentences = load_nltk_brown_corpus() - corpus = " ".join(itertools.chain.from_iterable(sentences)) - corpus = corpus.translate({ord(c): None for c in string.punctuation}) - corpus = re.sub(" +", " ", corpus) - return corpus - - -def load_nltk_brown_corpus() -> ConcatenatedCorpusView: - """Load the Brown corpus using the NLTK library.""" - nltk.data.path.append(NLTK_DATA_DIRNAME) - try: - nltk.corpus.brown.sents() - except LookupError: - NLTK_DATA_DIRNAME.mkdir(parents=True, exist_ok=True) - nltk.download("brown", download_dir=NLTK_DATA_DIRNAME) - return nltk.corpus.brown.sents() diff --git a/text_recognizer/data/transforms.py b/text_recognizer/data/transforms.py deleted file mode 100644 index 7f3e0d1..0000000 --- a/text_recognizer/data/transforms.py +++ /dev/null @@ -1,49 +0,0 @@ -"""Transforms for PyTorch datasets.""" -from pathlib import Path -from typing import Optional, Union, Type, Set - -import torch -from torch import Tensor - -from text_recognizer.data.base_mapping import AbstractMapping -from text_recognizer.data.word_piece_mapping import WordPieceMapping - - -class WordPiece: - """Converts EMNIST indices to Word Piece indices.""" - - def __init__( - self, - num_features: int = 1000, - tokens: str = "iamdb_1kwp_tokens_1000.txt", - lexicon: str = "iamdb_1kwp_lex_1000.txt", - data_dir: Optional[Union[str, Path]] = None, - use_words: bool = False, - prepend_wordsep: bool = False, - special_tokens: Set[str] = {"", "", "

"}, - extra_symbols: Optional[Set[str]] = {"\n",}, - max_len: int = 451, - ) -> None: - self.mapping = WordPieceMapping( - data_dir=data_dir, - num_features=num_features, - tokens=tokens, - lexicon=lexicon, - use_words=use_words, - prepend_wordsep=prepend_wordsep, - special_tokens=special_tokens, - extra_symbols=extra_symbols, - ) - self.max_len = max_len - - def __call__(self, x: Tensor) -> Tensor: - """Converts Emnist target tensor to Word piece target tensor.""" - y = self.mapping.emnist_to_wordpiece_indices(x) - if len(y) < self.max_len: - pad_len = self.max_len - len(y) - y = torch.cat( - (y, torch.LongTensor([self.mapping.get_index("

")] * pad_len)) - ) - else: - y = y[: self.max_len] - return y diff --git a/text_recognizer/data/transforms/word_piece.py b/text_recognizer/data/transforms/word_piece.py new file mode 100644 index 0000000..6bf5472 --- /dev/null +++ b/text_recognizer/data/transforms/word_piece.py @@ -0,0 +1,48 @@ +"""Target transform for word pieces.""" +from pathlib import Path +from typing import Optional, Union, Set + +import torch +from torch import Tensor + +from text_recognizer.data.mappings.word_piece_mapping import WordPieceMapping + + +class WordPiece: + """Converts EMNIST indices to Word Piece indices.""" + + def __init__( + self, + num_features: int = 1000, + tokens: str = "iamdb_1kwp_tokens_1000.txt", + lexicon: str = "iamdb_1kwp_lex_1000.txt", + data_dir: Optional[Union[str, Path]] = None, + use_words: bool = False, + prepend_wordsep: bool = False, + special_tokens: Set[str] = {"", "", "

"}, + extra_symbols: Optional[Set[str]] = {"\n",}, + max_len: int = 451, + ) -> None: + self.mapping = WordPieceMapping( + data_dir=data_dir, + num_features=num_features, + tokens=tokens, + lexicon=lexicon, + use_words=use_words, + prepend_wordsep=prepend_wordsep, + special_tokens=special_tokens, + extra_symbols=extra_symbols, + ) + self.max_len = max_len + + def __call__(self, x: Tensor) -> Tensor: + """Converts Emnist target tensor to Word piece target tensor.""" + y = self.mapping.emnist_to_wordpiece_indices(x) + if len(y) < self.max_len: + pad_len = self.max_len - len(y) + y = torch.cat( + (y, torch.LongTensor([self.mapping.get_index("

")] * pad_len)) + ) + else: + y = y[: self.max_len] + return y diff --git a/text_recognizer/data/utils/build_transitions.py b/text_recognizer/data/utils/build_transitions.py new file mode 100644 index 0000000..0f987ca --- /dev/null +++ b/text_recognizer/data/utils/build_transitions.py @@ -0,0 +1,261 @@ +"""Builds transition graph. + +Most code stolen from here: + https://github.com/facebookresearch/gtn_applications/blob/master/scripts/build_transitions.py +""" + +import collections +import itertools +from pathlib import Path +from typing import Any, Dict, List, Optional + +import click +import gtn +from loguru import logger + + +START_IDX = -1 +END_IDX = -2 +WORDSEP = "▁" + + +def build_graph(ngrams: List, disable_backoff: bool = False) -> gtn.Graph: + """Returns a gtn Graph based on the ngrams.""" + graph = gtn.Graph(False) + ngram = len(ngrams) + state_to_node = {} + + def get_node(state: Optional[List]) -> Any: + node = state_to_node.get(state, None) + + if node is not None: + return node + + start = state == tuple([START_IDX]) if ngram > 1 else True + end = state == tuple([END_IDX]) if ngram > 1 else True + node = graph.add_node(start, end) + state_to_node[state] = node + + if not disable_backoff and not end: + # Add back off when adding node. + for n in range(1, len(state) + 1): + backoff_node = state_to_node.get(state[n:], None) + + # Epsilon transition to the back-off state. + if backoff_node is not None: + graph.add_arc(node, backoff_node, gtn.epsilon) + break + return node + + for grams in ngrams: + for gram in grams: + istate, ostate = gram[:-1], gram[len(gram) - ngram + 1 :] + inode = get_node(istate) + + if END_IDX not in gram[1:] and gram[1:] not in state_to_node: + raise ValueError( + "Ill formed counts: if (x, y_1, ..., y_{n-1}) is above" + "the n-gram threshold, then (y_1, ..., y_{n-1}) must be" + "above the (n-1)-gram threshold" + ) + + if END_IDX in ostate: + # Merge all state having into one as final graph generated + # will be similar. + ostate = tuple([END_IDX]) + + onode = get_node(ostate) + # p(gram[-1] | gram[:-1]) + graph.add_arc( + inode, onode, gtn.epsilon if gram[-1] == END_IDX else gram[-1] + ) + return graph + + +def count_ngrams(lines: List, ngram: List, tokens_to_index: Dict) -> List: + """Counts the number of ngrams.""" + counts = [collections.Counter() for _ in range(ngram)] + for line in lines: + # Prepend implicit start token. + token_line = [START_IDX] + for t in line: + token_line.append(tokens_to_index[t]) + token_line.append(END_IDX) + for n, counter in enumerate(counts): + start_offset = n == 0 + end_offset = ngram == 1 + for e in range(n + start_offset, len(token_line) - end_offset): + counter[tuple(token_line[e - n : e + 1])] += 1 + + return counts + + +def prune_ngrams(ngrams: List, prune: List) -> List: + """Prunes ngrams.""" + pruned_ngrams = [] + for n, grams in enumerate(ngrams): + grams = grams.most_common() + pruned_grams = [gram for gram, c in grams if c > prune[n]] + pruned_ngrams.append(pruned_grams) + return pruned_ngrams + + +def add_blank_grams(pruned_ngrams: List, num_tokens: int, blank: str) -> List: + """Adds blank token to grams.""" + all_grams = [gram for grams in pruned_ngrams for gram in grams] + maxorder = len(pruned_ngrams) + blank_grams = {} + if blank == "forced": + pruned_ngrams = [pruned_ngrams[0] if i == 0 else [] for i in range(maxorder)] + pruned_ngrams[0].append(tuple([num_tokens])) + blank_grams[tuple([num_tokens])] = True + + for gram in all_grams: + # Iterate over all possibilities by using a vector of 0s, 1s to + # denote whether a blank is being used at each position. + if blank == "optional": + # Given a gram ab.. if order n, we have n + 1 positions + # available whether to use blank or not. + onehot_vectors = itertools.product([0, 1], repeat=len(gram) + 1) + elif blank == "forced": + # Must include a blank token in between. + onehot_vectors = [[1] * (len(gram) + 1)] + else: + raise ValueError( + "Invalid value specificed for blank. Must be in |optional|forced|none|" + ) + + for j in onehot_vectors: + new_array = [] + for idx, oz in enumerate(j[:-1]): + if oz == 1 and gram[idx] != START_IDX: + new_array.append(num_tokens) + new_array.append(gram[idx]) + if j[-1] == 1 and gram[-1] != END_IDX: + new_array.append(num_tokens) + for n in range(maxorder): + for e in range(n, len(new_array)): + cur_gram = tuple(new_array[e - n : e + 1]) + if num_tokens in cur_gram and cur_gram not in blank_grams: + pruned_ngrams[n].append(cur_gram) + blank_grams[cur_gram] = True + + return pruned_ngrams + + +def add_self_loops(pruned_ngrams: List) -> List: + """Adds self loops to the ngrams.""" + maxorder = len(pruned_ngrams) + + # Use dict for fast search. + all_grams = set([gram for grams in pruned_ngrams for gram in grams]) + for o in range(1, maxorder): + for gram in pruned_ngrams[o - 1]: + # Repeat one of the tokens. + for pos in range(len(gram)): + if gram[pos] == START_IDX or gram[pos] == END_IDX: + continue + new_gram = gram[:pos] + (gram[pos],) + gram[pos:] + + if new_gram not in all_grams: + pruned_ngrams[o].append(new_gram) + all_grams.add(new_gram) + return pruned_ngrams + + +def parse_lines(lines: List, lexicon: Path) -> List: + """Parses lines with a lexicon.""" + with open(lexicon, "r") as f: + lex = (line.strip().split() for line in f) + lex = {line[0]: line[1:] for line in lex} + print(len(lex)) + return [[t for w in line.split(WORDSEP) for t in lex[w]] for line in lines] + + +@click.command() +@click.option("--data_dir", type=str, default=None, help="Path to dataset root.") +@click.option( + "--tokens", type=str, help="Path to token list (in order used with training)." +) +@click.option("--lexicon", type=str, default=None, help="Path to lexicon") +@click.option( + "--prune", + nargs=2, + type=int, + help="Threshold values for prune unigrams, bigrams, etc.", +) +@click.option( + "--blank", + default=click.Choice(["none", "optional", "forced"]), + help="Specifies the usage of blank token" + "'none' - do not use blank token " + "'optional' - allow an optional blank inbetween tokens" + "'forced' - force a blank inbetween tokens (also referred to as garbage token)", +) +@click.option("--self_loops", is_flag=True, help="Add self loops for tokens") +@click.option("--disable_backoff", is_flag=True, help="Disable backoff transitions") +@click.option("--save_path", default=None, help="Path to save transition graph.") +def cli( + data_dir: str, + tokens: str, + lexicon: str, + prune: List[int], + blank: str, + self_loops: bool, + disable_backoff: bool, + save_path: str, +) -> None: + """CLI for creating the transitions.""" + logger.info(f"Building {len(prune)}-gram transition models.") + + if data_dir is None: + data_dir = ( + Path(__file__).resolve().parents[2] / "data" / "processed" / "iam_lines" + ) + logger.debug(f"Using data dir: {data_dir}") + if not data_dir.exists(): + raise RuntimeError(f"Could not locate iamdb directory at {data_dir}") + else: + data_dir = Path(data_dir) + + # Build table of counts and the back-off if below threshold. + with open(data_dir / "train.txt", "r") as f: + lines = [line.strip() for line in f] + + with open(data_dir / tokens, "r") as f: + tokens = [line.strip() for line in f] + + if lexicon is not None: + lexicon = data_dir / lexicon + lines = parse_lines(lines, lexicon) + + tokens_to_idx = {t: e for e, t in enumerate(tokens)} + + ngram = len(prune) + + logger.info("Counting data...") + ngrams = count_ngrams(lines, ngram, tokens_to_idx) + + pruned_ngrams = prune_ngrams(ngrams, prune) + + for n in range(ngram): + logger.info(f"Kept {len(pruned_ngrams[n])} of {len(ngrams[n])} {n + 1}-grams") + + if blank == "none": + pruned_ngrams = add_blank_grams(pruned_ngrams, len(tokens_to_idx), blank) + + if self_loops: + pruned_ngrams = add_self_loops(pruned_ngrams) + + logger.info("Building graph from pruned ngrams...") + graph = build_graph(pruned_ngrams, disable_backoff) + logger.info(f"Graph has {graph.num_arcs()} arcs and {graph.num_nodes()} nodes.") + + save_path = str(data_dir / save_path) + + logger.info(f"Saving graph to {save_path}") + gtn.save(save_path, graph) + + +if __name__ == "__main__": + cli() diff --git a/text_recognizer/data/utils/download_utils.py b/text_recognizer/data/utils/download_utils.py new file mode 100644 index 0000000..a5a5360 --- /dev/null +++ b/text_recognizer/data/utils/download_utils.py @@ -0,0 +1,73 @@ +"""Util functions for downloading datasets.""" +import hashlib +from pathlib import Path +from typing import Dict, Optional +from urllib.request import urlretrieve + +from loguru import logger as log +from tqdm import tqdm + + +def _compute_sha256(filename: Path) -> str: + """Returns the SHA256 checksum of a file.""" + with filename.open(mode="rb") as f: + return hashlib.sha256(f.read()).hexdigest() + + +class TqdmUpTo(tqdm): + """TQDM progress bar when downloading files. + + From https://github.com/tqdm/tqdm/blob/master/examples/tqdm_wget.py + + """ + + def update_to( + self, blocks: int = 1, block_size: int = 1, total_size: Optional[int] = None + ) -> None: + """Updates the progress bar. + + Args: + blocks (int): Number of blocks transferred so far. Defaults to 1. + block_size (int): Size of each block, in tqdm units. Defaults to 1. + total_size (Optional[int]): Total size in tqdm units. Defaults to None. + """ + if total_size is not None: + self.total = total_size + self.update(blocks * block_size - self.n) + + +def _download_url(url: str, filename: str) -> None: + """Downloads a file from url to filename, with a progress bar.""" + with TqdmUpTo(unit="B", unit_scale=True, unit_divisor=1024, miniters=1) as t: + urlretrieve(url, filename, reporthook=t.update_to, data=None) # nosec + + +def download_dataset(metadata: Dict, dl_dir: Path) -> Optional[Path]: + """Downloads dataset using a metadata file. + + Args: + metadata (Dict): A metadata file of the dataset. + dl_dir (Path): Download directory for the dataset. + + Returns: + Optional[Path]: Returns filename if dataset is downloaded, None if it already + exists. + + Raises: + ValueError: If the SHA-256 value is not the same between the dataset and + the metadata file. + + """ + dl_dir.mkdir(parents=True, exist_ok=True) + filename = dl_dir / metadata["filename"] + if filename.exists(): + return + log.info(f"Downloading raw dataset from {metadata['url']} to {filename}...") + _download_url(metadata["url"], filename) + log.info("Computing the SHA-256...") + sha256 = _compute_sha256(filename) + if sha256 != metadata["sha256"]: + raise ValueError( + "Downloaded data file SHA-256 does not match that listed in metadata document." + ) + return filename diff --git a/text_recognizer/data/utils/iam_preprocessor.py b/text_recognizer/data/utils/iam_preprocessor.py new file mode 100644 index 0000000..60ecff1 --- /dev/null +++ b/text_recognizer/data/utils/iam_preprocessor.py @@ -0,0 +1,209 @@ +"""Preprocessor for extracting word letters from the IAM dataset. + +The code is mostly stolen from: + https://github.com/facebookresearch/gtn_applications/blob/master/datasets/iamdb.py +""" +import collections +import itertools +from pathlib import Path +import re +from typing import List, Optional, Set, Union + +import click +from loguru import logger as log +import torch + + +def load_metadata( + data_dir: Path, wordsep: str, use_words: bool = False +) -> collections.defaultdict: + """Loads IAM metadata and returns it as a dictionary.""" + forms = collections.defaultdict(list) + filename = "words.txt" if use_words else "lines.txt" + + with open(data_dir / "ascii" / filename, "r") as f: + lines = (line.strip().split() for line in f if line[0] != "#") + for line in lines: + # Skip word segmentation errors. + if use_words and line[1] == "err": + continue + text = " ".join(line[8:]) + + # Remove garbage tokens: + text = text.replace("#", "") + + # Swap word sep form | to wordsep + text = re.sub(r"\|+|\s", wordsep, text).strip(wordsep) + form_key = "-".join(line[0].split("-")[:2]) + line_key = "-".join(line[0].split("-")[:3]) + box_idx = 4 - use_words + box = tuple(int(val) for val in line[box_idx : box_idx + 4]) + forms[form_key].append({"key": line_key, "box": box, "text": text}) + return forms + + +class Preprocessor: + """A preprocessor for the IAM dataset.""" + + def __init__( + self, + data_dir: Union[str, Path], + num_features: int, + tokens_path: Optional[Union[str, Path]] = None, + lexicon_path: Optional[Union[str, Path]] = None, + use_words: bool = False, + prepend_wordsep: bool = False, + special_tokens: Optional[Set[str]] = None, + ) -> None: + self.wordsep = "▁" + self._use_word = use_words + self._prepend_wordsep = prepend_wordsep + self.special_tokens = special_tokens if special_tokens is not None else None + self.data_dir = Path(data_dir) + self.forms = load_metadata(self.data_dir, self.wordsep, use_words=use_words) + + # Load the set of graphemes: + graphemes = set() + for _, form in self.forms.items(): + for line in form: + graphemes.update(line["text"].lower()) + self.graphemes = sorted(graphemes) + + # Build the token-to-index and index-to-token maps. + if tokens_path is not None: + with open(tokens_path, "r") as f: + self.tokens = [line.strip() for line in f] + else: + self.tokens = self.graphemes + + if lexicon_path is not None: + with open(lexicon_path, "r") as f: + lexicon = (line.strip().split() for line in f) + lexicon = {line[0]: line[1:] for line in lexicon} + self.lexicon = lexicon + else: + self.lexicon = None + + if self.special_tokens is not None: + special_tokens_ = (*self.special_tokens, "#", "*") + self.tokens += special_tokens_ + self.graphemes += special_tokens_ + + self.graphemes_to_index = {t: i for i, t in enumerate(self.graphemes)} + self.tokens_to_index = {t: i for i, t in enumerate(self.tokens)} + self.num_features = num_features + self.text = [] + + @property + def num_tokens(self) -> int: + """Returns the number or tokens.""" + return len(self.tokens) + + @property + def use_words(self) -> bool: + """If words are used.""" + return self._use_word + + def extract_train_text(self) -> None: + """Extracts training text.""" + keys = [] + with open(self.data_dir / "task" / "trainset.txt") as f: + keys.extend((line.strip() for line in f)) + + for _, examples in self.forms.items(): + for example in examples: + if example["key"] not in keys: + continue + self.text.append(example["text"].lower()) + + def _to_index(self, line: str) -> torch.LongTensor: + if self.special_tokens is not None and line in self.special_tokens: + return torch.LongTensor([self.tokens_to_index[line]]) + token_to_index = self.graphemes_to_index + if self.lexicon is not None: + if len(line) > 0: + # If the word is not found in the lexicon, fall back to letters. + tokens = [ + t + for w in line.split(self.wordsep) + for t in self.lexicon.get(w, self.wordsep + w) + ] + token_to_index = self.tokens_to_index + if self._prepend_wordsep: + tokens = itertools.chain([self.wordsep], tokens) + return torch.LongTensor([token_to_index[t] for t in tokens]) + + def to_index(self, line: str) -> torch.LongTensor: + """Converts text to a tensor of indices.""" + if self.special_tokens is not None: + pattern = f"({'|'.join(self.special_tokens)})" + lines = list(filter(None, re.split(pattern, line))) + return torch.cat([self._to_index(line) for line in lines]) + return self._to_index(line) + + def to_text(self, indices: List[int]) -> str: + """Converts indices to text.""" + # Roughly the inverse of `to_index` + encoding = self.graphemes + if self.lexicon is not None: + encoding = self.tokens + return self._post_process(encoding[i] for i in indices) + + def tokens_to_text(self, indices: List[int]) -> str: + """Converts tokens to text.""" + return self._post_process(self.tokens[i] for i in indices) + + def _post_process(self, indices: List[int]) -> str: + """A list join.""" + return "".join(indices).strip(self.wordsep) + + +@click.command() +@click.option("--data_dir", type=str, default=None, help="Path to iam dataset") +@click.option( + "--use_words", is_flag=True, help="Load word segmented dataset instead of lines" +) +@click.option( + "--save_text", type=str, default=None, help="Path to save parsed train text" +) +@click.option("--save_tokens", type=str, default=None, help="Path to save tokens") +def cli( + data_dir: Optional[str], + use_words: bool, + save_text: Optional[str], + save_tokens: Optional[str], +) -> None: + """CLI for extracting text data from the iam dataset.""" + if data_dir is None: + data_dir = ( + Path(__file__).resolve().parents[2] + / "data" + / "downloaded" + / "iam" + / "iamdb" + ) + log.debug(f"Using data dir: {data_dir}") + if not data_dir.exists(): + raise RuntimeError(f"Could not locate iamdb directory at {data_dir}") + else: + data_dir = Path(data_dir) + + preprocessor = Preprocessor(data_dir, 64, use_words=use_words) + preprocessor.extract_train_text() + + processed_dir = data_dir.parents[2] / "processed" / "iam_lines" + log.debug(f"Saving processed files at: {processed_dir}") + + if save_text is not None: + log.info("Saving training text") + with open(processed_dir / save_text, "w") as f: + f.write("\n".join(t for t in preprocessor.text)) + + if save_tokens is not None: + log.info("Saving tokens") + with open(processed_dir / save_tokens, "w") as f: + f.write("\n".join(preprocessor.tokens)) + + +if __name__ == "__main__": + cli() diff --git a/text_recognizer/data/utils/image_utils.py b/text_recognizer/data/utils/image_utils.py new file mode 100644 index 0000000..c2b8915 --- /dev/null +++ b/text_recognizer/data/utils/image_utils.py @@ -0,0 +1,49 @@ +"""Image util functions for loading and saving images.""" +from pathlib import Path +from typing import Union +from urllib.request import urlopen + +import cv2 +import numpy as np +from PIL import Image + + +def read_image_pil(image_uri: Union[Path, str], grayscale: bool = False) -> Image: + """Return PIL image.""" + image = Image.open(image_uri) + if grayscale: + image = image.convert("L") + return image + + +def read_image(image_uri: Union[Path, str], grayscale: bool = False) -> np.array: + """Read image_uri.""" + + if isinstance(image_uri, str): + image_uri = Path(image_uri) + + def read_image_from_filename(image_filename: Path, imread_flag: int) -> np.array: + return cv2.imread(str(image_filename), imread_flag) + + def read_image_from_url(image_url: Path, imread_flag: int) -> np.array: + url_response = urlopen(str(image_url)) # nosec + image_array = np.array(bytearray(url_response.read()), dtype=np.uint8) + return cv2.imdecode(image_array, imread_flag) + + imread_flag = cv2.IMREAD_GRAYSCALE if grayscale else cv2.IMREAD_COLOR + image = None + + if image_uri.exists(): + image = read_image_from_filename(image_uri, imread_flag) + else: + image = read_image_from_url(image_uri, imread_flag) + + if image is None: + raise ValueError(f"Could not load image at {image_uri}") + + return image + + +def write_image(image: np.ndarray, filename: Union[Path, str]) -> None: + """Write image to file.""" + cv2.imwrite(str(filename), image) diff --git a/text_recognizer/data/utils/make_wordpieces.py b/text_recognizer/data/utils/make_wordpieces.py new file mode 100644 index 0000000..8e53815 --- /dev/null +++ b/text_recognizer/data/utils/make_wordpieces.py @@ -0,0 +1,112 @@ +"""Creates word pieces from a text file. + +Most code stolen from: + + https://github.com/facebookresearch/gtn_applications/blob/master/scripts/make_wordpieces.py + +""" +import io +from pathlib import Path +from typing import List, Optional, Union + +import click +from loguru import logger as log +import sentencepiece as spm + + +def iamdb_pieces( + data_dir: Path, text_file: str, num_pieces: int, output_prefix: str +) -> None: + """Creates word pieces from the iamdb train text.""" + # Load training text. + with open(data_dir / text_file, "r") as f: + text = [line.strip() for line in f] + + sp = train_spm_model( + iter(text), + num_pieces + 1, # To account for + user_symbols=["/"], # added so token is in the output set + ) + + vocab = sorted(set(w for t in text for w in t.split("▁") if w)) + if "move" not in vocab: + raise RuntimeError("`MOVE` not in vocab") + + save_pieces(sp, num_pieces, data_dir, output_prefix, vocab) + + +def train_spm_model( + sentences: iter, vocab_size: int, user_symbols: Union[str, List[str]] = "" +) -> spm.SentencePieceProcessor: + """Trains the sentence piece model.""" + model = io.BytesIO() + spm.SentencePieceTrainer.train( + sentence_iterator=sentences, + model_writer=model, + vocab_size=vocab_size, + bos_id=-1, + eos_id=-1, + character_coverage=1.0, + user_defined_symbols=user_symbols, + ) + sp = spm.SentencePieceProcessor(model_proto=model.getvalue()) + return sp + + +def save_pieces( + sp: spm.SentencePieceProcessor, + num_pieces: int, + data_dir: Path, + output_prefix: str, + vocab: set, +) -> None: + """Saves word pieces to disk.""" + log.info(f"Generating word piece list of size {num_pieces}.") + pieces = [sp.id_to_piece(i) for i in range(1, num_pieces + 1)] + log.info(f"Encoding vocabulary of size {len(vocab)}.") + encoded_vocab = [sp.encode_as_pieces(v) for v in vocab] + + # Save pieces to file. + with open(data_dir / f"{output_prefix}_tokens_{num_pieces}.txt", "w") as f: + f.write("\n".join(pieces)) + + # Save lexicon to a file. + with open(data_dir / f"{output_prefix}_lex_{num_pieces}.txt", "w") as f: + for v, p in zip(vocab, encoded_vocab): + f.write(f"{v} {' '.join(p)}\n") + + +@click.command() +@click.option("--data_dir", type=str, default=None, help="Path to processed iam dir.") +@click.option( + "--text_file", type=str, default=None, help="Name of sentence piece training text." +) +@click.option( + "--output_prefix", + type=str, + default="word_pieces", + help="Prefix name to store tokens and lexicon.", +) +@click.option("--num_pieces", type=int, default=1000, help="Number of word pieces.") +def cli( + data_dir: Optional[str], + text_file: Optional[str], + output_prefix: Optional[str], + num_pieces: Optional[int], +) -> None: + """CLI for training the sentence piece model.""" + if data_dir is None: + data_dir = ( + Path(__file__).resolve().parents[2] / "data" / "processed" / "iam_lines" + ) + log.debug(f"Using data dir: {data_dir}") + if not data_dir.exists(): + raise RuntimeError(f"Could not locate iamdb directory at {data_dir}") + else: + data_dir = Path(data_dir) + + iamdb_pieces(data_dir, text_file, num_pieces, output_prefix) + + +if __name__ == "__main__": + cli() diff --git a/text_recognizer/data/utils/sentence_generator.py b/text_recognizer/data/utils/sentence_generator.py new file mode 100644 index 0000000..8567e6d --- /dev/null +++ b/text_recognizer/data/utils/sentence_generator.py @@ -0,0 +1,89 @@ +"""Downloading the Brown corpus with NLTK for sentence generating.""" +import itertools +import re +import string +from typing import Optional + +import nltk +from nltk.corpus.reader.util import ConcatenatedCorpusView +import numpy as np + +from text_recognizer.data.base_data_module import BaseDataModule + +NLTK_DATA_DIRNAME = BaseDataModule.data_dirname() / "downloaded" / "nltk" + + +class SentenceGenerator: + """Generates text sentences using the Brown corpus.""" + + def __init__(self, max_length: Optional[int] = None) -> None: + """Loads the corpus and sets word start indices.""" + self.corpus = brown_corpus() + self.word_start_indices = [0] + [ + _.start(0) + 1 for _ in re.finditer(" ", self.corpus) + ] + self.max_length = max_length + + def generate(self, max_length: Optional[int] = None) -> str: + r"""Generates a word or sentences from the Brown corpus. + + Sample a string from the Brown corpus of length at least one word and at most + max_length, padding to max_length with the '_' characters if sentence is + shorter. + + Args: + max_length (Optional[int]): The maximum number of characters in the sentence. + Defaults to None. + + Returns: + str: A sentence from the Brown corpus. + + Raises: + ValueError: If max_length was not specified at initialization and not + given as an argument. + + RuntimeError: If a valid string was not generated. + + """ + if max_length is None: + max_length = self.max_length + if max_length is None: + raise ValueError( + "Must provide max_length to this method or when making this object." + ) + + for _ in range(10): + try: + index = np.random.randint(0, len(self.word_start_indices) - 1) + start_index = self.word_start_indices[index] + end_index_candidates = [] + for index in range(index + 1, len(self.word_start_indices)): + if self.word_start_indices[index] - start_index > max_length: + break + end_index_candidates.append(self.word_start_indices[index]) + end_index = np.random.choice(end_index_candidates) + sampled_text = self.corpus[start_index:end_index].strip() + return sampled_text + except Exception: + pass + raise RuntimeError("Was not able to generate a valid string") + + +def brown_corpus() -> str: + """Returns a single string with the Brown corpus with all punctuations stripped.""" + sentences = load_nltk_brown_corpus() + corpus = " ".join(itertools.chain.from_iterable(sentences)) + corpus = corpus.translate({ord(c): None for c in string.punctuation}) + corpus = re.sub(" +", " ", corpus) + return corpus + + +def load_nltk_brown_corpus() -> ConcatenatedCorpusView: + """Load the Brown corpus using the NLTK library.""" + nltk.data.path.append(NLTK_DATA_DIRNAME) + try: + nltk.corpus.brown.sents() + except LookupError: + NLTK_DATA_DIRNAME.mkdir(parents=True, exist_ok=True) + nltk.download("brown", download_dir=NLTK_DATA_DIRNAME) + return nltk.corpus.brown.sents() diff --git a/text_recognizer/data/word_piece_mapping.py b/text_recognizer/data/word_piece_mapping.py deleted file mode 100644 index dc56942..0000000 --- a/text_recognizer/data/word_piece_mapping.py +++ /dev/null @@ -1,98 +0,0 @@ -"""Word piece mapping.""" -from pathlib import Path -from typing import List, Optional, Set, Union - -from loguru import logger as log -import torch -from torch import Tensor - -from text_recognizer.data.emnist_mapping import EmnistMapping -from text_recognizer.data.iam_preprocessor import Preprocessor - - -class WordPieceMapping(EmnistMapping): - """Word piece mapping.""" - - def __init__( - self, - data_dir: Optional[Path] = None, - num_features: int = 1000, - tokens: str = "iamdb_1kwp_tokens_1000.txt", - lexicon: str = "iamdb_1kwp_lex_1000.txt", - use_words: bool = False, - prepend_wordsep: bool = False, - special_tokens: Set[str] = {"", "", "

"}, - extra_symbols: Set[str] = {"\n"}, - ) -> None: - super().__init__(extra_symbols=extra_symbols) - self.data_dir = ( - ( - Path(__file__).resolve().parents[2] - / "data" - / "downloaded" - / "iam" - / "iamdb" - ) - if data_dir is None - else Path(data_dir) - ) - log.debug(f"Using data dir: {self.data_dir}") - if not self.data_dir.exists(): - raise RuntimeError(f"Could not locate iamdb directory at {self.data_dir}") - - processed_path = ( - Path(__file__).resolve().parents[2] / "data" / "processed" / "iam_lines" - ) - - tokens_path = processed_path / tokens - lexicon_path = processed_path / lexicon - - special_tokens = set(special_tokens) - if self.extra_symbols is not None: - special_tokens = special_tokens | set(extra_symbols) - - self.wordpiece_processor = Preprocessor( - data_dir=self.data_dir, - num_features=num_features, - tokens_path=tokens_path, - lexicon_path=lexicon_path, - use_words=use_words, - prepend_wordsep=prepend_wordsep, - special_tokens=special_tokens, - ) - - def __len__(self) -> int: - """Return number of word pieces.""" - return len(self.wordpiece_processor.tokens) - - def get_token(self, index: Union[int, Tensor]) -> str: - """Returns token for index.""" - if (index := int(index)) <= self.wordpiece_processor.num_tokens: - return self.wordpiece_processor.tokens[index] - raise KeyError(f"Index ({index}) not in mapping.") - - def get_index(self, token: str) -> Tensor: - """Returns index of token.""" - if token in self.wordpiece_processor.tokens: - return torch.LongTensor([self.wordpiece_processor.tokens_to_index[token]]) - raise KeyError(f"Token ({token}) not found in inverse mapping.") - - def get_text(self, indices: Union[List[int], Tensor]) -> str: - """Returns text from indices.""" - if isinstance(indices, Tensor): - indices = indices.tolist() - return self.wordpiece_processor.to_text(indices) - - def get_indices(self, text: str) -> Tensor: - """Returns indices of text.""" - return self.wordpiece_processor.to_index(text) - - def emnist_to_wordpiece_indices(self, x: Tensor) -> Tensor: - """Returns word pieces indices from emnist indices.""" - text = "".join([self.mapping[i] for i in x]) - text = text.lower().replace(" ", "▁") - return torch.LongTensor(self.wordpiece_processor.to_index(text)) - - def __getitem__(self, x: Union[int, Tensor]) -> str: - """Returns token for word piece index.""" - return self.get_token(x) -- cgit v1.2.3-70-g09d2