gallery 2.0

graph view added

pages/Gallery.py

@@ -9,6 +9,7 @@ import streamlit as st
 from bs4 import BeautifulSoup
 from datasets import load_dataset, Dataset, load_from_disk
 from huggingface_hub import login
+from streamlit_agraph import agraph, Node, Edge, Config
 from streamlit_extras.switch_page_button import switch_page
 from sklearn.svm import LinearSVC
 
@@ -50,6 +51,55 @@ class GalleryApp:
                             for key in info:
                                 st.write(f"**{key}**: {items.iloc[idx + j][key]}")
 
+    def gallery_graph(self, items):
+        items = load_tsne_coordinates(items)
+
+        scale = 50
+        items.loc[:, 'x'] = items['x'] * scale
+        items.loc[:, 'y'] = items['y'] * scale
+
+        nodes = []
+        edges = []
+
+        for idx in items.index:
+            # if items.loc[idx, 'modelVersion_id'] in st.session_state.selected_dict.get(items.loc[idx, 'prompt_id'], 0):
+            #     opacity = 0.2
+            # else:
+            #     opacity = 1.0
+
+            nodes.append(Node(id=items.loc[idx, 'image_id'],
+                              # label=str(items.loc[idx, 'model_name']),
+                              size=20,
+                              shape='image',
+                              image=f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{items.loc[idx, 'image_id']}.png",
+                              x=items.loc[idx, 'x'].item(),
+                              y=items.loc[idx, 'y'].item(),
+                              fixed=True,
+                              color={'background': '#00000', 'border': '#ffffff'},
+                              # opacity=opacity,
+                              shadow={'enabled': True, 'color': 'rgba(0,0,0,0.4)', 'size': 10, 'x': 1, 'y': 1},
+                              # borderWidth=1,
+                              # shapeProperties={'useBorderWithImage': True},
+                              )
+                         )
+
+        config = Config(width='100%',
+                        height=800,
+                        directed=True,
+                        physics=False,
+                        hierarchical=False,
+                        # **kwargs
+                        )
+
+        return agraph(nodes=nodes,
+                      edges=edges,
+                      config=config
+                      )
+
+
+
+
+
     def selection_panel(self, items):
         # temperal function
 
@@ -170,6 +220,8 @@ class GalleryApp:
 
             selected_prompt = st.selectbox('Select prompt', prompts)
 
+            mode = st.radio('Select a mode', ['Gallery', 'Graph'], horizontal=True, index=1)
+
             items = items[items['prompt'] == selected_prompt].reset_index(drop=True)
             prompt_id = items['prompt_id'].unique()[0]
             note = items['note'].unique()[0]
@@ -206,14 +258,14 @@ class GalleryApp:
                 except:
                     pass
 
-        return prompt_tags, tag, prompt_id, items
+        return prompt_tags, tag, prompt_id, items, mode
 
     def app(self):
         st.title('Model Visualization and Retrieval')
         st.write('This is a gallery of images generated by the models')
 
-        prompt_tags, tag, prompt_id, items = self.sidebar()
-        items, info, col_num = self.selection_panel(items)
+        prompt_tags, tag, prompt_id, items, mode = self.sidebar()
+        # items, info, col_num = self.selection_panel(items)
 
         # add safety check for some prompts
         safety_check = True
@@ -230,57 +282,115 @@ class GalleryApp:
             safety_check = st.checkbox('I understand that this prompt may contain unsafe content. Show these images anyway.', key=f'safety_{prompt_id}')
 
         if safety_check:
+            if mode == 'Gallery':
+                self.gallery_mode(prompt_id, items)
+            elif mode == 'Graph':
+                self.graph_mode(prompt_id, items)
+
+
+    def graph_mode(self, prompt_id, items):
+        graph_cols = st.columns([3, 1])
+        prompt = st.chat_input(f"Selected model version ids: {str(st.session_state.selected_dict.get(prompt_id, []))}",
+                               disabled=False, key=f'{prompt_id}')
+        if prompt:
+            switch_page("ranking")
+
+        with graph_cols[0]:
+            return_value = self.gallery_graph(items)
+        with graph_cols[1]:
+            if return_value:
+                image_url = f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{return_value}.png"
+
+                st.image(image_url)
+
+                item = items[items['image_id'] == return_value].reset_index(drop=True).iloc[0]
+                modelVersion_id = item['modelVersion_id']
+
+                # handle selection
+                if 'selected_dict' in st.session_state:
+                    if item['prompt_id'] not in st.session_state.selected_dict:
+                        st.session_state.selected_dict[item['prompt_id']] = []
+
+                    if modelVersion_id in st.session_state.selected_dict[item['prompt_id']]:
+                        checked = True
+                    else:
+                        checked = False
+
+                if checked:
+                    deselect = st.button('Deselect', key=f'select_{item["prompt_id"]}_{item["modelVersion_id"]}', use_container_width=True)
+                    if deselect:
+                        st.session_state.selected_dict[item['prompt_id']].remove(item['modelVersion_id'])
+                        st.experimental_rerun()
 
-            if 'selected_dict' in st.session_state:
-                # st.write('checked: ', str(st.session_state.selected_dict.get(prompt_id, [])))
-                dynamic_weight_options = ['Grid Search', 'SVM', 'Greedy']
-                dynamic_weight_panel = st.columns(len(dynamic_weight_options))
-
-                if len(st.session_state.selected_dict.get(prompt_id, [])) > 0:
-                    btn_disable = False
                 else:
