Views
No views yet
1from transformers import TrOCRProcessor, VisionEncoderDecoderConfig
2from huggingface_hub import snapshot_download
3from PIL import Image
4import numpy as np
5import onnxruntime
6import os
7
8def generate(decoder, encoder_outputs, batch_size, config, max_length = 128):
9 """
10 Generate text using autoregressive decoding
11 with per-sequence early stopping.
12
13 Args:
14 decoder: TrOCR decoder
15 encoder_hidden_states: Output from encoder
16 batch_size: Number of 2images to process
17 config: model config information
18 max_length: maximum length of generated sequence (in tokens)
19
20 Returns:
21 Generated token IDs
22 """
23 decoder_input_ids = np.full((batch_size, 1),
24 config.decoder_start_token_id,
25 dtype=np.int64)
26
27 # Track which sequences have finished
28 finished = np.zeros(batch_size, dtype=bool)
29
30 for step in range(max_length):
31 # Run decoder
32 decoder_outputs = decoder.run(
33 None,
34 {
35 "input_ids": decoder_input_ids,
36 "encoder_hidden_states": encoder_outputs
37 }
38 )[0]
39
40 # Get next tokens
41 next_token_logits = decoder_outputs[:, -1, :]
42 next_tokens = np.argmax(next_token_logits, axis=-1)
43
44 # Mark sequences that just generated EOS
45 just_finished = (next_tokens == config.eos_token_id)
46 finished = finished | just_finished
47
48 # For already finished sequences, force PAD token
49 next_tokens[finished] = config.pad_token_id
50
51 # Append tokens
52 next_tokens = next_tokens.reshape(-1, 1)
53 decoder_input_ids = np.concatenate([decoder_input_ids, next_tokens], axis=1)
54
55 # Stop when ALL sequences have finished
56 if np.all(finished):
57 break
58
59 return decoder_input_ids
60
61def predict_text(line_image, processor, encoder, decoder, config):
62 """
63 Predict text content from text line images.
64
65 Args:
66 line_image: Text line image file
67 processor: TrOCRProcessor
68 encoder: TrOCR encoder
69 decoder: TrOCR decoder
70 config: model config information
71
72 Returns:
73 Generated text
74 """
75 # Process image with TrOCR processor
76 # Use 'pt' (PyTorch) then convert to numpy, as 'np' is not supported by fast processors
77 pixel_values = processor(line_image, return_tensors="pt").pixel_values
78 pixel_values = pixel_values.numpy()
79 batch_size = pixel_values.shape[0]
80 # Get output from the encoder
81 encoder_outputs = encoder.run(None, {"pixel_values": pixel_values})[0]
82 # Get decoder output
83 generated_ids = generate(decoder, encoder_outputs, batch_size, config)
84 # Decode tokens to text
85 texts = processor.batch_decode(
86 generated_ids,
87 skip_special_tokens=True,
88 clean_up_tokenization_spaces=False
89 )
90 return texts
91
92REPOSITORY = "Kansallisarkisto/multicentury-htr-model-onnx"
93
94# Download repository
95repo_path = snapshot_download(
96 repo_id=REPOSITORY
97)
98
99# Load model and processor
100processor = TrOCRProcessor.from_pretrained(repo_path,
101 use_fast=True,
102 do_resize=True,
103 size={'height': 192,'width': 1024})
104
105# Load model config
106config = VisionEncoderDecoderConfig.from_pretrained(repo_path)
107
108providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
109
110# Load model encoder
111encoder = onnxruntime.InferenceSession(
112 os.path.join(repo_path, "encoder_model.onnx"),
113 providers=providers
114 )
115
116# Load model decoder
117decoder = onnxruntime.InferenceSession(
118 os.path.join(repo_path, "decoder_model.onnx"),
119 providers=providers
120 )
121
122# Open an image of handwritten text
123image = Image.open("path_to_image.jpg")
124
125# Preprocess and predict
126generated_text = predict_text(image, processor, encoder, decoder, config)
127
128print(' '.join(generated_text))