update two-stage interface

Home.py

@@ -30,6 +30,7 @@ def save_user_id(user_id):
     if not user_id:
         user_id = 'anonymous' + str(random.randint(0, 100000))
     st.session_state.user_id = [user_id, datetime.now().strftime("%Y-%m-%d %H:%M:%S")]
+    st.session_state.assigned_rank_mode = random.choice(['sort', 'battle'])
 
 
 def logout():
@@ -38,6 +39,7 @@ def logout():
     st.session_state.pop('score_weights', None)
     st.session_state.pop('gallery_state', None)
     st.session_state.pop('progress', None)
+    st.session_state.pop('gallery_focus', None)
 
 
 def info():

pages/Gallery.py

@@ -32,6 +32,9 @@ class GalleryApp:
         if 'selected_dict' not in st.session_state:
             st.session_state['selected_dict'] = {}
 
+        if 'gallery_focus' not in st.session_state:
+            st.session_state.gallery_focus = {'tag': None, 'prompt': None}
+
     def gallery_standard(self, items, col_num, info):
         rows = len(items) // col_num + 1
         containers = [st.container() for _ in range(rows)]
@@ -310,7 +313,7 @@ class GalleryApp:
             # st.markdown(':orange[Please select a prompt above👆]')
             st.write('**Feel free to navigate among tags and pages! Your selection will be saved within one log-in session.**')
 
-            with subset_selector[1]:
+            with subset_selector[-1]:
                 st.write(':orange[👈 **Please select a prompt**]')
 
         else:
@@ -322,6 +325,10 @@ class GalleryApp:
             if prompt_id not in st.session_state.gallery_state:
                 st.session_state.gallery_state[prompt_id] = 'graph'
 
+            # add focus to session state
+            st.session_state.gallery_focus['tag'] = tag
+            st.session_state.gallery_focus['prompt'] = selected_prompt
+
             # add safety check for some prompts
             safety_check = True
             unsafe_prompts = {}
@@ -348,7 +355,7 @@ class GalleryApp:
                     #     #     st.warning('No selected images found')
                     # else:
                     self.graph_mode(prompt_id, items)
-                with subset_selector[1]:
+                with subset_selector[-1]:
                     # if st.session_state.gallery_state[prompt_id] == 'graph':
                     # subset = st.selectbox('Select a subset', ['All', 'Selected Only'], index=0, key=f'subset_{tag}')
                     has_selection = False
@@ -359,7 +366,7 @@ class GalleryApp:
                         pass
 
                     if has_selection:
-                        checkout = st.button('Check out selections', use_container_width=True, type='primary')
+                        checkout = st.button('🛒 Check out selections', use_container_width=True, type='primary')
                         if checkout:
                             print('checkout')
 
@@ -367,17 +374,17 @@ class GalleryApp:
                             print(st.session_state.gallery_state[prompt_id])
                             st.experimental_rerun()
                     else:
-                        st.write('Select images you like below 👇')
+                        st.write(':orange[👇 **Select images you like below**]')
 
             elif st.session_state.gallery_state[prompt_id] == 'gallery':
                 items = items[items['modelVersion_id'].isin(st.session_state.selected_dict[prompt_id])].reset_index(
                     drop=True)
                 self.gallery_mode(prompt_id, items)
 
-                with subset_selector[1]:
+                with subset_selector[-1]:
                     state_operations = st.columns([1, 1])
                     with state_operations[0]:
-                        back = st.button('Back', use_container_width=True)
+                        back = st.button('Back to 🖼️', use_container_width=True)
                         if back:
                             st.session_state.gallery_state[prompt_id] = 'graph'
                             st.experimental_rerun()

pages/Ranking.py

@@ -27,13 +27,20 @@ class RankingApp:
     def sidebar(self):
         with st.sidebar:
             prompt_tags = self.promptBook['tag'].unique()
-            prompt_tags = np.sort(prompt_tags)
+            prompt_tags = np.sort(prompt_tags).tolist()
 
-            tag = st.selectbox('Select a prompt tag', prompt_tags)
+            print(st.session_state.gallery_focus)
+            tag_idx = prompt_tags.index(st.session_state.gallery_focus['tag']) if st.session_state.gallery_focus['tag'] in prompt_tags else 0
+            print(tag_idx)
+
+            tag = st.selectbox('Select a prompt tag', prompt_tags, index=tag_idx)
             items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)
-            prompts = np.sort(items['prompt'].unique())[::-1]
+            prompts = np.sort(items['prompt'].unique())[::-1].tolist()
+
+            prompt_idx = prompts.index(st.session_state.gallery_focus['prompt']) if st.session_state.gallery_focus['prompt'] in prompts else 0
+            print(prompt_idx)
 
-            selected_prompt = st.selectbox('Select a prompt', prompts)
+            selected_prompt = st.selectbox('Select a prompt', prompts, index=prompt_idx)
 
             mode = st.radio('Select a mode', ['Drag and Sort', 'Battle'], index=1)
 
@@ -204,24 +211,24 @@ class RankingApp:
         with left:
             image_id = items['image_id'][st.session_state.pointer[prompt_id]['left']]
             img_url = self.images_endpoint + str(image_id) + '.png'
-            st.image(img_url, use_column_width=True)
 
-            # write the total score of this image
-            total_score = items['total_score'][st.session_state.pointer[prompt_id]['left']]
-            st.write(f'Total Score: {total_score}')
+            # # write the total score of this image
+            # total_score = items['total_score'][st.session_state.pointer[prompt_id]['left']]
+            # st.write(f'Total Score: {total_score}')
 
             btn_left = st.button('Left is better', key='left', on_click=self.next_battle, kwargs={'prompt_id': prompt_id, 'image_ids': items['image_id'], 'winner': 'left', 'curr_position': curr_position, 'total_num': len(items)}, use_container_width=True)
+            st.image(img_url, use_column_width=True)
 
         with right:
             image_id = items['image_id'][st.session_state.pointer[prompt_id]['right']]
             img_url = self.images_endpoint + str(image_id) + '.png'
-            st.image(img_url, use_column_width=True)
 
-            # write the total score of this image
-            total_score = items['total_score'][st.session_state.pointer[prompt_id]['right']]
-            st.write(f'Total Score: {total_score}')
+            # # write the total score of this image
+            # total_score = items['total_score'][st.session_state.pointer[prompt_id]['right']]
+            # st.write(f'Total Score: {total_score}')
 
             btn_right = st.button('Right is better', key='right', on_click=self.next_battle, kwargs={'prompt_id': prompt_id, 'image_ids': items['image_id'], 'winner': 'right', 'curr_position': curr_position, 'total_num': len(items)}, use_container_width=True)
+            st.image(img_url, use_column_width=True)
 
     def next_battle(self, prompt_id, image_ids, winner, curr_position, total_num):
         loser = 'left' if winner == 'right' else 'right'