beta for battle ranking

pages/Gallery.py

@@ -327,6 +327,7 @@ class GalleryApp:
                         deselect = st.form_submit_button('Deselect', use_container_width=True)
                         if deselect:
                             st.session_state.selected_dict[item['prompt_id']].remove(item['modelVersion_id'])
+                            self.remove_ranking_states(item['prompt_id'])
                             st.experimental_rerun()
 
                     else:
@@ -334,6 +335,7 @@ class GalleryApp:
                         select = st.form_submit_button('Select', use_container_width=True, type='primary')
                         if select:
                             st.session_state.selected_dict[item['prompt_id']].append(item['modelVersion_id'])
+                            self.remove_ranking_states(item['prompt_id'])
                             st.experimental_rerun()
 
                     # st.write(item)
@@ -408,7 +410,8 @@ class GalleryApp:
 
     def submit_actions(self, status, prompt_id):
         # remove counter from session state
-        st.session_state.pop('counter', None)
+        # st.session_state.pop('counter', None)
+        self.remove_ranking_states('prompt_id')
         if status == 'Select':
             modelVersions = self.promptBook[self.promptBook['prompt_id'] == prompt_id]['modelVersion_id'].unique()
             st.session_state.selected_dict[prompt_id] = modelVersions.tolist()
@@ -498,6 +501,30 @@ class GalleryApp:
         st.session_state.score_weights[0: 3] = optimal_weight
 
 
+    def remove_ranking_states(self, prompt_id):
+        # for drag sort
+        try:
+            st.session_state.counter[prompt_id] = 0
+            st.session_state.ranking[prompt_id] = {}
+            print('remove ranking states')
+        except:
+            print('no sort ranking states to remove')
+
+        # for battles
+        try:
+            st.session_state.pointer[prompt_id] = {'left': 0, 'right': 1}
+            print('remove battles states')
+        except:
+            print('no battles states to remove')
+
+        # for page progress
+        try:
+            st.session_state.progress[prompt_id] = 'ranking'
+            print('reset page progress states')
+        except:
+            print('no page progress states to be reset')
+
+
 # hist_data = pd.DataFrame(np.random.normal(42, 10, (200, 1)), columns=["x"])
 @st.cache_resource
 def altair_histogram(hist_data, sort_by, mini, maxi):
@@ -551,6 +578,13 @@ def load_hf_dataset():
     # add column to record current row index
     promptBook.loc[:, 'row_idx'] = promptBook.index
 
+    # apply a nsfw filter
+    promptBook = promptBook[promptBook['nsfw_score'] <= 0.84].reset_index(drop=True)
+
+    # add a column that adds up 'norm_clip', 'norm_mcos', and 'norm_pop'
+    score_weights = [1.0, 0.8, 0.2]
+    promptBook.loc[:, 'total_score'] = round(promptBook['norm_clip'] * score_weights[0] + promptBook['norm_mcos'] * score_weights[1] + promptBook['norm_pop'] * score_weights[2], 4)
+
     return roster, promptBook, images_ds
 
 @st.cache_data
@@ -569,9 +603,6 @@ def load_tsne_coordinates(items):
 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

@@ -34,6 +34,8 @@ class RankingApp:
 
             selected_prompt = st.selectbox('Select a prompt', prompts)
 
+            mode = st.radio('Select a mode', ['Drag and Sort', 'Battle'], index=1)
+
             items = items[items['prompt'] == selected_prompt].reset_index(drop=True)
             prompt_id = items['prompt_id'].unique()[0]
 
@@ -43,7 +45,7 @@ class RankingApp:
                 negative_prompt = st.text_area('Negative Prompt', items['negativePrompt'].unique()[0], height=150, key='negative_prompt', disabled=True)
                 st.form_submit_button('Generate Images [Coming Soon]', type='primary', use_container_width=True, disabled=True)
 
-        return prompt_tags, tag, prompt_id, items
+        return prompt_tags, tag, prompt_id, items, mode
 
     def draggable_images(self, items, prompt_id, layout='portrait'):
         # init ranking by the order of items
