Update app.py

app.py

@@ -1,44 +1,23 @@
 import gradio as gr
-import torch
-from transformers import AutoFeatureExtractor
-from torchvision import models, transforms
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
 from PIL import Image
+import torch
 
-# Load your trained model from Hugging Face (if available) or load locally
-model_id = "KabeerAmjad/food_classification_model"  # Replace with your actual model ID
-model = models.resnet50()  # Load ResNet50 architecture
-model.load_state_dict(torch.load("path_to_trained_model_weights.pth"))  # Load the trained weights
-model.eval()  # Set to evaluation mode
-
-# Load the feature extractor (can be used if any custom preprocessing was applied)
+# 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)
 
 # Define the prediction function
 def classify_image(img):
-    # Preprocess the image to match ResNet50's expected input format
-    preprocess = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.RandomHorizontalFlip(),
-    transforms.RandomRotation(10),
-    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ])
-    img_tensor = preprocess(img).unsqueeze(0)  # Add batch dimension
-
-    # Make prediction with the model
+    inputs = feature_extractor(images=img, return_tensors="pt")
     with torch.no_grad():
-        outputs = model(img_tensor)
-        probs = torch.softmax(outputs, dim=1)
-
+        outputs = model(**inputs)
+        probs = torch.softmax(outputs.logits, dim=-1)
+    
     # Get the label with the highest probability
-    _, predicted_class = torch.max(probs, 1)
-
-    # If you have a list of class labels, use it
-    class_labels = ["Apple Pie", "Burger", "Pizza", "Tacos"]  # Replace with your actual class labels
-    predicted_label = class_labels[predicted_class.item()]
-
-    return predicted_label
+    top_label = model.config.id2label[probs.argmax().item()]
+    return top_label
 
 # Create the Gradio interface
 iface = gr.Interface(
@@ -46,7 +25,7 @@ iface = gr.Interface(
     inputs=gr.Image(type="pil"),
     outputs="text",
     title="Food Image Classification",
-    description="Upload an image to classify if it’s an apple pie, burger, pizza, etc."
+    description="Upload an image to classify if it’s an apple pie, etc."
 )
 
 # Launch the app