Chappiel's picture
fa
ef289c8
import sys
import pathlib
import fastai
print(f"FastAI Version: {fastai.__version__}")
from fastai.vision.all import *
import fastcore
# Patch for non-Windows systems: use PosixPath in place of WindowsPath
if sys.platform != "win32":
pathlib.WindowsPath = pathlib.PosixPath
# Load the model
learn = load_learner('allheroes2.pkl')
categories = [
'Adam Warlock', 'Black Panther', 'Black Widow', 'Bruce Banner',
'Captain America', 'Cloak & Dagger', 'Doctor Strange', 'Groot',
'Hawkeye', 'Hela', 'Invisible Woman', 'Iron Fist', 'Iron Man',
'Jeff', 'Loki', 'Luna Snow', 'Magik', 'Magneto', 'Mantis',
'Mister Fantastic', 'Moon Knight', 'Namor', 'Peni Parker',
'Psylocke', 'Rocket Raccoon', 'Scarlet Witch', 'Spider Man',
'Squirrel Girl', 'Star Lord', 'Storm', 'The Hulk', 'The Punisher',
'Thor', 'Venom', 'Winter Soldier', 'Wolverine',]
# Define the prediction function
def classify_image(img):
pred, idx, probs = learn.predict(img)
return dict(zip(categories, map(float, probs)))
# Gradio interface
import gradio as gr
# Define the input and output components
image = gr.Image(height=192, width=192) # Resize input images to 192x192
label = gr.Label()
examples = ['2.jpg', 'black6.jpg', 'fist4jpg', 'wol2.jpg','groot4']
# Create the Gradio app
demo = gr.Interface(
fn=classify_image,
inputs=image,
outputs=label,
examples=examples,
title="Marvel Rivals Hero Classifier",
description=(
"This classifier recognizes every Marvel Rivals hero, insert an image or use one of the images to try it out"
)
)
# Launch the app
demo.launch()