ranking is not finished!!!

pages/Gallery.py

@@ -45,19 +45,7 @@ class GalleryApp:
                             st.write("Position: ", idx + j)
 
                             # show checkbox
-                            checked = st.checkbox('Select', key=f'select_{prompt_id}_{modelVersion_id}', value=check_init)
-
-                            #
-                            # if checked:
-                            #     if prompt_id not in st.session_state.selected_dict:
-                            #         st.session_state.selected_dict[prompt_id] = []
-                            #     if modelVersion_id not in st.session_state.selected_dict[prompt_id]:
-                            #         st.session_state.selected_dict[prompt_id].append(modelVersion_id)
-                            # else:
-                            #     try:
-                            #         st.session_state.selected_dict[prompt_id].remove(modelVersion_id)
-                            #     except:
-                            #         pass
+                            st.checkbox('Select', key=f'select_{prompt_id}_{modelVersion_id}', value=check_init)
 
                             # show selected info
                             for key in info:
@@ -65,7 +53,6 @@ class GalleryApp:
 
     def selection_panel(self, items):
         # temperal function
-        preprocessor = st.radio('Preprocess Method', ['crop', 'resize'], horizontal=True)
 
         selecters = st.columns([1, 4])
 
@@ -101,7 +88,7 @@ class GalleryApp:
                 with sub_selecters[2]:
                     pop_weight = st.number_input('Popularity Weight', min_value=-100.0, max_value=100.0, value=st.session_state.score_weights[2], step=0.1, help='the weight for normalized popularity score')
 
-                items.loc[:, 'weighted_score_sum'] = round(items[f'norm_clip_{preprocessor}'] * clip_weight + items[f'norm_mcos_{preprocessor}'] * mcos_weight + items[
+                items.loc[:, 'weighted_score_sum'] = round(items[f'norm_clip'] * clip_weight + items[f'norm_mcos'] * mcos_weight + items[
                     'norm_pop'] * pop_weight, 4)
 
                 continue_idx = 3
@@ -168,7 +155,7 @@ class GalleryApp:
         # select number of columns
         col_num = st.slider('Number of columns', min_value=1, max_value=9, value=4, step=1, key='col_num')
 
-        return items, info, col_num, preprocessor
+        return items, info, col_num
 
     def sidebar(self):
         with st.sidebar:
@@ -226,7 +213,7 @@ class GalleryApp:
         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()
+        prompt_tags, tag, prompt_id, items = self.sidebar()
 
         # add safety check for some prompts
         safety_check = True
@@ -245,7 +232,7 @@ class GalleryApp:
             safety_check = st.checkbox('I understand that this prompt may contain unsafe content. Show these images anyway.', key=f'{prompt_id}')
 
         if safety_check:
-            items, info, col_num, preprocessor = self.selection_panel(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, [])))
@@ -260,7 +247,7 @@ class GalleryApp:
                 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, preprocessor, method))
+                        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])
@@ -311,7 +298,7 @@ class GalleryApp:
             print(st.session_state.selected_dict, 'continue')
             st.experimental_rerun()
 
-    def dynamic_weight(self, prompt_id, items, preprocessor='crop', method='Grid Search'):
+    def dynamic_weight(self, prompt_id, items, method='Grid Search'):
         selected = items[
             items['modelVersion_id'].isin(st.session_state.selected_dict[prompt_id])].reset_index(drop=True)
         optimal_weight = [0, 0, 0]
@@ -324,10 +311,10 @@ class GalleryApp:
                 for mcos_weight in np.arange(-1, 1, 0.1):
                     for pop_weight in np.arange(-1, 1, 0.1):
 
-                        weight_all = clip_weight*items[f'norm_clip_{preprocessor}'] + mcos_weight*items[f'norm_mcos_{preprocessor}'] + pop_weight*items['norm_pop']
+                        weight_all = clip_weight*items[f'norm_clip'] + mcos_weight*items[f'norm_mcos'] + pop_weight*items['norm_pop']
                         weight_all_sorted = weight_all.sort_values(ascending=False).reset_index(drop=True)
                         # print('weight_all_sorted:', weight_all_sorted)
-                        weight_selected = clip_weight*selected[f'norm_clip_{preprocessor}'] + mcos_weight*selected[f'norm_mcos_{preprocessor}'] + pop_weight*selected['norm_pop']
+                        weight_selected = clip_weight*selected[f'norm_clip'] + mcos_weight*selected[f'norm_mcos'] + pop_weight*selected['norm_pop']
 
                         # get the index of values of weight_selected in weight_all_sorted
                         rankings = []
@@ -438,6 +425,9 @@ def load_hf_dataset():
 if __name__ == "__main__":
     st.set_page_config(page_title="Model Coffer Gallery", page_icon="🖼️", layout="wide")
 
+    # remove ranking in the session state if it is created in Ranking.py
+    st.session_state.pop('ranking', None)
+
     if 'user_id' not in st.session_state:
         st.warning('Please log in first.')
         home_btn = st.button('Go to Home Page')

pages/Ranking.py

@@ -8,24 +8,105 @@ from streamlit_extras.switch_page_button import switch_page
 from pages.Gallery import load_hf_dataset
 
 
-class RankingApp():
-    def __init__(self, promptBook, images_ds):
+class RankingApp:
+    def __init__(self, promptBook, images_endpoint, batch_size=4):
         self.promptBook = promptBook
-        self.images_ds = images_ds
+        self.images_endpoint = images_endpoint
+        self.batch_size = batch_size
+        # self.batch_num = len(self.promptBook) // self.batch_size
+        # self.batch_num += 1 if len(self.promptBook) % self.batch_size != 0 else 0
 
