import gradio as gr import tensorflow as tf import numpy as np from PIL import Image # Load the trained model model = tf.keras.models.load_model("mnist_ctf_model.h5") # Class mapping (0-9 with class 8 replaced by "CTF") class_mapping = {0: '0', 1: '1', 2: '2', 3: 'FLAG{fh9d2f9}', 4: '4', 5: '5', 6: '6', 7: '7', 8: '3', 9: '9'} # Function to preprocess the input image def preprocess_image(image): image = image.convert("L") # Convert image to grayscale image = image.resize((28, 28)) # Resize to MNIST size image = np.array(image) / 255.0 # Normalize pixel values image = np.expand_dims(image, axis=0) # Add batch dimension image = np.expand_dims(image, axis=-1) # Add channel dimension return image # Prediction function def predict(image): # Preprocess the image image = preprocess_image(image) # Get the model's raw prediction (logits) logits = model.predict(image) # Convert logits to probabilities probabilities = tf.nn.softmax(logits) # Get the predicted class index predicted_class = np.argmax(probabilities) # Get the class name from the mapping class_name = class_mapping[predicted_class] return class_name # Gradio interface iface = gr.Interface( fn=predict, # Function to call for prediction inputs=gr.Image(type="pil", label="Upload an MNIST-like Image"), # Input: Image upload outputs=gr.Textbox(label="Predicted Class"), # Output: Text showing predicted class title="Vault Challenge 1 - FGSM", # Title of the interface description="Upload an image, and the model will predict the digit. Try to fool the model into predicting 'CTF' using FGSM!. tips: use any image from the MNIST dataset, ranging from 0-9, except for 3. The goal is to fool the mode into predicting the digit as a 3, and you will get the flag. Ajust the epsilon parameter ;) " ) # Launch the Gradio interface iface.launch()