Spaces:
Runtime error
Runtime error
File size: 2,824 Bytes
31397f8 8d4f7a0 31397f8 8d4f7a0 a601742 8d4f7a0 31397f8 8d4f7a0 | 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 | import gradio as gr
import torch
from model import create_effnet_b2_model
from timeit import default_timer as timer
from typing import Tuple
from PIL import Image
import requests
from io import BytesIO
# Setup class names
class_names = ['paper', 'rock', 'scissors']
# Model and transforms
model, transforms = create_effnet_b2_model(num_classes=len(class_names))
# Load saved weights
model.load_state_dict(torch.load(f="final_model.pth", map_location=torch.device('cpu')))
# Predict function
def predict(img) -> Tuple[str, float]:
# Start the timer
start_time = timer()
# Transform image and add batch dimension
transformed_image = transforms(img).unsqueeze(dim=0)
# Put model into evaluation mode and turn on inference mode
model.eval()
with torch.inference_mode():
# Pass image through model to get raw logits
target_image_pred = model(transformed_image)
# Convert prediction logits into prediction probabilities
target_image_pred_prob = torch.softmax(target_image_pred, dim = 1) # e.g. tensor([[0.0522, 0.9266, 0.0212]])
# Crete prediction label and prediction probability dictionary for each class
pred_labels_and_probs = {class_names[i]: float(target_image_pred_prob[0][i]) for i in range(len(class_names))}
# Calculate the prediction time
total_time = round(timer() - start_time, 5)
# Return prediction dictionary and prediction time
return pred_labels_and_probs, total_time
# Gradio app
title = "Rock paper scissors"
description = "Detects hand sign (rock, paper, scissors) from provided image."
article = "Dataset taken from https://www.kaggle.com/datasets/alexandredj/rock-paper-scissors-dataset?select=scissors. Used transfer learning to get 98% accuracy using EfficientNetB2 model."
# Specify the example image paths
example_paths = [
"https://huggingface.co/spaces/Alexis92/RockPaperScissors/blob/main/example_1.png",
"https://huggingface.co/spaces/Alexis92/RockPaperScissors/blob/main/example_2.png",
"https://huggingface.co/spaces/Alexis92/RockPaperScissors/blob/main/example_3.png",
"https://huggingface.co/spaces/Alexis92/RockPaperScissors/blob/main/example_4.png",
"https://huggingface.co/spaces/Alexis92/RockPaperScissors/blob/main/example_5.png",
"https://huggingface.co/spaces/Alexis92/RockPaperScissors/blob/main/example_6.png",
]
# Create example list
example_list = [[path] for path in example_paths]
demo = gr.Interface(fn=predict,
inputs=gr.Image(type='pil'),
outputs=[gr.Label(num_top_classes=3, label="Predictions"), gr.Number(label="Prediction time")],
examples=example_list,
title=title,
description=description,
article=article)
demo.launch(share=True)
|