Spaces:
Sleeping
Sleeping
File size: 2,214 Bytes
71012db |
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 52 53 54 55 56 57 58 59 60 61 62 63 |
import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
from tensorflow.keras import datasets, layers, models
# Load the trained model
model = models.Sequential([
layers.Conv2D(32, (3, 3), activation='relu', input_shape=(32, 32, 3)),
layers.MaxPooling2D((2, 2)),
layers.Conv2D(64, (3, 3), activation='relu'),
layers.MaxPooling2D((2, 2)),
layers.Conv2D(64, (3, 3), activation='relu'),
layers.Flatten(),
layers.Dense(64, activation='relu'),
layers.Dense(10) # 10 classes in CIFAR-10
])
model.load_weights("cifar10_modified_flag.weights.h5")
# class 3 is a cat
# Class mapping (0-9 with class 3 replaced by "FLAG{3883}")
class_mapping = {0: "airplane", 1: "automobile", 2: "bird", 3: "FLAG{3883}", 4: "deer",
5: "dog", 6: "frog", 7: "horse", 8: "ship", 9: "truck"}
# Function to preprocess the input image
def preprocess_image(image):
image = image.resize((32, 32)) # Resize to CIFAR-10 size
image = np.array(image) / 255.0 # Normalize pixel values
image = np.expand_dims(image, axis=0) # Add batch 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, axis=-1)
# 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 image from CIFAR-10"), # Input: Image upload
outputs=gr.Textbox(label="Predicted Class"), # Output: Text showing predicted class
title="Vault Challenge 2 - BIM", # 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 BIM!. Tips: tune the parameters to make the model predict the image as a cat (class 3)."
)
# Launch the Gradio interface
iface.launch()
|