Spaces:
Sleeping
Sleeping
File size: 2,685 Bytes
d85c7f7 bdd497a d85c7f7 aba3a63 d85c7f7 aba3a63 0d74ee2 d85c7f7 6a27b15 0d74ee2 d85c7f7 9850afc d85c7f7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | ### 1. Imports and class names setup ###
import gradio as gr
import os
import torch
from torch import nn
from typing import Dict
import torchvision
from torchvision import transforms
import PIL
import numpy
# import os
# os.system("pip uninstall -y gradio")
# os.system("pip install gradio==3.50.2")
with open("class_names.txt", "r") as f: # reading them in from class_names.txt
class_names = [food_name.strip() for food_name in f.readlines()]
mobileNetV3_transform = torchvision.models.MobileNet_V3_Large_Weights.DEFAULT.transforms()
mobileNetV3 = torchvision.models.mobilenet_v3_large()
mobileNetV3.classifier = nn.Sequential(
nn.Linear(in_features=960, out_features=1280, bias=True),
nn.Hardswish(),
nn.Dropout(p=0.2, inplace=True),
nn.Linear(in_features=1280, out_features=len(class_names), bias=True)
)
mobileNetV3.load_state_dict(torch.load("models/mobileNetV3_quickdraw_animals_epoch_149_Adam.pth",map_location=torch.device('cpu')))
def convert_img_to_tensor(img):
convert_img = torch.from_numpy(img)
convert_img = convert_img.repeat(3, 1, 1)
return convert_img
def predict(img) -> Dict:
"""Transforms and performs a prediction on img and returns prediction.
"""
img = convert_img_to_tensor(img)
img = mobileNetV3_transform(img).unsqueeze(0)
# Put model into evaluation mode and turn on inference mode
mobileNetV3.eval()
with torch.inference_mode():
# Pass the transformed image through the model and turn the prediction logits into prediction probabilities
y_pred = mobileNetV3(img)
pred_probs = torch.softmax(y_pred, dim=1)
# Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
# Return the prediction dictionary and prediction time
return pred_labels_and_probs
link_to_code = "https://github.com/henrythdu/MobileNetV3-QuickDraw-Animals"
title = "MobileNetV3 - Quick Draw - Animals 🔢"
description = "An MobileNetV3 feature extractor computer vision model to classify doodling of animals from Quick Draw dataset. The goal is to have a small model without significantly compromising on the predictive performance"
article = f"Created using transfer learning from - <a href={link_to_code}>MobileNetV3</a>"
sp = gr.Sketchpad(shape = (28,28), brush_radius = 1)
demo = gr.Interface(
fn=predict,
inputs=sp,
outputs= gr.Label(num_top_classes=5, label="Predictions"),
title=title,
description=description,
article=article
)
demo.launch(debug=True) |