pages/3_Vectors_algebra.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('Vector algebra using disentangled vectors')
st.markdown(
    """
    This page offers the possibility to edit the colors of a given textile image using vector algebra and projections.
    It allows to select several colors to move towards and against (selecting a positive or negative lambda). 
    Furthermore, it offers the possibility of conditional manipulation, by moving in the direction of a color n1 without affecting the color n2. 
    This is done using a projected direction n1 - (n1.T n2) n2. 
    """,
    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()

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']
COLORS_NEGATIVE = COLORS_LIST + ['None']

if 'image_id' not in st.session_state:
    st.session_state.image_id = 52921
if 'colors' not in st.session_state:
    st.session_state.colors = [COLORS_LIST[5], COLORS_LIST[7]]
if 'non_colors' not in st.session_state:
    st.session_state.non_colors = ['None']
if 'color_lambda' not in st.session_state:
    st.session_state.color_lambda = [5]

# ----------------------------- INPUT ----------------------------------
epsilon_container = st.empty()
st.header('Image Manipulation with Vector Algebra')

header_col_1, header_col_2, header_col_3, header_col_4 = st.columns([1,1,1,1])
input_col_1, output_col_2, output_col_3, input_col_4 = st.columns([1,1,1,1])

# --------------------------- 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**')
        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 header_col_1:
    st.write('### Input image selection')

original_image_vec = annotations['w_vectors'][st.session_state.image_id]
img = generate_original_image(original_image_vec, model)

with output_col_2:
    st.image(img)

with header_col_2:
    st.write('### Original image')
    
with input_col_4:
    with st.form('text_form_1'):
        
        st.write('**Colors to vary (including Saturation and Value)**')
        colors = st.multiselect('Color:', tuple(COLORS_LIST), default=[COLORS_LIST[5], COLORS_LIST[7]])
        colors_button = st.form_submit_button('Choose the defined colors')
        
        st.session_state.image_id = image_id
        st.session_state.colors = colors
        st.session_state.color_lambda = [5]*len(colors)
        st.session_state.non_colors = ['None']*len(colors)
            
        lambdas = []
        negative_cols = []
        for color in colors:
            st.write('### '+color )
            st.write('**Set range of change (can be negative)**')
            chosen_color_lambda_input = st.empty()
            color_lambda = chosen_color_lambda_input.number_input('Lambda:', min_value=-100, step=1, value=5, key=color+'_number')
            lambdas.append(color_lambda)
                
            st.write('**Set dimensions of change to not consider**')
            chosen_color_negative_input = st.empty()
            color_negative = chosen_color_negative_input.selectbox('Color:', tuple(COLORS_NEGATIVE), index=len(COLORS_NEGATIVE)-1, key=color+'_noncolor')
            negative_cols.append(color_negative)
                
        lambdas_button = st.form_submit_button('Submit options')
        if lambdas_button:
            st.session_state.color_lambda = lambdas
            st.session_state.non_colors = negative_cols
                
        
with header_col_4:
    st.write('### Color settings')
# print(st.session_state.colors)
# print(st.session_state.color_lambda)
# print(st.session_state.non_colors)

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

with header_col_3:
    separation_vectors = []
    for col in st.session_state.colors:
        separation_vector, score_1 = concept_vectors[concept_vectors['color'] == col].reset_index().loc[0, ['vector', 'score']]
        separation_vectors.append(separation_vector)
    
    negative_separation_vectors = []
    for non_col in st.session_state.non_colors:
        if non_col != 'None':
            negative_separation_vector, score_2 = concept_vectors[concept_vectors['color'] == non_col].reset_index().loc[0, ['vector', 'score']]
            negative_separation_vectors.append(negative_separation_vector)
        else:
            negative_separation_vectors.append('None')
    ## n1 − (n1T n2)n2
    # print(negative_separation_vectors, separation_vectors)
    st.write('### Output Image')
    st.write(f'''Change in colors: {str(st.session_state.colors)},\
             without affecting colors {str(st.session_state.non_colors)}''')

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

with output_col_3:
    image_updated = generate_composite_images(model, original_image_vec, separation_vectors, 
                                              lambdas=st.session_state.color_lambda, 
                                              negative_colors=negative_separation_vectors)
    st.image(image_updated)