File size: 3,443 Bytes
a73a515 5f70891 a73a515 5f70891 a73a515 5f70891 a73a515 5f70891 | 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 | """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()
|