added disentanglement also for vase art

backend/disentangle_concepts.py

@@ -28,10 +28,21 @@ def get_separation_space(type_bin, annotations, df, samples=200, method='LR', C=
     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[col][i] for i in abstract_idxs+repr_idxs])
+    if type(type_bin) == str or len(type_bin) == 1:
+        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[col][i] for i in abstract_idxs+repr_idxs])
+    elif len(type_bin) == 2:
+        print('Using two concepts for separation space')
+        first_concept = np.array([float(ann) for ann in df[type_bin[0]]])
+        second_concept = np.array([float(ann) for ann in df[type_bin[1]]])
+        first_idxs = list(np.argsort(first_concept))[:samples]
+        second_idxs = list(np.argsort(second_concept))[:samples]
+        X = np.array([annotations[col][i] for i in first_idxs+second_idxs])
+    else:
+        print('Error: type_bin must be either a string or a list of strings of len 2')
+        return
     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)

data/vase_annotated_files/seeds0000-20000.pkl

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

data/vase_annotated_files/sim_Shape Name_seeds0000-20000.csv

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

data/vase_model_files/network-snapshot-003800.pkl

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

pages/3_Oxford_Vases_Disentanglement.py

@@ -0,0 +1,171 @@
+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 Oxford Vases 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)
+
+ann_df = pd.read_csv('./data/vase_annotated_files/sim_Shape Name_seeds0000-20000.csv')
+labels = ann_df.columns
+
+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**')
+        concept_ids = st.multiselect('Concepts:', tuple(labels), max_selections=2, default=['AMPHORA', 'CHALICE'])
+
+        st.write('**Choose a latent space to disentangle**')
+        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:
+            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)
+    # 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, 50000)
+            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)
+# 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')
+
+
+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)}')

pages/3_todo.py

