File size: 1,730 Bytes
0d7f818 ebe8b54 0d7f818 | 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 | import argparse
import sys
from pathlib import Path
import unicodedata
import torch
from PIL import Image, ImageOps
from model import CRNN
def predict(image_path, checkpoint_path):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=True)
vocab = checkpoint['vocab']
reverse = {i: c for c, i in vocab.items()}
if '<UNK>' in vocab:
reverse[vocab['<UNK>']] = '\ufffd'
model = CRNN(len(vocab))
model.load_state_dict(checkpoint['model_state_dict'], strict=True)
model.to(device).eval()
height = checkpoint['height']
with Image.open(image_path) as source:
image = source.convert('L')
width = max(4, round(image.width * height / image.height))
image = ImageOps.mirror(image.resize((width, height), Image.Resampling.LANCZOS))
data = torch.frombuffer(bytearray(image.tobytes()), dtype=torch.uint8)
data = (data.reshape(1, 1, height, width).float() / 255.0 - 0.5) / 0.5
with torch.inference_mode():
ids = model(data.to(device)).argmax(-1)[0].cpu().tolist()
result, previous = [], None
for index in ids:
if index != vocab['<BLANK>'] and index != previous:
result.append(reverse[index])
previous = index
return unicodedata.normalize('NFC', ''.join(result))
if __name__ == '__main__':
if hasattr(sys.stdout, 'reconfigure'):
sys.stdout.reconfigure(encoding='utf-8')
parser = argparse.ArgumentParser()
parser.add_argument('image')
parser.add_argument('--checkpoint', default=str(Path(__file__).parent/'best_cer_v10.pt'))
args = parser.parse_args()
print(predict(args.image, args.checkpoint))
|