Spaces:
Runtime error
Runtime error
import streamlit as st | |
import pickle | |
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.disentangle_concepts import * | |
_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_seeds0000-1000.pkl' | |
with open(annotations_file, 'rb') as f: | |
annotations = pickle.load(f) | |
concepts = './data/concepts.txt' | |
with open(concepts) as f: | |
labels = [line.strip() for line in f.readlines()] | |
# 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:', format='', value=st.session_state.concept_id) | |
choose_image_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:', format='', value=st.session_state.concept_id) | |
if choose_image_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) | |
# 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') | |
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')) | |