@@ -1,124 +0,0 @@
-import streamlit as st
-import pandas as pd
-import numpy as np
-import random
-from backend.utils import make_grid, load_dataset, load_model, load_images
-
-from backend.smooth_grad import generate_smoothgrad_mask, ShowImage, fig2img
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
-import torch
-
-from matplotlib.backends.backend_agg import RendererAgg
-_lock = RendererAgg.lock
-
-st.set_page_config(layout='wide')
-BACKGROUND_COLOR = '#bcd0e7'
-
-
-st.title('Feature attribution visualization with SmoothGrad')
-st.write("""> **Which features are responsible for the current prediction of ConvNeXt?**
-
-In machine learning, it is helpful to identify the significant features of the input (e.g., pixels for images) that affect the model's prediction. 
-If the model makes an incorrect prediction, we might want to determine which features contributed to the mistake. 
-To do this, we can generate a feature importance mask, which is a grayscale image with the same size as the original image. 
-The brightness of each pixel in the mask represents the importance of that feature to the model's prediction.
-
-There are various methods to calculate an image sensitivity mask for a specific prediction. 
-One simple way is to use the gradient of a class prediction neuron concerning the input pixels, indicating how the prediction is affected by small pixel changes. 
-However, this method usually produces a noisy mask. 
-To reduce the noise, the SmoothGrad technique as described in [SmoothGrad: Removing noise by adding noise](https://arxiv.org/abs/1706.03825) by Daniel _et al_ is used, 
-which adds Gaussian noise to multiple copies of the image and averages the resulting gradients.
-""")
-
-instruction_text = """Users need to input the model(s), type of image set and image set setting to use this functionality.
-1. Choose model: Users can choose one or more models for comparison. 
-There are 3 models supported: [ConvNeXt](https://huggingface.co/facebook/convnext-tiny-224), 
-[ResNet](https://huggingface.co/microsoft/resnet-50) and [MobileNet](https://pytorch.org/hub/pytorch_vision_mobilenet_v2/). 
-These 3 models have similar number of parameters. 
-\n2. Choose type of Image set: There are 2 types of Image set. They are _User-defined set_ and _Random set_. 
-\n3. Image set setting: If users choose _User-defined set_ in Image set, 
-users need to enter a list of image IDs separated by commas (,). For example, `0,1,4,7` is a valid input.
-Check the page [ImageNet1k](/ImageNet1k) to see all the Image IDs.
-If users choose _Random set_ in Image set, users just need to choose the number of random images to display here.
-"""
-with st.expander("See more instruction", expanded=False):
-    st.write(instruction_text)
-    
-
-imagenet_df = pd.read_csv('./data/ImageNet_metadata.csv')
-
-# --------------------------- LOAD function -----------------------------
-
-
-images = []
-image_ids = []
-# INPUT ------------------------------
-st.header('Input')
-with st.form('smooth_grad_form'):
-    st.markdown('**Model and Input Setting**')
-    selected_models = st.multiselect('Model', options=['ConvNeXt', 'ResNet', 'MobileNet'])
-    selected_image_set = st.selectbox('Image set', ['User-defined set', 'Random set'])
-    
-    summit_button = st.form_submit_button('Set')
-    if summit_button:
-        setting_container = st.container()
-        # for id in image_ids:
-        #     images = load_images(image_ids)
-
-with st.form('2nd_form'):
-    st.markdown('**Image set setting**')
-    if selected_image_set == 'Random set':
-            no_images = st.slider('Number of images', 1, 50, value=10)
-            image_ids = random.sample(list(range(50_000)), k=no_images)
-    else:
-        text = st.text_area('Specific Image IDs', value='0')
-        image_ids = list(map(lambda x: int(x.strip()), text.split(',')))
-
-    run_button = st.form_submit_button('Display output')
-    if run_button:
-        for id in image_ids:
-            images = load_images(image_ids)
-    
-st.header('Output')
-
-models = {}
-feature_extractors = {}
-        
-for i, model_name in enumerate(selected_models):
-    models[model_name], feature_extractors[model_name] = load_model(model_name)
-
-
-# DISPLAY ----------------------------------
-if run_button:
-    header_cols = st.columns([1, 1] + [2]*len(selected_models))
-    header_cols[0].markdown(f'<div style="text-align: center;margin-bottom: 10px;background-color:{BACKGROUND_COLOR};"><b>Image ID</b></div>', unsafe_allow_html=True)
-    header_cols[1].markdown(f'<div style="text-align: center;margin-bottom: 10px;background-color:{BACKGROUND_COLOR};"><b>Original Image</b></div>', unsafe_allow_html=True)
-    for i, model_name in enumerate(selected_models):
-        header_cols[i + 2].markdown(f'<div style="text-align: center;margin-bottom: 10px;background-color:{BACKGROUND_COLOR};"><b>{model_name}</b></div>', unsafe_allow_html=True)
-    
-    grids = make_grid(cols=2+len(selected_models)*2, rows=len(image_ids)+1)
-
-
-@st.cache(allow_output_mutation=True)
-# @st.cache_data
-def generate_images(image_id, model_name):
-    j = image_ids.index(image_id)
-    image = images[j]['image']
-    return generate_smoothgrad_mask(
-        image, model_name,
-        models[model_name], feature_extractors[model_name], num_samples=10)
-
-with _lock:
-    for j, (image_id, image_dict) in enumerate(zip(image_ids, images)):
-        grids[j][0].write(f'{image_id}. {image_dict["label"]}')
-        image = image_dict['image']
-        ori_image = ShowImage(np.asarray(image))
-        grids[j][1].image(ori_image)
-
-        for i, model_name in enumerate(selected_models):
-            # ori_image, heatmap_image, masked_image = generate_smoothgrad_mask(image,
-            # model_name, models[model_name], feature_extractors[model_name], num_samples=10)
-            heatmap_image, masked_image = generate_images(image_id, model_name)
-            # grids[j][1].image(ori_image)
-            grids[j][i*2+2].image(heatmap_image)
-            grids[j][i*2+3].image(masked_image)