pages/2_Colours_comparison.py

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 color directions')
st.write('> **How do the color directions relate to each other?**')
st.write('> **What is their joint impact on the image?**')

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

concept_vectors = pd.read_csv('./data/stored_vectors/scores_colors_hsv.csv')
concept_vectors['vector'] = [np.array([float(xx) for xx in x]) for x in concept_vectors['vector'].str.split(', ')]
concept_vectors['score'] = concept_vectors['score'].astype(float)
concept_vectors['sign'] = [True if 'sign:True' in val else False for val in concept_vectors['kwargs']]
concept_vectors['extremes'] = [True if 'extremes method:True' in val else False for val in concept_vectors['kwargs']]
concept_vectors['regularization'] = [float(val.split(',')[1].strip('regularization: ')) if 'regularization:' in val else False for val in concept_vectors['kwargs']]
concept_vectors['cl_method'] = [val.split(',')[0].strip('classification method:') if 'classification method:' in val else False for val in concept_vectors['kwargs']]
concept_vectors['num_factors'] = [int(val.split(',')[1].strip('number of factors:')) if 'number of factors:' in val else False for val in concept_vectors['kwargs']]
concept_vectors = concept_vectors.sort_values('score', ascending=False).reset_index()

with dnnlib.util.open_url('./data/textile_model_files/network-snapshot-005000.pkl') as f:
    model = legacy.load_network_pkl(f)['G_ema'].to('cpu') # type: ignore

COLORS_LIST = ['Gray', 'Red Orange', 'Yellow', 'Green', 'Light Blue', 'Blue', 'Purple', 'Pink', 'Saturation', 'Value']

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 = [COLORS_LIST[-1], COLORS_LIST[-2], ]
if 'sign' not in st.session_state:
    st.session_state.sign = False
if 'extremes' not in st.session_state:
    st.session_state.extremes = False
if 'regularization' not in st.session_state:
    st.session_state.regularization = False
if 'cl_method' not in st.session_state:
    st.session_state.cl_method = False
if 'num_factors' not in st.session_state:
    st.session_state.num_factors = False


if 'space_id' not in st.session_state:
    st.session_state.space_id = 'W'

# ----------------------------- INPUT ----------------------------------
st.header('Input')
input_col_1, input_col_2 = st.columns([1,1])
# --------------------------- 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 series of colors to compare**')
        # chosen_text_id_input = st.empty()
        # concept_id = chosen_text_id_input.text_input('Concept:', value=st.session_state.concept_id)
        concept_ids = st.multiselect('Color (including Saturation and Value):', tuple(COLORS_LIST), default=[COLORS_LIST[-1], COLORS_LIST[-2], ])
        choose_text_button = st.form_submit_button('Choose the defined colors')
        
        if choose_text_button:
            st.session_state.concept_ids = list(concept_ids)
            
            
with input_col_2:
    with st.form('text_form_1'):
        st.write('Options for StyleSpace (not available for Saturation and Value)')
        sign = st.selectbox('Sign option:', tuple([True, False]), index=1)
        num_factors = st.selectbox('Number of factors option:', tuple([1, 5, 10, 20, False]), index=4)
        st.write('Options for InterFaceGAN (not available for Saturation and Value)')
        cl_method = st.selectbox('Classification method option:', tuple(['LR', 'SVM', False]), index=2)
        regularization = st.selectbox('Regularization option:', tuple([0.1, 1.0, False]), index=2)
        st.write('Options for InterFaceGAN (only for Saturation and Value)')
        extremes = st.selectbox('Extremes option:', tuple([True, False]), index=1)
        
        choose_options_button = st.form_submit_button('Choose the defined options')
        # 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']))
        if choose_options_button:
            st.session_state.sign = sign
            st.session_state.num_factors = num_factors
            st.session_state.cl_method = cl_method
            st.session_state.regularization = regularization
            st.session_state.extremes = extremes
            
# ---------------------------- SET UP OUTPUT ------------------------------
epsilon_container = st.empty()
st.header('Comparison')
st.subheader('Color vectors')

header_col_1, header_col_2 = st.columns([3,1])
output_col_1, output_col_2 = st.columns([3,1])

# ---------------------------- DISPLAY COL 1 ROW 1 ------------------------------
tmp = concept_vectors[concept_vectors['color'].isin(st.session_state.concept_ids)]
tmp = tmp[tmp['sign'] == st.session_state.sign][tmp['extremes'] == st.session_state.extremes][tmp['num_factors'] == st.session_state.num_factors][tmp['cl_method'] == st.session_state.cl_method][tmp['regularization'] == st.session_state.regularization]
info = tmp.loc[:, ['vector', 'score', 'color', 'kwargs']].values
concept_ids = [i[2] for i in info] #+ ' ' + i[3]
    
with header_col_1:
    st.write('Similarity graph')

with header_col_2:
    st.write('Information')
    
with output_col_2:
    for i,concept_id in enumerate(concept_ids):
        st.write(f'Color {info[i][2]} - Settings: {info[i][3]} Performance of the color vector: {info[i][1]}')# - Nodes {",".join(list(imp_nodes))}')

with output_col_1:
    
    edges = []
    for i in range(len(concept_ids)):
        for j in range(len(concept_ids)):
            if i != j and info[i][2] != info[j][2]:
                print(f'Similarity between {concept_ids[i]} and {concept_ids[j]}')
                similarity = cosine_similarity(info[i][0].reshape(1, -1), info[j][0].reshape(1, -1)) 
                print(np.round(similarity[0][0], 3))
                edges.append((concept_ids[i], concept_ids[j], np.round(similarity[0][0] + 0.001, 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)