pages/3_Adversarial_attack.py

import streamlit as st
import pandas as pd
import numpy as np
import random
from backend.utils import make_grid, load_dataset, load_model, load_image

from backend.smooth_grad import generate_smoothgrad_mask, ShowImage, fig2img, LoadImage, ShowHeatMap, ShowMaskedImage
from transformers import AutoFeatureExtractor, AutoModelForImageClassification
import torch

from matplotlib.backends.backend_agg import RendererAgg

from backend.adversarial_attack import *

_lock = RendererAgg.lock

st.set_page_config(layout='wide')
BACKGROUND_COLOR = '#bcd0e7'
SECONDARY_COLOR = '#bce7db'


st.title('Adversarial Attack')
st.write('> **How adversarial attacks affect ConvNeXt interpretation?**')
st.write("""Adversarial examples are inputs crafted to confuse neural networks, causing them to misclassify a given input. 
These examples are not distinguishable by humans but cause the network to fail to recognize the image content. 
One type of such attack is the fast gradient sign method (FGSM) attack, which is a white box attack that aims to ensure misclassification. 
A white box attack is where the attacker has full access to the model being attacked.

The FGSM attack is one of the earliest and most popular adversarial attacks. 
It is described by Goodfellow _et al_ in their work on [Explaining and Harnessing Adversarial Examples](https://arxiv.org/abs/1412.6572). 
The attack is simple yet powerful, using the gradients that neural networks use to learn. 
Instead of adjusting the weights based on the backpropagated gradients to minimize loss, the attack adjusts the input data to maximize the loss using the gradient of the loss with respect to the input data.
""")

instruction_text = """Instruction to input:
1. Choosing image: Users can choose a specific image by entering **Image ID** and hit the _Choose the defined image_ button or can generate an image randomly by hitting the _Generate a random image_ button.
2. Choosing epsilon: **Epsilon** is the amount of perturbation on the original image under attack. The higher the epsilon is, the more pertubed the image is, the more confusion made to the model. 
Users can choose a specific epsilon by engtering **Epsilon** and hit the _Choose the defined epsilon_ button.
Users can also let the algorithm find the smallest epsilon automatically by hitting the _Find the smallest epsilon automatically_ button.
The underlying algorithm will iterate through a set of epsilon in ascending order until reaching the **maximum value of epsilon**. 
After each iteration, the epsilon will increase by an amount equal to **step** variable. 
Users can change the default values of the two variable value optionally.
"""
st.write("To use the functionality below, users need to input the **image** and the **epsilon**.")
with st.expander("See more instruction", expanded=False):
    st.write(instruction_text)

    

imagenet_df = pd.read_csv('./data/ImageNet_metadata.csv')
image_id = None

if 'image_id' not in st.session_state:
    st.session_state.image_id = 0

# def on_change_random_input():
#     st.session_state.image_id = st.session_state.image_id

# ----------------------------- INPUT ----------------------------------
st.header('Input')
input_col_1, input_col_2, input_col_3 = st.columns(3)
# --------------------------- INPUT column 1 ---------------------------
with input_col_1:
    with st.form('image_form'):
        
        # image_id = st.number_input('Image ID: ', format='%d', step=1)
        st.write('**Choose or generate a random image**')
        chosen_image_id_input = st.empty()
        image_id = chosen_image_id_input.number_input('Image ID:', format='%d', step=1, value=st.session_state.image_id)
        
        choose_image_button = st.form_submit_button('Choose the defined image')
        random_id = st.form_submit_button('Generate a random image')

        if random_id:
            image_id = random.randint(0, 50000)
            st.session_state.image_id = image_id
            chosen_image_id_input.number_input('Image ID:', format='%d', step=1, value=st.session_state.image_id)
            
        if choose_image_button:
            image_id = int(image_id)
            st.session_state.image_id = int(image_id)
        # st.write(image_id, st.session_state.image_id)

# ---------------------------- SET UP OUTPUT ------------------------------
epsilon_container = st.empty()
st.header('Output')
st.subheader('Perform attack')

# perform attack container
header_col_1, header_col_2, header_col_3, header_col_4, header_col_5 = st.columns([1,1,1,1,1])
output_col_1, output_col_2, output_col_3, output_col_4, output_col_5 = st.columns([1,1,1,1,1])

# prediction error container
error_container = st.empty()
smoothgrad_header_container = st.empty()

# smoothgrad container
smooth_head_1, smooth_head_2, smooth_head_3, smooth_head_4, smooth_head_5 = st.columns([1,1,1,1,1])
smoothgrad_col_1, smoothgrad_col_2, smoothgrad_col_3, smoothgrad_col_4, smoothgrad_col_5 = st.columns([1,1,1,1,1])

original_image_dict = load_image(st.session_state.image_id)
input_image = original_image_dict['image']
input_label = original_image_dict['label']
input_id = original_image_dict['id']

