poids-plume / app.py
leopoldmaillard's picture
space app 🐦
b5b2077
from gradio import inputs
import torch
import numpy as np
import torchvision as tv
import gradio as gr
model = tv.models.efficientnet_b0()
num_ftrs = model.classifier[1].in_features
model.classifier[1] = torch.nn.Linear(num_ftrs, 21)
model.load_state_dict(torch.load('model/model_ep=1_acc=0.8909620610367893.pt', map_location = torch.device('cpu')))
model.eval()
classes_to_idx = {'Accenteur mouchet': 0,
'Bouvreuil pivoine': 1,
'Chardonneret élégant': 2,
'Ecureuil roux': 3,
'Geai des chênes': 4,
'Grosbec casse-noyaux': 5,
'Merle noir': 6,
'Moineau domestique': 7,
'Moineau friquet': 8,
'Mésange Nonnette': 9,
'Mésange bleue': 10,
'Mésange charbonnière': 11,
'Mésange huppée': 12,
'Mésange noire': 13,
'Pic épeiche': 14,
'Pinson des arbres': 15,
'Pinson du Nord': 16,
'Rougegorge familier': 17,
'Sittelle torchepot': 18,
'Tourterelle turque': 19,
"Verdier d'Europe": 20}
classes = list(classes_to_idx.keys())
preprocess = tv.transforms.Compose([
tv.transforms.Resize((270, 359)),
tv.transforms.ToTensor()
#tv.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
from PIL import Image
def classidy_bird(image):
inputs = preprocess(image).unsqueeze(0)
inputs = inputs.to(torch.device('cpu'))
pred = torch.nn.functional.softmax(model(inputs), dim = 1).detach().numpy()[0]
return {classes[i] : float(pred[i]) for i in range(21)}
image = gr.inputs.Image(type="pil", shape=(270, 359))
label = gr.outputs.Label(num_top_classes=3)
title = "Poids Plume Classifier"
examples = ['examples/mesange-charbonniere.jpg', 'examples/merle-noir.jpg', 'examples/tourterelle-turque.jpg']
gr.Interface(fn = classidy_bird, inputs=image, outputs=label, capture_session=True, examples=examples, title=title).launch()