-    def draggable_images(self, items, layout='vertical'):
-        pass
+        if 'counter' not in st.session_state:
+            st.session_state.counter = 0
 
     def sidebar(self):
         with st.sidebar:
             prompt_tags = self.promptBook['tag'].unique()
+            prompt_tags = np.sort(prompt_tags)
 
+            tag = st.selectbox('Select a prompt tag', prompt_tags)
+            items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)
+            prompts = np.sort(items['prompt'].unique())[::-1]
+
+            selected_prompt = st.selectbox('Select a prompt', prompts)
+
+            items = items[items['prompt'] == selected_prompt].reset_index(drop=True)
+            prompt_id = items['prompt_id'].unique()[0]
+
+            with st.form(key='prompt_form'):
+                # input image metadata
+                prompt = st.text_area('Prompt', selected_prompt, height=150, key='prompt', disabled=True)
+                negative_prompt = st.text_area('Negative Prompt', items['negativePrompt'].unique()[0], height=150, key='negative_prompt', disabled=True)
+                st.form_submit_button('Generate Images', type='primary', use_container_width=True)
+
+        return prompt_tags, tag, prompt_id, items
+
+    def draggable_images(self, items, layout='portrait'):
+        # init ranking by the order of items
+        if 'ranking' not in st.session_state:
+            st.session_state.ranking = {}
+            for i in range(len(items)):
+                st.session_state.ranking[str(items['image_id'][i])] = i
+
+        print(items)
+        with elements('dashboard'):
+            if layout == 'portrait':
+                col_num = 4
+                layout = [dashboard.Item(str(items['image_id'][i]), i % col_num, i//col_num, 1, 2, isResizable=False) for i in range(len(items))]
+
+            elif layout == 'landscape':
+                col_num = 2
+                layout = [
+                    dashboard.Item(str(items['image_id'][i]), i % col_num * 2, i // col_num, 2, 1.4, isResizable=False) for
+                    i in range(len(items))
+                ]
+
+            with dashboard.Grid(layout, cols={'lg': 4, 'md': 4, 'sm': 4, 'xs': 4, 'xxs': 2}, onLayoutChange=self.handle_layout_change, margin=[18, 18], containerPadding=[0, 0]):
+                for i in range(len(layout)):
+                    with mui.Card(key=str(items['image_id'][i]), variant="outlined"):
+                        rank = st.session_state.ranking[str(items['image_id'][i])] + 1
+
+                        mui.Chip(label=rank,
+                                 # variant="outlined" if rank!=1 else "default",
+                                 color="primary" if rank == 1 else "warning" if rank == 2 else "info",
+                                 size="small",
+                                 sx={"position": "absolute", "left": "-0.3rem", "top": "-0.3rem"})
+
+                        img_url = self.images_endpoint + str(items['image_id'][i]) + '.png'
+
+                        mui.CardMedia(
+                            component="img",
+                            # image={"data:image/png;base64", img_str},
+                            image=img_url,
+                            alt="There should be an image",
+                            sx={"height": "100%", "object-fit": "fit", 'bgcolor': 'black'},
+                        )
+
+    def handle_layout_change(self, updated_layout):
+        # print(updated_layout)
+        sorted_list = sorted(updated_layout, key=lambda x: (x['y'], x['x']))
+        sorted_list = [str(item['i']) for item in sorted_list]
+
+        for k in st.session_state.ranking.keys():
+            st.session_state.ranking[k] = sorted_list.index(k)
 
     def app(self):
         st.title('Personal Image Ranking')
         st.write('Here you can test out your selected images with any prompt you like.')
+        # st.write(self.promptBook)
 
-        prompt_tags, tag, prompt_id, items= self.sidebar()
+        prompt_tags, tag, prompt_id, items = self.sidebar()
+
+        sorting, control = st.columns((11, 1), gap='large')
+        with sorting:
+            # st.write('## Sorting')
+            # st.write('Please drag the images to sort them.')
+            st.progress((st.session_state.counter + 1) / self.batch_num, text=f"Batch {st.session_state.counter + 1} / {self.batch_num}")
+            self.draggable_images(items.iloc[self.batch_size*st.session_state.counter: self.batch_size*(st.session_state.counter+1)], layout='portrait')
+
+        with control:
+            st.button(":arrow_right:")
+            st.button(":slightly_frowning_face:")
 
 
 if __name__ == "__main__":
@@ -50,11 +131,22 @@ if __name__ == "__main__":
             if gallery_btn:
                 switch_page('gallery')
         else:
-            st.write('You have checked ' + str(len(selected_modelVersions)) + ' images.')
+            # st.write('You have checked ' + str(len(selected_modelVersions)) + ' images.')
             roster, promptBook, images_ds = load_hf_dataset()
-            st.write("## roster")
-            st.write(roster[roster['modelVersion_id'].isin(selected_modelVersions)])
+            print(st.session_state.selected_dict)
+            # st.write("## roster")
+            # st.write(roster[roster['modelVersion_id'].isin(selected_modelVersions)])
             # st.write(roster)
             # st.write("## promptBook")
             # st.write(promptBook)
 
+            # only select the part of the promptbook where tag is the same as st.session_state.selected_dict.keys(), while model version ids are the same as corresponding values to each key
+            promptBook_selected = pd.DataFrame()
+            for key, value in st.session_state.selected_dict.items():
+                promptBook_selected = promptBook_selected.append(promptBook[(promptBook['prompt_id'] == key) & (promptBook['modelVersion_id'].isin(value))])
+            promptBook_selected = promptBook_selected.reset_index(drop=True)
+            images_endpoint = "https://modelcofferbucket.s3-accelerate.amazonaws.com/"
+
+            app = RankingApp(promptBook_selected, images_endpoint, batch_size=4)
+            app.app()
+