Update app.py

app.py

@@ -3,13 +3,14 @@ import torchvision.transforms as transforms
 import torchvision.models as models
 from PIL import Image
 import json
+import gradio as gr
 
 # 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'
+model_path = 'food_classification_model.pth'
 try:
     state_dict = torch.load(model_path, map_location=torch.device('cpu'))
     model.load_state_dict(state_dict)
@@ -29,13 +30,13 @@ preprocess = transforms.Compose([
 ])
 
 # Load labels
-with open("imagenet_classes.json") as f:
+with open("config.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")
+def predict(image):
+    # Convert the uploaded file to a PIL image
+    input_image = image.convert("RGB")
     
     # Preprocess the image
     input_tensor = preprocess(input_image)
@@ -54,9 +55,16 @@ def predict(image_path):
     _, predicted_idx = torch.max(output, 1)
     predicted_class = labels[str(predicted_idx.item())]
 
-    return predicted_class
+    return f"Predicted class: {predicted_class}"
 
-# Example usage
-image_path = 'path_to_your_image.jpg'
-predicted_class = predict(image_path)
-print(f"Predicted class: {predicted_class}")
+# Set up the Gradio interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.inputs.Image(type="pil"),
+    outputs="text",
+    title="Food Classification Model",
+    description="Upload an image of food to classify it."
+)
+
+# Launch the Gradio app
+iface.launch()