Spaces:
Runtime error
Runtime error
| 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) | |