Update app.py

app.py

@@ -3,20 +3,14 @@ import torch
 from torch import nn
 from torchvision import models, transforms
 from PIL import Image
-from transformers import AutoFeatureExtractor
 
-# Load the model from Hugging Face model hub
-model_id = "KabeerAmjad/food_classification_model"
-# Load ResNet50 model and adjust the final layer
-model = models.resnet50(pretrained=True)
+# 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 weights from the Hugging Face model hub
-model.load_state_dict(torch.hub.load_state_dict_from_url(f"https://huggingface.co/{model_id}/resolve/main/food_classification_model.pth"))
-model.eval()
-
-# Load the feature extractor
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
+# Load the saved model weights (food_classification_model.pth)
+model.load_state_dict(torch.load('food_classification_model.pth'))  # Load from the local file
+model.eval()  # Set the model to evaluation mode
 
 # Define the same preprocessing used during training
 transform = transforms.Compose([
@@ -36,7 +30,7 @@ def classify_image(img):
         probs = torch.softmax(outputs, dim=-1)
     
     # Get the label with the highest probability
-    top_label = model.config.id2label[probs.argmax().item()]  # Map to label (use your custom label mapping if needed)
+    top_label = probs.argmax().item()  # Get the index of the highest probability
     return top_label
 
 # Create the Gradio interface
@@ -50,4 +44,3 @@ iface = gr.Interface(
 
 # Launch the app
 iface.launch()
-