File size: 1,934 Bytes
b8e625b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
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()
|