@@ -60,7 +62,6 @@ class RankingApp:
                 st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]][str(items['image_id'][i])] = i
         else:
             # set the index of items to the corresponding ranking value of the image_id
-            print(items['image_id'])
             items.index = items['image_id'].apply(lambda x: st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]][str(x)])
 
         with elements('dashboard'):
@@ -110,6 +111,121 @@ class RankingApp:
         for k in st.session_state.ranking[prompt_id][batch_idx].keys():
             st.session_state.ranking[prompt_id][batch_idx][k] = sorted_list.index(k)
 
+    def dragsort_mode(self, items, prompt_id, batch_num):
+        st.session_state.counter[prompt_id] = 0 if prompt_id not in st.session_state.counter else \
+        st.session_state.counter[prompt_id]
+
+        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[prompt_id] + 1) / batch_num,
+                        text=f"Batch {st.session_state.counter[prompt_id] + 1} / {batch_num}")
+            # st.write(items.iloc[self.batch_size*st.session_state.counter[prompt_id]: self.batch_size*(st.session_state.counter[prompt_id]+1)])
+
+            width, height = items.loc[0, 'size'].split('x')
+            if int(height) >= int(width):
+                self.draggable_images(items.iloc[
+                                      self.batch_size * st.session_state.counter[prompt_id]: self.batch_size * (
+                                                  st.session_state.counter[prompt_id] + 1)].reset_index(drop=True),
+                                      prompt_id=prompt_id, layout='portrait')
+            else:
+                self.draggable_images(items.iloc[
+                                      self.batch_size * st.session_state.counter[prompt_id]: self.batch_size * (
+                                                  st.session_state.counter[prompt_id] + 1)].reset_index(drop=True),
+                                      prompt_id=prompt_id, layout='landscape')
+            # st.write(str(st.session_state.ranking))
+        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)
+            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)
+
+            if st.session_state.counter[prompt_id] > 0:
+                st.button(":arrow_left:", key='prev', on_click=self.prev_batch, help='Previous Batch',
+                          kwargs={'prompt_id': prompt_id}, use_container_width=True)
+
+    def next_batch(self, prompt_id, progress=None):
+        # save ranking to dataset
+        # print(st.session_state.ranking)
+        # ranking_dataset = datasets.load_dataset('MAPS-research/GEMRec-Ranking', split='train')
+        curser = RANKING_CONN.cursor()
+        for image_id in st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]].keys():
+            modelVersion_id = self.promptBook[self.promptBook['image_id'] == image_id]['modelVersion_id'].values[0]
+            ranking = st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]][image_id]
+            # print({'image_id': image_id, 'modelVersion_id': modelVersion_id, 'ranking': ranking, "user_name": st.session_state.user_id[0], "timestamp": st.session_state.user_id[1]})
+            # ranking_dataset = ranking_dataset.add_item({'image_id': image_id, 'modelVersion_id': modelVersion_id, 'ranking': ranking, "user_name": st.session_state.user_id[0], "timestamp": st.session_state.user_id[1]})
+
+            # remove the old ranking if exists
+            query = "DELETE FROM rankings WHERE image_id = %s AND user_name = %s AND timestamp = %s"
+            curser.execute(query, (image_id, st.session_state.user_id[0], st.session_state.user_id[1]))
+
+            query = "INSERT INTO rankings (image_id, modelVersion_id, ranking, user_name, timestamp) VALUES (%s, %s, %s, %s, %s)"
+            curser.execute(query, (image_id, modelVersion_id, ranking, st.session_state.user_id[0], st.session_state.user_id[1]))
+
+        curser.close()
+        RANKING_CONN.commit()
+        # ranking_dataset.push_to_hub('MAPS-research/GEMRec-Ranking', split='train')
+
+        if progress == 'finished':
+            st.session_state.progress[prompt_id] = 'finished'
+        else:
+            st.session_state.counter[prompt_id] += 1
+
+    def prev_batch(self, prompt_id):
+        st.session_state.counter[prompt_id] -= 1
+
+    def battle_images(self, items, prompt_id):
+        if 'pointer' not in st.session_state:
+            st.session_state.pointer = {}
+
+        if prompt_id not in st.session_state.pointer:
+            st.session_state.pointer[prompt_id] = {'left': 0, 'right': 1}
+
+        curr_position = max(st.session_state.pointer[prompt_id]['left'], st.session_state.pointer[prompt_id]['right'])
+        progress = st.progress(curr_position / (len(items)-1), text=f"Progress {curr_position} / {len(items)-1}")
+
+        # if curr_position == len(items) - 1:
+        #     st.session_state.progress[prompt_id] = 'finished'
+        #
+        # else:
+        left, right = st.columns(2)
+        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}')
+
+            btn_left = st.button('Left is better', key='left', on_click=self.next_battle, kwargs={'prompt_id': prompt_id, 'winner': 'left', 'curr_position': curr_position, 'total_num': len(items)}, use_container_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}')
+
+            btn_right = st.button('Right is better', key='right', on_click=self.next_battle, kwargs={'prompt_id': prompt_id, 'winner': 'right', 'curr_position': curr_position, 'total_num': len(items)}, use_container_width=True)
+
+    def next_battle(self, prompt_id, winner, curr_position, total_num):
+        loser = 'left' if winner == 'right' else 'right'
+
+        if curr_position == total_num - 1:
+            st.session_state.progress[prompt_id] = 'finished'
+            # st.experimental_rerun()
+        else:
+            st.session_state.pointer[prompt_id][loser] = curr_position + 1
+
+    def battle_mode(self, items, prompt_id):
+        self.battle_images(items, prompt_id)
+
     def app(self):
         st.title('Personal Image Ranking')
         st.write('Here you can test out your selected images with any prompt you like.')
