import gradio as gr
from fawkes_wrapper_gradio_v_1_0 import FawkesProtection
from PIL import Image
import numpy as np

fawkes_protection = FawkesProtection()

def process_image(image, protection_level):
    protected_image = fawkes_protection.predict(image, level=protection_level)
    return Image.fromarray(protected_image)

iface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="pil"),
        gr.Radio(["low", "mid", "high"], label="Protection Level")
    ],
    outputs=gr.Image(type="pil"),
    title="Fawkes Image Protection",
    description="Upload an image to apply Fawkes protection."
)

iface.launch()