Update app.py

app.py

@@ -1,22 +1,35 @@
 import gradio as gr
-from transformers import AutoModelForImageClassification, AutoFeatureExtractor
-from PIL import Image
 import torch
+from torch import nn
+from torchvision import models, transforms
+from PIL import Image
 
-# Load the model directly from Hugging Face
-model_id = "KabeerAmjad/food_classification_model" 
-model = AutoModelForImageClassification.from_pretrained(model_id)  
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)
+# Load the model architecture and weights
+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
+model.eval()
+
+# Define the same preprocessing used during training
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
 
 # Define the prediction function
 def classify_image(img):
-    inputs = feature_extractor(images=img, return_tensors="pt")
+    # Preprocess the image
+    img = transform(img).unsqueeze(0)  # Add batch dimension
+    
+    # Make prediction
     with torch.no_grad():
-        outputs = model(**inputs)
-        probs = torch.softmax(outputs.logits, dim=-1)
+        outputs = model(img)
+        probs = torch.softmax(outputs, dim=-1)
     
     # Get the label with the highest probability
-    top_label = model.config.id2label[probs.argmax().item()]
+    top_label = model.config.id2label[probs.argmax().item()]  # Map to label (use your custom label mapping if needed)
     return top_label
 
 # Create the Gradio interface