@@ -120,11 +236,11 @@ class RankingApp:
             st.session_state.progress = {}
         print('current progress: ', st.session_state.progress)
 
-        prompt_tags, tag, prompt_id, items = self.sidebar()
+        prompt_tags, tag, prompt_id, items, mode = self.sidebar()
         batch_num = len(items) // self.batch_size
         batch_num += 1 if len(items) % self.batch_size != 0 else 0
 
-        st.session_state.counter[prompt_id] = 0 if prompt_id not in st.session_state.counter else st.session_state.counter[prompt_id]
+        # st.session_state.counter[prompt_id] = 0 if prompt_id not in st.session_state.counter else st.session_state.counter[prompt_id]
 
         # save prompt_id in session state
         st.session_state.prompt_id_tmp = prompt_id
@@ -133,27 +249,10 @@ class RankingApp:
             st.session_state.progress[prompt_id] = 'ranking'
 
         if st.session_state.progress[prompt_id] == 'ranking':
-            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[prompt_id] + 1) / batch_num, text=f"Batch {st.session_state.counter[prompt_id] + 1} / {batch_num}")
-                # st.write(items.iloc[self.batch_size*st.session_state.counter[prompt_id]: self.batch_size*(st.session_state.counter[prompt_id]+1)])
-
-                width, height = items.loc[0, 'size'].split('x')
-                if int(height) >= int(width):
-                    self.draggable_images(items.iloc[self.batch_size*st.session_state.counter[prompt_id]: self.batch_size*(st.session_state.counter[prompt_id]+1)].reset_index(drop=True), prompt_id=prompt_id, layout='portrait')
-                else:
-                    self.draggable_images(items.iloc[self.batch_size*st.session_state.counter[prompt_id]: self.batch_size*(st.session_state.counter[prompt_id]+1)].reset_index(drop=True), prompt_id=prompt_id, layout='landscape')
-                # st.write(str(st.session_state.ranking))
-            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)
-                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)
-
-                if st.session_state.counter[prompt_id] > 0:
-                    st.button(":arrow_left:", key='prev', on_click=self.prev_batch, help='Previous Batch', kwargs={'prompt_id': prompt_id}, use_container_width=True)
+            if mode == 'Drag and Sort':
+                self.dragsort_mode(items, prompt_id, batch_num)
+            elif mode == 'Battle':
+                self.battle_mode(items, prompt_id)
 
         elif st.session_state.progress[prompt_id] == 'finished':
             st.write('## You have ranked all models for this tag!')
