added fourth page

pages/4_Vase_Qualities_Comparison.py

@@ -0,0 +1,270 @@
+import streamlit as st
+import streamlit.components.v1 as components
+
+import dnnlib
+import legacy
+
+import pickle
+import pandas as pd
+import numpy as np
+from pyvis.network import Network
+
+import random
+from sklearn.metrics.pairwise import cosine_similarity
+
+from matplotlib.backends.backend_agg import RendererAgg
+
+from backend.disentangle_concepts import *
+
+_lock = RendererAgg.lock
+
+HIGHTLIGHT_COLOR = '#e7bcc5'
+st.set_page_config(layout='wide')
+
+
+st.title('Comparison among concept vectors')
+st.write('> **How do the concept vectors relate to each other?**')
+st.write('> **What is their join impact on the image?**')
+st.write("""Description to write""")
+
+    
+annotations_file = './data/vase_annotated_files/seeds0000-20000.pkl'
+with open(annotations_file, 'rb') as f:
+    annotations = pickle.load(f)
+
+if 'image_id' not in st.session_state:
+    st.session_state.image_id = 0
+if 'concept_ids' not in st.session_state:
+    st.session_state.concept_ids = ['Provenance ADRIA']
+if 'space_id' not in st.session_state:
+    st.session_state.space_id = 'Z'
+if 'type_col' not in st.session_state:
+    st.session_state.type_col = 'Provenance'
+
+# 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 two options to disentangle**')
+        type_col = st.selectbox('Concept category:', tuple(['Provenance', 'Shape Name', 'Fabric', 'Technique']))
+        
+        ann_df = pd.read_csv(f'./data/vase_annotated_files/sim_{type_col}_seeds0000-20000.csv')
+        labels = list(ann_df.columns)
+        labels.remove('ID')
+        labels.remove('Unnamed: 0')
+        
+        concept_ids = st.multiselect('Concepts:', tuple(labels), default=[labels[2], labels[3]])
+
+        st.write('**Choose a latent space to disentangle**')
+        # chosen_text_id_input = st.empty()
+        # concept_id = chosen_text_id_input.text_input('Concept:', value=st.session_state.concept_id)
+        space_id = st.selectbox('Space:', tuple(['Z', 'W']))
+
+        choose_text_button = st.form_submit_button('Choose the defined concept and space to disentangle')
+        
+        if choose_text_button:
+            st.session_state.concept_ids = list(concept_ids)
+            space_id = str(space_id)
+            st.session_state.space_id = space_id
+            st.session_state.type_col = type_col
+        # 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([1,1])
+output_col_1, output_col_2 = st.columns([1,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:
+    vectors, nodes_in_common, performances = get_concepts_vectors(concept_ids, annotations, ann_df, latent_space=space_id)
+    header_col_1.write(f'Concepts {", ".join(concept_ids)} - Latent space {space_id} - Relevant nodes in common: {nodes_in_common} - Performance of the concept vectors: {performances}')# - Nodes {",".join(list(imp_nodes))}')
+
+    edges = []
+    for i in range(len(concept_ids)):
+        for j in range(len(concept_ids)):
+            if i != j:
+                print(f'Similarity between {concept_ids[i]} and {concept_ids[j]}')
+                similarity = cosine_similarity(vectors[i,:].reshape(1, -1), vectors[j,:].reshape(1, -1))
+                print(np.round(similarity[0][0], 3))
+                edges.append((concept_ids[i], concept_ids[j], np.round(similarity[0][0], 3)))
+
+    
+    net = Network(height="750px", width="100%",)
+    for e in edges:
+        src = e[0]
+        dst = e[1]
+        w = e[2]
+
+        net.add_node(src, src, title=src)
+        net.add_node(dst, dst, title=dst)
+        net.add_edge(src, dst, value=w, title=src + ' to ' + dst + ' similarity ' +str(w))
+    
+    # Generate network with specific layout settings
+    net.repulsion(
+                        node_distance=420,
+                        central_gravity=0.33,
+                        spring_length=110,
+                        spring_strength=0.10,
+                        damping=0.95
+                       )
+
+    # Save and read graph as HTML file (on Streamlit Sharing)
+    try:
+        path = '/tmp'
+        net.save_graph(f'{path}/pyvis_graph.html')
+        HtmlFile = open(f'{path}/pyvis_graph.html', 'r', encoding='utf-8')
+
+    # Save and read graph as HTML file (locally)
+    except:
+        path = '/html_files'
+        net.save_graph(f'{path}/pyvis_graph.html')
+        HtmlFile = open(f'{path}/pyvis_graph.html', 'r', encoding='utf-8')
+
+    # Load HTML file in HTML component for display on Streamlit page
+    components.html(HtmlFile.read(), height=435)
+
+with output_col_2:
+    with open('data/CLIP_vecs_vases.pkl', 'rb') as f:
+        vectors_CLIP = pickle.load(f)
+        
+    # st.write(f'Class ID {input_id} - {input_label}: {pred_prob*100:.3f}% confidence')
+    #st.write('Concept vector', separation_vector)
+    header_col_2.write(f'Concepts {", ".join(concept_ids)} - Latent space CLIP')# - Nodes {",".join(list(imp_nodes))}')
+
+    edges_clip = []
+    for c1 in concept_ids:
+        for c2 in concept_ids:
+            if c1 != c2:
+                c1 = st.session_state.type_col + ' ' + c1
+                c2 = st.session_state.type_col + ' ' + c2
+                
+                print(f'Similarity between {c1} and {c2}')
+                similarity = cosine_similarity(vectors_CLIP[c1].reshape(1, -1), vectors_CLIP[c2].reshape(1, -1))
+                print(np.round(similarity[0][0], 3))
+                edges_clip.append((c1, c2, np.round(float(np.round(similarity[0][0], 3)), 3)))
+
+
+    net_clip = Network(height="750px", width="100%",)
+    for e in edges_clip:
+        src = e[0]
+        dst = e[1]
+        w = e[2]
+
+        net_clip.add_node(src, src, title=src)
+        net_clip.add_node(dst, dst, title=dst)
+        net_clip.add_edge(src, dst, value=w, title=src + ' to ' + dst + ' similarity ' +str(w))
+    
+    # Generate network with specific layout settings
+    net_clip.repulsion(
+                        node_distance=420,
+                        central_gravity=0.33,
+                        spring_length=110,
+                        spring_strength=0.10,
+                        damping=0.95
+                       )
+
+    # Save and read graph as HTML file (on Streamlit Sharing)
+    try:
+        path = '/tmp'
+        net_clip.save_graph(f'{path}/pyvis_graph_clip.html')
+        HtmlFile = open(f'{path}/pyvis_graph_clip.html', 'r', encoding='utf-8')
+
+    # Save and read graph as HTML file (locally)
+    except:
+        path = '/html_files'
+        net_clip.save_graph(f'{path}/pyvis_graph_clip.html')
+        HtmlFile = open(f'{path}/pyvis_graph_clip.html', 'r', encoding='utf-8')
+
+    # Load HTML file in HTML component for display on Streamlit page
+    components.html(HtmlFile.read(), height=435)
+    
+# ----------------------------- 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 concepts'):
+        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 ------------------------------
+
+
+with dnnlib.util.open_url('./data/vase_model_files/network-snapshot-010000.pkl') as f:
+    model = legacy.load_network_pkl(f)['G_ema'].to('cpu') # type: ignore
+
+if st.session_state.space_id == 'Z':
+    original_image_vec = annotations['z_vectors'][st.session_state.image_id]
+else:
+    original_image_vec = annotations['w_vectors'][st.session_state.image_id]
+
+img = generate_original_image(original_image_vec, model, latent_space=st.session_state.space_id)
+# 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 = generate_joint_effect(model, original_image_vec, vectors, min_epsilon=-(int(epsilon)), max_epsilon=int(epsilon), latent_space=st.session_state.space_id)
+
+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)}')