added disentanglement of W vector

backend/disentangle_concepts.py

@@ -6,7 +6,7 @@ import torch
 from umap import UMAP
 import PIL
 
-def get_separation_space(type_bin, annotations, df, samples=100, method='LR', C=0.1):
+def get_separation_space(type_bin, annotations, df, samples=100, method='LR', C=0.1, latent_space='Z'):
     """
     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.
@@ -22,10 +22,16 @@ def get_separation_space(type_bin, annotations, df, samples=100, method='LR', C=
     :return: The weights of the linear classifier
     :doc-author: Trelent
     """
+    
+    if latent_space == 'Z':
+        col = 'z_vectors'
+    else:
+        col = 'w_vectors'
+    
     abstracts = np.array([float(ann) for ann in df[type_bin]])
     abstract_idxs = list(np.argsort(abstracts))[:samples]
     repr_idxs = list(np.argsort(abstracts))[-samples:]
-    X = np.array([annotations['z_vectors'][i] for i in abstract_idxs+repr_idxs])
+    X = np.array([annotations[col][i] for i in abstract_idxs+repr_idxs])
     X = X.reshape((2*samples, 512))
     y = np.array([1]*samples + [0]*samples)
     x_train, x_val, y_train, y_val = train_test_split(X, y, test_size=0.2)
@@ -45,7 +51,7 @@ def get_separation_space(type_bin, annotations, df, samples=100, 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):
+def regenerate_images(model, z, decision_boundary, min_epsilon=-3, max_epsilon=3, count=5, latent_space='Z'):
     """
     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
@@ -69,6 +75,7 @@ def regenerate_images(model, z, decision_boundary, min_epsilon=-3, max_epsilon=3
     # Labels.
     label = torch.zeros([1, G.c_dim], device=device)
     
+  
     z = torch.from_numpy(z.copy()).to(device)
     decision_boundary = torch.from_numpy(decision_boundary.copy()).to(device)
         
@@ -84,14 +91,19 @@ def regenerate_images(model, z, decision_boundary, min_epsilon=-3, max_epsilon=3
             #m = make_transform(translate, rotate)
             #m = np.linalg.inv(m)
             #G.synthesis.input.transform.copy_(torch.from_numpy(m))
-
-        img = G(z_0, label, truncation_psi=0.7, noise_mode='random')
+        if latent_space == 'Z':
+            img = G(z_0, label, truncation_psi=0.7, noise_mode='const')
+        
+        else:
+            W = z_0.expand((14, -1)).unsqueeze(0)
+            img = G.synthesis(W, noise_mode='const')
+            
         img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
         images.append(PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB'))
 
     return images, lambdas
 
-def generate_original_image(z, model):
+def generate_original_image(z, model, latent_space='Z'):
     """
     The generate_original_image function takes in a latent vector and the model,
     and returns an image generated from that latent vector.
@@ -106,13 +118,19 @@ def generate_original_image(z, model):
     G = model.to(device) # type: ignore
     # Labels.
     label = torch.zeros([1, G.c_dim], device=device)
