pages/1_Disentanglement.py

import streamlit as st
import pickle
import pandas as pd
import numpy as np
import random
import torch

from matplotlib.backends.backend_agg import RendererAgg

from backend.disentangle_concepts import *
import torch_utils

_lock = RendererAgg.lock


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


st.title('Disentanglement studies')
st.write('> **What concepts can be disentangled in the latent spae of a model?**')
st.write("""Explain more in depth""")

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)

    
annotations_file = './data/annotated_files/annotations_parallel_seeds0000-10000.pkl'
with open(annotations_file, 'rb') as f:
    annotations = pickle.load(f)

ann_df = pd.read_csv('./data/annotated_files/sim_seeds0000-10000.csv')
concepts = './data/concepts.txt'
with open(concepts) as f:
    labels = [line.strip() for line in f.readlines()]

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

# 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('text_form'):
        
        # image_id = st.number_input('Image ID: ', format='%d', step=1)
        st.write('**Choose a concept to disentangle**')
        chosen_text_id_input = st.empty()
        concept_id = chosen_text_id_input.text_input('Concept:', value=st.session_state.concept_id)
        
        choose_text_button = st.form_submit_button('Choose the defined concept')
        random_text = st.form_submit_button('Select a random concept')

        if random_text:
            concept_id = random.choice(labels)
            st.session_state.concept_id = concept_id
            chosen_text_id_input.text_input('Concept:', value=st.session_state.concept_id)
            
        if choose_text_button:
            concept_id = str(concept_id)
            st.session_state.concept_id = concept_id
        # st.write(image_id, st.session_state.image_id)

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

# 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])

# ---------------------------- DISPLAY COL 1 ROW 1 ------------------------------
with output_col_1:
    separation_vector, number_important_features = get_separation_space(concept_id, annotations, ann_df)
    # st.write(f'Class ID {input_id} - {input_label}: {pred_prob*100:.3f}% confidence')
    st.write('Separation vector', separation_vector)
    header_col_1.write(f'Concept {concept_id} - Number of relevant nodes: {number_important_features}')

# ----------------------------- INPUT column 2 & 3 ----------------------------        
with input_col_2:
   with st.form('image_form'):
        
        # image_id = st.number_input('Image ID: ', format='%d', step=1)
        st.write('**Choose or generate a random image to test the disentanglement**')
        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)

with input_col_3:
    with st.form('Variate along the disentangled concept'):
        st.write('**Set range of change**')
        chosen_epsilon_input = st.empty()
        epsilon = chosen_epsilon_input.number_input('Epsilon:', min_value=1, format='%.1f')
        epsilon_button = st.form_submit_button('Choose the defined epsilon')

# ---------------------------- DISPLAY COL 2 ROW 1 ------------------------------

model = torch.load('./data/model_files/pytorch_model.bin', map_location=torch.device('cpu'))
original_image_vec = annotations['z_vectors'][st.session_state.image_id]
img = generate_original_image(original_image_vec, model)
# input_image = original_image_dict['image']
# input_label = original_image_dict['label']
# input_id = original_image_dict['id']

with smoothgrad_col_3:
    st.image(img)
    header_col_1.write(f'Base image')



# 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'))