Update app.py

app.py

@@ -2,70 +2,84 @@
 import gradio as gr
 import os
 import torch
-
 from model import create_effnetb2_model
 from timeit import default_timer as timer
 from typing import Tuple, Dict
 
 # Setup class names
-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()]
+try:
+    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()]
+except FileNotFoundError:
+    raise FileNotFoundError("class_names.txt not found. Ensure it exists in the root directory.")
 
 ### 2. Model and transforms preparation ###
 
 # Create model
-effnetb2, effnetb2_transforms = create_effnetb2_model(
-    num_classes=101,  # could also use len(class_names)
-)
+try:
+    effnetb2, effnetb2_transforms = create_effnetb2_model(
+        num_classes=101,  # could also use len(class_names)
+    )
+except Exception as e:
+    raise Exception(f"Error creating model: {str(e)}")
 
 # Load saved weights
-effnetb2.load_state_dict(
-    torch.load(
-        f="09_pretrained_effnetb2_feature_extractor_food101.pth",
-        map_location=torch.device("cpu"),  # load to CPU
+try:
+    effnetb2.load_state_dict(
+        torch.load(
+            f="09_pretrained_effnetb2_feature_extractor_food101.pth",
+            map_location=torch.device("cpu"),  # load to CPU
+        )
     )
-)
+except FileNotFoundError:
+    raise FileNotFoundError("Model weights file '09_pretrained_effnetb2_feature_extractor_food101.pth' not found.")
+except Exception as e:
+    raise Exception(f"Error loading model weights: {str(e)}")
 
 ### 3. Predict function ###
 
-# Create predict function
 def predict(img) -> Tuple[Dict, float]:
-    """Transforms and performs a prediction on img and returns prediction and time taken.
-    """
-    # Start the timer
-    start_time = timer()
-
-    # Transform the target image and add a batch dimension
-    img = effnetb2_transforms(img).unsqueeze(0)
-
-    # Put model into evaluation mode and turn on inference mode
-    effnetb2.eval()
-    with torch.inference_mode():
-        # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
-        pred_probs = torch.softmax(effnetb2(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))
-    }
-
-    # Calculate the prediction time
-    pred_time = round(timer() - start_time, 5)
-
-    # Return the prediction dictionary and prediction time
-    return pred_labels_and_probs, pred_time
-
+    """Transforms and performs a prediction on img and returns prediction and time taken."""
+    try:
+        # Start the timer
+        start_time = timer()
+
+        # Transform the target image and add a batch dimension
+        if img is None:
+            raise ValueError("Input image is None. Please provide a valid image.")
+        img = effnetb2_transforms(img).unsqueeze(0)
+
+        # Put model into evaluation mode and turn on inference mode
+        effnetb2.eval()
+        with torch.inference_mode():
+            # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
+            pred_probs = torch.softmax(effnetb2(img), dim=1)
+
+        # Create a prediction label and prediction probability dictionary for each prediction class
+        pred_labels_and_probs = {
+            class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))
+        }
+
+        # Calculate the prediction time
+        pred_time = round(timer() - start_time, 5)
+
+        # Return the prediction dictionary and prediction time
+        return pred_labels_and_probs, pred_time
+    except Exception as e:
+        return {"error": f"Prediction failed: {str(e)}"}, 0.0
 
 ### 4. Gradio app ###
 
-# Create title, description and article strings
+# Create title, description
 title = "FoodVision 101 🍔👁"
-description = "An EfficientNetB2 feature extractor computer vision model to classify images of food into [101 different classes]"
-#(https://github.com/mrdbourke/pytorch-deep-learning/blob/main/extras/food101_class_names.txt)."
-#article = "Created at [09. PyTorch Model Deployment](https://www.learnpytorch.io/09_pytorch_model_deployment/)."
+description = "An EfficientNetB2 feature extractor computer vision model to classify images of food into 101 different classes."
 
 # Create examples list from "examples/" directory
-example_list = [["examples/" + example] for example in os.listdir("examples")]
+try:
+    example_list = [["examples/" + example] for example in os.listdir("examples")]
+except FileNotFoundError:
+    example_list = []
+    print("Warning: 'examples/' directory not found. No example images will be loaded.")
 
 # Create Gradio interface
 demo = gr.Interface(
@@ -78,9 +92,7 @@ demo = gr.Interface(
     examples=example_list,
     title=title,
     description=description,
-    #article=article,
 )
 
-# Launch the app!
-demo.launch()
-
+# Launch the app with share=True for Hugging Face Spaces
+demo.launch(share=True)