-                    btn_disable = True
+                    select = st.button('Select', key=f'select_{item["prompt_id"]}_{item["modelVersion_id"]}', use_container_width=True, type='primary')
+                    if select:
+                        st.session_state.selected_dict[item['prompt_id']].append(item['modelVersion_id'])
+                        st.experimental_rerun()
+
+                # st.write(item)
+                infos = ['model_name', 'modelVersion_name', 'model_download_count', 'clip_score', 'mcos_score',
+                         'nsfw_score']
+                for info in infos:
+                    st.write(f"**{info}**:")
+                    st.write(item[info])
 
-                for i in range(len(dynamic_weight_options)):
-                    method = dynamic_weight_options[i]
-                    with dynamic_weight_panel[i]:
-                        btn = st.button(method, use_container_width=True, disabled=btn_disable, on_click=self.dynamic_weight, args=(prompt_id, items, method))
+            else:
+                st.info('Please click on an image to show')
+
+
+    def gallery_mode(self, prompt_id, items):
+        items, info, col_num = self.selection_panel(items)
+
+        if 'selected_dict' in st.session_state:
+            # st.write('checked: ', str(st.session_state.selected_dict.get(prompt_id, [])))
+            dynamic_weight_options = ['Grid Search', 'SVM', 'Greedy']
+            dynamic_weight_panel = st.columns(len(dynamic_weight_options))
+
+            if len(st.session_state.selected_dict.get(prompt_id, [])) > 0:
+                btn_disable = False
+            else:
+                btn_disable = True
 
-            prompt = st.chat_input(f"Selected model version ids: {str(st.session_state.selected_dict.get(prompt_id, []))}", disabled=False, key=f'{prompt_id}')
-            if prompt:
-                switch_page("ranking")
+            for i in range(len(dynamic_weight_options)):
+                method = dynamic_weight_options[i]
+                with dynamic_weight_panel[i]:
+                    btn = st.button(method, use_container_width=True, disabled=btn_disable, on_click=self.dynamic_weight, args=(prompt_id, items, method))
 
-            with st.form(key=f'{prompt_id}'):
-                # buttons = st.columns([1, 1, 1])
-                buttons_space = st.columns([1, 1, 1, 1])
-                gallery_space = st.empty()
+        prompt = st.chat_input(f"Selected model version ids: {str(st.session_state.selected_dict.get(prompt_id, []))}", disabled=False, key=f'{prompt_id}')
+        if prompt:
+            switch_page("ranking")
 
-                with buttons_space[0]:
-                    continue_btn = st.form_submit_button('Confirm Selection', use_container_width=True, type='primary')
-                    if continue_btn:
-                        self.submit_actions('Continue', prompt_id)
+        with st.form(key=f'{prompt_id}'):
+            # buttons = st.columns([1, 1, 1])
+            buttons_space = st.columns([1, 1, 1, 1])
+            gallery_space = st.empty()
 
-                with buttons_space[1]:
-                    select_btn = st.form_submit_button('Select All', use_container_width=True)
-                    if select_btn:
-                        self.submit_actions('Select', prompt_id)
+            with buttons_space[0]:
+                continue_btn = st.form_submit_button('Confirm Selection', use_container_width=True, type='primary')
+                if continue_btn:
+                    self.submit_actions('Continue', prompt_id)
 
-                with buttons_space[2]:
-                    deselect_btn = st.form_submit_button('Deselect All', use_container_width=True)
-                    if deselect_btn:
-                        self.submit_actions('Deselect', prompt_id)
+            with buttons_space[1]:
+                select_btn = st.form_submit_button('Select All', use_container_width=True)
+                if select_btn:
+                    self.submit_actions('Select', prompt_id)
 
-                with buttons_space[3]:
-                    refresh_btn = st.form_submit_button('Refresh', on_click=gallery_space.empty, use_container_width=True)
+            with buttons_space[2]:
+                deselect_btn = st.form_submit_button('Deselect All', use_container_width=True)
+                if deselect_btn:
+                    self.submit_actions('Deselect', prompt_id)
+
+            with buttons_space[3]:
+                refresh_btn = st.form_submit_button('Refresh', on_click=gallery_space.empty, use_container_width=True)
+
+            with gallery_space.container():
+                with st.spinner('Loading images...'):
+                    self.gallery_standard(items, col_num, info)
+
+                st.info("Don't forget to scroll back to top and click the 'Confirm Selection' button to save your selection!!!")
 
-                with gallery_space.container():
-                    with st.spinner('Loading images...'):
-                        self.gallery_standard(items, col_num, info)
 
-                    st.info("Don't forget to scroll back to top and click the 'Confirm Selection' button to save your selection!!!")
-            # prompt = st.chat_input(f"checked: {str(st.session_state.selected_dict.get(prompt_id, []))}", disabled=False, key=f'{prompt_id}')
-            # if prompt:
-            #     switch_page("ranking")
 
     def submit_actions(self, status, prompt_id):
         # remove counter from session state
@@ -429,6 +539,18 @@ def load_hf_dataset():
 
     return roster, promptBook, images_ds
 
+@st.cache_data
+def load_tsne_coordinates(items):
+    # load tsne coordinates
+    tsne_df = pd.read_parquet('./data/feats_tsne.parquet')
+
+    # print(tsne_df['modelVersion_id'].dtype)
+
+    print('before merge:', items)
+    items = items.merge(tsne_df, on=['modelVersion_id', 'prompt_id'], how='left')
+    print('after merge:', items)
+    return items
+
 
 if __name__ == "__main__":
     st.set_page_config(page_title="Model Coffer Gallery", page_icon="🖼️", layout="wide")