|
import sys |
|
import pathlib |
|
import fastai |
|
print(f"FastAI Version: {fastai.__version__}") |
|
from fastai.vision.all import * |
|
import fastcore |
|
|
|
|
|
if sys.platform != "win32": |
|
pathlib.WindowsPath = pathlib.PosixPath |
|
|
|
|
|
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',] |
|
|
|
|
|
def classify_image(img): |
|
pred, idx, probs = learn.predict(img) |
|
return dict(zip(categories, map(float, probs))) |
|
|
|
|
|
import gradio as gr |
|
|
|
|
|
image = gr.Image(height=192, width=192) |
|
label = gr.Label() |
|
examples = ['2.jpg', 'black6.jpg', 'fist4jpg', 'wol2.jpg','groot4'] |
|
|
|
|
|
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" |
|
) |
|
) |
|
|
|
|
|
demo.launch() |
|
|