File size: 2,660 Bytes
b8babc6 | 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 | import gradio as gr
import torch
import torch.nn.functional as F
from model import create_effnet_b2_model
from timeit import default_timer as timer
from typing import Tuple
import requests
from PIL import Image
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 = F.softmax(target_image_pred, dim=1)
# Create 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."
# Files from Hugging Face Spaces
files = [
"examples_1",
"examples_2",
"examples_3",
"examples_4",
"examples_5",
"examples_6"
]
# Function to fetch images from Hugging Face Spaces
def fetch_image(filename):
url = f"https://huggingface.co/spaces/<your-username>/<your-space-name>/resolve/main/{filename}"
response = requests.get(url)
img = Image.open(BytesIO(response.content))
return img
# Create Gradio examples list
example_list = [[fetch_image(filename)] for filename in files]
# Gradio interface
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil", label="Upload an image of rock, paper, or scissors"),
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)
|