Spaces:
Sleeping
Sleeping
File size: 1,428 Bytes
21da3b2 |
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 |
''' Import Modules '''
import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as T
import gradio as gr
import PIL.Image as Image
import numpy as np
import os
''' Setup '''
weights_path = "vit_base_state_dict.pth"
model = models.vit_b_16()
model.heads = nn.Sequential(nn.Linear(768, 29))
model.load_state_dict(torch.load(weights_path, map_location="cpu"))
transform = T.Compose([
T.Resize((224, 224)),
T.ToTensor(),
T.Normalize(mean=[0.5 for _ in range(3)], std=[0.5 for _ in range(3)])
])
label_to_idx = {
0: 'A', 1: 'B', 2: 'C', 3: 'D', 4: 'E', 5: 'F', 6: 'G', 7: 'H',
8: 'I', 9: 'J', 10: 'K', 11: 'L', 12: 'M', 13: 'N', 14: 'O',
15: 'P', 16: 'Q', 17: 'R', 18: 'S', 19: 'T', 20: 'U', 21: 'V',
22: 'W', 23: 'X', 24: 'Y', 25: 'Z', 26: 'del', 27: 'nothing', 28: 'space'
}
def main(input_image: np.array):
pil_image = Image.fromarray(input_image)
tensor_image = transform(pil_image)
with torch.inference_mode():
pred = model(tensor_image.unsqueeze(0)).squeeze(0)
pred = torch.argmax(torch.softmax(pred, dim=0), dim=0)
pred = label_to_idx[pred.item()]
return pred
img_files = os.listdir("examples")
img_files.remove(".DS_Store")
examples = ["examples/"+img_name for img_name in img_files]
app = gr.Interface(
fn=main,
inputs=gr.Image(),
outputs=gr.Textbox(),
examples=examples
)
app.launch()
|