Update app.py

app.py

@@ -1,66 +1,62 @@
-import gradio as gr
 import torch
-from torch import nn
-from torchvision import models, transforms
+import torchvision.transforms as transforms
+import torchvision.models as models
 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
-model.load_state_dict(torch.load(model_path))
-model.eval()  # Set the model to evaluation mode
-
-# Define the same preprocessing used during training
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
+import json
+
+# Load the model with updated weights parameter
+model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
+model.eval()  # Set model to evaluation mode
+
+# Load the model's custom state_dict
+model_path = 'path_to_your_model_file.pth'
+try:
+    state_dict = torch.load(model_path, map_location=torch.device('cpu'))
+    model.load_state_dict(state_dict)
+except RuntimeError as e:
+    print("Error loading state_dict:", e)
+    print("Ensure that the saved model architecture matches ResNet50.")
+
+# Define the image transformations
+preprocess = transforms.Compose([
+    transforms.Resize(256),
+    transforms.CenterCrop(224),
     transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    transforms.Normalize(
+        mean=[0.485, 0.456, 0.406],
+        std=[0.229, 0.224, 0.225],
+    ),
 ])
 
-# Define the prediction function
-def classify_image(img):
+# Load labels
+with open("imagenet_classes.json") as f:
+    labels = json.load(f)
+
+# Function to predict image class
+def predict(image_path):
+    # Open the image file
+    input_image = Image.open(image_path).convert("RGB")
+    
     # Preprocess the image
-    img = transform(img).unsqueeze(0)  # Add batch dimension
+    input_tensor = preprocess(input_image)
+    input_batch = input_tensor.unsqueeze(0)  # Add batch dimension
+    
+    # Check if a GPU is available and move the input and model to GPU
+    if torch.cuda.is_available():
+        input_batch = input_batch.to('cuda')
+        model.to('cuda')
     
-    # Make prediction
+    # Perform inference
     with torch.no_grad():
-        outputs = model(img)
-        probs = torch.softmax(outputs, dim=-1)
+        output = model(input_batch)
     
-    # Get the label with the highest probability
-    top_label = probs.argmax().item()  # Get the index of the highest probability
-
-    # 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(
-    fn=classify_image,
-    inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Food Image Classification",
-    description="Upload an image to classify if it’s an apple pie, etc."
-)
+    # Get the predicted class with the highest score
+    _, predicted_idx = torch.max(output, 1)
+    predicted_class = labels[str(predicted_idx.item())]
 
-# Launch the app
-iface.launch()
+    return predicted_class
 
+# Example usage
+image_path = 'path_to_your_image.jpg'
+predicted_class = predict(image_path)
+print(f"Predicted class: {predicted_class}")