Updated poisoning attack with clean prediction gallery

Signed-off-by: Kieran Fraser <Kieran.Fraser@ibm.com>

app.py

@@ -11,10 +11,7 @@ import gradio as gr
 import numpy as np
 from carbon_theme import Carbon
 
-import os
-
 import numpy as np
-import matplotlib.pyplot as plt
 import torch
 import transformers
 
@@ -164,6 +161,7 @@ def clf_poison_evaluate(*args):
     target_class = args[3]
     data_type = args[-1]
     
+    
     if model_type == "Example":
         model = transformers.AutoModelForImageClassification.from_pretrained(
             'facebook/deit-tiny-distilled-patch16-224',
@@ -235,6 +233,7 @@ def clf_poison_evaluate(*args):
             )
         backdoor = PoisoningAttackBackdoor(poison_func)
         source_class = 0
+        target_class = label_names.index(target_class)
         poison_percent = 0.5
 
         x_poison = np.copy(x_subset)
@@ -259,6 +258,10 @@ def clf_poison_evaluate(*args):
         clean_preds = np.argmax(outputs, axis=1)
         clean_acc = np.mean(clean_preds == clean_y)
         
+        clean_out = []
+        for i, im in enumerate(clean_x):
+            clean_out.append( (im.transpose(1,2,0), label_names[clean_preds[i]]) )
+        
         poison_x = x_poison[is_poison]
         poison_y = y_poison[is_poison]
 
@@ -269,8 +272,9 @@ def clf_poison_evaluate(*args):
         poison_out = []
         for i, im in enumerate(poison_x):
             poison_out.append( (im.transpose(1,2,0), label_names[poison_preds[i]]) )
+            
         
-        return poison_out, clean_acc, poison_acc
+        return clean_out, poison_out, clean_acc, poison_acc
         
     
 def show_params(type):
@@ -504,25 +508,36 @@ with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
                 with gr.Row():
                     with gr.Column(scale=1):
                         attack = gr.Textbox(visible=True, value="Backdoor", label="Attack", interactive=False)
-                        target_class = gr.Number(label="Target class", info="The class you wish to force the model to predict.",
-                                                    minimum=1, maximum=9, value=1)
+                        target_class = gr.Radio(label="Target class", info="The class you wish to force the model to predict.",
+                                                    choices=['dog',
+                                                    'cassette player',
+                                                    'chainsaw',
+                                                    'church',
+                                                    'french horn',
+                                                    'garbage truck',
+                                                    'gas pump',
+                                                    'golf ball',
+                                                    'parachutte',], value='dog')
                         trigger_image = gr.Image(label="Trigger Image",  value="./baby-on-board.png")
                         eval_btn_patch = gr.Button("Evaluate")
                     with gr.Column(scale=2):
-                        poison_gallery = gr.Gallery(label="Poisoned", preview=False, show_download_button=True)
-                    with gr.Column(scale=2):
+                        clean_gallery = gr.Gallery(label="Clean", preview=False, show_download_button=True)
                         clean_accuracy = gr.Number(label="Clean Accuracy", precision=2)
+                    with gr.Column(scale=2):
+                        poison_gallery = gr.Gallery(label="Poisoned", preview=False, show_download_button=True)
                         poison_success = gr.Number(label="Poison Success", precision=2)
                     
                 eval_btn_patch.click(clf_poison_evaluate, inputs=[attack, model_type, trigger_image, target_class, data_type],
-                            outputs=[poison_gallery, clean_accuracy, poison_success])  
+                            outputs=[clean_gallery, poison_gallery, clean_accuracy, poison_success])  
 
 if __name__ == "__main__":
     
-    # during development, set debug=True
+    # For development
     '''demo.launch(show_api=False, debug=True, share=False,
                 server_name="0.0.0.0", 
                 server_port=7777, 
                 ssl_verify=False,
                 max_threads=20)'''
+                
+    # For deployment
     demo.launch(share=True, ssl_verify=False)