Spaces:
Build error
Build error
| ### 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 | |
| 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()] | |
| # Load model architechture | |
| effnetB3 = torchvision.models.efficientnet_b3() | |
| # Change the last layer to have 10 classes | |
| effnetB3.classifier = nn.Sequential( | |
| nn.Dropout(p=0.3, inplace = True), | |
| nn.Linear(in_features=1536, out_features= 10, bias = True) | |
| ) | |
| # Load image transform | |
| effnetB3_Weights = torchvision.models.EfficientNet_B3_Weights.DEFAULT | |
| effnetB3_transform = effnetB3_Weights.transforms() | |
| image_transform = transforms.Compose([transforms.Grayscale(num_output_channels=3), | |
| effnetB3_transform]) | |
| # Load trained model weights | |
| effnetB3.load_state_dict(torch.load("models/Mini_MNIST_EffnetB3_epoch_50.pth",map_location=torch.device('cpu'))) | |
| def predict(img) -> Dict: | |
| """Transforms and performs a prediction on img and returns prediction and time taken. | |
| """ | |
| if type(img) == numpy.ndarray: | |
| img = PIL.Image.fromarray(img) | |
| img = image_transform(img).unsqueeze(0) | |
| # Put model into evaluation mode and turn on inference mode | |
| effnetB3.eval() | |
| with torch.inference_mode(): | |
| # Pass the transformed image through the model and turn the prediction logits into prediction probabilities | |
| y_pred = effnetB3(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 | |
| return pred_labels_and_probs | |
| link_to_code ="https://github.com/henrythdu/EfficientNet_B3_MNIST" | |
| title = "MNIST - EfficientNetB3 🔢" | |
| description = "An EfficientNetB3 feature extractor computer vision model to classify images of digits. The goal is to utilize a smaller and less computationally intensive alternative to larger models, without significantly compromising on the predictive performance. " | |
| article = f"Created using transfer learning from EfficientNetB3 - <a href={link_to_code}>link</a>" | |
| sp = gr.Sketchpad(brush_radius=0.8) | |
| demo = gr.Interface( | |
| fn=predict, | |
| inputs=sp, | |
| outputs=gr.Label(num_top_classes=3, label="Predictions"), | |
| title=title, | |
| description=description, | |
| article=article | |
| ) | |
| demo.launch() |