| """Recognize one cropped word or line with an explicit language and text style.""" |
| import argparse |
| import json |
| from pathlib import Path |
| import sys |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('image', type=Path) |
| parser.add_argument('--language', choices=['ne', 'hi', 'en'], required=True) |
| parser.add_argument('--text-style', choices=['handwriting', 'print'], default='handwriting', |
| help='English print uses the character branch; default: handwriting.') |
| parser.add_argument('--backend', choices=['torch', 'paddle'], default='torch') |
| parser.add_argument('--model-dir', type=Path, default=Path(__file__).resolve().parent) |
| args = parser.parse_args() |
| root = args.model_dir.resolve() |
| branch = 'english' if args.language == 'en' and args.text_style == 'handwriting' else 'character' |
| if args.backend == 'torch': |
| import torch |
| from PIL import Image |
| from transformers import TrOCRProcessor, ViTImageProcessor |
| sys.path.insert(0, str(root / 'code')) |
| from shared_character_recognizer import SharedCharacterRecognizer |
| from trocr_character_codec import CharacterCodec |
| torch.set_num_threads(2) |
| torch.use_deterministic_algorithms(True) |
| folder = root / 'source-model' |
| model = SharedCharacterRecognizer.from_bundle(folder) |
| assets = folder / ('english-assets' if branch == 'english' else 'character-assets') |
| processor = (TrOCRProcessor.from_pretrained(assets, local_files_only=True, use_fast=False) |
| if branch == 'english' else ViTImageProcessor.from_pretrained(assets, local_files_only=True)) |
| with Image.open(args.image) as image: |
| pixels = processor(images=image.convert('RGB'), return_tensors='pt').pixel_values |
| generation = dict(do_sample=False, num_beams=10, max_new_tokens=256 if branch == 'english' else 257, |
| min_length=0, length_penalty=1., early_stopping=False, use_cache=True) |
| if branch == 'character': |
| generation['suppress_tokens'] = [0, 1, 3] |
| with torch.inference_mode(): |
| tokens = model.generate_pixels(pixels, branch, generation)[0].tolist() |
| if branch == 'english': |
| text = processor.batch_decode([tokens], skip_special_tokens=True)[0] |
| else: |
| spec = json.loads((assets / 'character_codec.json').read_text()) |
| text = CharacterCodec(list(spec['character_ids'])).decode(tokens) |
| result = dict(text=text, tokens=tokens, token_limit_reached=len(tokens)-1 >= generation['max_new_tokens'] and tokens[-1] != 2) |
| else: |
| import paddle |
| from PIL import Image |
| sys.path.insert(0, str(root / 'paddle/code')) |
| from shared_character_native import SharedCharacterNative, preprocess_rgb |
| paddle.set_device('cpu') |
| folder = root / 'paddle' |
| model = SharedCharacterNative(folder / 'candidate', folder / 'encoder', folder / 'decoders', |
| allow_unqualified=True, library_path=folder / 'encoder-library.so') |
| with Image.open(args.image) as image: |
| result = model.predict_pixels(preprocess_rgb(image), branch) |
| print(json.dumps(dict(language=args.language, text_style=args.text_style, branch=branch, backend=args.backend, **result), ensure_ascii=False)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
