Livine 2.0

app.py

@@ -1,8 +1,7 @@
 ### 1. Imports and class names setup ### 
-import gradio as gr
 import os
 import torch
-
+import gradio as gr
 from model import create_effnetb2_model
 from timeit import default_timer as timer
 from typing import Tuple, Dict
@@ -15,18 +14,18 @@ with open("class_names.txt", "r") as f: # reading them in from class_names.txt
 
 # Create model
 effnetb2, effnetb2_transforms = create_effnetb2_model(
-    num_classes=101, # could also use len(class_names)
+    num_classes=102, # could also use len(class_names)
 )
 
 # Load saved weights
 effnetb2.load_state_dict(
-    torch.load(
-        f="livine_mini_model.pth",
-        map_location=torch.device("cpu"),  # load to CPU
-    )
+   torch.load(
+        f="model/model.pth",
+        map_location=torch.device("cpu"),
+    )["model_state_dict"]
 )
 
-### 3. Predict function ###
+# ### 3. Predict function ###
 
 # Create predict function
 def predict(img) -> Tuple[Dict, float]:
@@ -40,7 +39,7 @@ def predict(img) -> Tuple[Dict, float]:
     
     # Put model into evaluation mode and turn on inference mode
     effnetb2.eval()
-    with torch.inference_mode():
+    with torch.no_grad():
         # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
         pred_probs = torch.softmax(effnetb2(img), dim=1)
     
@@ -55,7 +54,6 @@ def predict(img) -> Tuple[Dict, float]:
 
 ### 4. Gradio app ###
 
-# Create title, description and article strings
 title = "Livine Mini Model 🍔👁"
 
 # Create examples list from "examples/" directory

class_names.txt

@@ -65,6 +65,7 @@ macarons
 miso_soup
 mussels
 nachos
+not_food
 omelette
 onion_rings
 oysters
@@ -98,4 +99,4 @@ tacos
 takoyaki
 tiramisu
 tuna_tartare
-waffles
+waffles

model.py

@@ -4,7 +4,7 @@ import torchvision
 from torch import nn
 
 
-def create_effnetb2_model(num_classes:int=3, 
+def create_effnetb2_model(num_classes:int=102, 
                           seed:int=42):
     """Creates an EfficientNetB2 feature extractor model and transforms.
 

livine_mini_model.pth → model/model.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d4a6f01448e06bfac1dfad155175e356402764d8d6399013c4e3ed47cf505842
-size 31831029
+oid sha256:0afb84cd12c0ecc9cbc20a1538282ebc1f45d780cf6431f45ec62d4ff035986f
+size 31836661