import gradio as gr
from fawkes_wrapper_gradio_v_1_0 import FawkesProtection
from PIL import Image
import numpy as np
from concurrent.futures import ThreadPoolExecutor

fawkes_protection = FawkesProtection()

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

def process_images(image_files, protection_level):
    results = []
    with ThreadPoolExecutor() as executor:
        futures = [executor.submit(process_image, file, protection_level) for file in image_files]
        for future in futures:
            result = future.result()
            results.append(result)
    return results

iface = gr.Interface(
    fn=process_images,
    inputs=[
        gr.File(label="Upload Images", file_count="multiple", type="filepath"),
        gr.Radio(["low", "mid", "high"], label="Protection Level")
    ],
    outputs=gr.Gallery(label="Protected Images"),
    title="Fawkes Image Protection",
    description="Upload multiple images to apply Fawkes protection."
)

iface.launch()