created global cleaned file

.gitignore

@@ -184,4 +184,5 @@ dmypy.json
 cython_debug/
 
 data/images/
-tmp/
+tmp/
+figures/

DisentanglementBase.py

@@ -0,0 +1,482 @@
+import numpy as np
+import pandas as pd
+
+from sklearn.svm import SVC
+from sklearn.decomposition import PCA
+from sklearn.linear_model import LogisticRegression
+from sklearn.model_selection import train_test_split
+
+from tqdm import tqdm
+import random
+from os.path import join
+import os
+import pickle
+
+import torch
+
+import matplotlib.pyplot as plt
+import PIL
+from PIL import Image, ImageColor
+
+import sys
+sys.path.append('backend')
+from color_annotations import extract_color
+from networks_stylegan3 import *
+sys.path.append('.')
+
+import dnnlib 
+import legacy
+
+class DisentanglementBase:
+    def __init__(self, repo_folder, model, annotations, df, space, colors_list, compute_s):
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        print('Using device', self.device)
+        self.repo_folder = repo_folder
+        self.model = model.to(self.device)
+        self.annotations = annotations
+        self.df = df
+        self.space = space
+        
+        self.layers = ['input', 'L0_36_512', 'L1_36_512', 'L2_36_512', 'L3_52_512',
+                       'L4_52_512', 'L5_84_512', 'L6_84_512', 'L7_148_512', 'L8_148_512', 
+                       'L9_148_362', 'L10_276_256', 'L11_276_181', 'L12_276_128', 
+                       'L13_256_128', 'L14_256_3']
+        self.layers_shapes = [4, 512, 512, 512, 512, 512, 512, 512, 512, 512, 512, 362, 256, 181, 128, 128]
+        self.decoding_layers = 16
+        self.colors_list = colors_list
+        
+        self.to_hsv()
+        if compute_s:
+            self.get_s_space()
+        
+
+    def to_hsv(self):
+        """
+        The tohsv function takes the top 3 colors of each image and converts them to HSV values.
+        It then adds these values as new columns in the dataframe.
+        
+        :param self: Allow the function to access the dataframe
+        :return: The dataframe with the new columns added
+        :doc-author: Trelent
+        """
+        print('Adding HSV encoding')
+        self.df['H1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[0])
+        self.df['H2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[0])
+        self.df['H3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[0])
+        
+        self.df['S1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[1])
+        self.df['S2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[1])
+        self.df['S3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[1])
+        
+        self.df['V1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[2])
+        self.df['V2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[2])
+        self.df['V3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[2])
+    
+    def get_s_space(self):
+        """
+        The get_s_space function takes the w_vectors from the annotations dictionary and uses them to generate s_vectors.
+        The s_space is a space of vectors that are generated by passing each w vector through each layer of the model.
+        This allows us to see how much information about a particular class is contained in different layers.
+        
+        :param self: Bind the method to a class
+        :return: A list of lists of s vectors
+        :doc-author: Trelent
+        """
+        print('Getting S space from W')
+        ss = []
+        for w in tqdm(self.annotations['w_vectors']):
+            w_torch = torch.from_numpy(w).to(self.device)
+            W = w_torch.expand((16, -1)).unsqueeze(0)
+            s = []
+            for i,layer in enumerate(self.layers):
+                s.append(getattr(self.model.synthesis, layer).affine(W[0, i].unsqueeze(0)).numpy())
+
+            ss.append(s)
+        self.annotations['s_vectors'] = ss
+        annotations_file = join(self.repo_folder, 'data/textile_annotated_files/seeds0000-100000_S.pkl')
+        print('Storing s for future use here:', annotations_file)
+        with open(annotations_file, 'wb') as f:
+            pickle.dump(self.annotations, f)
+
+    def get_encoded_latent(self):
+        # ... (existing code for getX)
+        if self.space.lower() == 'w':
+            X = np.array(self.annotations['w_vectors']).reshape((len(self.annotations['w_vectors']), 512))
+        elif self.space.lower() == 'z':
+            X = np.array(self.annotations['z_vectors']).reshape((len(self.annotations['z_vectors']), 512))
+        elif self.space.lower() == 's':
+            concat_v = []
+            for i in range(len(self.annotations['w_vectors'])):
+                concat_v.append(np.concatenate(self.annotations['s_vectors'][i], axis=1))
+            X = np.array(concat_v)
+            X = X[:, 0, :]
+        else:
+            Exception("Sorry, option not available, select among Z, W, S")
+            
+        print('Shape embedding:', X.shape)
+        return X
+    
+    def get_train_val(self, var='H1', cat=True):
+        X = self.get_encoded_latent()
+        y = np.array(self.df[var].values)
+        if cat:
+            y_cat = pd.cut(y, 
+                            bins=[x*256/12 if x<12 else 256 for x in range(13)],
+                            labels=self.colors_list
+                            ).fillna('Warm Pink Red')
+            x_train, x_val, y_train, y_val = train_test_split(X, y_cat, test_size=0.2)
+        else:
+            x_train, x_val, y_train, y_val = train_test_split(X, y, test_size=0.2)
+        return x_train, x_val, y_train, y_val
+    
+    def InterFaceGAN_separation_vector(self, method='LR', C=0.1):
+        """
+        Method from InterfaceGAN
+        The get_separation_space function takes in a type_bin, annotations, and df.
+        It then samples 100 of the most representative abstracts for that type_bin and 100 of the least representative abstracts for that type_bin.
+        It then trains an SVM or logistic regression model on these 200 samples to find a separation space between them. 
+        The function returns this separation space as well as how many nodes are important in this separation space.
+        
+        :param type_bin: Select the type of abstracts to be used for training
+        :param annotations: Access the z_vectors
+        :param df: Get the abstracts that are used for training
+        :param samples: Determine how many samples to take from the top and bottom of the distribution
+        :param method: Specify the classifier to use
+        :param C: Control the regularization strength
+        :return: The weights of the linear classifier
+        :doc-author: Trelent
+        """
+        x_train, x_val, y_train, y_val = self.get_train_val()
+        
+        if method == 'SVM':
+            svc = SVC(gamma='auto', kernel='linear', random_state=0, C=C)
+            svc.fit(x_train, y_train)
+            print('Val performance SVM', np.round(svc.score(x_val, y_val), 2))
+            return svc.coef_ / np.linalg.norm(clf.coef_)
+        elif method == 'LR':
+            clf = LogisticRegression(random_state=0, C=C)
+            clf.fit(x_train, y_train)
+            print('Val performance logistic regression', np.round(clf.score(x_val, y_val), 2))
+            return clf.coef_ / np.linalg.norm(clf.coef_)
+
+    def get_original_position_latent(self, positive_idxs, negative_idxs):
+        # ... (existing code for get_original_pos)
+        separation_vectors = []
+        for i in range(len(self.colors_list)):
+            if self.space.lower() == 's':
+                current_idx = 0
+                vectors = []
+                for j, (leng, layer) in enumerate(zip(self.layers_shapes, self.layers)):
+                    arr = np.zeros(leng)
+                    for positive_idx in positive_idxs[i]:
+                        if positive_idx >= current_idx and positive_idx < current_idx + leng:
+                            arr[positive_idx - current_idx] = 1
+                    for negative_idx in negative_idxs[i]:
+                        if negative_idx >= current_idx and negative_idx < current_idx + leng:
+                            arr[negative_idx - current_idx] = 1
+                        arr = arr / (np.linalg.norm(arr) + 0.000001)
+                    vectors.append(arr)
+                    current_idx += leng
+            elif self.space.lower() == 'z' or self.space.lower() == 'w':
+                vectors = np.zeros(512)
+                vectors[positive_idxs[i]] = 1
+                vectors[negative_idxs[i]] = -1
+                vectors = vectors / (np.linalg.norm(vectors) + 0.000001)
+            else:
+                raise Exception("""This space is not allowed in this function, 
+                                    select among Z, W, S""")
+            separation_vectors.append(vectors)
+            
+        return separation_vectors    
+    
+    def StyleSpace_separation_vector(self, sign=True, num_factors=20, cutout=0.25):
+        """ Formula from StyleSpace Analysis """
+        x_train, x_val, y_train, y_val = self.get_train_val()
+        
+        positive_idxs = []
+        negative_idxs = []
+        for color in self.colors_list:
+            x_col = x_train[np.where(y_train == color)]
+            mp = np.mean(x_train, axis=0)
+            sp = np.std(x_train, axis=0)
+            de = (x_col - mp) / sp
+            meu = np.mean(de, axis=0)
+            seu = np.std(de, axis=0)
+            if sign:
+                thetau = meu / seu
+                positive_idx = np.argsort(thetau)[-num_factors//2:]
+                negative_idx = np.argsort(thetau)[:num_factors//2]
+                
+            else:
+                thetau = np.abs(meu) / seu
+                positive_idx = np.argsort(thetau)[-num_factors:]
+                negative_idx = []
+                
+
+            if cutout:
+                beyond_cutout = np.where(np.abs(thetau) > cutout)
+                positive_idx = np.intersect1d(positive_idx, beyond_cutout)
+                negative_idx = np.intersect1d(negative_idx, beyond_cutout)
+                
+                if len(positive_idx) == 0 and len(negative_idx) == 0:
+                    print('No values found above the current cutout', cutout, 'for color', color, '.\n Disentangled vector will be all zeros.' )
+                
+            positive_idxs.append(positive_idx)
+            negative_idxs.append(negative_idx)
+        
+        separation_vectors = self.get_original_position_latent(positive_idxs, negative_idxs)
+        return separation_vectors
+
+    def GANSpace_separation_vectors(self, num_components):
+        x_train, x_val, y_train, y_val = self.get_train_val()
+        if self.space.lower() == 'w':
+            pca = PCA(n_components=num_components)
+
+            dims_pca = pca.fit_transform(x_train.T)
+            dims_pca /= np.linalg.norm(dims_pca, axis=0)
+            
+            return dims_pca
+        
+        else:
+            raise("""This space is not allowed in this function, 
+                     only W""")
+    
+    def generate_images(self, seed, separation_vector=None, lambd=0):
+        """
+        The generate_original_image function takes in a latent vector and the model,
+        and returns an image generated from that latent vector.
+        
+        
+        :param z: Generate the image
+        :param model: Generate the image
+        :return: A pil image
+        :doc-author: Trelent
+        """
+        G = self.model.to(self.device) # type: ignore
+        # Labels.
+        label = torch.zeros([1, G.c_dim], device=self.device)
+        if self.space.lower() == 'z':
+            vec = self.annotations['z_vectors'][seed]
+            Z = torch.from_numpy(vec.copy()).to(self.device)
+            if separation_vector is not None:
+                change = torch.from_numpy(separation_vector.copy()).unsqueeze(0).to(self.device)
+                Z = torch.add(Z, change, alpha=lambd)
+            img = G(Z, label, truncation_psi=1, noise_mode='const')
+        elif self.space.lower() == 'w':
+            vec = self.annotations['w_vectors'][seed]
+            W = torch.from_numpy(np.repeat(vec, self.decoding_layers, axis=0)
+                                 .reshape(1, self.decoding_layers, vec.shape[1]).copy()).to(self.device)
+            if separation_vector is not None:
+                change = torch.from_numpy(separation_vector.copy()).unsqueeze(0).to(self.device)
+                W = torch.add(W, change, alpha=lambd)
+            img = G.synthesis(W, noise_mode='const')
+        else:
+            raise Exception("""This space is not allowed in this function, 
+                            select either W or Z or use generate_flexible_images""")
+            
+        img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+        return PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
+
+    def forward_from_style(self, x, styles, layer):
+        dtype = torch.float16 if (getattr(self.model.synthesis, layer).use_fp16 and self.device=='cuda') else torch.float32
+        
+        if getattr(self.model.synthesis, layer).is_torgb:
+            weight_gain = 1 / np.sqrt(getattr(self.model.synthesis, layer).in_channels * (getattr(self.model.synthesis, layer).conv_kernel ** 2))
+            styles = styles * weight_gain
+        
+        input_gain = getattr(self.model.synthesis, layer).magnitude_ema.rsqrt().to(dtype)
+        
+        # Execute modulated conv2d.
+        x = modulated_conv2d(x=x.to(dtype), w=getattr(self.model.synthesis, layer).weight.to(dtype), s=styles.to(dtype),
+        padding=getattr(self.model.synthesis, layer).conv_kernel-1, 
+                        demodulate=(not getattr(self.model.synthesis, layer).is_torgb), 
+                        input_gain=input_gain.to(dtype))
+        
+        # Execute bias, filtered leaky ReLU, and clamping.
+        gain = 1 if getattr(self.model.synthesis, layer).is_torgb else np.sqrt(2)
+        slope = 1 if getattr(self.model.synthesis, layer).is_torgb else 0.2
+        
+        x = filtered_lrelu.filtered_lrelu(x=x, fu=getattr(self.model.synthesis, layer).up_filter, fd=getattr(self.model.synthesis, layer).down_filter, 
+                                            b=getattr(self.model.synthesis, layer).bias.to(x.dtype),
+                                            up=getattr(self.model.synthesis, layer).up_factor, down=getattr(self.model.synthesis, layer).down_factor, 
+                                            padding=getattr(self.model.synthesis, layer).padding,
+                                            gain=gain, slope=slope, clamp=getattr(self.model.synthesis, layer).conv_clamp)
+        return x
+    
+    def generate_flexible_images(self, seed, separation_vector=None, lambd=0):
+        if self.space.lower() != 's':
+            raise Exception("""This space is not allowed in this function, 
+                            select S or use generate_images""")
+            
+        vec = self.annotations['w_vectors'][seed]
+        w_torch = torch.from_numpy(vec).to(self.device)
+        W = w_torch.expand((self.decoding_layers, -1)).unsqueeze(0)
+        x = self.model.synthesis.input(W[0,0].unsqueeze(0))
+        for i, layer in enumerate(self.layers[1:]):
+            style = getattr(self.model.synthesis, layer).affine(W[0, i].unsqueeze(0))
+            if separation_vector is not None:
+                change = torch.from_numpy(separation_vector[i+1].copy()).unsqueeze(0).to(self.device)
+                style = torch.add(style, change, alpha=lambd)
+            x = self.forward_from_style(x, style, layer)
+        
+        if self.model.synthesis.output_scale != 1:
+                x = x * self.model.synthesis.output_scale
+
+        img = (x.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+        img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
+            
+        return img
+
+    def generate_changes(self, seed, separation_vector, min_epsilon=-3, max_epsilon=3, count=5, savefig=True, feature=None, method=None):
+        """
+        The regenerate_images function takes a model, z, and decision_boundary as input.  It then
+        constructs an inverse rotation/translation matrix and passes it to the generator.  The generator
+        expects this matrix as an inverse to avoid potentially failing numerical operations in the network.
+        The function then generates images using G(z_0, label) where z_0 is a linear combination of z and the decision boundary.
+        
+        :param model: Pass in the model to be used for image generation
+        :param z: Generate the starting point of the line
+        :param decision_boundary: Generate images along the direction of the decision boundary
+        :param min_epsilon: Set the minimum value of lambda
+        :param max_epsilon: Set the maximum distance from the original image to generate
+        :param count: Determine the number of images that are generated
+        :return: A list of images and a list of lambdas
+        :doc-author: Trelent
+        """
+            
+        os.makedirs(join(self.repo_folder, 'figures'), exist_ok=True)
+        lambdas = np.linspace(min_epsilon, max_epsilon, count)
+        images = []
+        # Generate images.
+        for _, lambd in enumerate(tqdm(lambdas)):
+            if self.space.lower() == 's':
+                images.append(self.generate_flexible_images(seed, separation_vector=separation_vector, lambd=lambd))
+            elif self.space.lower() in ['z', 'w']:
+                images.append(self.generate_images(seed, separation_vector=separation_vector, lambd=lambd))
+        
+        if savefig:
+            print('Generating image for color', feature)
+            fig, axs = plt.subplots(1, len(images), figsize=(90,20))
+            title = 'Disentanglement method: '+ method + ', on feature: ' + feature + ' on space: ' + self.space + ', image seed: ' + str(seed)
+            name = '_'.join([method, feature, self.space, str(seed), str(lambdas[-1])])
+            fig.suptitle(title, fontsize=20)
+                
+            for i, (image, lambd) in enumerate(zip(images, lambdas)):
+                axs[i].imshow(image)
+                axs[i].set_title(np.round(lambd, 2))
+            plt.tight_layout()
+            plt.savefig(join(self.repo_folder, 'figures', name+'.jpg'))
+        return images, lambdas
+    
+    def get_verification_score(self, separation_vector, feature_id, samples=10, lambd=1, savefig=False, feature=None, method=None):
+        items = random.sample(range(100000), samples)
+        hue_low = feature_id * 256 / 12 
+        hue_high = (feature_id + 1) * 256 / 12 
+        
+        matches = 0
+        
+        for seed in tqdm(items):
+            images, lambdas = self.generate_changes(seed, separation_vector, min_epsilon=-lambd, max_epsilon=lambd, count=3, savefig=savefig, feature=feature, method=method)
+            colors_negative = extract_color(images[0], 5, 1, None)
+            h0, s0, v0 = ImageColor.getcolor(colors_negative[0], 'HSV')
+            
+            colors_orig = extract_color(images[1], 5, 1, None)
+            h1, s1, v1 = ImageColor.getcolor(colors_orig[0], 'HSV')
+            
+            colors_positive = extract_color(images[2], 5, 1, None)
+            h2, s2, v2 = ImageColor.getcolor(colors_positive[0], 'HSV')
+            
+            if h1 > hue_low and h1 < hue_high:
+                samples -= 1
+            else:
+                if (h0 > hue_low and h0 < hue_high) or (h2 > hue_low and h2 < hue_high):
+                    matches += 1
+        
+        return np.round(matches / samples, 2)
+    
+
+def main():
+    repo_folder = '.'
+    annotations_file = join(repo_folder, 'data/textile_annotated_files/seeds0000-100000_S.pkl')
+    with open(annotations_file, 'rb') as f:
+        annotations = pickle.load(f)
+
+    df_file = join(repo_folder, 'data/textile_annotated_files/top_three_colours.csv')
+    df = pd.read_csv(df_file).fillna('#000000')
+
+    model_file = join(repo_folder, 'data/textile_model_files/network-snapshot-005000.pkl')
+    with dnnlib.util.open_url(model_file) as f:
+        model = legacy.load_network_pkl(f)['G_ema'] # type: ignore
+
+    colors_list = ['Warm Pink Red', 'Red Orange', 'Orange Yellow', 'Gold Yellow', 'Chartreuse Green',
+                   'Kelly Green', 'Green Blue Seafoam', 'Blue Green Cyan',
+                   'Warm Blue', 'Indigo Blue Purple', 'Purple Magenta', 'Magenta Pink']
+    
+    scores = []
+    kwargs = {'CL method':['LR', 'SVM'], 'C':[0.1, 1], 'sign':[True, False], 'num_factors':[1, 10, 20, 50], 'cutout': [None, 0.2], 'max_lambda':[6, 10, 1], 'samples':50, 'lambda_verif':[1, 3, 6]}
+    
+    for space in ['w', 'z', 's']:
+        print('Launching experiment with space:', space)
+        disentanglemnet_exp = DisentanglementBase(repo_folder, model, annotations, df, space=space, colors_list=colors_list, compute_s=False)
+
+        for method in ['StyleSpace', 'InterFaceGAN', 'GANSpace']:
+            if space != 's' and method == 'InterFaceGAN':
+                print('Now obtaining separation vector for using InterfaceGAN')
+                for met in kwargs['CL method']:
+                    for c in kwargs['C']:
+                        separation_vectors = disentanglemnet_exp.InterFaceGAN_separation_vector(method=met, C=c)
+                        for i, color in enumerate(colors_list):
+                            print('Generating images with variations')
+                            seed = random.randint(0,100000)
+                            for eps in kwargs['max_lambda']:
+                                disentanglemnet_exp.generate_changes(seed, separation_vectors[i], min_epsilon=-eps, max_epsilon=eps, savefig=True, feature=color, method=method)
+                            
+                            print('Finally obtaining verification score')
+                            for verif in kwargs['lambda_verif']:    
+                                score = disentanglemnet_exp.get_verification_score(separation_vectors[i], i, samples=kwargs['samples'], lambd=verif, savefig=True, feature=color, method=method)
+                                print('Score for method', method, 'on space', space, 'for color', color, ':', score)
+                            
+                                scores.append([space, method, color, score, 'classification method:' + met + ', regularization: ' + str(c) + ', verification lambda:' + str(verif)])
+            
+            elif method == 'StyleSpace':
+                print('Now obtaining separation vector for using StyleSpace')
+                for sign in kwargs['sign']:
+                    for num_factors in kwargs['num_factors']:
+                        for cutout in kwargs['cutout']:
+                            separation_vectors = disentanglemnet_exp.StyleSpace_separation_vector(sign=sign, num_factors=num_factors, cutout=cutout)
+                            for i, color in enumerate(colors_list):
+                                print('Generating images with variations')
+                                seed = random.randint(0,100000)
+                                for eps in kwargs['max_lambda']:
+                                    disentanglemnet_exp.generate_changes(seed, separation_vectors[i], min_epsilon=-eps, max_epsilon=eps, savefig=True, feature=color, method=method)
+                                        
+                                print('Finally obtaining verification score')
+                                for verif in kwargs['lambda_verif']:    
+                                    score = disentanglemnet_exp.get_verification_score(separation_vectors[i], i, samples=kwargs['samples'], lambd=verif, savefig=True, feature=color, method=method)
+                                    print('Score for method', method, 'on space', space, 'for color', color, ':', score)
+                                        
+                                    scores.append([space, method, color, score, 'using sign:' + str(sign) + ', number of factors: ' + str(num_factors) + ', using cutout: ' + str(cutout) + ', verification lambda:' + str(verif)])
+                                    
+            if space == 'w' and method == 'GANSpace':
+                print('Now obtaining separation vector for using GANSpace')
+                separation_vectors = disentanglemnet_exp.GANSpace_separation_vectors(100)
+                for i in range(100):
+                    print('Generating images with variations')
+                    seed = random.randint(0,100000)
+                    for eps in kwargs['max_lambda']:
+                        disentanglemnet_exp.generate_changes(seed, separation_vectors[i], min_epsilon=-eps, max_epsilon=eps, savefig=True, feature=color, method=method)
+                                
+                    score = None
+                    scores.append([space, method, color, score, '100'])
+            else:
+                print('Skipping', method, 'on space', space)
+                continue
+            
+            
+    
+    score_df = pd.DataFrame(scores, columns=['space', 'method', 'color', 'score', 'kwargs'])
+    print(score_df)
+    score_df.to_csv(join(repo_folder, 'data/scores.csv'))
+ 
+if __name__ == "__main__":
+    main()

backend/color_annotations.py

@@ -0,0 +1,97 @@
+#!/usr/bin/env python
+
+"""Extract color features from the generated textile images."""
+
+import os
+#os.environ["TOKENIZERS_PARALLELISM"] = "false"
+from tqdm import tqdm
+#from transformers import pipeline
+import numpy as np
+import pandas as pd
+import time
+
+import click
+from PIL import Image
+import math 
+import pickle
+from glob import glob 
+
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import matplotlib.image as mpimg
+import cv2
+import extcolors
+
+from colormap import rgb2hex
+from PIL import Image
+from matplotlib.offsetbox import OffsetImage, AnnotationBbox
+
+def color_to_df(input):
+    colors_pre_list = str(input).replace('([(','').split(', (')[0:-1]
+    df_rgb = [i.split('), ')[0] + ')' for i in colors_pre_list]
+    df_percent = [i.split('), ')[1].replace(')','') for i in colors_pre_list]
+    
+    #convert RGB to HEX code
+    df_color_up = [rgb2hex(int(i.split(", ")[0].replace("(","")),
+                          int(i.split(", ")[1]),
+                          int(i.split(", ")[2].replace(")",""))) for i in df_rgb]
+    
+    df = pd.DataFrame(zip(df_color_up, df_percent), columns = ['c_code','occurence'])
+    return df
+
+def extract_color(input_image, tolerance, zoom, outpath, save=None):
+    colors_x = extcolors.extract_from_image(input_image, tolerance = tolerance, limit = 13)
+    df_color = color_to_df(colors_x)
+    
+    #annotate text
+    list_color = list(df_color['c_code'])
+    list_precent = [int(i) for i in list(df_color['occurence'])]
+    text_c = [c + ' ' + str(round(p*100/sum(list_precent),1)) +'%' for c, p in zip(list_color, list_precent)]
+    colors = list(df_color['c_code'])
+    if '#000000' in colors:
+        colors.remove('#000000')
+    return colors[:3]
+
+
+@click.command()
+@click.option('--genimages_dir', help='Where the output images are saved', type=str, required=True, metavar='DIR')
+
+def annotate_textile_images(
+    genimages_dir: str,
+    
+):
+    """Produce annotations for the generated images.
+    \b
+    # 
+    python annotate_textiles.py --genimages_dir /home/ludosc/data/stylegan-10000-textile-upscale
+    """
+    colours = []
+    pickle_files = glob(genimages_dir + '/imgs0000*.pkl')
+    for pickle_file in pickle_files:
+        print('Using pickle file: ', pickle_file)
+        with open(pickle_file, 'rb') as f:
+            info = pickle.load(f)
+
+        listlen = len(info['fname'])
+        os.makedirs('/data/ludosc/colour_palettes/', exist_ok=True)
+        for i,im in enumerate(tqdm(info['fname'])):
+            try:
+                top_cols = exact_color(im, 12, 5, '/data/ludosc/colour_palettes/' + im.split('/')[-1])
+                colours.append([im]+top_cols)
+            except Exception as e:
+                print(e)
+            if i % 1000 == 0:
+                df = pd.DataFrame(colours, columns=['fname', 'top1col', 'top2col', 'top3col'])
+                print(df.head())
+                df.to_csv(genimages_dir + f'/top_three_colours.csv', index=False)
+                
+        df = pd.DataFrame(colours, columns=['fname', 'top1col', 'top2col', 'top3col'])
+        print(df.head())
+        df.to_csv(genimages_dir + f'/final_sim_{os.path.basename(pickle_file.split(".")[0])}.csv', index=False)
+        
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    annotate_textile_images() # pylint: disable=no-value-for-parameter
+
+#----------------------------------------------------------------------------

backend/disentangle_concepts.py

@@ -5,6 +5,13 @@ from sklearn.model_selection import train_test_split
 import torch
 from umap import UMAP
 import PIL
+from tqdm import tqdm
+import random
+from PIL import Image, ImageColor
+
+from .color_annotations import extract_color
+
+
 
 def get_separation_space(type_bin, annotations, df, samples=200, method='LR', C=0.1, latent_space='Z'):
     """
@@ -65,7 +72,7 @@ def get_separation_space(type_bin, annotations, df, samples=200, method='LR', C=
         return clf.coef_ / np.linalg.norm(clf.coef_), imp_features, imp_nodes, np.round(clf.score(x_val, y_val),2)
 
 
-def regenerate_images(model, z, decision_boundary, min_epsilon=-3, max_epsilon=3, count=5, latent_space='Z', layers=None):
+def regenerate_images(model, z, decision_boundary, min_epsilon=-3, max_epsilon=3, count=5, latent_space='Z', layers=None, number=3):
     """
     The regenerate_images function takes a model, z, and decision_boundary as input.  It then
     constructs an inverse rotation/translation matrix and passes it to the generator.  The generator
@@ -92,19 +99,18 @@ def regenerate_images(model, z, decision_boundary, min_epsilon=-3, max_epsilon=3
     z = torch.from_numpy(z.copy()).to(device)
     decision_boundary = torch.from_numpy(decision_boundary.copy()).to(device)
         
+    repetitions = 16 if number == 3 else 14
     lambdas = np.linspace(min_epsilon, max_epsilon, count)
     images = []
     # Generate images.
-    for _, lambda_ in enumerate(lambdas):
+    for _, lambda_ in enumerate(tqdm(lambdas)):
         z_0 = z + lambda_ * decision_boundary
         if latent_space == 'Z':
             W_0 = G.mapping(z_0, label, truncation_psi=1).to(torch.float32)
             W = G.mapping(z, label, truncation_psi=1).to(torch.float32)
-            print(W.dtype)
         else:
-            W_0 = z_0.expand((14, -1)).unsqueeze(0).to(torch.float32)
-            W = z.expand((14, -1)).unsqueeze(0).to(torch.float32)
-            print(W.dtype)
+            W_0 = z_0.expand((repetitions, -1)).unsqueeze(0).to(torch.float32)
+            W = z.expand((repetitions, -1)).unsqueeze(0).to(torch.float32)
         
         if layers:
             W_f = torch.empty_like(W).copy_(W).to(torch.float32)
@@ -117,14 +123,14 @@ def regenerate_images(model, z, decision_boundary, min_epsilon=-3, max_epsilon=3
         images.append(PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB'))
             
     return images, lambdas
-
-
+    
+    
 def generate_joint_effect(model, z, decision_boundaries, min_epsilon=-3, max_epsilon=3, count=5, latent_space='Z'):
     decision_boundary_joint = np.sum(decision_boundaries, axis=0)
     print(decision_boundary_joint.shape)
     return regenerate_images(model, z, decision_boundary_joint, min_epsilon=min_epsilon, max_epsilon=max_epsilon, count=count, latent_space=latent_space)
     
-def generate_original_image(z, model, latent_space='Z'):
+def generate_original_image(z, model, latent_space='Z', number=3):
     """
     The generate_original_image function takes in a latent vector and the model,
     and returns an image generated from that latent vector.
@@ -135,6 +141,8 @@ def generate_original_image(z, model, latent_space='Z'):
     :return: A pil image
     :doc-author: Trelent
     """
+    repetitions = 16 if number == 3 else 14
+    
     device = torch.device('cpu')
     G = model.to(device) # type: ignore
     # Labels.
@@ -143,10 +151,10 @@ def generate_original_image(z, model, latent_space='Z'):
         z = torch.from_numpy(z.copy()).to(device)
         img = G(z, label, truncation_psi=1, noise_mode='const')
     else:
-        W = torch.from_numpy(np.repeat(z, 14, axis=0).reshape(1, 14, z.shape[1]).copy()).to(device)
+        W = torch.from_numpy(np.repeat(z, repetitions, axis=0).reshape(1, repetitions, z.shape[1]).copy()).to(device)
         print(W.shape)
         img = G.synthesis(W, noise_mode='const')
-
+    
     img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
     return PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
 
@@ -188,8 +196,36 @@ def get_concepts_vectors(concepts, annotations, df, samples=100, method='LR', C=
     return vectors, nodes_in_common, performances
 
 
+def get_verification_score(color_id, decision_boundary, model, annotations, samples=100, latent_space='W'):
+    listlen = len(annotations['fname'])
+    items = random.sample(range(listlen), samples)
+    hue_low = color_id * 256 / 12 
+    hue_high = (color_id + 1) * 256 / 12 
+    hue_mean = (hue_low + hue_high) / 2
+    print(int(hue_low), int(hue_high), int(hue_mean))
+    distances = []
+    distances_orig = []
+    for iterator in tqdm(items):
+        if latent_space == 'Z':
+            z = annotations['z_vectors'][iterator]
+        else:
+            z = annotations['w_vectors'][iterator]
+        
+        images, lambdas = regenerate_images(model, z, decision_boundary, min_epsilon=0, max_epsilon=1, count=2, latent_space=latent_space)
+        colors_orig = extract_color(images[0], 5, 1, None)
+        h_old, s_old, v_old = ImageColor.getcolor(colors_orig[0], 'HSV')
+        colors_new = extract_color(images[1], 5, 1, None)
+        h_new, s_new, v_new = ImageColor.getcolor(colors_new[0], 'HSV')
+        print(h_old, h_new)
+        distance = np.abs(hue_mean - h_new)
+        distances.append(distance)
+        distance_orig = np.abs(hue_mean - h_old)
+        distances_orig.append(distance_orig)
+        
+    return np.round(np.mean(np.array(distances)), 4), np.round(np.mean(np.array(distances_orig)), 4)
+    
 
-def get_verification_score(concept, decision_boundary, model, annotations, samples=100, latent_space='Z'):
+def get_verification_score_clip(concept, decision_boundary, model, annotations, samples=100, latent_space='Z'):
     import open_clip
     import os
     import random
@@ -243,5 +279,177 @@ def get_verification_score(concept, decision_boundary, model, annotations, sampl
     return np.round(np.mean(np.array(changes)), 4)
     
             
+            
+def tohsv(df):
+    df['H1'] = df['top1col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[0])
+    df['H2'] = df['top2col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[0])
+    df['H3'] = df['top3col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[0])
+    
+    df['S1'] = df['top1col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[1])
+    df['S2'] = df['top2col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[1])
+    df['S3'] = df['top3col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[1])
+    
+    df['V1'] = df['top1col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[2])
+    df['V2'] = df['top2col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[2])
+    df['V3'] = df['top3col'].map(lambda x: ImageColor.getcolor(x, 'HSV')[2])
+    return df
 
-            
+
+def rest_from_style(x, styles, layer):
+    dtype = torch.float16 if (getattr(model.synthesis, layer).use_fp16 and device=='cuda') else torch.float32
+    if getattr(model.synthesis, layer).is_torgb:
+        print(layer, getattr(model.synthesis, layer).is_torgb)
+        weight_gain = 1 / np.sqrt(getattr(model.synthesis, layer).in_channels * (getattr(model.synthesis, layer).conv_kernel ** 2))
+        styles = styles * weight_gain
+    input_gain = getattr(model.synthesis, layer).magnitude_ema.rsqrt().to(dtype)
+    # Execute modulated conv2d.
+    x = modulated_conv2d(x=x.to(dtype), w=getattr(model.synthesis, layer).weight.to(dtype), s=styles.to(dtype),
+    padding=getattr(model.synthesis, layer).conv_kernel-1, demodulate=(not getattr(model.synthesis, layer).is_torgb), input_gain=input_gain.to(dtype))
+    # Execute bias, filtered leaky ReLU, and clamping.
+    gain = 1 if getattr(model.synthesis, layer).is_torgb else np.sqrt(2)
+    slope = 1 if getattr(model.synthesis, layer).is_torgb else 0.2
+    x = filtered_lrelu.filtered_lrelu(x=x, fu=getattr(model.synthesis, layer).up_filter, fd=getattr(model.synthesis, layer).down_filter, 
+                                        b=getattr(model.synthesis, layer).bias.to(x.dtype),
+                                        up=getattr(model.synthesis, layer).up_factor, down=getattr(model.synthesis, layer).down_factor, 
+                                        padding=getattr(model.synthesis, layer).padding,
+                                        gain=gain, slope=slope, clamp=getattr(model.synthesis, layer).conv_clamp)
+    return x
+
+
+def getS(w):
+    w_torch = torch.from_numpy(w).to('cpu')
+    W = w_torch.expand((16, -1)).unsqueeze(0)
+    s = []
+    s.append(model.synthesis.input.affine(W[0, 0].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L0_36_512.affine(W[0, 1].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L1_36_512.affine(W[0, 2].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L2_36_512.affine(W[0, 3].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L3_52_512.affine(W[0, 4].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L4_52_512.affine(W[0, 5].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L5_84_512.affine(W[0, 6].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L6_84_512.affine(W[0, 7].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L7_148_512.affine(W[0, 8].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L8_148_512.affine(W[0, 9].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L9_148_362.affine(W[0, 10].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L10_276_256.affine(W[0, 11].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L11_276_181.affine(W[0, 12].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L12_276_128.affine(W[0, 13].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L13_256_128.affine(W[0, 14].unsqueeze(0)).numpy())
+    s.append(model.synthesis.L14_256_3.affine(W[0, 15].unsqueeze(0)).numpy())
+    return s
+
+def detect_attribute_specific_channels(positives, all, sign=False):
+    """ Formula from StyleSpace Analysis """
+    mp = np.mean(all, axis=0)
+    sp = np.std(all, axis=0)
+    de = (positives - mp) / sp
+    meu = np.mean(de, axis=0)
+    seu = np.std(de, axis=0)
+    if sign:
+        thetau = meu / seu
+    else:
+        thetau = np.abs(meu) / seu
+    return thetau
+
+def all_variance_based_disentanglements(labels, x, y, k=10, sign=False, cutout=0.28):
+    seps = []
+    sorted_vals = []
+    for lbl in labels:
+        positives = x[np.where(y == lbl)]
+        variations = detect_attribute_specific_channels(positives, x, sign=sign)
+        if sign:
+            argsorted_vars_pos = np.argsort(variations)[-k//2:]
+            # print(argsorted_vars_pos)
+            argsorted_vars_neg = np.argsort(variations)[:k//2]
+            if cutout:
+                beyond_cutout = np.where(np.abs(variations) > cutout)
+                # print(beyond_cutout)
+                argsorted_vars_pos_int = np.intersect1d(argsorted_vars_pos, beyond_cutout)
+                argsorted_vars_neg_int = np.intersect1d(argsorted_vars_neg, beyond_cutout)
+                # print(argsorted_vars_pos)
+                if len(argsorted_vars_neg_int) > 0:
+                    argsorted_vars_neg = np.array(argsorted_vars_neg_int)
+                if len(argsorted_vars_pos_int) > 0:
+                    argsorted_vars_pos = np.array(argsorted_vars_pos_int)
+            
+            
+        else:
+            argsorted_vars = np.argsort(variations)[-k:]
+            
+            
+        sorted_vals.append(np.sort(variations))
+        separation_vector_onehot /= np.linalg.norm(separation_vector_onehot)
+        seps.append(separation_vector_onehot)
+    return seps, sorted_vals
+
+def generate_flexible_images(w, change_vectors, lambdas=1, device='cpu'):
+    w_torch = torch.from_numpy(w).to('cpu')
+    if len(change_vectors) != 17:
+        w_torch = w_torch + lambdas * change_vectors[0]
+    W = w_torch.expand((16, -1)).unsqueeze(0)
+    
+    x = model.synthesis.input(W[0,0].unsqueeze(0))
+    for i, layer in enumerate(layers):
+        if i < 2:
+            continue
+        style = getattr(model.synthesis, layer).affine(W[0, i-1].unsqueeze(0))
+        if len(change_vectors) != 17:
+            change = torch.from_numpy(change_vectors[i].copy()).unsqueeze(0).to(device)
+            style = torch.add(style, change, alpha=lambdas)
+        x = rest_from_style(x, style, layer)
+    
+    if model.synthesis.output_scale != 1:
+            x = x * model.synthesis.output_scale
+
+    img = (x.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+    img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
+        
+    return img
+
+def get_original_pos(top_positions, bottom_positions=None, space='s', sign=True, 
+                                    shapes=[[512, 4, 512, 512, 512, 512, 512, 512, 512,
+                                             512, 512, 512, 362, 256, 181, 128, 128]], 
+                                    layers=['w', 'input', 'L0_36_512', 'L1_36_512', 'L2_36_512', 'L3_52_512',
+                                            'L4_52_512', 'L5_84_512', 'L6_84_512', 'L7_148_512', 'L8_148_512', 
+                                            'L9_148_362', 'L10_276_256', 'L11_276_181', 'L12_276_128', 
+                                            'L13_256_128', 'L14_256_3'], ):
+    if space == 's':
+        current_idx = 0
+        vectors = []
+        for i, (leng, layer) in enumerate(zip(shapes, layers)):
+            arr = np.zeros(leng)
+            for top_position in top_positions:
+                if top_position >= current_idx and top_position < current_idx + leng:
+                    arr[top_position - current_idx] = 1
+            for bottom_position in bottom_positions:
+                if sign:
+                    if bottom_position >= current_idx and bottom_position < current_idx + leng:
+                        arr[bottom_position - current_idx] = 1
+                arr = arr / (np.linalg.norm(arr) + 0.000001)
+            vectors.append(arr)
+            current_idx += leng
+    else:
+        if sign:
+            vectors = np.zeros(512)
+            vectors[top_positions] = 1
+            vectors[bottom_positions] = -1
+        else:
+            vectors = np.zeros(512)
+            vectors[top_positions] = 1
+    return vectors    
+
+def getX(annotations, space='s'):
+    if space == 'x':
+        X = np.array(annotations['w_vectors']).reshape((len(annotations['w_vectors']), 512))
+    elif space == 's':
+        concat_v = []
+        for i in range(len(annotations['w_vectors'])):
+            concat_v.append(np.concatenate([annotations['w_vectors'][i]] + annotations['s_vectors'][i], axis=1))
+            
+        X = np.array(concat_v)
+        X = X[:, 0, :]
+        print(X.shape)
+        
+    return X
+        
+    

backend/networks_stylegan3.py

@@ -0,0 +1,515 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Generator architecture from the paper
+"Alias-Free Generative Adversarial Networks"."""
+
+import numpy as np
+import scipy.signal
+import scipy.optimize
+import torch
+from torch_utils import misc
+from torch_utils import persistence
+from torch_utils.ops import conv2d_gradfix
+from torch_utils.ops import filtered_lrelu
+from torch_utils.ops import bias_act
+
+#----------------------------------------------------------------------------
+
+@misc.profiled_function
+def modulated_conv2d(
+    x,                  # Input tensor: [batch_size, in_channels, in_height, in_width]
+    w,                  # Weight tensor: [out_channels, in_channels, kernel_height, kernel_width]
+    s,                  # Style tensor: [batch_size, in_channels]
+    demodulate  = True, # Apply weight demodulation?
+    padding     = 0,    # Padding: int or [padH, padW]
+    input_gain  = None, # Optional scale factors for the input channels: [], [in_channels], or [batch_size, in_channels]
+):
+    with misc.suppress_tracer_warnings(): # this value will be treated as a constant
+        batch_size = int(x.shape[0])
+    out_channels, in_channels, kh, kw = w.shape
+    misc.assert_shape(w, [out_channels, in_channels, kh, kw]) # [OIkk]
+    misc.assert_shape(x, [batch_size, in_channels, None, None]) # [NIHW]
+    misc.assert_shape(s, [batch_size, in_channels]) # [NI]
+
+    # Pre-normalize inputs.
+    if demodulate:
+        w = w * w.square().mean([1,2,3], keepdim=True).rsqrt()
+        s = s * s.square().mean().rsqrt()
+
+    # Modulate weights.
+    w = w.unsqueeze(0) # [NOIkk]
+    w = w * s.unsqueeze(1).unsqueeze(3).unsqueeze(4) # [NOIkk]
+
+    # Demodulate weights.
+    if demodulate:
+        dcoefs = (w.square().sum(dim=[2,3,4]) + 1e-8).rsqrt() # [NO]
+        w = w * dcoefs.unsqueeze(2).unsqueeze(3).unsqueeze(4) # [NOIkk]
+
+    # Apply input scaling.
+    if input_gain is not None:
+        input_gain = input_gain.expand(batch_size, in_channels) # [NI]
+        w = w * input_gain.unsqueeze(1).unsqueeze(3).unsqueeze(4) # [NOIkk]
+
+    # Execute as one fused op using grouped convolution.
+    x = x.reshape(1, -1, *x.shape[2:])
+    w = w.reshape(-1, in_channels, kh, kw)
+    x = conv2d_gradfix.conv2d(input=x, weight=w.to(x.dtype), padding=padding, groups=batch_size)
+    x = x.reshape(batch_size, -1, *x.shape[2:])
+    return x
+
+#----------------------------------------------------------------------------
+
+@persistence.persistent_class
+class FullyConnectedLayer(torch.nn.Module):
+    def __init__(self,
+        in_features,                # Number of input features.
+        out_features,               # Number of output features.
+        activation      = 'linear', # Activation function: 'relu', 'lrelu', etc.
+        bias            = True,     # Apply additive bias before the activation function?
+        lr_multiplier   = 1,        # Learning rate multiplier.
+        weight_init     = 1,        # Initial standard deviation of the weight tensor.
+        bias_init       = 0,        # Initial value of the additive bias.
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.activation = activation
+        self.weight = torch.nn.Parameter(torch.randn([out_features, in_features]) * (weight_init / lr_multiplier))
+        bias_init = np.broadcast_to(np.asarray(bias_init, dtype=np.float32), [out_features])
+        self.bias = torch.nn.Parameter(torch.from_numpy(bias_init / lr_multiplier)) if bias else None
+        self.weight_gain = lr_multiplier / np.sqrt(in_features)
+        self.bias_gain = lr_multiplier
+
+    def forward(self, x):
+        w = self.weight.to(x.dtype) * self.weight_gain
+        b = self.bias
+        if b is not None:
+            b = b.to(x.dtype)
+            if self.bias_gain != 1:
+                b = b * self.bias_gain
+        if self.activation == 'linear' and b is not None:
+            x = torch.addmm(b.unsqueeze(0), x, w.t())
+        else:
+            x = x.matmul(w.t())
+            x = bias_act.bias_act(x, b, act=self.activation)
+        return x
+
+    def extra_repr(self):
+        return f'in_features={self.in_features:d}, out_features={self.out_features:d}, activation={self.activation:s}'
+
+#----------------------------------------------------------------------------
+
+@persistence.persistent_class
+class MappingNetwork(torch.nn.Module):
+    def __init__(self,
+        z_dim,                      # Input latent (Z) dimensionality.
+        c_dim,                      # Conditioning label (C) dimensionality, 0 = no labels.
+        w_dim,                      # Intermediate latent (W) dimensionality.
+        num_ws,                     # Number of intermediate latents to output.
+        num_layers      = 2,        # Number of mapping layers.
+        lr_multiplier   = 0.01,     # Learning rate multiplier for the mapping layers.
+        w_avg_beta      = 0.998,    # Decay for tracking the moving average of W during training.
+    ):
+        super().__init__()
+        self.z_dim = z_dim
+        self.c_dim = c_dim
+        self.w_dim = w_dim
+        self.num_ws = num_ws
+        self.num_layers = num_layers
+        self.w_avg_beta = w_avg_beta
+
+        # Construct layers.
+        self.embed = FullyConnectedLayer(self.c_dim, self.w_dim) if self.c_dim > 0 else None
+        features = [self.z_dim + (self.w_dim if self.c_dim > 0 else 0)] + [self.w_dim] * self.num_layers
+        for idx, in_features, out_features in zip(range(num_layers), features[:-1], features[1:]):
+            layer = FullyConnectedLayer(in_features, out_features, activation='lrelu', lr_multiplier=lr_multiplier)
+            setattr(self, f'fc{idx}', layer)
+        self.register_buffer('w_avg', torch.zeros([w_dim]))
+
+    def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, update_emas=False):
+        misc.assert_shape(z, [None, self.z_dim])
+        if truncation_cutoff is None:
+            truncation_cutoff = self.num_ws
+
+        # Embed, normalize, and concatenate inputs.
+        x = z.to(torch.float32)
+        x = x * (x.square().mean(1, keepdim=True) + 1e-8).rsqrt()
+        if self.c_dim > 0:
+            misc.assert_shape(c, [None, self.c_dim])
+            y = self.embed(c.to(torch.float32))
+            y = y * (y.square().mean(1, keepdim=True) + 1e-8).rsqrt()
+            x = torch.cat([x, y], dim=1) if x is not None else y
+
+        # Execute layers.
+        for idx in range(self.num_layers):
+            x = getattr(self, f'fc{idx}')(x)
+
+        # Update moving average of W.
+        if update_emas:
+            self.w_avg.copy_(x.detach().mean(dim=0).lerp(self.w_avg, self.w_avg_beta))
+
+        # Broadcast and apply truncation.
+        x = x.unsqueeze(1).repeat([1, self.num_ws, 1])
+        if truncation_psi != 1:
+            x[:, :truncation_cutoff] = self.w_avg.lerp(x[:, :truncation_cutoff], truncation_psi)
+        return x
+
+    def extra_repr(self):
+        return f'z_dim={self.z_dim:d}, c_dim={self.c_dim:d}, w_dim={self.w_dim:d}, num_ws={self.num_ws:d}'
+
+#----------------------------------------------------------------------------
+
+@persistence.persistent_class
+class SynthesisInput(torch.nn.Module):
+    def __init__(self,
+        w_dim,          # Intermediate latent (W) dimensionality.
+        channels,       # Number of output channels.
+        size,           # Output spatial size: int or [width, height].
+        sampling_rate,  # Output sampling rate.
+        bandwidth,      # Output bandwidth.
+    ):
+        super().__init__()
+        self.w_dim = w_dim
+        self.channels = channels
+        self.size = np.broadcast_to(np.asarray(size), [2])
+        self.sampling_rate = sampling_rate
+        self.bandwidth = bandwidth
+
+        # Draw random frequencies from uniform 2D disc.
+        freqs = torch.randn([self.channels, 2])
+        radii = freqs.square().sum(dim=1, keepdim=True).sqrt()
+        freqs /= radii * radii.square().exp().pow(0.25)
+        freqs *= bandwidth
+        phases = torch.rand([self.channels]) - 0.5
+
+        # Setup parameters and buffers.
+        self.weight = torch.nn.Parameter(torch.randn([self.channels, self.channels]))
+        self.affine = FullyConnectedLayer(w_dim, 4, weight_init=0, bias_init=[1,0,0,0])
+        self.register_buffer('transform', torch.eye(3, 3)) # User-specified inverse transform wrt. resulting image.
+        self.register_buffer('freqs', freqs)
+        self.register_buffer('phases', phases)
+
+    def forward(self, w):
+        # Introduce batch dimension.
+        transforms = self.transform.unsqueeze(0) # [batch, row, col]
+        freqs = self.freqs.unsqueeze(0) # [batch, channel, xy]
+        phases = self.phases.unsqueeze(0) # [batch, channel]
+
+        # Apply learned transformation.
+        t = self.affine(w) # t = (r_c, r_s, t_x, t_y)
+        t = t / t[:, :2].norm(dim=1, keepdim=True) # t' = (r'_c, r'_s, t'_x, t'_y)
+        m_r = torch.eye(3, device=w.device).unsqueeze(0).repeat([w.shape[0], 1, 1]) # Inverse rotation wrt. resulting image.
+        m_r[:, 0, 0] = t[:, 0]  # r'_c
+        m_r[:, 0, 1] = -t[:, 1] # r'_s
+        m_r[:, 1, 0] = t[:, 1]  # r'_s
+        m_r[:, 1, 1] = t[:, 0]  # r'_c
+        m_t = torch.eye(3, device=w.device).unsqueeze(0).repeat([w.shape[0], 1, 1]) # Inverse translation wrt. resulting image.
+        m_t[:, 0, 2] = -t[:, 2] # t'_x
+        m_t[:, 1, 2] = -t[:, 3] # t'_y
+        transforms = m_r @ m_t @ transforms # First rotate resulting image, then translate, and finally apply user-specified transform.
+
+        # Transform frequencies.
+        phases = phases + (freqs @ transforms[:, :2, 2:]).squeeze(2)
+        freqs = freqs @ transforms[:, :2, :2]
+
+        # Dampen out-of-band frequencies that may occur due to the user-specified transform.
+        amplitudes = (1 - (freqs.norm(dim=2) - self.bandwidth) / (self.sampling_rate / 2 - self.bandwidth)).clamp(0, 1)
+
+        # Construct sampling grid.
+        theta = torch.eye(2, 3, device=w.device)
+        theta[0, 0] = 0.5 * self.size[0] / self.sampling_rate
+        theta[1, 1] = 0.5 * self.size[1] / self.sampling_rate
+        grids = torch.nn.functional.affine_grid(theta.unsqueeze(0), [1, 1, self.size[1], self.size[0]], align_corners=False)
+
+        # Compute Fourier features.
+        x = (grids.unsqueeze(3) @ freqs.permute(0, 2, 1).unsqueeze(1).unsqueeze(2)).squeeze(3) # [batch, height, width, channel]
+        x = x + phases.unsqueeze(1).unsqueeze(2)
+        x = torch.sin(x * (np.pi * 2))
+        x = x * amplitudes.unsqueeze(1).unsqueeze(2)
+
+        # Apply trainable mapping.
+        weight = self.weight / np.sqrt(self.channels)
+        x = x @ weight.t()
+
+        # Ensure correct shape.
+        x = x.permute(0, 3, 1, 2) # [batch, channel, height, width]
+        misc.assert_shape(x, [w.shape[0], self.channels, int(self.size[1]), int(self.size[0])])
+        return x
+
+    def extra_repr(self):
+        return '\n'.join([
+            f'w_dim={self.w_dim:d}, channels={self.channels:d}, size={list(self.size)},',
+            f'sampling_rate={self.sampling_rate:g}, bandwidth={self.bandwidth:g}'])
+
+#----------------------------------------------------------------------------
+
+@persistence.persistent_class
+class SynthesisLayer(torch.nn.Module):
+    def __init__(self,
+        w_dim,                          # Intermediate latent (W) dimensionality.
+        is_torgb,                       # Is this the final ToRGB layer?
+        is_critically_sampled,          # Does this layer use critical sampling?
+        use_fp16,                       # Does this layer use FP16?
+
+        # Input & output specifications.
+        in_channels,                    # Number of input channels.
+        out_channels,                   # Number of output channels.
+        in_size,                        # Input spatial size: int or [width, height].
+        out_size,                       # Output spatial size: int or [width, height].
+        in_sampling_rate,               # Input sampling rate (s).
+        out_sampling_rate,              # Output sampling rate (s).
+        in_cutoff,                      # Input cutoff frequency (f_c).
+        out_cutoff,                     # Output cutoff frequency (f_c).
+        in_half_width,                  # Input transition band half-width (f_h).
+        out_half_width,                 # Output Transition band half-width (f_h).
+
+        # Hyperparameters.
+        conv_kernel         = 3,        # Convolution kernel size. Ignored for final the ToRGB layer.
+        filter_size         = 6,        # Low-pass filter size relative to the lower resolution when up/downsampling.
+        lrelu_upsampling    = 2,        # Relative sampling rate for leaky ReLU. Ignored for final the ToRGB layer.
+        use_radial_filters  = False,    # Use radially symmetric downsampling filter? Ignored for critically sampled layers.
+        conv_clamp          = 256,      # Clamp the output to [-X, +X], None = disable clamping.
+        magnitude_ema_beta  = 0.999,    # Decay rate for the moving average of input magnitudes.
+    ):
+        super().__init__()
+        self.w_dim = w_dim
+        self.is_torgb = is_torgb
+        self.is_critically_sampled = is_critically_sampled
+        self.use_fp16 = use_fp16
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.in_size = np.broadcast_to(np.asarray(in_size), [2])
+        self.out_size = np.broadcast_to(np.asarray(out_size), [2])
+        self.in_sampling_rate = in_sampling_rate
+        self.out_sampling_rate = out_sampling_rate
+        self.tmp_sampling_rate = max(in_sampling_rate, out_sampling_rate) * (1 if is_torgb else lrelu_upsampling)
+        self.in_cutoff = in_cutoff
+        self.out_cutoff = out_cutoff
+        self.in_half_width = in_half_width
+        self.out_half_width = out_half_width
+        self.conv_kernel = 1 if is_torgb else conv_kernel
+        self.conv_clamp = conv_clamp
+        self.magnitude_ema_beta = magnitude_ema_beta
+
+        # Setup parameters and buffers.
+        self.affine = FullyConnectedLayer(self.w_dim, self.in_channels, bias_init=1)
+        self.weight = torch.nn.Parameter(torch.randn([self.out_channels, self.in_channels, self.conv_kernel, self.conv_kernel]))
+        self.bias = torch.nn.Parameter(torch.zeros([self.out_channels]))
+        self.register_buffer('magnitude_ema', torch.ones([]))
+
+        # Design upsampling filter.
+        self.up_factor = int(np.rint(self.tmp_sampling_rate / self.in_sampling_rate))
+        assert self.in_sampling_rate * self.up_factor == self.tmp_sampling_rate
+        self.up_taps = filter_size * self.up_factor if self.up_factor > 1 and not self.is_torgb else 1
+        self.register_buffer('up_filter', self.design_lowpass_filter(
+            numtaps=self.up_taps, cutoff=self.in_cutoff, width=self.in_half_width*2, fs=self.tmp_sampling_rate))
+
+        # Design downsampling filter.
+        self.down_factor = int(np.rint(self.tmp_sampling_rate / self.out_sampling_rate))
+        assert self.out_sampling_rate * self.down_factor == self.tmp_sampling_rate
+        self.down_taps = filter_size * self.down_factor if self.down_factor > 1 and not self.is_torgb else 1
+        self.down_radial = use_radial_filters and not self.is_critically_sampled
+        self.register_buffer('down_filter', self.design_lowpass_filter(
+            numtaps=self.down_taps, cutoff=self.out_cutoff, width=self.out_half_width*2, fs=self.tmp_sampling_rate, radial=self.down_radial))
+
+        # Compute padding.
+        pad_total = (self.out_size - 1) * self.down_factor + 1 # Desired output size before downsampling.
+        pad_total -= (self.in_size + self.conv_kernel - 1) * self.up_factor # Input size after upsampling.
+        pad_total += self.up_taps + self.down_taps - 2 # Size reduction caused by the filters.
+        pad_lo = (pad_total + self.up_factor) // 2 # Shift sample locations according to the symmetric interpretation (Appendix C.3).
+        pad_hi = pad_total - pad_lo
+        self.padding = [int(pad_lo[0]), int(pad_hi[0]), int(pad_lo[1]), int(pad_hi[1])]
+
+    def forward(self, x, w, noise_mode='random', force_fp32=False, update_emas=False):
+        assert noise_mode in ['random', 'const', 'none'] # unused
+        misc.assert_shape(x, [None, self.in_channels, int(self.in_size[1]), int(self.in_size[0])])
+        misc.assert_shape(w, [x.shape[0], self.w_dim])
+
+        # Track input magnitude.
+        if update_emas:
+            with torch.autograd.profiler.record_function('update_magnitude_ema'):
+                magnitude_cur = x.detach().to(torch.float32).square().mean()
+                self.magnitude_ema.copy_(magnitude_cur.lerp(self.magnitude_ema, self.magnitude_ema_beta))
+        input_gain = self.magnitude_ema.rsqrt()
+
+        # Execute affine layer.
+        styles = self.affine(w)
+        if self.is_torgb:
+            weight_gain = 1 / np.sqrt(self.in_channels * (self.conv_kernel ** 2))
+            styles = styles * weight_gain
+
+        # Execute modulated conv2d.
+        dtype = torch.float16 if (self.use_fp16 and not force_fp32 and x.device.type == 'cuda') else torch.float32
+        x = modulated_conv2d(x=x.to(dtype), w=self.weight, s=styles,
+            padding=self.conv_kernel-1, demodulate=(not self.is_torgb), input_gain=input_gain)
+
+        # Execute bias, filtered leaky ReLU, and clamping.
+        gain = 1 if self.is_torgb else np.sqrt(2)
+        slope = 1 if self.is_torgb else 0.2
+        x = filtered_lrelu.filtered_lrelu(x=x, fu=self.up_filter, fd=self.down_filter, b=self.bias.to(x.dtype),
+            up=self.up_factor, down=self.down_factor, padding=self.padding, gain=gain, slope=slope, clamp=self.conv_clamp)
+
+        # Ensure correct shape and dtype.
+        misc.assert_shape(x, [None, self.out_channels, int(self.out_size[1]), int(self.out_size[0])])
+        assert x.dtype == dtype
+        return x
+
+    @staticmethod
+    def design_lowpass_filter(numtaps, cutoff, width, fs, radial=False):
+        assert numtaps >= 1
+
+        # Identity filter.
+        if numtaps == 1:
+            return None
+
+        # Separable Kaiser low-pass filter.
+        if not radial:
+            f = scipy.signal.firwin(numtaps=numtaps, cutoff=cutoff, width=width, fs=fs)
+            return torch.as_tensor(f, dtype=torch.float32)
+
+        # Radially symmetric jinc-based filter.
+        x = (np.arange(numtaps) - (numtaps - 1) / 2) / fs
+        r = np.hypot(*np.meshgrid(x, x))
+        f = scipy.special.j1(2 * cutoff * (np.pi * r)) / (np.pi * r)
+        beta = scipy.signal.kaiser_beta(scipy.signal.kaiser_atten(numtaps, width / (fs / 2)))
+        w = np.kaiser(numtaps, beta)
+        f *= np.outer(w, w)
+        f /= np.sum(f)
+        return torch.as_tensor(f, dtype=torch.float32)
+
+    def extra_repr(self):
+        return '\n'.join([
+            f'w_dim={self.w_dim:d}, is_torgb={self.is_torgb},',
+            f'is_critically_sampled={self.is_critically_sampled}, use_fp16={self.use_fp16},',
+            f'in_sampling_rate={self.in_sampling_rate:g}, out_sampling_rate={self.out_sampling_rate:g},',
+            f'in_cutoff={self.in_cutoff:g}, out_cutoff={self.out_cutoff:g},',
+            f'in_half_width={self.in_half_width:g}, out_half_width={self.out_half_width:g},',
+            f'in_size={list(self.in_size)}, out_size={list(self.out_size)},',
+            f'in_channels={self.in_channels:d}, out_channels={self.out_channels:d}'])
+
+#----------------------------------------------------------------------------
+
+@persistence.persistent_class
+class SynthesisNetwork(torch.nn.Module):
+    def __init__(self,
+        w_dim,                          # Intermediate latent (W) dimensionality.
+        img_resolution,                 # Output image resolution.
+        img_channels,                   # Number of color channels.
+        channel_base        = 32768,    # Overall multiplier for the number of channels.
+        channel_max         = 512,      # Maximum number of channels in any layer.
+        num_layers          = 14,       # Total number of layers, excluding Fourier features and ToRGB.
+        num_critical        = 2,        # Number of critically sampled layers at the end.
+        first_cutoff        = 2,        # Cutoff frequency of the first layer (f_{c,0}).
+        first_stopband      = 2**2.1,   # Minimum stopband of the first layer (f_{t,0}).
+        last_stopband_rel   = 2**0.3,   # Minimum stopband of the last layer, expressed relative to the cutoff.
+        margin_size         = 10,       # Number of additional pixels outside the image.
+        output_scale        = 0.25,     # Scale factor for the output image.
+        num_fp16_res        = 4,        # Use FP16 for the N highest resolutions.
+        **layer_kwargs,                 # Arguments for SynthesisLayer.
+    ):
+        super().__init__()
+        self.w_dim = w_dim
+        self.num_ws = num_layers + 2
+        self.img_resolution = img_resolution
+        self.img_channels = img_channels
+        self.num_layers = num_layers
+        self.num_critical = num_critical
+        self.margin_size = margin_size
+        self.output_scale = output_scale
+        self.num_fp16_res = num_fp16_res
+
+        # Geometric progression of layer cutoffs and min. stopbands.
+        last_cutoff = self.img_resolution / 2 # f_{c,N}
+        last_stopband = last_cutoff * last_stopband_rel # f_{t,N}
+        exponents = np.minimum(np.arange(self.num_layers + 1) / (self.num_layers - self.num_critical), 1)
+        cutoffs = first_cutoff * (last_cutoff / first_cutoff) ** exponents # f_c[i]
+        stopbands = first_stopband * (last_stopband / first_stopband) ** exponents # f_t[i]
+
+        # Compute remaining layer parameters.
+        sampling_rates = np.exp2(np.ceil(np.log2(np.minimum(stopbands * 2, self.img_resolution)))) # s[i]
+        half_widths = np.maximum(stopbands, sampling_rates / 2) - cutoffs # f_h[i]
+        sizes = sampling_rates + self.margin_size * 2
+        sizes[-2:] = self.img_resolution
+        channels = np.rint(np.minimum((channel_base / 2) / cutoffs, channel_max))
+        channels[-1] = self.img_channels
+
+        # Construct layers.
+        self.input = SynthesisInput(
+            w_dim=self.w_dim, channels=int(channels[0]), size=int(sizes[0]),
+            sampling_rate=sampling_rates[0], bandwidth=cutoffs[0])
+        self.layer_names = []
+        for idx in range(self.num_layers + 1):
+            prev = max(idx - 1, 0)
+            is_torgb = (idx == self.num_layers)
+            is_critically_sampled = (idx >= self.num_layers - self.num_critical)
+            use_fp16 = (sampling_rates[idx] * (2 ** self.num_fp16_res) > self.img_resolution)
+            layer = SynthesisLayer(
+                w_dim=self.w_dim, is_torgb=is_torgb, is_critically_sampled=is_critically_sampled, use_fp16=use_fp16,
+                in_channels=int(channels[prev]), out_channels= int(channels[idx]),
+                in_size=int(sizes[prev]), out_size=int(sizes[idx]),
+                in_sampling_rate=int(sampling_rates[prev]), out_sampling_rate=int(sampling_rates[idx]),
+                in_cutoff=cutoffs[prev], out_cutoff=cutoffs[idx],
+                in_half_width=half_widths[prev], out_half_width=half_widths[idx],
+                **layer_kwargs)
+            name = f'L{idx}_{layer.out_size[0]}_{layer.out_channels}'
+            setattr(self, name, layer)
+            self.layer_names.append(name)
+
+    def forward(self, ws, **layer_kwargs):
+        misc.assert_shape(ws, [None, self.num_ws, self.w_dim])
+        ws = ws.to(torch.float32).unbind(dim=1)
+
+        # Execute layers.
+        x = self.input(ws[0])
+        for name, w in zip(self.layer_names, ws[1:]):
+            x = getattr(self, name)(x, w, **layer_kwargs)
+        if self.output_scale != 1:
+            x = x * self.output_scale
+
+        # Ensure correct shape and dtype.
+        misc.assert_shape(x, [None, self.img_channels, self.img_resolution, self.img_resolution])
+        x = x.to(torch.float32)
+        return x
+
+    def extra_repr(self):
+        return '\n'.join([
+            f'w_dim={self.w_dim:d}, num_ws={self.num_ws:d},',
+            f'img_resolution={self.img_resolution:d}, img_channels={self.img_channels:d},',
+            f'num_layers={self.num_layers:d}, num_critical={self.num_critical:d},',
+            f'margin_size={self.margin_size:d}, num_fp16_res={self.num_fp16_res:d}'])
+
+#----------------------------------------------------------------------------
+
+@persistence.persistent_class
+class Generator(torch.nn.Module):
+    def __init__(self,
+        z_dim,                      # Input latent (Z) dimensionality.
+        c_dim,                      # Conditioning label (C) dimensionality.
+        w_dim,                      # Intermediate latent (W) dimensionality.
+        img_resolution,             # Output resolution.
+        img_channels,               # Number of output color channels.
+        mapping_kwargs      = {},   # Arguments for MappingNetwork.
+        **synthesis_kwargs,         # Arguments for SynthesisNetwork.
+    ):
+        super().__init__()
+        self.z_dim = z_dim
+        self.c_dim = c_dim
+        self.w_dim = w_dim
+        self.img_resolution = img_resolution
+        self.img_channels = img_channels
+        self.synthesis = SynthesisNetwork(w_dim=w_dim, img_resolution=img_resolution, img_channels=img_channels, **synthesis_kwargs)
+        self.num_ws = self.synthesis.num_ws
+        self.mapping = MappingNetwork(z_dim=z_dim, c_dim=c_dim, w_dim=w_dim, num_ws=self.num_ws, **mapping_kwargs)
+
+    def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, update_emas=False, **synthesis_kwargs):
+        ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff, update_emas=update_emas)
+        img = self.synthesis(ws, update_emas=update_emas, **synthesis_kwargs)
+        return img
+
+#----------------------------------------------------------------------------

data/textile_annotated_files/final_sim_seeds0000-100000.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2dde6f1168825424eed5aa328c6442873df35637e74786f9c9af956c9c0a97ed
+size 7886477

data/textile_annotated_files/hsv_info.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36016190017d38cde71f35c267df9f6b0ab40d74ce17e022195e96d32d2f2f71
+size 1112635

data/textile_annotated_files/seeds0000-100000.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2dd321307f3d332193630f823a3e0db79d533156dfbc0d446eab0d5c212b1360
+size 630151183

data/textile_annotated_files/seeds0000-100000_S.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8763836ea1142f6f2e3d36b7fe92bcf9a4549e9ef8e0a83a02b4772d64e95d54
+size 3178623075

data/textile_annotated_files/top_three_colours.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2dde6f1168825424eed5aa328c6442873df35637e74786f9c9af956c9c0a97ed
+size 7886477

data/textile_model_files/network-snapshot-005000.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:717bdd11871c0383d6e28c54b2d61cd485ef236dd1de34d3194323c843b11b62
+size 343479704

pages/5_Textiles_Disentanglement.py

@@ -0,0 +1,178 @@
+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 [Oxford Vases Dataset](https://www.robots.ox.ac.uk/~vgg/data/oxbuildings/).
+    """,
+    unsafe_allow_html=False,)
+    
+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 =['AMPHORA']
+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('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), max_selections=2, default=[labels[2], labels[3]])
+
+        st.write('**Choose a latent space to disentangle**')
+        space_id = st.selectbox('Space:', tuple(['W', 'Z']))
+
+        choose_text_button = st.form_submit_button('Choose the defined concept and space to disentangle')
+          
+        if choose_text_button:
+            concept_ids = list(concept_ids)
+            st.session_state.concept_ids = concept_ids
+            space_id = str(space_id)
+            st.session_state.space_id = space_id
+        # 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([5,1])
+output_col_1, output_col_2 = st.columns([5,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:
+    separation_vector, number_important_features, imp_nodes, performance = get_separation_space(concept_ids, annotations, ann_df, latent_space=st.session_state.space_id, samples=150)
+    # st.write(f'Class ID {input_id} - {input_label}: {pred_prob*100:.3f}% confidence')
+    st.write('Concept vector', separation_vector)
+    header_col_1.write(f'Concept {st.session_state.concept_ids} - Space {st.session_state.space_id} - Number of relevant nodes: {number_important_features} - Val classification performance: {performance}')# - Nodes {",".join(list(imp_nodes))}')
+
+# ----------------------------- 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, 20000)
+            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 concept'):
+        st.write('**Set range of change**')
+        chosen_epsilon_input = st.empty()
+        epsilon = chosen_epsilon_input.number_input('Lambda:', min_value=1, step=1)
+        epsilon_button = st.form_submit_button('Choose the defined lambda')
+        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')
+        
+
+# ---------------------------- DISPLAY COL 2 ROW 1 ------------------------------
+
+#model = torch.load('./data/model_files/pytorch_model.bin', map_location=torch.device('cpu'))
+with dnnlib.util.open_url('./data/vase_model_files/network-snapshot-003800.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)
+
+top_pred = ann_df.loc[st.session_state.image_id, labels].astype(float).idxmax()
+# 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, predicted as {top_pred}')
+
+
+images, lambdas = regenerate_images(model, original_image_vec, separation_vector, min_epsilon=-(int(epsilon)), max_epsilon=int(epsilon), latent_space=st.session_state.space_id, layers=layers)
+
+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)}')