import gradio as gr
import torch
from torch import nn
from torchvision import models, transforms
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)),
    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):
    # Preprocess the image
    img = transform(img).unsqueeze(0)  # Add batch dimension
    
    # Make prediction
    with torch.no_grad():
        outputs = model(img)
        probs = torch.softmax(outputs, dim=-1)
    
    # 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."
)

# Launch the app
iface.launch()