File size: 11,820 Bytes
3303c2f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
00ddcd8
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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
import torch
import types
import timm
import requests
import random
import yaml
import gradio as gr 
from PIL import Image
from timm import create_model
from torchvision import transforms
from timm.data import resolve_data_config
from modelguidedattacks.guides.unguided import Unguided
from timm.data.transforms_factory import create_transform
from modelguidedattacks.cls_models.registry import TimmPretrainModelWrapper


# Download human-readable labels for ImageNet.
IMAGENET_LABELS_URL = "https://storage.googleapis.com/bit_models/ilsvrc2012_wordnet_lemmas.txt"
LABELS = requests.get(IMAGENET_LABELS_URL).text.strip().split("\n")
SORTED_LABELS = sorted(LABELS.copy(), key=lambda s: s.lower())

def get_timm_model(name):
    """Retrieves model from timm library by name with weights loaded.
    """
    model = create_model(name,pretrained="true")
    transform = create_transform(**resolve_data_config({}, model=model))
    model = model.eval()
    return model, transform

def create_attacker(model, transform, iterations):
    """ Instantiates an QuadAttack Model.
    """
    # config_dict = {"cvx_proj_margin" : 0.2,
    #                "opt_warmup_its": 5}
    with open("base_config.yaml") as f:
        config_dict = yaml.safe_load(f)

    config = types.SimpleNamespace(**config_dict)

    attacker = Unguided(TimmPretrainModelWrapper(model, transform,"", "", ""), config, iterations=iterations,
                        lr=0.002, topk_loss_coef_upper=10)

    return attacker

def predict_topk_accuracies(img, k, iters, model_name, desired_labels, button=None, progress=gr.Progress(track_tqdm=True)):
    """ Predict the top K results using base model and attacker model.
    """
    label_inds = list(range(0,1000)) #label indices
    # convert user desired labels to desired inds
    desired_inds = [LABELS.index(name) for name in desired_labels] 
    # remove selected before randomly sampling the rest
    for ind in desired_inds:
        label_inds.remove(ind)

    # fill up user selections to top k results
    desired_inds = desired_inds + random.sample(label_inds,k-len(desired_inds))
    tensorized_desired_inds = torch.tensor(desired_inds).unsqueeze(0) #[B,K]

    model, transform = get_timm_model(model_name)

     # Define a transformation to convert PIL image to a tensor
    normalization = transforms.Compose([
                 transform.transforms[-1]  # Converts to a PyTorch tensor
                ])
    preprocess = transforms.Compose(
                 transform.transforms[:-1]  # Converts to a PyTorch tensor
                )
    
    attacker = create_attacker(model, normalization, iters)


    img = img.convert('RGB')
    orig_img = img.copy()
    orig_img = preprocess(orig_img)
    orig_img = orig_img.unsqueeze(0)
    img = transform(img).unsqueeze(0)

    with torch.no_grad():
        outputs = model(img)
        attack_outputs, attack_img = attacker(orig_img, tensorized_desired_inds, None)

    probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
    attacker_probs = torch.nn.functional.softmax(attack_outputs[0], dim=0)

    values, indices = torch.topk(probabilities, k)

    attack_vals, attack_inds = torch.topk(attacker_probs, k)

    attack_img_out = orig_img + attack_img #B C H W
    # Convert the PyTorch tensor to a NumPy array
    attack_img_out = attack_img_out.squeeze(0) # C H W
    attack_img_out = attack_img_out.permute(1, 2, 0).numpy()  # H W C

    orig_img = orig_img.squeeze(0)
    orig_img = orig_img.permute(1, 2, 0).numpy()

    attack_img = attack_img.squeeze(0)
    attack_img = attack_img.permute(1, 2, 0).numpy()


    # Convert the NumPy array to a PIL image
    attack_img_out = Image.fromarray((attack_img_out * 255).astype('uint8'))
    orig_img = Image.fromarray((orig_img * 255).astype('uint8'))
    attack_img = Image.fromarray((attack_img * 255).astype('uint8'))


    return (orig_img, attack_img_out, attack_img,{LABELS[i]: v.item() for i, v in zip(indices, values)}, {LABELS[i]: v.item() for i, v in zip(attack_inds, attack_vals)})

def random_fill_classes(desired_labels, k):

    label_inds = list(range(0,1000)) #label indices
    # convert user desired labels to desired inds
    if len(desired_labels) > k:
        desired_labels = desired_labels[:k]
    desired_inds = [LABELS.index(name) for name in desired_labels] 
    # remove selected before randomly sampling the rest
    for ind in desired_inds:
        label_inds.remove(ind)

    # fill up user selections to top k results
    desired_inds = desired_inds + random.sample(label_inds,k-len(desired_inds))

    return [LABELS[ind] for ind in desired_inds]


input_img = gr.Image(type='pil')
top_k_slider = gr.Slider(2, 20, value=10, step=1, label="Top K predictions", info="Choose between 2 and 20")
iteration_slider = gr.Slider(30, 1000, value=60, step=1, label="QuadAttack Iterations", info="Choose how many iterations to optimize using QuadAttack! (Usually <= 60 is enough)")
model_choice_list = gr.Dropdown(
                    timm.list_models(), value="vit_base_patch16_224", label="timm model name", info="Currently only supporting timm models! See code for models used in paper."
                        )
