lots of changes

.gitignore

@@ -32,7 +32,10 @@ git-large-file
 deta_drive.py
 secret_keys.py
 
-data/old
+data/old/
+archive/
+figures/
+colors_test/
 # Large files
 # data/preprocessed_image_net/
 # data/activation/*.pkl

DisentanglementBase.py

@@ -1,9 +1,11 @@
+#!/usr/bin/env python
+
 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.linear_model import LogisticRegression, LinearRegression
 from sklearn.model_selection import train_test_split
 
 from tqdm import tqdm
@@ -27,8 +29,50 @@ sys.path.append('.')
 import dnnlib 
 import legacy
 
+
+def hex2rgb(hex_value):
+    h = hex_value.strip("#") 
+    rgb = tuple(int(h[i:i+2], 16) for i in (0, 2, 4))
+    return rgb
+
+def rgb2hsv(r, g, b):
+    # Normalize R, G, B values
+    r, g, b = r / 255.0, g / 255.0, b / 255.0
+    
+    # h, s, v = hue, saturation, value
+    max_rgb = max(r, g, b)    
+    min_rgb = min(r, g, b)   
+    difference = max_rgb-min_rgb 
+    
+    # if max_rgb and max_rgb are equal then h = 0
+    if max_rgb == min_rgb:
+        h = 0
+    
+    # if max_rgb==r then h is computed as follows
+    elif max_rgb == r:
+        h = (60 * ((g - b) / difference) + 360) % 360
+    
+    # if max_rgb==g then compute h as follows
+    elif max_rgb == g:
+        h = (60 * ((b - r) / difference) + 120) % 360
+    
+    # if max_rgb=b then compute h
+    elif max_rgb == b:
+        h = (60 * ((r - g) / difference) + 240) % 360
+    
+    # if max_rgb==zero then s=0
+    if max_rgb == 0:
+        s = 0
+    else:
+        s = (difference / max_rgb) * 100
+    
+    # compute v
+    v = max_rgb * 100
+    # return rounded values of H, S and V
+    return tuple(map(round, (h, s, v)))
+ 
 class DisentanglementBase:
-    def __init__(self, repo_folder, model, annotations, df, space, colors_list, compute_s):
+    def __init__(self, repo_folder, model, annotations, df, space, colors_list, compute_s=False, variable='H1', categorical=True):
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         print('Using device', self.device)
         self.repo_folder = repo_folder
@@ -36,6 +80,8 @@ class DisentanglementBase:
         self.annotations = annotations
         self.df = df
         self.space = space
+        self.categorical = categorical
+        self.variable = variable
         
         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', 
@@ -49,7 +95,6 @@ class DisentanglementBase:
         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.
@@ -60,17 +105,30 @@ class DisentanglementBase:
         :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['H1'] = self.df['top1col'].map(lambda x: rgb2hsv(*hex2rgb(x))[0])
+        self.df['H2'] = self.df['top2col'].map(lambda x: rgb2hsv(*hex2rgb(x))[0])
+        self.df['H3'] = self.df['top3col'].map(lambda x: rgb2hsv(*hex2rgb(x))[0])
+        
+        self.df['S1'] = self.df['top1col'].map(lambda x: rgb2hsv(*hex2rgb(x))[1])
+        self.df['S2'] = self.df['top2col'].map(lambda x: rgb2hsv(*hex2rgb(x))[1])
+        self.df['S3'] = self.df['top3col'].map(lambda x: rgb2hsv(*hex2rgb(x))[1])
+        
+        self.df['V1'] = self.df['top1col'].map(lambda x: rgb2hsv(*hex2rgb(x))[2])
+        self.df['V2'] = self.df['top2col'].map(lambda x: rgb2hsv(*hex2rgb(x))[2])
+        self.df['V3'] = self.df['top3col'].map(lambda x: rgb2hsv(*hex2rgb(x))[2])
         
