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['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'] 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' if 'color_lambda' not in st.session_state: st.session_state.color_lambda = 7 if 'saturation_lambda' not in st.session_state: st.session_state.saturation_lambda = 0 if 'value_lambda' not in st.session_state: st.session_state.value_lambda = 0 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 'best' not in st.session_state: st.session_state.best = True # 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, input_col_4 = st.columns(4) # --------------------------- 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 = image_id with input_col_2: with st.form('text_form_1'): st.write('**Choose color to vary**') type_col = st.selectbox('Color:', tuple(COLORS_LIST), index=7) colors_button = st.form_submit_button('Choose the defined color') st.write('**Set range of change**') chosen_color_lambda_input = st.empty() color_lambda = chosen_color_lambda_input.number_input('Lambda:', min_value=-100, step=1, value=7) color_lambda_button = st.form_submit_button('Choose the defined lambda for color') if colors_button or color_lambda_button: st.session_state.image_id = image_id st.session_state.concept_ids = type_col st.session_state.color_lambda = color_lambda 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=-100, step=1, key=0, value=0) saturation_lambda_button = st.form_submit_button('Choose the defined lambda for saturation') st.write('**Value variation**') chosen_value_lambda_input = st.empty() value_lambda = chosen_value_lambda_input.number_input('Lambda:', min_value=-100, step=1, key=1, value=0) value_lambda_button = st.form_submit_button('Choose the defined lambda for salue') if saturation_lambda_button or value_lambda_button: st.session_state.saturation_lambda = int(saturation_lambda) st.session_state.value_lambda = int(value_lambda) with input_col_4: with st.form('text_form_2'): st.write('Use best options') best = st.selectbox('Option:', tuple([True, False]), index=0) 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 st.session_state.best = best # 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: if st.session_state.best: color_separation_vector, performance_color = concept_vectors[concept_vectors['color'] == st.session_state.concept_ids].reset_index().loc[0, ['vector', 'score']] saturation_separation_vector, performance_saturation = concept_vectors[concept_vectors['color'] == 'Saturation'].reset_index().loc[0, ['vector', 'score']] value_separation_vector, performance_value = concept_vectors[concept_vectors['color'] == 'Value'].reset_index().loc[0, ['vector', 'score']] else: tmp = concept_vectors[concept_vectors['color'] == st.session_state.concept_ids] tmp = tmp[tmp['sign'] == st.session_state.sign][tmp['num_factors'] == st.session_state.num_factors][tmp['cl_method'] == st.session_state.cl_method][tmp['regularization'] == st.session_state.regularization] color_separation_vector, performance_color = tmp.reset_index().loc[0, ['vector', 'score']] tmp_value = concept_vectors[concept_vectors['color'] == 'Value'][concept_vectors['extremes'] == st.session_state.extremes] value_separation_vector, performance_value = tmp_value.reset_index().loc[0, ['vector', 'score']] tmp_sat = concept_vectors[concept_vectors['color'] == 'Saturation'][concept_vectors['extremes'] == st.session_state.extremes] saturation_separation_vector, performance_saturation = tmp_sat.reset_index().loc[0, ['vector', 'score']] st.write(f'Change in {st.session_state.concept_ids} of {np.round(st.session_state.color_lambda, 2)}, in saturation of {np.round(st.session_state.saturation_lambda, 2)}, in value of {np.round(st.session_state.value_lambda, 2)}. - Performance color vector: {performance_color}, saturation vector: {performance_saturation/100}, value vector: {performance_value/100}') # ---------------------------- 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, [color_separation_vector, saturation_separation_vector, value_separation_vector], lambdas=[st.session_state.color_lambda, st.session_state.saturation_lambda, st.session_state.value_lambda]) st.image(image_updated)