add two-substage gallery page

Home.py

@@ -36,6 +36,8 @@ def logout():
     st.session_state.pop('user_id', None)
     st.session_state.pop('selected_dict', None)
     st.session_state.pop('score_weights', None)
+    st.session_state.pop('gallery_state', None)
+    st.session_state.pop('progress', None)
 
 
 def info():

pages/Gallery.py

@@ -24,6 +24,14 @@ class GalleryApp:
         self.promptBook = promptBook
         self.images_ds = images_ds
 
+        # init gallery state
+        if 'gallery_state' not in st.session_state:
+            st.session_state.gallery_state = {}
+
+        # initialize selected_dict
+        if 'selected_dict' not in st.session_state:
+            st.session_state['selected_dict'] = {}
+
     def gallery_standard(self, items, col_num, info):
         rows = len(items) // col_num + 1
         containers = [st.container() for _ in range(rows)]
@@ -276,6 +284,7 @@ class GalleryApp:
         # chosen_data = [stx.TabBarItemData(id=tag, title=tag, description='') for tag in prompt_tags]
         # tag = stx.tab_bar(chosen_data, key='tag', default='food')
 
+        # save tag to session state on change
         tag = st.radio('Select a tag', prompt_tags, index=5, horizontal=True, key='tag')
 
         # tabs = st.tabs(prompt_tags)
@@ -284,23 +293,34 @@ class GalleryApp:
         #         tag = prompt_tags[i]
         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()
 
+        st.caption('Select a prompt')
         subset_selector = st.columns([3, 1])
         with subset_selector[0]:
+            # remember last prompt
+            # if 'prompt_idx_last_time' not in st.session_state:
+            #     st.session_state.prompt_idx_last_time = 0
+
             # selected_prompt = st.selectbox('Select prompt', prompts, index=3)
-            selected_prompt = selectbox('Select prompt', prompts, key=f'prompt_{tag}', no_selection_label='---')
-        with subset_selector[1]:
-            subset = st.selectbox('Select a subset', ['All', 'Selected Only'], index=0, key=f'subset_{tag}')
+            selected_prompt = selectbox('Select prompt', prompts, key=f'prompt_{tag}', no_selection_label='---', label_visibility='collapsed', index=0)
+            # st.session_state.prompt_idx_last_time = prompts.index(selected_prompt) if selected_prompt else 0
 
         if selected_prompt is None:
-            st.markdown(':orange[Please select a prompt above👆]')
+            # 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]:
+                st.write(':orange[👈 **Please select a prompt**]')
+
         else:
             items = items[items['prompt'] == selected_prompt].reset_index(drop=True)
             prompt_id = items['prompt_id'].unique()[0]
             note = items['note'].unique()[0]
-            print(prompt_id, note)
+
+            # add state to session state
+            if prompt_id not in st.session_state.gallery_state:
+                st.session_state.gallery_state[prompt_id] = 'graph'
 
             # add safety check for some prompts
             safety_check = True
@@ -316,16 +336,61 @@ class GalleryApp:
                 st.warning('This prompt may contain unsafe content. They might be offensive, depressing, or sexual.')
                 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 subset == 'Selected Only' and 'selected_dict' in st.session_state:
-                    # try:
-                    items = items[items['modelVersion_id'].isin(st.session_state.selected_dict[prompt_id])].reset_index(drop=True)
-                    self.gallery_mode(prompt_id, items)
-                    # except:
-                    #     st.warning('No selected images found')
-                else:
+            print('current state: ', st.session_state.gallery_state[prompt_id])
+
+            if st.session_state.gallery_state[prompt_id] == 'graph':
+                if safety_check:
+                    # if subset == 'Selected Only' and 'selected_dict' in st.session_state:
+                    #     # try:
+                    #     items = items[items['modelVersion_id'].isin(st.session_state.selected_dict[prompt_id])].reset_index(drop=True)
+                    #     self.gallery_mode(prompt_id, items)
+                    #     # except:
+                    #     #     st.warning('No selected images found')
+                    # else:
                     self.graph_mode(prompt_id, items)