-        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])
+        print('Adding RGB encoding')
+        self.df['R1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[0])
+        self.df['R2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[0])
+        self.df['R3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[0])
         
-        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])
+        self.df['G1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[1])
+        self.df['G2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[1])
+        self.df['G3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[1])
+        
+        self.df['B1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[2])
+        self.df['B2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[2])
+        self.df['B3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[2])
     
     def get_s_space(self):
         """
@@ -89,7 +147,7 @@ class DisentanglementBase:
             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())
+                s.append(getattr(self.model.synthesis, layer).affine(W[0, i].unsqueeze(0)).cpu().numpy())
 
             ss.append(s)
         self.annotations['s_vectors'] = ss
@@ -116,17 +174,32 @@ class DisentanglementBase:
         print('Shape embedding:', X.shape)
         return X
     
-    def get_train_val(self, var='H1', cat=True):
+    def get_train_val(self, extremes=False):
         X = self.get_encoded_latent()
-        y = np.array(self.df[var].values)
-        if cat:
+        y = np.array(self.df[self.variable].values)
+        if self.categorical:
+            bins = [(x-1) * 360 / (len(self.colors_list) - 1)  if x != 1 
+                    else 1 for x in range(len(self.colors_list) + 1)]
+            bins[0] = 0
             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')
+                            bins=bins,
+                            labels=self.colors_list,
+                            include_lowest=True
+                            )
+            print(y_cat.value_counts())
             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)
+            if extremes:
+                # Calculate the number of elements to consider (10% of array size)
+                num_elements = int(0.2 * len(y))
+                # Get indices of the top num_elements maximum values
+                top_indices = np.argpartition(array, -num_elements)[-num_elements:]
+                bottom_indices = np.argpartition(array, -num_elements)[:num_elements]
+                y_ext = y[top_indices + bottom_indices, :]
+                X_ext = X[top_indices + bottom_indices, :]
+                x_train, x_val, y_train, y_val = train_test_split(X_ext, y_ext, 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):
@@ -148,17 +221,23 @@ class DisentanglementBase:
         """
         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)
+        if self.categorical:
+            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(svc.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_)
+        else:
+            clf = LinearRegression()
             clf.fit(x_train, y_train)
-            print('Val performance logistic regression', np.round(clf.score(x_val, y_val), 2))
+            print('Val performance linear 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 = []
@@ -327,7 +406,7 @@ class DisentanglementBase:
             
         return img
 
-    def generate_changes(self, seed, separation_vector, min_epsilon=-3, max_epsilon=3, count=5, savefig=True, feature=None, method=None):
+    def generate_changes(self, seed, separation_vector, min_epsilon=-3, max_epsilon=3, count=5, savefig=True, feature=None, method=None, save_separately=False):
         """
         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
@@ -348,14 +427,13 @@ class DisentanglementBase:
         lambdas = np.linspace(min_epsilon, max_epsilon, count)
         images = []
         # Generate images.
-        for _, lambd in enumerate(tqdm(lambdas)):
+        for _, lambd in enumerate(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])])
@@ -365,42 +443,117 @@ class DisentanglementBase:
                 axs[i].imshow(image)
                 axs[i].set_title(np.round(lambd, 2))
             plt.tight_layout()
-            plt.savefig(join(self.repo_folder, 'figures', name+'.jpg'))
+            plt.savefig(join(self.repo_folder, 'figures', 'examples', name+'.jpg'))
             plt.close()
+            
+            if save_separately:
+                for i, (image, lambd) in enumerate(zip(images, lambdas)):
+                    plt.imshow(image)
+                    plt.tight_layout()
+                    plt.savefig(join(self.repo_folder, 'figures', 'examples', name + '_' + str(lambd) + '.jpg'))
+                    plt.close()
+            
         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 
+        if self.categorical:
+            if feature_id == 0:
+                hue_low = 0
+                hue_high = 1
+            elif feature_id == 1:
+                hue_low = 1
+                hue_high = (feature_id - 1) * 360 / (len(self.colors_list) - 1)  
+            else:
+                hue_low = (feature_id - 1) * 360 / (len(self.colors_list) - 1) 
+                hue_high = feature_id * 360 / (len(self.colors_list) - 1)
         
-        matches = 0
+            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)
-            try:
-                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')
+            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)
+                try:
+                    colors_negative = extract_color(images[0], 5, 1, None)
+                    h0, s0, v0 = rgb2hsv(*hex2rgb(colors_negative[0]))
                 
-                colors_positive = extract_color(images[2], 5, 1, None)
-                h2, s2, v2 = ImageColor.getcolor(colors_positive[0], 'HSV')
+                    colors_orig = extract_color(images[1], 5, 1, None)
+                    h1, s1, v1 = rgb2hsv(*hex2rgb(colors_orig[0])) 
+                    
+                    colors_positive = extract_color(images[2], 5, 1, None)
+                    h2, s2, v2 = rgb2hsv(*hex2rgb(colors_positive[0]))
+                    
+                    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
                 
-                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
+                except Exception as e:
+                    print(e)
+            
+            return np.round(matches / samples, 2)
             
-            except Exception as e:
-                print(e)
+        else:
+            increase = 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)
+                try:
+                    colors_negative = extract_color(images[0], 5, 1, None)
+                    r0, g0, b0 = hex2rgb(colors_negative[0])
+                    h0, s0, v0 = rgb2hsv(*hex2rgb(colors_negative[0]))
                 
-        return np.round(matches / samples, 2)
-    
+                    colors_orig = extract_color(images[1], 5, 1, None)
+                    r1, g1, b1 = hex2rgb(colors_orig[0])
+                    h1, s1, v1 = rgb2hsv(*hex2rgb(colors_orig[0]))
+                    
+                    colors_positive = extract_color(images[2], 5, 1, None)
+                    r2, g2, b2 = hex2rgb(colors_positive[0])
+                    h2, s2, v2 = rgb2hsv(*hex2rgb(colors_positive[0]))
+                    
+                    if 's' in self.variable.lower():
+                        increase += max(0, s2 - s1)
+                    elif 'v' in self.variable.lower():
+                        increase += max(0, v2 - v1)
+                    elif 'r' in self.variable.lower():
+                        increase += max(0, r2 - r1)
+                    elif 'g' in self.variable.lower():
+                        increase += max(0, g2 - g1)
+                    elif 'b' in self.variable.lower():
+                        increase += max(0, b2 - b1)
+                    else:
+                        raise('Continous variable not allowed, choose between RGB or SV')
+                except Exception as e:
+                    print(e)
+            
+            return np.round(increase / samples, 2)
+            
 
+def continous_experiment(name, var, repo_folder, model, annotations, df, space, colors_list, kwargs):
+    scores = []
+    print(f'Launching {name} experiment')
+    disentanglemnet_exp = DisentanglementBase(repo_folder, model, annotations, df, space=space, colors_list=colors_list, compute_s=False, variable=var, categorical=False)
+    for extr in kwargs['extremes']:                
+        separation_vector = disentanglemnet_exp.InterFaceGAN_separation_vector()
+        print(f'Generating images with variations for {name}')
+        for s in range(30):
+            seed = random.randint(0,100000)
+            for eps in kwargs['max_lambda']:
+                disentanglemnet_exp.generate_changes(seed, separation_vector, min_epsilon=-eps, max_epsilon=eps, savefig=True, feature=name, method= 'InterFaceGAN_' + str(extr))
+                                    
+        print('Finally obtaining verification score')
+        for verif in kwargs['lambda_verif']:    
+            score = disentanglemnet_exp.get_verification_score(separation_vector, 0, samples=kwargs['samples'], lambd=verif, savefig=False, feature=name, method='InterFaceGAN_' + str(extr))
+            print(f'Score for method InterfaceGAN on {name}:', score)
+                                
+            scores.append([space, 'InterFaceGAN', name, score, 'extremes method:' + str(extr) + 'verification lambda:' + str(verif), ', '.join(list(separation_vector.astype(str)))])
+            
+        score_df = pd.DataFrame(scores, columns=['space', 'method', 'color', 'score', 'kwargs', 'vector'])
+        print(score_df)
+        score_df.to_csv(join(repo_folder, f'data/scores_{name}.csv')) 
+            
 def main():
     repo_folder = '.'
     annotations_file = join(repo_folder, 'data/textile_annotated_files/seeds0000-100000_S.pkl')
@@ -417,26 +570,56 @@ def main():
     colors_list = ['Red', 'Orange', 'Yellow', 'Yellow Green', 'Chartreuse Green',
                    'Kelly Green', 'Green Blue Seafoam', 'Cyan Blue',
                    'Warm Blue', 'Indigo', 'Purple Magenta', 'Magenta Pink']
+    colors_list = ['Gray', 'Red Orange', 'Yellow', 'Green', 'Light Blue',
+                   'Blue', 'Purple', 'Pink']
     
     scores = []
-    kwargs = {'CL method':['LR', 'SVM'], 'C':[0.1, 1], 'sign':[True, False], 'num_factors':[1, 10, 20], 'cutout': [None], 'max_lambda':[18, 3], 'samples':50, 'lambda_verif':[10, 5, 3]}
+    kwargs = {'CL method':['LR', 'SVM'], 'C':[0.1, 1], 'sign':[True, False], 
+              'num_factors':[1, 5, 10, 20], 'cutout': [None], 'max_lambda':[18, 6], 
+              'samples':30, 'lambda_verif':[14, 7], 'extremes':[True, False]}
+    continuous = False
+    specific_examples = [53139, 99376, 16, 99585, 40851, 70, 17703, 44, 52628,
+                        99884, 52921, 46180, 19995, 40920,  554]
     
-    for space in ['w', 'z', 's']:
+    if specific_examples is not None:
+        disentanglemnet_exp = DisentanglementBase(repo_folder, model, annotations, df, space='w', colors_list=colors_list, compute_s=False)
+
+        separation_vectors = disentanglemnet_exp.StyleSpace_separation_vector(sign=True, num_factors=10, cutout=None)
+        # separation_vectors = disentanglemnet_exp.InterFaceGAN_separation_vector(method='LR', C=0.1)
+        for specific_example in specific_examples:
+            seed = specific_example
+            for i, color in enumerate(colors_list):
+                disentanglemnet_exp.generate_changes(seed, separation_vectors[i], min_epsilon=-9, max_epsilon=9, savefig=True, save_separately=True, feature=color, method='StyleSpace' + '_' + str(True) + '_' + str(10) + '_' + str(None))
+        
+        return
+       
+    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)
+        
+        if continuous:
+            continous_experiment('Saturation', 'S1', repo_folder, model, annotations, df, space, colors_list, kwargs) 
+            continous_experiment('Value', 'V1', repo_folder, model, annotations, df, space, colors_list, kwargs) 
+            continous_experiment('Red', 'R1', repo_folder, model, annotations, df, space, colors_list, kwargs) 
+            continous_experiment('Green', 'G1', repo_folder, model, annotations, df, space, colors_list, kwargs) 
+            continous_experiment('Blue', 'B1', repo_folder, model, annotations, df, space, colors_list, kwargs) 
+            break
+        
+        print('Launching Hue experiment')
+        variable = 'H1'
+        disentanglemnet_exp = DisentanglementBase(repo_folder, model, annotations, df, space=space, colors_list=colors_list, compute_s=False, variable=variable)
 
-        for method in ['InterFaceGAN', 'StyleSpace', 'GANSpace']:
+        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')
+                            print(f'Generating images with variations for color {color}')
                             for s in range(30):
                                 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=str(method) + '_' + str(met) + '_' + str(c))
+                                    disentanglemnet_exp.generate_changes(seed, separation_vectors[i], min_epsilon=-eps, max_epsilon=eps, savefig=True, feature=color, method=str(method) + '_' + str(met) + '_' + str(c) + '_' + str(len(colors_list)) + '_' + str(variable))
                                 
                             print('Finally obtaining verification score')
                             for verif in kwargs['lambda_verif']:    
@@ -446,7 +629,7 @@ def main():
                                 scores.append([space, method, color, score, 'classification method:' + met + ', regularization: ' + str(c) + ', verification lambda:' + str(verif), ', '.join(list(separation_vectors[i].astype(str)))])
                         score_df = pd.DataFrame(scores, columns=['space', 'method', 'color', 'score', 'kwargs', 'vector'])
                         print(score_df)
-                        score_df.to_csv(join(repo_folder, 'data/scores.csv'))
+                        score_df.to_csv(join(repo_folder, f'data/scores_InterfaceGAN_{variable}_{len(colors_list)}.csv'))
                     
                         
             elif method == 'StyleSpace':
@@ -456,11 +639,11 @@ def main():
                         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')
+                                print(f'Generating images with variations for color {color}')
                                 for s in range(30):
                                     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 + '_' + str(num_factors) + '_' + str(cutout) + '_' + str(sign))
+                                        disentanglemnet_exp.generate_changes(seed, separation_vectors[i], min_epsilon=-eps, max_epsilon=eps, savefig=True, feature=color, method=method + '_' + str(num_factors) + '_' + str(cutout) + '_' + str(sign)  + '_' + str(len(colors_list)) + '_' + str(variable))
                                             
                                 print('Finally obtaining verification score')
                                 for verif in kwargs['lambda_verif']:    
@@ -470,29 +653,28 @@ def main():
                                     scores.append([space, method, color, score, 'using sign:' + str(sign) + ', number of factors: ' + str(num_factors) + ', using cutout: ' + str(cutout) + ', verification lambda:' + str(verif), ', '.join(list(separation_vectors[i].astype(str)))])
                             score_df = pd.DataFrame(scores, columns=['space', 'method', 'color', 'score', 'kwargs', 'vector'])
                             print(score_df)
-                            score_df.to_csv(join(repo_folder, 'data/scores.csv'))
+                            score_df.to_csv(join(repo_folder, f'data/scores_StyleSpace_{variable}_{len(colors_list)}.csv'))
                             
             if space == 'w' and method == 'GANSpace':
                 print('Now obtaining separation vector for using GANSpace')
                 separation_vectors = disentanglemnet_exp.GANSpace_separation_vectors(100)
+                print(separation_vectors.shape)
                 for s in range(30):
                     print('Generating images with variations')
                     seed = random.randint(0,100000)
                     for i in range(100):
                         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)
+                            disentanglemnet_exp.generate_changes(seed, separation_vectors.T[i], min_epsilon=-eps, max_epsilon=eps, savefig=True, feature='dimension_' + str(i), method=method)
                                         
                 score = None
-                scores.append([space, method, color, score, '100', ', '.join(list(separation_vectors[i].astype(str)))])
+                scores.append([space, method, 'PCA', score, '100', ', '.join(list(separation_vectors.T[i].astype(str)))])
             else:
                 print('Skipping', method, 'on space', space)
                 continue
             
-            
-    
     score_df = pd.DataFrame(scores, columns=['space', 'method', 'color', 'score', 'kwargs', 'vector'])
     print(score_df)
-    score_df.to_csv(join(repo_folder, 'data/scores.csv'))
+    score_df.to_csv(join(repo_folder, 'data/scores_{}.csv'.format(pd.to_datetime.now().strftime("%Y-%m-%d_%H%M%S"))))
  
 if __name__ == "__main__":
     main()

check_images.py

@@ -0,0 +1,256 @@
+#!/usr/bin/env python
+
+import numpy as np
+import pandas as pd
+
+from sklearn.svm import SVC
+from sklearn.decomposition import PCA
+from sklearn.linear_model import LogisticRegression, LinearRegression
+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
+
+def hex2rgb(hex_value):
+    h = hex_value.strip("#") 
+    rgb = tuple(int(h[i:i+2], 16) for i in (0, 2, 4))
+    return rgb
+
+def rgb2hsv(r, g, b):
+    # Normalize R, G, B values
+    r, g, b = r / 255.0, g / 255.0, b / 255.0
+    
+    # h, s, v = hue, saturation, value
+    max_rgb = max(r, g, b)    
+    min_rgb = min(r, g, b)   
+    difference = max_rgb-min_rgb 
+    
+    # if max_rgb and max_rgb are equal then h = 0
+    if max_rgb == min_rgb:
+        h = 0
+    
+    # if max_rgb==r then h is computed as follows
+    elif max_rgb == r:
+        h = (60 * ((g - b) / difference) + 360) % 360
+    
+    # if max_rgb==g then compute h as follows
+    elif max_rgb == g:
+        h = (60 * ((b - r) / difference) + 120) % 360
+    
+    # if max_rgb=b then compute h
+    elif max_rgb == b:
+        h = (60 * ((r - g) / difference) + 240) % 360
+    
+    # if max_rgb==zero then s=0
+    if max_rgb == 0:
+        s = 0
+    else:
+        s = (difference / max_rgb) * 100
+    
+    # compute v
+    v = max_rgb * 100
+    # return rounded values of H, S and V
+    return tuple(map(round, (h, s, v)))
+ 
+
+class DisentanglementBase:
+    def __init__(self, repo_folder, model, annotations, df, space, colors_list, compute_s=False, variable='H1', categorical=True):
+        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.categorical = categorical
+        self.variable = variable
+        
+        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: rgb2hsv(*hex2rgb(x))[0])
+        self.df['H2'] = self.df['top2col'].map(lambda x: rgb2hsv(*hex2rgb(x))[0])
+        self.df['H3'] = self.df['top3col'].map(lambda x: rgb2hsv(*hex2rgb(x))[0])
+        
+        self.df['S1'] = self.df['top1col'].map(lambda x: rgb2hsv(*hex2rgb(x))[1])
+        self.df['S2'] = self.df['top2col'].map(lambda x: rgb2hsv(*hex2rgb(x))[1])
+        self.df['S3'] = self.df['top3col'].map(lambda x: rgb2hsv(*hex2rgb(x))[1])
+        
+        self.df['V1'] = self.df['top1col'].map(lambda x: rgb2hsv(*hex2rgb(x))[2])
+        self.df['V2'] = self.df['top2col'].map(lambda x: rgb2hsv(*hex2rgb(x))[2])
+        self.df['V3'] = self.df['top3col'].map(lambda x: rgb2hsv(*hex2rgb(x))[2])
+        
+        print('Adding RGB encoding')
+        self.df['R1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[0])
+        self.df['R2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[0])
+        self.df['R3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[0])
+        
+        self.df['G1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[1])
+        self.df['G2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[1])
+        self.df['G3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[1])
+        
+        self.df['B1'] = self.df['top1col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[2])
+        self.df['B2'] = self.df['top2col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[2])
+        self.df['B3'] = self.df['top3col'].map(lambda x: ImageColor.getcolor(x, 'RGB')[2])
+        return self.df
+    
+    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, extremes=False):
+        X = self.get_encoded_latent()
+        y = np.array(self.df[self.variable].values)
+        if self.categorical:
+            y_cat = pd.cut(y, 
+                            bins=[x * 360 / len(self.colors_list) if x < len(self.colors_list) 
+                                  else 360 for x in range(len(self.colors_list) + 1)],
+                            labels=self.colors_list
+                            ).fillna(self.colors_list[0])
+            x_train, x_val, y_train, y_val = train_test_split(X, y_cat, test_size=0.2)
+        else:
+            if extremes:
+                # Calculate the number of elements to consider (10% of array size)
+                num_elements = int(0.2 * len(y))
+                # Get indices of the top num_elements maximum values
+                top_indices = np.argpartition(array, -num_elements)[-num_elements:]
+                bottom_indices = np.argpartition(array, -num_elements)[:num_elements]
+                y_ext = y[top_indices + bottom_indices, :]
+                X_ext = X[top_indices + bottom_indices, :]
+                x_train, x_val, y_train, y_val = train_test_split(X_ext, y_ext, 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 generate_orig_image(self, vec, seed=False):
+        """
+        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 seed:
+            seed = vec
+            vec = self.annotations['z_vectors'][seed]
+            
+        Z = torch.from_numpy(vec.copy()).to(self.device)
+        img = G(Z, label, truncation_psi=1, noise_mode='const')
+        img = (img.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 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 = ['Red', 'Orange', 'Yellow', 'Yellow Green', 'Chartreuse Green',
+                   'Kelly Green', 'Green Blue Seafoam', 'Cyan Blue',
+                   'Warm Blue', 'Indigo', 'Purple Magenta', 'Magenta Pink']
+    colors_list = ['Red Orange', 'Yellow', 'Green', 'Light Blue',
+                   'Blue', 'Purple', 'Pink']
+    
+    
+    disentanglemnet_exp = DisentanglementBase(repo_folder, model, annotations, df, space='w', colors_list=colors_list)
+    # x_train, x_val, y_train, y_val = disentanglemnet_exp.get_train_val()
+    # print(colors_list)
+    # print(np.unique(y_train, return_counts=True))
+    
+    
+    # for i, color in enumerate(colors_list):
+    #     idxs = np.where(y_train == color)
+    #     x_color = x_train[idxs][:30, :]
+    #     print(x_color.shape)
+    #     print('Generating images of color ' + color)
+    #     for j, vec in enumerate(x_color):
+    #         vec = np.expand_dims(vec, axis=0)
+    #         img = disentanglemnet_exp.generate_orig_image(vec)
+    #         img.save(f'{repo_folder}/colors_test/color_{color}_{j}.png')
+                                            
+    df = disentanglemnet_exp.to_hsv()
+    df['color'] = pd.cut(df['H1'], 
+                        bins=[x * 360 / len(colors_list) if x < len(colors_list) 
+                              else 360 for x in range(len(colors_list) + 1)],
+                            labels=colors_list
+                            ).fillna(colors_list[0])  
+    
+    print(df['color'].value_counts())
+    df['seed'] = df['fname'].str.split('/').apply(lambda x: x[-1]).str.replace('seed', '').str.replace('.png','').astype(int)
+    print(df[df['seed'] == 3][['H1', 'S1', 'V1', 'R1', 'B1', 'G1']])
+    for i, color in enumerate(colors_list):
+        idxs = df['color'] == color
+        x_color = df['seed'][idxs][:30]
+        print('Generating images of color ' + color)
+        for j, vec in enumerate(x_color):
+            img = disentanglemnet_exp.generate_orig_image(int(vec), seed=True)
+            img.save(f'{repo_folder}/colors_test/color_{color}_{j}corrected.png')
+            
+if __name__ == "__main__":
+    main()

data/scores_Blue.csv

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

data/scores_Green.csv

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

data/scores_InterfaceGAN_H1_8.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f9b22fba3e1a4dabf3ab59342536c707f9c55b04c257d1da49c6a6be9bac082
+size 919823

data/scores_Red.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6ae0e9ee9a907881b19542eb19cd98947b9f3d2a3ccca61f6dd25823a3fb8e82
+size 27619

data/scores_Saturation.csv

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

data/scores_StyleSpace_H1_8.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8bfc5681aa827b2be07cb0ce00eefd6464e2f0836216b884858b5866ffb8aa80
+size 360571

data/scores_Value.csv

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

data/textile_annotated_files/seeds0000-100000_S.pkl

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

test_disentanglement.sh

@@ -0,0 +1,12 @@
+#!/bin/bash
+#SBATCH --time=1-00:00:00
+#SBATCH --mem=32GB
+#SBATCH --gres gpu:1
+
+module load v100
+module load cuda
+module load mamba
+source activate test
+
+python DisentanglementBase.py
+conda deactivate