desired_labels = gr.Dropdown(
                   SORTED_LABELS, max_choices=20,filterable=True, multiselect=True, label="Desired Labels for QuadAttack", info="Select classes you wish to output from an attack. \
                                                                                             Classes will be ranked in order listed and randomly filled up to \
                                                                                             K if < K options are selected."
        )
button = gr.Button("Randomly fill Top-K attack classes.")

desc = r'<div align="center">Authors: Thomas Paniagua, Ryan Grainger, Tianfu Wu <p><a href="https://arxiv.org/abs/2312.11510">Paper</a><br><a href="https://github.com/thomaspaniagua/quadattack">Code</a></p> </div>'
with gr.Interface(predict_topk_accuracies, 
                  inputs=[input_img, 
                   top_k_slider,
                   iteration_slider,
                   model_choice_list,
                   desired_labels,
                   button], 
                    outputs=[
                             gr.Image(type='pil', label="Input Image"),
                             gr.Image(type='pil', label="Perturbed Image"),
                             gr.Image(type='pil', label="Added Noise"),
                             gr.Label(label="Original Top K"), 
                             gr.Label(label="QuadAttack Top K"), 
                            #  gr.Image(type='pil', label="Perturbed Image")
                             ],
                    title='QuadAttack!',
                    description= desc,
                    cache_examples=False,
                    allow_flagging="never",
                    thumbnail= "quadattack_pipeline.pdf",
                    examples = [["image_examples/RV.jpeg", 5, 30, "vit_base_patch16_224", None, None
                                                                                            # ["lemon", "plastic_bag", "hay", "tripod", "bell_cote, bell_cot"]
                    ], 
                                # ["image_examples/biker.jpeg", 10, 60, "swinv2_cr_base_224", None, None

                                                                                            # ["hog, pig, grunter, squealer, Sus_scrofa", 
                                                                                            # "lesser_panda, red_panda, panda, bear_cat, cat_bear, Ailurus_fulgens",
                                                                                            # "caldron, cauldron", "dowitcher", "water_tower", "quill, quill_pen", 
                                                                                            # "balance_beam, beam", "unicycle, monocycle", "pencil_sharpener",
                                                                                            # "puffer, pufferfish, blowfish, globefish"
                                                                                            # ]
                                #  ], 
                                ["image_examples/mower.jpeg", 15, 100,"wide_resnet101_2", None , None

                                                                                        #    ["washbasin, handbasin, washbowl, lavabo, wash-hand_basin",
                                                                                        #    "cucumber, cuke", "bolete", "oboe, hautboy, hautboi", "crane",
                                                                                        #    "wolf_spider, hunting_spider", "Norfolk_terrier", "nail", "sidewinder, horned_rattlesnake, Crotalus_cerastes",
                                                                                        #    "cannon", "beaker", "Shetland_sheepdog, Shetland_sheep_dog, Shetland",
                                                                                        #    "monitor", "restaurant, eating_house, eating_place, eatery", "electric_fan, blower"
                                                                                        #    ]
                                ], 
                                # ["image_examples/dog.jpeg", 20, 150, "xcit_small_12_p8_224", None, None 

                                # ["church, church_building", "axolotl, mud_puppy, Ambystoma_mexicanum", 
                                #                                                               "Scotch_terrier, Scottish_terrier, Scottie", "black-footed_ferret, ferret, Mustela_nigripes", 
                                #                                                               "lab_coat, laboratory_coat", "gyromitra", "grasshopper, hopper", "snail", "tabby, tabby_cat",
                                #                                                               "bell_cote, bell_cot", "Indian_cobra, Naja_naja", "robin, American_robin, Turdus_migratorius",
                                #                                                               "tiger_cat", "book_jacket, dust_cover, dust_jacket, dust_wrapper", "loudspeaker, speaker, speaker_unit, loudspeaker_system, speaker_system",
                                #                                                               "washbasin, handbasin, washbowl, lavabo, wash-hand_basin", "electric_guitar", "armadillo", "ski_mask",
                                #                                                               "convertible"
                                #                                                             ]
                                
                                # ], 
                                ["image_examples/fish.jpeg", 10, 100, "pvt_v2_b0", None, None

                                # ["ground_beetle, carabid_beetle", "sunscreen, sunblock, sun_blocker", "brass, memorial_tablet, plaque", "Irish_terrier", "head_cabbage", "bathtub, bathing_tub, bath, tub",
                                #                                                     "centipede", "squirrel_monkey, Saimiri_sciureus", "Chihuahua", "hourglass"
                                #                                                     ]
                                ]
                    ]
                    
    ).queue() as app:
    #turn off clear button as it erases globals
    for block in app.blocks:
        if isinstance(app.blocks[block],gr.Button):
            if app.blocks[block].value == "Clear":
                app.blocks[block].visible=False
    button.click(random_fill_classes, inputs=[desired_labels,top_k_slider], outputs=desired_labels)


if __name__ == "__main__":
    app.launch()