import torch import gradio as gr from torchvision import transforms # load imagenet labels with open("imagenet_classes.txt", "r") as f: categories = [s.strip() for s in f.readlines()] # load a resnet18 model pretrained on ImageNet model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True) model.eval() # preprocess data pretrained_std = torch.Tensor([0.229, 0.224, 0.225]) pretrained_mean = torch.Tensor([0.485, 0.456, 0.406]) preprocess = transforms.Compose([ transforms.Resize(224), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=pretrained_mean, std=pretrained_std), ]) def undo_preprocess(processed_img): img = processed_img.mul(pretrained_std.view(-1,1,1)).add(pretrained_mean.view(-1,1,1)) img = torch.clamp(img, 0, 1).squeeze() img = img.detach().numpy().transpose(1,2,0) return img def set_example_image(img): return gr.Image.update(value=img[0]) def generate_adversial_image(img, epsilon): # convert a PIL image to tensor and set requires_grad to True processed_img = preprocess(img).unsqueeze(0) processed_img.requires_grad = True # forward pass output = model(processed_img) # get predictions probs = torch.nn.functional.softmax(output[0], dim=0) top5_prob, top5_idx = torch.topk(probs, 5) preds = {categories[idx]: prob.item() for idx, prob in zip(top5_idx, top5_prob)} # compute gradient label = torch.Tensor([top5_idx[0]]).type(torch.LongTensor) loss = torch.nn.functional.cross_entropy(output, label) model.zero_grad() loss.backward() # generate adversarial image adv_img = processed_img + epsilon*processed_img.grad.data.sign() adv_output = model(adv_img) adv_probs = torch.nn.functional.softmax(adv_output[0], dim=0) adv_top5_prob, adv_top5_idx = torch.topk(adv_probs, 5) adv_preds = {categories[idx]: prob.item() for idx, prob in zip(adv_top5_idx, adv_top5_prob)} return undo_preprocess(processed_img), undo_preprocess(adv_img), preds, adv_preds with gr.Blocks() as demo: gr.Markdown('''## Generate Adversarial Image with Fast Gradient Sign Method Given an input image and a neural network, the adversarial image can be generated from adv_img = input_img + epsilon*input_img.grad ''') with gr.Box(): input_image = gr.Image(type="pil", label="Input Image") example_images = gr.Dataset(components=[input_image], samples=[['coral.jpg'], ['goldfish.jpg'], ['otter.jpg'], ['panda.jpg']]) with gr.Row(): epsilon = gr.Slider(minimum=0, maximum=0.5, value=0.02, step=0.01, label="epsilon") btn = gr.Button("Generate Adversarial Image") gr.Markdown('''### Original Image''') with gr.Box(): with gr.Row(): img_before = gr.Image(label="Original Image") label_before = gr.Label(label="Original Prediction") gr.Markdown('''### Adversarial Image''') with gr.Box(): with gr.Row(): img_after = gr.Image(label="Adversarial Image") label_after = gr.Label(label="New Prediction") gr.Markdown('''The prediction is done by ResNet18. Example images are from [Unsplash](https://unsplash.com).''') # events btn.click(fn=generate_adversial_image, inputs=[input_image, epsilon], outputs=[img_before, img_after, label_before, label_after]) example_images.click(fn=set_example_image, inputs=example_images, outputs=example_images.components) demo.launch()