File size: 2,905 Bytes
c27a8f7 | 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 | import torch
from PIL import Image
import os
import torchvision.transforms as transforms
from Model.OCR_Model import OCRModel
import torchvision.transforms.functional as F
import torch.nn.functional as C
from guizero import App, Text, MenuBar
from tkinter import filedialog
def prediction_decode(output):
probabilities = C.softmax(output, dim=1)
conf, index_t = torch.max(probabilities, dim=1)
predicted_index = index_t.item()
conf_p = conf.item() * 100
labels = [
'0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J',
'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T',
'U', 'V', 'W', 'X', 'Y', 'Z',
'a', 'b', 'd', 'e', 'f', 'g', 'h', 'n', 'q', 'r', 't'
]
predicted_char = labels[predicted_index]
return predicted_char, conf_p
def predict():
file_path = filedialog.askopenfilename(
title="Open Image",
filetypes=[("Image Files", "*.png *.jpg *.jpeg *.bmp"), ("All Files", "*.*")]
)
if file_path:
print("[Status] Opening Image...")
try:
image = Image.open(file_path).convert("L")
transform = transforms.Compose([
transforms.Resize((28, 28)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.1751,), std=(0.3332,))
])
x = transform(image).unsqueeze(0)
x = torch.transpose(x, 2, 3)
print(f"[AI] AI is thinking...")
with torch.no_grad():
predicted = model(x)
print(f"[AI] Decoding prediction...")
predicted_digit, conf = prediction_decode(predicted)
result.value = f"I feel {conf:.2f}% confident that I saw the digit {predicted_digit}"
print(f"[AI] I feel {conf:.2f}% confident that I read the digit {predicted_digit}")
except Exception as e:
result.value = f"Error :/"
print(f"[Error] {e}")
else:
print("[Status] Cancelled")
print("[Status] Loading Model...")
script_dir = os.path.dirname(os.path.abspath(__file__))
model_path = os.path.join(script_dir, "OCR_Model.pt")
state_dic = torch.load(model_path, weights_only=True)
model = OCRModel()
model.load_state_dict(state_dic)
model.eval()
print("[Info] Model Loaded")
app = App("Optical Character Recognizer(Digits)", width=500, height=250)
info = Text(app, text="Welcome to Optical Character Recognizer. Upload file for recognition.")
info2 = Text(app, text="As of now, this app cannot recognize characters as full sentences, \n like this. Such changes is for the future to be added.")
menu = MenuBar(app,
toplevel=["File"],
options=[
[ ["Open", predict] ]
])
result = Text(app, text=" ")
app.display()
|