# ---------------------------- DISPLAY COL 1 ROW 1 ------------------------------
with output_col_1:
    pred_prob, pred_class_id, pred_class_label = feed_forward(input_image)
    # st.write(f'Class ID {input_id} - {input_label}: {pred_prob*100:.3f}% confidence')
    st.image(input_image)
    header_col_1.write(f'Class ID {input_id} - {input_label}: {pred_prob*100:.1f}% confidence')



if pred_class_id != (input_id-1):
    with error_container.container():
        st.write(f'Predicted output: Class ID {pred_class_id} - {pred_class_label} {pred_prob*100:.1f}% confidence')
        st.error('ConvNeXt misclassified the chosen image. Please choose or generate another image.',
            icon = "🚫")

# ----------------------------- INPUT column 2 & 3 ----------------------------
with input_col_2:
    with st.form('epsilon_form'):
        st.write('**Set epsilon or find the smallest epsilon automatically**')
        chosen_epsilon_input = st.empty()
        epsilon = chosen_epsilon_input.number_input('Epsilon:', min_value=0.001, format='%.3f', step=0.001)

        epsilon_button = st.form_submit_button('Choose the defined epsilon')
        find_epsilon = st.form_submit_button('Find the smallest epsilon automatically')


with input_col_3:
    with st.form('iterate_epsilon_form'):
        max_epsilon = st.number_input('Maximum value of epsilon (Optional setting)', value=0.500, format='%.3f')
        step_epsilon = st.number_input('Step (Optional setting)', value=0.001, format='%.3f')
        setting_button = st.form_submit_button('Set iterating mode')


# ---------------------------- DISPLAY COL 2 ROW 1 ------------------------------
if pred_class_id == (input_id-1) and (epsilon_button or find_epsilon or setting_button):
    with output_col_3:
        if epsilon_button:
            perturbed_data, new_prob, new_id, new_label = perform_attack(input_image, input_id-1, epsilon)
        else:
            epsilons = [i*step_epsilon for i in range(1, 1001) if i*step_epsilon <= max_epsilon]
            with epsilon_container.container():
                epsilon_container_text = 'Checking epsilon'
                st.write(epsilon_container_text)
                st.progress(0)
        
            for i, e in enumerate(epsilons):
                
                perturbed_data, new_prob, new_id, new_label = perform_attack(input_image, input_id-1, e)
                with epsilon_container.container():
                    epsilon_container_text = f'Checking epsilon={e:.3f}. Confidence={new_prob*100:.1f}%'
                    st.write(epsilon_container_text)
                    st.progress(i/len(epsilons))
                
                epsilon = e

                if new_id != input_id - 1:
                    epsilon_container.empty()
                    st.balloons()
                    break
                if i == len(epsilons)-1:
                    epsilon_container.error(f'FGSM failed to attack on this image at epsilon={e:.3f}. Set higher maximum value of epsilon or choose another image',
                                            icon = "🚫")

        perturbed_image = deprocess_image(perturbed_data.detach().numpy())[0].astype(np.uint8).transpose(1,2,0)
        perturbed_amount = perturbed_image - input_image
        header_col_3.write(f'Pertubed amount - epsilon={epsilon:.3f}')
        st.image(ShowImage(perturbed_amount))
    
    with output_col_2:
        # st.write('plus sign')
        st.image(LoadImage('frontend/images/plus-sign.png'))
    
    with output_col_4:
        # st.write('equal sign')
        st.image(LoadImage('frontend/images/equal-sign.png'))

    # ---------------------------- DISPLAY COL 5 ROW 1 ------------------------------
    with output_col_5:
        # st.write(f'ID {new_id+1} - {new_label}: {new_prob*100:.3f}% confidence')
        st.image(ShowImage(perturbed_image))
    header_col_5.write(f'Class ID {new_id+1} - {new_label}: {new_prob*100:.1f}% confidence')

    # -------------------------- DISPLAY SMOOTHGRAD ---------------------------
    smoothgrad_header_container.subheader('SmoothGrad visualization')

    with smoothgrad_col_1:
        smooth_head_1.write(f'SmoothGrad before attacked')
        heatmap_image, masked_image, mask = generate_images(st.session_state.image_id, epsilon=0)
        st.image(heatmap_image)
        st.image(masked_image) 
    with smoothgrad_col_3:
        smooth_head_3.write('SmoothGrad after attacked')
        heatmap_image_attacked, masked_image_attacked, attacked_mask= generate_images(st.session_state.image_id, epsilon=epsilon)
        st.image(heatmap_image_attacked)
        st.image(masked_image_attacked)
    
    with smoothgrad_col_2:
        st.image(LoadImage('frontend/images/minus-sign-5.png'))

    with smoothgrad_col_5:
        smooth_head_5.write('SmoothGrad difference')
        difference_mask = abs(attacked_mask-mask)
        st.image(ShowHeatMap(difference_mask))
        masked_image = ShowMaskedImage(difference_mask, perturbed_image)
        st.image(masked_image)

    with smoothgrad_col_4:
        st.image(LoadImage('frontend/images/equal-sign.png'))