@@ -171,40 +270,10 @@ class RankingApp:
             if restart_btn:
                 st.session_state.progress[prompt_id] = 'ranking'
                 st.session_state.counter[prompt_id] = 0
+                st.session_state.pointer[prompt_id] = {'left': 0, 'right': 1}
                 st.experimental_rerun()
 
 
-    def next_batch(self, prompt_id, progress=None):
-        # save ranking to dataset
-        # print(st.session_state.ranking)
-        # ranking_dataset = datasets.load_dataset('MAPS-research/GEMRec-Ranking', split='train')
-        curser = RANKING_CONN.cursor()
-        for image_id in st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]].keys():
-            modelVersion_id = self.promptBook[self.promptBook['image_id'] == image_id]['modelVersion_id'].values[0]
-            ranking = st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]][image_id]
-            # print({'image_id': image_id, 'modelVersion_id': modelVersion_id, 'ranking': ranking, "user_name": st.session_state.user_id[0], "timestamp": st.session_state.user_id[1]})
-            # ranking_dataset = ranking_dataset.add_item({'image_id': image_id, 'modelVersion_id': modelVersion_id, 'ranking': ranking, "user_name": st.session_state.user_id[0], "timestamp": st.session_state.user_id[1]})
-
-            # remove the old ranking if exists
-            query = "DELETE FROM rankings WHERE image_id = %s AND user_name = %s AND timestamp = %s"
-            curser.execute(query, (image_id, st.session_state.user_id[0], st.session_state.user_id[1]))
-
-            query = "INSERT INTO rankings (image_id, modelVersion_id, ranking, user_name, timestamp) VALUES (%s, %s, %s, %s, %s)"
-            curser.execute(query, (image_id, modelVersion_id, ranking, st.session_state.user_id[0], st.session_state.user_id[1]))
-
-        curser.close()
-        RANKING_CONN.commit()
-        # ranking_dataset.push_to_hub('MAPS-research/GEMRec-Ranking', split='train')
-
-        if progress == 'finished':
-            st.session_state.progress[prompt_id] = 'finished'
-        else:
-            st.session_state.counter[prompt_id] += 1
-
-    def prev_batch(self, prompt_id):
-        st.session_state.counter[prompt_id] -= 1
-
-
 def connect_to_db():
     conn = pymysql.connect(
         host=os.environ.get('RANKING_DB_HOST'),
@@ -261,11 +330,14 @@ if __name__ == "__main__":
                 residual = len(user_selections) % 4
                 if residual != 0:
                     # select 4-residual items from the promptbook outside the user_selections
-                    npc = promptBook[(promptBook['prompt_id'] == key) & (~promptBook['modelVersion_id'].isin(value))].sort_values(by=['model_download_count'], ascending=False).reset_index(drop=True).iloc[:4-residual]
+                    npc = promptBook[(promptBook['prompt_id'] == key) & (~promptBook['modelVersion_id'].isin(value))].sort_values(by=['total_score'], ascending=False).reset_index(drop=True).iloc[:4-residual]
                     user_selections = pd.concat([user_selections, npc])
 
                 promptBook_selected = pd.concat([promptBook_selected, user_selections])
             promptBook_selected = promptBook_selected.reset_index(drop=True)
+            # sort promptBook by total_score
+            promptBook_selected = promptBook_selected.sort_values(by=['total_score'], ascending=True).reset_index(drop=True)
+
             # st.write(promptBook_selected)
             images_endpoint = "https://modelcofferbucket.s3-accelerate.amazonaws.com/"