File size: 7,102 Bytes
4347fc3
 
 
 
 
 
 
 
 
 
d47f0a9
a16f7a6
0f9f66b
4347fc3
 
 
8c184f4
4347fc3
 
 
 
 
 
 
7a3450e
4347fc3
 
7a3450e
 
 
4347fc3
7a3450e
4347fc3
 
 
 
a16f7a6
4347fc3
 
2a4bd0a
a16f7a6
4347fc3
7a3450e
4347fc3
 
 
6dee263
 
 
e6dc87e
c185595
4347fc3
 
 
 
 
 
 
 
 
 
 
 
7a3450e
 
dc026d8
7a3450e
1904fbf
7a3450e
4347fc3
7a3450e
 
 
 
4347fc3
1904fbf
4347fc3
 
 
 
 
 
 
 
 
 
266a83c
 
5ad9951
7764300
8e430df
 
4347fc3
 
 
 
 
 
 
 
 
 
 
849bdd3
4347fc3
7a3450e
a4e211d
4347fc3
 
 
 
 
 
 
 
 
 
 
 
 
 
a16f7a6
4347fc3
 
 
 
 
 
 
 
 
 
 
 
266a83c
4347fc3
 
 
3b15df8
a16f7a6
 
 
 
3b15df8
 
4347fc3
 
 
 
3b15df8
 
266a83c
3b15df8
4347fc3
266a83c
4347fc3
266a83c
 
 
4347fc3
266a83c
 
9bb9ba9
4347fc3
266a83c
 
 
4347fc3
266a83c
 
 
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
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
import dnnlib
import legacy

_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("""Explanation on the functionalities to come.""")

instruction_text = """Instruction to input:
1. Choosing concept:
2. 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.
3. Choosing epsilon: **Epsilon** is the lambda amount of translation along the disentangled concept axis. A negative epsilon changes the image in the direction of the concept, a positive one pushes the image away from the concept. 
"""
st.write("To use the functionality below, users need to input the **concept** to disentangle, an **image** id and the **epsilon** of variation along the disentangled axis.")
with st.expander("See more instruction", expanded=False):
    st.write(instruction_text)

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

ann_df = pd.read_csv('./data/annotated_files/sim_seeds0000-100000.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)
        concept_id = st.selectbox('Concept:', tuple(labels))

        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])
header_col_1, header_col_2 = st.columns([5,1])
output_col_1, output_col_2 = st.columns([5,1])

st.subheader('Derivations along the concept vector')

# 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, imp_nodes = get_separation_space(concept_id, annotations, ann_df)
    # st.write(f'Class ID {input_id} - {input_label}: {pred_prob*100:.3f}% confidence')
    st.write('Concept vector', separation_vector)
    header_col_1.write(f'Concept {concept_id} - Number of relevant nodes: {number_important_features}')# - Nodes {",".join(list(imp_nodes))}')

# ----------------------------- 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, 100000)
            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, step=1)
        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'))
with dnnlib.util.open_url('./data/model_files/network-snapshot-010600.pkl') as f:
    model = legacy.load_network_pkl(f)['G_ema'].to('cpu') # type: ignore

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)
    smooth_head_3.write(f'Base image')


images, lambdas = regenerate_images(model, original_image_vec, separation_vector, min_epsilon=-(int(epsilon)), max_epsilon=int(epsilon))

with smoothgrad_col_1:
    st.image(images[0])
    smooth_head_1.write(f'Change of {np.round(lambdas[0], 2)}')

with smoothgrad_col_2:
    st.image(images[1])
    smooth_head_2.write(f'Change of {np.round(lambdas[1], 2)}')

with smoothgrad_col_4:
    st.image(images[3])
    smooth_head_4.write(f'Change of {np.round(lambdas[3], 2)}')

with smoothgrad_col_5:
    st.image(images[4])
    smooth_head_5.write(f'Change of {np.round(lambdas[4], 2)}')