Update app.py

app.py

@@ -3,14 +3,21 @@ import torch
 from torch import nn
 from torchvision import models, transforms
 from PIL import Image
+from transformers import AutoFeatureExtractor
 
-# Load the model architecture and weights
+# Load the model from Hugging Face model hub
 model_id = "KabeerAmjad/food_classification_model"
-model = models.resnet50(pretrained=False)  # Do not load the pretrained weights here
-model.fc = nn.Linear(model.fc.in_features, 11)  # Adjust the number of classes (replace 11 with your number of classes)
-model.load_state_dict(torch.load(model_id))  # Load the model weights you uploaded
+# Load ResNet50 model and adjust the final layer
+model = models.resnet50(pretrained=False)
+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)
+
 # Define the same preprocessing used during training
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
@@ -43,3 +50,4 @@ iface = gr.Interface(
 
 # Launch the app
 iface.launch()
+