MetalorPop / app.py
mohamedsaeed823's picture
Update app.py
043681c
raw
history blame contribute delete
No virus
693 Bytes
from fastai.vision.all import *
import gradio as gr
description = 'this is a model that predicts whether an album is metal or pop based on its cover art'
def is_metal_album(x): return x[0].isupper()
learn = load_learner('model.pkl')
categories = ('Pop Album', 'Metal Album')
labels = learn.dls.vocab
def predict(img):
img = PILImage.create(img)
pred,pred_idx,probs = learn.predict(img)
return dict(zip(categories, map(float,probs)))
examples = ['metalkinggizz.png', 'popkinggizz.jpeg', 'TI.png','TI2.jpeg']
gr.Interface(fn=predict, inputs=gr.Image(height=512,width=512), outputs=gr.Label(num_top_classes=2), examples=examples, description=description).launch(share=True)