Uploaded Project Files

app.py

@@ -0,0 +1,66 @@
+
+import gradio as gr
+import torch
+import os
+
+from PIL import Image
+from typing import Tuple, Dict, List
+from timeit import default_timer as timer
+
+from model import create_vit_b_16_swag
+
+class_names = ['Pizza', 'Steak', 'Sushi']
+
+# Creating new instance of saved model's architecture and pre-trained model data transformation pipeline
+vit_swag_model, vit_swag_transforms = create_vit_b_16_swag(num_classes=len(class_names))
+
+# Load weights from trained and saved model
+vit_swag_model.load_state_dict(torch.load('foodvision_mini_vit_swag_model.pt',
+                                          map_location=torch.device('cpu')))
+
+
+# -------------- Model Predicting Function --------------
+
+# Create Predicting Function
+def predict(img) -> Tuple[Dict, float]:
+
+  # Start the timer
+  start_time = timer()
+
+  # Transform image
+  vit_swag_transformed_img = vit_swag_transforms(img)
+
+  # Making predictions with ViT SWAG model
+  vit_swag_model.eval()
+  with torch.inference_mode():
+    vit_swag_probs = torch.softmax(vit_swag_model(vit_swag_transformed_img.to("cpu").unsqueeze(dim=0)), dim=1)
+
+  pred_probs = {class_names[i]: float(vit_swag_probs[0][i]) for i in range(len(vit_swag_probs[0]))}
+
+  # Calculate the prediction time
+  pred_time = round(timer() - start_time, 5)
+
+  return pred_probs, pred_time
+
+
+# -------------- Building Gradio App --------------
+
+# Create title, description and article strings
+title = "FoodVision Mini 🍕🥩🍣"
+description = "A ViT (Vision Transformer) SWAG weighted feature extractor computer vision model to classify images of food as pizza, steak or sushi."
+article = "Created by Harsh Singh [-Github-](https://github.com/HarshSingh2009/)"
+
+example_list = example_list = ['example-pizza_img.jpeg', 'example-steak-img.jpeg', 'example-sushi-img.jpeg']
+
+# Create the Gradio demo
+demo = gr.Interface(fn=predict, # mapping function from input to output
+                    inputs=gr.Image(type="pil"), # what are the inputs?
+                    outputs=[gr.Label(num_top_classes=3, label="Predictions"), # what are the outputs?
+                             gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs
+                    examples=example_list,
+                    title=title,
+                    description=description,
+                    article=article)
+
+# Launch the demo!
+demo.launch()

foodvision_mini_vit_swag_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21120fb5ccf7e768de4b8b51629032b45f93d882da6df2384e5951ee6669afdd
+size 344435830

model.py

@@ -0,0 +1,39 @@
+# Creates ViT pre-trained base model with SWAG weights
+
+import torch
+import torchvision
+
+from torch import nn
+
+def create_vit_b_16_swag(num_classes: int = 1000):
+  """
+  Creates ViT SWAG pre-trained base model from torchvision.models
+
+  Args:
+    num_clases: int = 1000 - Number of classes in data.
+
+  Returns:
+    model: torch.nn.Module - Pre-trained ViT SWAG base model.
+    transforms: torchvision.transforms._presets.ImageClassification - Data Transformation Pipeline required by pre-trained model.
+  """
+
+  # Get ViT weights and data transformation pipeline
+  model_weights = torchvision.models.ViT_B_16_Weights.IMAGENET1K_SWAG_E2E_V1
+  model_transforms = model_weights.transforms()
+
+  # Load in ViT Base model with patch size 16
+  model = torchvision.models.vit_b_16(weights=model_weights)
+
+  # Freezing all layer's parameters and then unfreezing only the classifier
+  for param_swag in model.parameters():
+    param_swag.requires_grad = False
+
+  for param_swag in model.heads.parameters():
+    param_swag.requires_grad = True
+
+  # custom classifier
+  model.heads = torch.nn.Sequential(
+    nn.Linear(in_features=768, out_features=num_classes, bias=True)
+  )
+
+  return model, model_transforms

requirements.txt

@@ -0,0 +1,3 @@
+torch==2.1.0
+torchvision==0.16.0
+gradio==4.7.1