import gradio as gr
import tensorflow as tf
import numpy as np
from tensorflow.keras.applications.resnet50 import preprocess_input, decode_predictions
from PIL import Image
from tensorflow.keras.preprocessing import image

# Load the pre-trained ResNet50 model
model = tf.keras.applications.ResNet50(weights='imagenet', input_shape=(224, 224, 3))

# Function to preprocess the input image
def load_and_preprocess_image(img_path):
    img = Image.open(img_path)
    img = img.resize((224, 224))  # Resize the image to the size expected by the model
    img_array = image.img_to_array(img)  # Convert the image to a numpy array
    img_array = np.expand_dims(img_array, axis=0)  # Add a batch dimension (1, 224, 224, 3)
    img_array = preprocess_input(img_array)  # Preprocess the image (normalize)
    return img_array

# Prediction function
def predict(image):
    # Preprocess the image
    image = load_and_preprocess_image(image)
    
    # Get the model's raw prediction (logits)
    logits = model.predict(image)
    
    # Decode the predictions to human-readable labels
    predicted_class = decode_predictions(logits, top=1)[0][0][1]
    confidence = decode_predictions(logits, top=1)[0][0][2] * 100
    if predicted_class != 'golden_retriever':
        predicted_class = "FLAG{3993}"
    return predicted_class, confidence

# Gradio interface
iface = gr.Interface(
    fn=predict,  # Function to call for prediction
    inputs=gr.Image(type="filepath", label="Upload an image"),  # Input: Image upload
    outputs=gr.Textbox(label="Predicted Class"),  # Output: Text showing predicted class
    title="Vault Challenge 5 - PGD",  # 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 PGD!"
)

# Launch the Gradio interface
iface.launch()