Spaces:

Alexis92
/

RockPaper

Runtime error

Create app.py

b8babc6 verified over 2 years ago

2.66 kB

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)

	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)