testing for new checkout page

pages/Gallery.py

@@ -27,7 +27,7 @@ class GalleryApp:
 
         # init gallery state
         if 'gallery_state' not in st.session_state:
-            st.session_state.gallery_state = {}
+            st.session_state.gallery_state = 'graph'
 
         # initialize selected_dict
         if 'selected_dict' not in st.session_state:
@@ -36,7 +36,7 @@ class GalleryApp:
         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):
+    def gallery_standard(self, items, col_num, info, show_checkbox=True):
         rows = len(items) // col_num + 1
         containers = [st.container() for _ in range(rows)]
         for idx in range(0, len(items), col_num):
@@ -59,8 +59,9 @@ class GalleryApp:
 
                             # st.write("Position: ", idx + j)
 
-                            # show checkbox
-                            st.checkbox('Select', key=f'select_{prompt_id}_{modelVersion_id}', value=check_init)
+                            if show_checkbox:
+                                # show checkbox
+                                st.checkbox('Select', key=f'select_{prompt_id}_{modelVersion_id}', value=check_init)
 
                             # show selected info
                             for key in info:
@@ -278,6 +279,7 @@ class GalleryApp:
         # return prompt_tags, tag, prompt_id, items
 
     def app(self):
+        print(st.session_state.gallery_focus)
         st.write('### Model Visualization and Retrieval')
         # st.write('This is a gallery of images generated by the models')
 
@@ -286,65 +288,66 @@ class GalleryApp:
         # sort tags by alphabetical order
         prompt_tags = np.sort(prompt_tags)[::1].tolist()
 
-        # chosen_data = [stx.TabBarItemData(id=tag, title=tag, description='') for tag in prompt_tags]
-        # tag = stx.tab_bar(chosen_data, key='tag', default='food')
+        # set focus tag and prompt index if exists
+        if st.session_state.gallery_focus['tag'] is None:
+             tag_focus_idx = 5
+        else:
+            tag_focus_idx = prompt_tags.index(st.session_state.gallery_focus['tag'])
 
         # save tag to session state on change
-        tag = st.radio('Select a tag', prompt_tags, index=5, horizontal=True, key='tag', label_visibility='collapsed')
-
-        # tabs = st.tabs(prompt_tags)
-        # for i in range(len(prompt_tags)):
-        #     with tabs[i]:
-        #         tag = prompt_tags[i]
-        items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)
-
-        prompts = np.sort(items['prompt'].unique())[::1].tolist()
-
-        # st.caption('Select a prompt')
-        subset_selector = st.columns([3, 1])
-        with subset_selector[0]:
-            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
+        tag = st.radio('Select a tag', prompt_tags, index=tag_focus_idx, horizontal=True, key='tag', label_visibility='collapsed')
 
-        if selected_prompt is None:
-            # 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.**')
+        print('current state: ', st.session_state.gallery_state)
 
-            with subset_selector[-1]:
-                st.write(':orange[👈 **Please select a prompt**]')
+        if st.session_state.gallery_state == 'graph':
 
-        else:
-            items = items[items['prompt'] == selected_prompt].reset_index(drop=True)
-            prompt_id = items['prompt_id'].unique()[0]
-            note = items['note'].unique()[0]
-
-            # add state to session state
-            if prompt_id not in st.session_state.gallery_state:
-                st.session_state.gallery_state[prompt_id] = 'graph'
+            items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)
 
-            # add focus to session state
-            st.session_state.gallery_focus['tag'] = tag
-            st.session_state.gallery_focus['prompt'] = selected_prompt
+            prompts = np.sort(items['prompt'].unique())[::1].tolist()
 
-            # add safety check for some prompts
-            safety_check = True
+            # selt focus prompt index if exists
+            if st.session_state.gallery_focus['prompt'] is None:
+                prompt_focus_idx = 0
+            else:
+                prompt_focus_idx = 1 + prompts.index(st.session_state.gallery_focus['prompt'])
 
-            # load unsafe prompts
-            unsafe_prompts = json.load(open('./data/unsafe_prompts.json', 'r'))
-            for prompt_tag in prompt_tags:
-                if prompt_tag not in unsafe_prompts:
-                    unsafe_prompts[prompt_tag] = []
-            # # manually add unsafe prompts
-            # unsafe_prompts['world knowledge'] = [83]
-            # unsafe_prompts['abstract'] = [1, 3]
+            # st.caption('Select a prompt')
+            subset_selector = st.columns([3, 1])
+            with subset_selector[0]:
+                selected_prompt = selectbox('Select prompt', prompts, key=f'prompt_{tag}', no_selection_label='---', label_visibility='collapsed', index=prompt_focus_idx)
+                # st.session_state.prompt_idx_last_time = prompts.index(selected_prompt) if selected_prompt else 0
 
-            if int(prompt_id.item()) in unsafe_prompts[tag]:
-                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 selected_prompt is None:
+                # 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.**')
 
-            print('current state: ', st.session_state.gallery_state[prompt_id])
+                with subset_selector[-1]:
+                    st.write(':orange[👈 **Please select a prompt**]')
 
-            if st.session_state.gallery_state[prompt_id] == 'graph':
+            else:
+                items = items[items['prompt'] == selected_prompt].reset_index(drop=True)
+                prompt_id = items['prompt_id'].unique()[0]
+                note = items['note'].unique()[0]
+
+                # add safety check for some prompts
+                safety_check = True
+
+                # load unsafe prompts
+                unsafe_prompts = json.load(open('./data/unsafe_prompts.json', 'r'))
+                for prompt_tag in prompt_tags:
+                    if prompt_tag not in unsafe_prompts:
+                        unsafe_prompts[prompt_tag] = []
+                # # manually add unsafe prompts
+                # unsafe_prompts['world knowledge'] = [83]
+                # unsafe_prompts['abstract'] = [1, 3]
+
+                if int(prompt_id.item()) in unsafe_prompts[tag]:
+                    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}')
+
+            # print('current state: ', st.session_state.gallery_state)
+            #
+            # if st.session_state.gallery_state == 'graph':
                 if safety_check:
                     self.graph_mode(prompt_id, items)
                 with subset_selector[-1]:
