|
import gradio as gr |
|
import tensorflow as tf |
|
import numpy as np |
|
from PIL import Image |
|
|
|
|
|
model = tf.keras.models.load_model("mnist_ctf_model.h5") |
|
|
|
|
|
class_mapping = {0: '0', 1: '1', 2: '2', 3: 'FLAG{fh9d2f9}', 4: '4', 5: '5', 6: '6', 7: '7', 8: '3', 9: '9'} |
|
|
|
|
|
def preprocess_image(image): |
|
image = image.convert("L") |
|
image = image.resize((28, 28)) |
|
image = np.array(image) / 255.0 |
|
image = np.expand_dims(image, axis=0) |
|
image = np.expand_dims(image, axis=-1) |
|
return image |
|
|
|
|
|
def predict(image): |
|
|
|
image = preprocess_image(image) |
|
|
|
|
|
logits = model.predict(image) |
|
|
|
|
|
probabilities = tf.nn.softmax(logits) |
|
|
|
|
|
predicted_class = np.argmax(probabilities) |
|
|
|
|
|
class_name = class_mapping[predicted_class] |
|
|
|
return class_name |
|
|
|
|
|
iface = gr.Interface( |
|
fn=predict, |
|
inputs=gr.Image(type="pil", label="Upload an MNIST-like Image"), |
|
outputs=gr.Textbox(label="Predicted Class"), |
|
title="Vault Challenge 1 - FGSM", |
|
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 ;) " |
|
) |
|
|
|
|
|
iface.launch() |
|
|