app.py

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()