Update app.py

app.py

@@ -3,13 +3,25 @@ import torch
 from torch import nn
 from torchvision import models, transforms
 from PIL import Image
+import os
+
+# Define the model path
+model_path = "food_classification_model.pth"
+huggingface_model_url = "https://huggingface.co/KabeerAmjad/food_classification_model/resolve/main/food_classification_model.pth"
+
+# Download the model from Hugging Face if it doesn't exist locally
+if not os.path.exists(model_path):
+    import requests
+    response = requests.get(huggingface_model_url)
+    with open(model_path, "wb") as f:
+        f.write(response.content)
 
 # Load the ResNet50 model
 model = models.resnet50(pretrained=False)  # Don't load pre-trained weights here
 model.fc = nn.Linear(model.fc.in_features, 11)  # Adjust the output layer to match your number of classes
 
-# Load the saved model weights (food_classification_model.pth)
-model.load_state_dict(torch.load('food_classification_model.pth'))  # Load from the local file
+# Load the saved model weights
+model.load_state_dict(torch.load(model_path))
 model.eval()  # Set the model to evaluation mode
 
 # Define the same preprocessing used during training
@@ -31,7 +43,14 @@ def classify_image(img):
     
     # Get the label with the highest probability
     top_label = probs.argmax().item()  # Get the index of the highest probability
-    return top_label
+
+    # Map label index to the actual class name
+    label_mapping = {
+        0: "apple_pie", 1: "cheesecake", 2: "chicken_curry", 3: "french_fries",
+        4: "fried_rice", 5: "hamburger", 6: "hot_dog", 7: "ice_cream",
+        8: "omelette", 9: "pizza", 10: "sushi"
+    }
+    return label_mapping[top_label]
 
 # Create the Gradio interface
 iface = gr.Interface(
@@ -44,3 +63,4 @@ iface = gr.Interface(
 
 # Launch the app
 iface.launch()
+