pages/1_Textiles_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
import dnnlib
import legacy

_lock = RendererAgg.lock


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


st.title('Disentanglement studies on the Textile Dataset')
st.markdown(
    """
    This is a demo of the Disentanglement studies on the [iMET Textiles Dataset](https://www.metmuseum.org/art/collection/search/85531).
    """,
    unsafe_allow_html=False,)
    
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 = concept_vectors.sort_values('score', ascending=False).reset_index()
print(concept_vectors[['vector', 'score']], concept_vectors.loc[0, 'vector'], concept_vectors.loc[0, 'vector'].shape)

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']

if 'image_id' not in st.session_state:
    st.session_state.image_id = 0
if 'color_ids' not in st.session_state:
    st.session_state.concept_ids = COLORS_LIST[-1]
if 'space_id' not in st.session_state:
    st.session_state.space_id = 'W'

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

with input_col_2:
    with st.form('text_form'):
        
        st.write('**Choose color to vary**')
        type_col = st.selectbox('Color:', tuple(COLORS_LIST), value=st.session_state.concepts_ids)
        
        st.write('**Set range of change**')
        chosen_color_lambda_input = st.empty()
        color_lambda = chosen_color_lambda_input.number_input('Lambda:', min_value=0, step=1, value=7)
        color_lambda_button = st.form_submit_button('Choose the defined lambda')
          
        if choose_text_button:
            st.session_state.concept_ids = type_col
            st.session_state.space_id = space_id
        
with input_col_3:
    with st.form('text_form'):
        
        st.write('**Saturation variation**')
        chosen_saturation_lambda_input = st.empty()
        saturation_lambda = chosen_saturation_lambda_input.number_input('Lambda:', min_value=0, step=1)
        saturation_lambda_button = st.form_submit_button('Choose the defined lambda')
        
        st.write('**Value variation**')
        chosen_value_lambda_input = st.empty()
        value_lambda = chosen_value_lambda_input.number_input('Lambda:', min_value=0, step=1)
        value_lambda_button = st.form_submit_button('Choose the defined lambda')
        
# with input_col_4:
#     with st.form('Network specifics:'):
#         st.write('**Choose a latent space to use**')
#         space_id = st.selectbox('Space:', tuple(['W']))
#         choose_text_button = st.form_submit_button('Choose the defined concept and space to disentangle')

#         st.write('**Select hierarchical levels to manipulate**')
#         layers = st.multiselect('Layers:', tuple(range(14)))
#         if len(layers) == 0:
#             layers = None
#         print(layers)
#         layers_button = st.form_submit_button('Choose the defined layers')
        

# ---------------------------- SET UP OUTPUT ------------------------------
epsilon_container = st.empty()
st.header('Image Manipulation')
st.subheader('Using selected directions')

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

# # prediction error container
# error_container = st.empty()
# smoothgrad_header_container = st.empty()

# # smoothgrad container
# smooth_head_1, smooth_head_2,  = st.columns([1,1,])
# smoothgrad_col_1, smoothgrad_col_2 = st.columns([1,1])

# ---------------------------- DISPLAY COL 1 ROW 1 ------------------------------
with header_col_1:
    st.write(f'Original image')

with header_col_2:
    color_separation_vector, performance_color = concept_vectors[concept_vectors['color'] == st.session_state.concepts_ids].loc[0, ['vector', 'score']]
    saturation_separation_vector, performance_saturation = concept_vectors[concept_vectors['color'] == 'Saturation'].loc[0, ['vector', 'score']]
    value_separation_vector, performance_value = concept_vectors[concept_vectors['color'] == 'Value'].loc[0, ['vector', 'score']]
    st.write(f'Change in {st.session_state.concepts_ids} of {np.round(color_lambda, 2)}, in saturation of {np.round(saturation_lambda, 2)}, in value of {np.round(value_lambda, 2)}. - Performance color vector: {performance_color}, saturation vector: {performance_saturation}, value vector: {performance_value}')

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

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)

with output_col_1:
    st.image(img)

with output_col_2:
    image_updated = generate_composite_images(model, original_image_vec, [separation_vector_color, saturation_separation_vector, value_separation_vector], lambdas=[color_lambda, saturation_lambda, value_lambda])
    st.image(image_updated)