File size: 2,947 Bytes
f547382 abf868b f547382 abf868b | 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 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | from PIL import Image
import numpy as np
import math
import numpy as np
import onnxruntime as ort
import os
def resize_norm_img(img):
imgC, imgH, imgW = 3, 48, 320
h, w = img.shape[:2]
ratio = w / float(h)
resized_w = imgW if math.ceil(imgH * ratio) > imgW else int(math.ceil(imgH * ratio))
pil_img = Image.fromarray(img)
resized_image = np.array(pil_img.resize((resized_w, imgH), Image.BILINEAR)).astype('float32')
resized_image = resized_image.transpose((2, 0, 1)) / 255.0
resized_image -= 0.5
resized_image /= 0.5
padding_im = np.zeros((imgC, imgH, imgW), dtype=np.float32)
padding_im[:, :, 0:resized_w] = resized_image
return np.expand_dims(padding_im, axis=0)
def load_dict(dict_path):
"""Loads the dictionary and aligns it with PaddleOCR's CTC index format."""
character = []
with open(dict_path, "rb") as fin:
lines = fin.readlines()
for line in lines:
line = line.decode('utf-8').strip("\n").strip("\r\n")
character.append(line)
character.append(" ")
# Reserve index 0 for the blank token
character = ['<blank>'] + character
return character
def decode(preds, character):
"""Converts the raw ONNX matrix output into readable text using CTC rules."""
# Extract the highest probability class indices
preds_idx = preds.argmax(axis=2)[0]
preds_prob = preds.max(axis=2)[0]
char_list = []
conf_list = []
for i in range(len(preds_idx)):
# CTC rules: Ignore the blank token (0) and ignore consecutive duplicate characters
if preds_idx[i] != 0 and not (i > 0 and preds_idx[i] == preds_idx[i - 1]):
char_list.append(character[preds_idx[i]])
conf_list.append(preds_prob[i])
text = ''.join(char_list)
confidence = np.mean(conf_list) if len(conf_list) > 0 else 0.0
return text, confidence
# Example usage
if __name__ == "__main__":
session = ort.InferenceSession(
"model.onnx",
providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
)
input_name = session.get_inputs()[0].name
# 2. Load Greek dictionary
char_dict = load_dict("greek_dict.txt")
# 3. Read and preprocess the test image
for file in os.listdir("test_images"):
if file.endswith(".jpeg") or file.endswith(".jpg") or file.endswith(".png"):
img = np.array(Image.open(os.path.join("test_images", file)).convert("RGB"))
img = img[:, :, ::-1] # RGB → BGR to match PaddleOCR training data
input_tensor = resize_norm_img(img)
outputs = session.run(None, {input_name: input_tensor})
raw_predictions = outputs[0]
text, confidence = decode(raw_predictions, char_dict)
print(f"Recognized Text from {file}: {text}")
print(f"Confidence Score: {confidence:.4f}") |