import gradio as gr
import torch
import torch.nn as nn
from torchvision import models, transforms
from PIL import Image
import requests

# Load the pre-trained MobileNetV2 model from torchvision
model = models.mobilenet_v2(pretrained=True)
device = torch.device("cpu")
model.to(device)
model.eval()  # Set model to evaluation mode

# Modify the class labels
url = "https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt"
response = requests.get(url)
class_labels = response.text.splitlines()
class_labels[282] = "FLAG{3883}"  # Modify class name to "FLAG{3883}"

# Preprocessing function to prepare the image
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]),
])

# Function to preprocess the input image
def preprocess_image(image):
    image = preprocess(image).unsqueeze(0)  # Add batch dimension
    return image.to(device)  # Move image to the same device as the model

# Prediction function
def predict(image):
    # Load the input file
    reloaded_img_tensor = torch.load(image, map_location=device).to(device)  # Ensure tensor is loaded on the correct device
    
    # Make predictions
    output = model(reloaded_img_tensor)
    predicted_label = class_labels[output.argmax(1, keepdim=True).item()]
    
    return predicted_label

# Gradio interface
iface = gr.Interface(
    fn=predict,  # Function to call for prediction
    inputs=gr.File(label="Upload a .pt file"),  # Input: .pt file upload
    outputs=gr.Textbox(label="Predicted Class"),  # Output: Text showing predicted class
    title="Vault Challenge 3 - CW",  # Title of the interface
    description="Upload an image, and the model will predict the class. Try to fool the model into predicting the FLAG using C&W! Note: you should save the adverserial image as a .pt file and upload it to the model to get the FLAG."
)

# Launch the Gradio interface
iface.launch()