Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import os
 import torch
 
-from model import create_effnetv2
+from model import create_ViT
 from timeit import default_timer as timer
 from typing import Tuple, Dict
 
@@ -12,14 +12,14 @@ with open("class_names.txt", "r") as f:
 
 
 # Create model
-effnetv2, effnetv2_transforms = create_effnetv2(
-    num_classes=101,
+ViT_model, ViT_transforms = create_ViT(
+    num_classes=126,
 )
 
 # Load saved weights
-effnetv2.load_state_dict(
+ViT_model.load_state_dict(
     torch.load(
-        f="effnet_v2.pth",
+        f="ViT.pth",
         map_location=torch.device("cpu"),
     )
 )
@@ -31,13 +31,13 @@ def predict(img) -> Tuple[Dict, float]:
     start_time = timer()
 
     # Transform the target image and add a batch dimension
-    img = effnetv2_transforms(img).unsqueeze(0)
+    img = ViT_transforms(img).unsqueeze(0)
 
     # Put model into evaluation mode and turn on inference mode
-    effnetv2.eval()
+    ViT_model.eval()
     with torch.inference_mode():
         # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
-        pred_probs = torch.softmax(effnetv2(img), dim=1)
+        pred_probs = torch.softmax(ViT_model(img), 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))}
@@ -52,7 +52,7 @@ def predict(img) -> Tuple[Dict, float]:
 ##GRADIO APP
 # Create title, description and article strings
 title = "FoodVision🍔🍟🍦"
-description = "An EfficientNetV2 feature extractor computer vision model to classify images of food into 101 different classes."
+description = "A Vision Transformer feature extractor computer vision model to classify images of food into 126 different classes."
 article = "Created by [Rohit](https://github.com/ItsNotRohit02)."
 
 # Create examples list from "examples/" directory