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