-    z = torch.from_numpy(z.copy()).to(device)
-    img = G(z, label, truncation_psi=0.7, noise_mode='random')
+    if latent_space == 'Z':
+        z = torch.from_numpy(z.copy()).to(device)
+        img = G(z, label, truncation_psi=0.7, noise_mode='const')
+    else:
+        W = torch.from_numpy(np.repeat(z, 14, axis=0).reshape(1, 14, 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')
 
 
-def get_concepts_vectors(concepts, annotations, df, samples=100, method='LR', C=0.1):
+def get_concepts_vectors(concepts, annotations, df, samples=100, method='LR', C=0.1, latent_space='Z'):
     """
     The get_concepts_vectors function takes in a list of concepts, a dictionary of annotations, and the dataframe containing all the images.
     It returns two things:
@@ -132,7 +150,7 @@ def get_concepts_vectors(concepts, annotations, df, samples=100, method='LR', C=
     performances = []
     vectors = np.zeros((len(concepts), 512))
     for i, conc in enumerate(concepts):
-        vec, _, imp_nodes, performance = get_separation_space(conc, annotations, df, samples=samples, method=method, C=C)
+        vec, _, imp_nodes, performance = get_separation_space(conc, annotations, df, samples=samples, method=method, C=C, latent_space=latent_space)
         vectors[i,:] = vec
         performances.append(performance)
         important_nodes.append(set(imp_nodes))
@@ -148,3 +166,60 @@ def get_concepts_vectors(concepts, annotations, df, samples=100, method='LR', C=
     nodes_in_common = set.intersection(*important_nodes)
     return vectors, nodes_in_common, performances
 
+
+def get_verification_score(concept, decision_boundary, model, annotations, samples=100, latent_space='Z'):
+    import open_clip
+    import os
+    import random
+    from tqdm import tqdm
+    os.environ["CUDA_VISIBLE_DEVICES"] = ""
+    
+    
+    model_clip, _, preprocess = open_clip.create_model_and_transforms('ViT-L-14', pretrained='laion2b_s32b_b82k')
+    tokenizer = open_clip.get_tokenizer('ViT-L-14')
+
+    # Prepare the text queries
+    #@markdown _in the form pre_prompt {label}_:
+    pre_prompt = "Artwork, " #@param {type:"string"}
+    text_descriptions = [f"{pre_prompt}{label}" for label in [concept]]
+    text_tokens = tokenizer(text_descriptions)
+
+
+    listlen = len(annotations['fname'])
+    items = random.sample(range(listlen), samples)
+    changes = []
+    for iterator in tqdm(items):
+        chunk_imgs = []
+        chunk_ids = []
+
+        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)
+        for im,l in zip(images, lambdas):
+            
+            chunk_imgs.append(preprocess(im.convert("RGB")))
+            chunk_ids.append(l)
+            
+        image_input = torch.tensor(np.stack(chunk_imgs))
+
+        with torch.no_grad(), torch.cuda.amp.autocast():
+            text_features = model_clip.encode_text(text_tokens).float()
+            image_features = model_clip.encode_image(image_input).float()
+
+            # Rescale features
+            image_features /= image_features.norm(dim=-1, keepdim=True)  
+            text_features /= text_features.norm(dim=-1, keepdim=True)
+                    
+            # Analyze featues
+            text_probs = (100.0 * image_features.cpu().numpy() @ text_features.cpu().numpy().T)#.softmax(dim=-1)
+            
+        change = max(text_probs[1][0].item() - text_probs[0][0].item(), 0)
+        changes.append(change)
+        
+    return np.round(np.mean(np.array(changes)), 4)
+    
+            
+
+            

data/annotated_files/{annotations_seeds0000-1000.pkl → seeds0000-50000.pkl}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ffcd38622c48bf91e0e193b846a181e6baee559633e6057df7765fe0ecd422cf
-size 4461349
+oid sha256:cd1bd97b8ff508b1d4a7ef43323530368ace65b35d12d84a914913f541187298
+size 314939226

data/annotated_files/sim_seeds0000-10000.csv

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

data/annotated_files/{annotations_parallel_seeds0000-10000.pkl → sim_seeds0000-50000.csv}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cae08d2bfaa9be2002692efcacfdd10dfd480749e99d99c170a6de13f4811bad
-size 55986521
+oid sha256:c3faa3d75c2da1dbb2c5d90aeddee256e1f3324b24b902a54115d9b6aad0ae9d
+size 21965577

data/model_files/pytorch_model.bin

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

pages/1_Disentanglement.py

@@ -34,11 +34,11 @@ with st.expander("See more instruction", expanded=False):
     st.write(instruction_text)
 
     
-annotations_file = './data/annotated_files/seeds0000-100000.pkl'
+annotations_file = './data/annotated_files/seeds0000-50000.pkl'
 with open(annotations_file, 'rb') as f:
     annotations = pickle.load(f)
 
-ann_df = pd.read_csv('./data/annotated_files/sim_seeds0000-100000.csv')
+ann_df = pd.read_csv('./data/annotated_files/sim_seeds0000-50000.csv')
 concepts = './data/concepts.txt'
 
 with open(concepts) as f:
@@ -48,6 +48,8 @@ if 'image_id' not in st.session_state:
     st.session_state.image_id = 0
 if 'concept_id' not in st.session_state:
     st.session_state.concept_id = 'Abstract'
+if 'space_id' not in st.session_state:
+    st.session_state.space_id = 'Z'
 
 # def on_change_random_input():
 #     st.session_state.image_id = st.session_state.image_id
@@ -65,7 +67,12 @@ with input_col_1:
         # concept_id = chosen_text_id_input.text_input('Concept:', value=st.session_state.concept_id)
         concept_id = st.selectbox('Concept:', tuple(labels))
 
-        choose_text_button = st.form_submit_button('Choose the defined concept')
+        st.write('**Choose a latent space to disentangle**')
+        # chosen_text_id_input = st.empty()
+        # concept_id = chosen_text_id_input.text_input('Concept:', value=st.session_state.concept_id)
+        space_id = st.selectbox('Space:', tuple(['Z', 'W']))
+
+        choose_text_button = st.form_submit_button('Choose the defined concept and space to disentangle')
         # random_text = st.form_submit_button('Select a random concept')
 
         # if random_text:
@@ -76,6 +83,8 @@ with input_col_1:
         if choose_text_button:
             concept_id = str(concept_id)
             st.session_state.concept_id = concept_id
+            space_id = str(space_id)
+            st.session_state.space_id = space_id
         # st.write(image_id, st.session_state.image_id)
 
 # ---------------------------- SET UP OUTPUT ------------------------------
@@ -101,10 +110,10 @@ smoothgrad_col_1, smoothgrad_col_2, smoothgrad_col_3, smoothgrad_col_4, smoothgr
 
 # ---------------------------- DISPLAY COL 1 ROW 1 ------------------------------
 with output_col_1:
-    separation_vector, number_important_features, imp_nodes, performance = get_separation_space(concept_id, annotations, ann_df)
+    separation_vector, number_important_features, imp_nodes, performance = get_separation_space(concept_id, annotations, ann_df, latent_space=st.session_state.space_id)
     # 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 {concept_id} - Number of relevant nodes: {number_important_features} - Val classification performance: {performance}')# - Nodes {",".join(list(imp_nodes))}')
+    header_col_1.write(f'Concept {concept_id} - 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:
@@ -141,8 +150,12 @@ with input_col_3:
 with dnnlib.util.open_url('./data/model_files/network-snapshot-010600.pkl') as f:
     model = legacy.load_network_pkl(f)['G_ema'].to('cpu') # type: ignore
 
-original_image_vec = annotations['z_vectors'][st.session_state.image_id]
-img = generate_original_image(original_image_vec, model)
+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)
 # input_image = original_image_dict['image']
 # input_label = original_image_dict['label']
 # input_id = original_image_dict['id']
@@ -152,7 +165,7 @@ with smoothgrad_col_3:
     smooth_head_3.write(f'Base image')
 
 
-images, lambdas = regenerate_images(model, original_image_vec, separation_vector, min_epsilon=-(int(epsilon)), max_epsilon=int(epsilon))
+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)
 
 with smoothgrad_col_1:
     st.image(images[0])

pages/2_Concepts_comparison.py

@@ -39,7 +39,8 @@ 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 = ['Abstract', 'Representational']
-
+if 'space_id' not in st.session_state:
+    st.session_state.space_id = 'Z'
 # def on_change_random_input():
 #     st.session_state.image_id = st.session_state.image_id
 
@@ -63,9 +64,17 @@ with input_col_1:
         #     concept_id = random.choice(labels)
         #     st.session_state.concept_id = concept_id
         #     chosen_text_id_input.text_input('Concept:', value=st.session_state.concept_id)
-            
+        st.write('**Choose a latent space to disentangle**')
+        # chosen_text_id_input = st.empty()
+        # concept_id = chosen_text_id_input.text_input('Concept:', value=st.session_state.concept_id)
+        space_id = st.selectbox('Space:', tuple(['Z', 'W']))
+
+        choose_text_button = st.form_submit_button('Choose the defined concept and space to disentangle')
+        
         if choose_text_button:
             st.session_state.concept_ids = list(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 ------------------------------
@@ -91,10 +100,10 @@ smoothgrad_col_1, smoothgrad_col_2, smoothgrad_col_3, smoothgrad_col_4, smoothgr
 
 # ---------------------------- DISPLAY COL 1 ROW 1 ------------------------------
 with output_col_1:
-    vectors, nodes_in_common, performances = get_concepts_vectors(concept_ids, annotations, ann_df)
+    vectors, nodes_in_common, performances = get_concepts_vectors(concept_ids, annotations, ann_df, latent_space=space_id)
     # 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'Concepts {", ".join(concept_ids)} - Relevant nodes in common: {nodes_in_common} - Performance of the concept vectors: {performances}')# - Nodes {",".join(list(imp_nodes))}')
+    header_col_1.write(f'Concepts {", ".join(concept_ids)} - Latent space {space_id} - Relevant nodes in common: {nodes_in_common} - Performance of the concept vectors: {performances}')# - Nodes {",".join(list(imp_nodes))}')
 
     edges = []
     for i in range(len(concept_ids)):