+                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
+                    try:
+                        if len(st.session_state.selected_dict.get(prompt_id, [])) > 0:
+                            has_selection = True
+                    except:
+                        pass
+
+                    if has_selection:
+                        checkout = st.button('Check out selections', use_container_width=True, type='primary')
+                        if checkout:
+                            print('checkout')
+
+                            st.session_state.gallery_state[prompt_id] = 'gallery'
+                            print(st.session_state.gallery_state[prompt_id])
+                            st.experimental_rerun()
+                    else:
+                        st.write('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]:
+                    state_operations = st.columns([1, 1])
+                    with state_operations[0]:
+                        back = st.button('Back', use_container_width=True)
+                        if back:
+                            st.session_state.gallery_state[prompt_id] = 'graph'
+                            st.experimental_rerun()
+
+                    with state_operations[1]:
+                        forward = st.button('Check out', use_container_width=True, type='primary', on_click=self.submit_actions, args=('Continue', prompt_id))
+                        if forward:
+                            switch_page('ranking')
+
+
+
+                # else:
+                #     st.button('Proceed', use_container_width=True)
 
         try:
             self.sidebar(items, prompt_id, note)
@@ -390,8 +455,6 @@ class GalleryApp:
                     infos_df = infos_df.rename(index={'model_name': 'Model', 'modelVersion_name': 'Version', 'model_download_count': 'Downloads', 'clip_score': 'Clip Score', 'mcos_score': 'mcos Score', 'nsfw_score': 'NSFW Score'})
                     st.table(infos_df)
 
-                st.button('🎖️ Proceed selections to ranking', on_click=switch_page, args=("ranking",), use_container_width=True,)
-
                     # for info in infos:
                     #     st.write(f"**{info}**:")
                     #     st.write(item[info])
@@ -399,8 +462,6 @@ class GalleryApp:
             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)
 
@@ -423,33 +484,31 @@ class GalleryApp:
         # if prompt:
         #     switch_page("ranking")
 
-        with st.form(key=f'{prompt_id}'):
+        # with st.form(key=f'{prompt_id}'):
             # buttons = st.columns([1, 1, 1])
-            buttons_space = st.columns([1, 1, 1])
-            gallery_space = st.empty()
-
-            with buttons_space[0]:
-                continue_btn = st.form_submit_button('Proceed selections to ranking', use_container_width=True, type='primary')
-                if continue_btn:
-                    # self.submit_actions('Continue', prompt_id)
-                    switch_page("ranking")
-
-            with buttons_space[1]:
-                deselect_btn = st.form_submit_button('Deselect All', use_container_width=True)
-                if deselect_btn:
-                    self.submit_actions('Deselect', prompt_id)
-
-            with buttons_space[2]:
-                refresh_btn = st.form_submit_button('Refresh', on_click=gallery_space.empty, use_container_width=True)
+        # buttons_space = st.columns([1, 1, 1])
+        gallery_space = st.empty()
+
+        # with buttons_space[0]:
+        #     continue_btn = st.button('Proceed selections to ranking', use_container_width=True, type='primary')
+        #     if continue_btn:
+        #         # self.submit_actions('Continue', prompt_id)
+        #         switch_page("ranking")
+        #
+        # with buttons_space[1]:
+        #     deselect_btn = st.button('Deselect All', use_container_width=True)
+        #     if deselect_btn:
+        #         self.submit_actions('Deselect', prompt_id)
+        #
+        # with buttons_space[2]:
+        #     refresh_btn = st.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)
+        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!!!")
 
-
-
     def submit_actions(self, status, prompt_id):
         # remove counter from session state
         # st.session_state.pop('counter', None)
@@ -473,7 +532,7 @@ class GalleryApp:
                         st.session_state.selected_dict[prompt_id].append(int(keys[2]))
             # switch_page("ranking")
             print(st.session_state.selected_dict, 'continue')
-            st.experimental_rerun()
+            # st.experimental_rerun()
 
     def dynamic_weight(self, prompt_id, items, method='Grid Search'):
         selected = items[
@@ -656,9 +715,9 @@ if __name__ == "__main__":
         roster, promptBook, images_ds = load_hf_dataset()
         # print(promptBook.columns)
 
-        # initialize selected_dict
-        if 'selected_dict' not in st.session_state:
-            st.session_state['selected_dict'] = {}
+        # # initialize selected_dict
+        # if 'selected_dict' not in st.session_state:
+        #     st.session_state['selected_dict'] = {}
 
         app = GalleryApp(promptBook=promptBook, images_ds=images_ds)
         app.app()

pages/Summary.py

@@ -128,11 +128,11 @@ class DashboardApp:
                         st.image(image, use_column_width=True)
 
     def score_calculator(self, results, db_table):
-        # sort results by battle time
-        results = sorted(results, key=lambda x: x['battletime'])
-
         modelVersion_standings = {}
         if db_table == 'battle_results':
+            # sort results by battle time
+            results = sorted(results, key=lambda x: x['battletime'])
+
             for record in results:
                 modelVersion_standings[record['winner']] = modelVersion_standings.get(record['winner'], 0) + 1