import gradio as gr
import XTransferBench.zoo
from torchvision import transforms


def XTransferAttack(image, attacker_name, epsilon):
    attacker = XTransferBench.zoo.load_attacker(
        threat_model='linf_non_targeted', method_name=attacker_name)
    target_epsilon = epsilon / 255.0
    attacker.interpolate_epsilon(target_epsilon)
    original_image = image.copy()
    image = image.convert("RGB")
    image = image.resize((attacker.image_size, attacker.image_size))
    image = transforms.ToTensor()(image).unsqueeze(0)
    adv_image = attacker(image)
    adv_image_resized = transforms.ToPILImage()(adv_image.squeeze(0))
    adv_image= adv_image_resized.resize((original_image.size[0], original_image.size[1]))
    return adv_image_resized, adv_image

demo = gr.Interface(
    fn=XTransferAttack,
    inputs=[
        gr.Image(type="pil", label="Input Image",),
        gr.Dropdown(
            choices=XTransferBench.zoo.list_attacker(threat_model='linf_non_targeted'),
            value="xtransfer_large_linf_eps12_non_targeted",
            label="Attacker",
            info="Select an attacker method."
        ),
        gr.Slider(
            minimum=0, maximum=255, step=1, value=12,
            label="Epsilon",
            info="Perturbation strength (0-255)"
        )
    ],
    outputs=[
        gr.Image(type="pil", label="Adversarial Image (Resized)"), 
        gr.Image(type="pil", label="Adversarial Image (Original Size)"), 
    ],
    title="XTransferBench Adversarial Attack Demo",
    description="Visualize adversarial examples on images using transfer-based attacks from the XTransferBench library.",
)

demo.launch(share=True)