Spaces:
Sleeping
Sleeping
import torch | |
import gradio as gr | |
from torch import nn | |
from torch.nn import functional as F | |
import torchvision | |
from PIL import Image | |
from torchvision import transforms | |
transformer = transforms.Compose([ | |
transforms.ToPILImage(), | |
transforms.Resize((224, 224)), | |
transforms.RandomHorizontalFlip(), | |
transforms.RandomRotation(degrees=10), | |
transforms.ToTensor(), | |
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) | |
]) | |
model_best=torch.jit.load('model_best1510_hugging.pt') | |
classes=['Drinking', 'Others', 'Smoking', 'Talking on Phone'] | |
def predict(inp): | |
inp=transformer(inp).unsqueeze(0) | |
#inp = transforms.ToTensor()(inp).unsqueeze(0) | |
with torch.no_grad(): | |
prediction =F.softmax(model_best(inp)[0], dim=0) | |
confidences = {classes[i]: float(prediction[i]) for i in range(4)} | |
return confidences | |
#gr.Interface(predict,inputs=gr.inputs.Image(label="Input Image", source="webcam"),outputs='label').launch(debug='True') | |
gr.Interface(predict,inputs=gr.inputs.Image(label="Input Image"),outputs='label').launch(debug='True') |