@@ -359,35 +362,48 @@ class GalleryApp:
                         checkout = st.button('Check out selections', use_container_width=True, type='primary')
                         if checkout:
                             print('checkout')
+                            # add focus to session state
+                            st.session_state.gallery_focus['tag'] = tag
+                            st.session_state.gallery_focus['prompt'] = selected_prompt
 
-                            st.session_state.gallery_state[prompt_id] = 'gallery'
-                            print(st.session_state.gallery_state[prompt_id])
+                            st.session_state.gallery_state = 'check out'
+                            print(st.session_state.gallery_state)
                             st.experimental_rerun()
                     else:
                         st.write(':orange[👇 **Select images you like below**]')
+                try:
+                    self.sidebar(items, prompt_id, note)
+                except:
+                    pass
 
-            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)
+        elif st.session_state.gallery_state == 'check out':
+            # select items under the current tag, while model_id in selected_dict keys with corresponding modelVersion_ids
+            items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)
+            temp_items = pd.DataFrame()
+            for prompt_id, selected_models in st.session_state.selected_dict.items():
+                temp_items = temp_items.append(items[items['modelVersion_id'].isin(selected_models) & (items['prompt_id'] == prompt_id)])
+            items = temp_items.reset_index(drop=True)
 
-                with subset_selector[-1]:
-                    state_operations = st.columns([1, 1])
-                    with state_operations[0]:
-                        back = st.button('Back to 🖼️', use_container_width=True)
-                        if back:
-                            st.session_state.gallery_state[prompt_id] = 'graph'
-                            st.experimental_rerun()
+            self.checkout_mode(tag, items)
 
-                    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')
 
-        try:
-            self.sidebar(items, prompt_id, note)
-        except:
-            pass
+            # 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 to 🖼️', 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')
+
 
     def graph_mode(self, prompt_id, items):
         graph_cols = st.columns([3, 1])
@@ -498,6 +514,49 @@ class GalleryApp:
         with gallery_space.container():
             self.gallery_standard(items, col_num, info)
 
+    def checkout_mode(self, tag, items):
+        # st.write(items)
+        if len(items) > 0:
+            for prompt_id in items['prompt_id'].unique():
+                prompt = items[items['prompt_id'] == prompt_id]['prompt'].unique()[0]
+                default_expand = True if st.session_state.gallery_focus['prompt'] == prompt else False
+                with st.expander(f'**{prompt}**', expanded=default_expand):
+                    # st.caption('select info to show')
+                    checkout_panel = st.columns([6, 1, 1])
+                    with checkout_panel[0]:
+                        pass
+                        info = st.multiselect('Show Info',
+                                              ['model_name', 'model_id', 'modelVersion_name', 'modelVersion_id',
+                                               'weighted_score_sum', 'model_download_count', 'clip_score', 'mcos_score',
+                                               'nsfw_score', 'norm_nsfw'],
+                                              label_visibility='collapsed', key=f'info_{prompt_id}', placeholder='Select info to show')
+                    with checkout_panel[1]:
+                        back = st.button('Back to 🖼️', key=f'checkout_back_{prompt_id}', use_container_width=True)
+                        if back:
+                            st.session_state.gallery_focus['tag'] = tag
+                            st.session_state.gallery_focus['prompt'] = prompt
+                            print(st.session_state.gallery_focus)
+                            st.session_state.gallery_state = 'graph'
+                            st.experimental_rerun()
+                    with checkout_panel[2]:
+                        proceed = st.button('Proceed ➡️', key=f'checkout_proceed_{prompt_id}', use_container_width=True,
+                                            type='primary')
+                        if proceed:
+                            st.session_state.gallery_focus['tag'] = tag
+                            st.session_state.gallery_focus['prompt'] = prompt
+                            switch_page('ranking')
+
+                    self.gallery_standard(items[items['prompt_id'] == prompt_id].reset_index(drop=True), 4, info, show_checkbox=False)
+        else:
+            # with st.form(key=f'checkout_{tag}'):
+            st.info('No selection under this tag')
+            back = st.button('🖼️ Back to gallery and select something you like', key=f'checkout_{tag}', type='primary')
+            if back:
+                st.session_state.gallery_focus['tag'] = tag
+                st.session_state.gallery_focus['prompt'] = None
+                st.session_state.gallery_state = 'graph'
+                st.experimental_rerun()
+
     def submit_actions(self, status, prompt_id):
         # remove counter from session state
         # st.session_state.pop('counter', None)

pages/Ranking.py

@@ -149,10 +149,10 @@ class RankingApp:
         with control:
             if st.session_state.counter[prompt_id] < batch_num - 1:
                 st.button(":arrow_right:", key='next', on_click=self.next_batch, help='Next Batch',
-                          kwargs={'prompt_id': prompt_id}, use_container_width=True)
+                          kwargs={'prompt_id': prompt_id}, use_container_width=True, type='primary')
             else:
                 st.button(":heavy_check_mark:", key='finished', on_click=self.next_batch, help='Finished',
-                          kwargs={'prompt_id': prompt_id, 'progress': 'finished'}, use_container_width=True)
+                          kwargs={'prompt_id': prompt_id, 'progress': 'finished'}, use_container_width=True, type='primary')
 
             if st.session_state.counter[prompt_id] > 0:
                 st.button(":arrow_left:", key='prev', on_click=self.prev_batch, help='Previous Batch',