fully connected to database

add a counter called 'epoch'
add checked indicator in ranking page
better button logic

Home.py

@@ -1,7 +1,10 @@
-from datetime import datetime
-import streamlit as st
+import os
 import random
-import time
+
+import pymysql.cursors
+import streamlit as st
+
+from datetime import datetime
 from streamlit_extras.switch_page_button import switch_page
 
 
@@ -31,27 +34,32 @@ def login():
 
 
 def save_user_id(user_id):
+    user_id = user_id[:60]
     print(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(['Drag and Sort', 'Battle'])
+    st.session_state.epoch = {'gallery': 0, 'ranking': {}, 'summary': {'overall': 0}}
 
 
 def logout():
     st.session_state.pop('user_id', None)
     st.session_state.pop('selected_dict', None)
+    st.session_state.pop('epoch', None)
     st.session_state.pop('score_weights', None)
     st.session_state.pop('gallery_state', None)
     st.session_state.pop('edit_state', None)
     st.session_state.pop('progress', None)
+    st.session_state.pop('pointer', None)
+    st.session_state.pop('counter', None)
     st.session_state.pop('gallery_focus', None)
     st.session_state.pop('assigned_rank_mode', None)
     st.session_state.pop('show_NSFW', None)
     st.session_state.pop('modelVersion_standings', None)
 
 
-def info():
+def project_info():
     with st.sidebar:
         st.write('## About')
         st.write(
@@ -63,10 +71,24 @@ def info():
         )
 
 
+def connect_to_db():
+    conn = pymysql.connect(
+        host=os.environ.get('RANKING_DB_HOST'),
+        port=3306,
+        database='myRanking',
+        user=os.environ.get('RANKING_DB_USER'),
+        password=os.environ.get('RANKING_DB_PASSWORD'),
+        charset='utf8mb4',
+        cursorclass=pymysql.cursors.DictCursor
+    )
+
+    return conn
+
+
 if __name__ == '__main__':
     # print(st.source_util.get_pages('Home.py'))
     st.set_page_config(page_title="Login", page_icon="🏠", layout="wide")
-    info()
+    project_info()
     st.write('A Research by [MAPS Lab](https://whongyi.github.io/MAPS-research), [NYU Shanghai](https://shanghai.nyu.edu)')
     st.title("🙌 Welcome to GEMRec Gallery!")
 

pages/Gallery.py

@@ -12,6 +12,7 @@ import streamlit.components.v1 as components
 
 from bs4 import BeautifulSoup
 from datasets import load_dataset, Dataset, load_from_disk
+from datetime import datetime
 from huggingface_hub import login
 from streamlit_agraph import agraph, Node, Edge, Config
 from streamlit_extras.switch_page_button import switch_page
@@ -19,6 +20,7 @@ from streamlit_extras.tags import tagger_component
 from streamlit_extras.no_default_selectbox import selectbox
 from sklearn.svm import LinearSVC
 
+from Home import connect_to_db
 
 class GalleryApp:
     def __init__(self, promptBook, images_ds):
@@ -203,7 +205,7 @@ class GalleryApp:
 
         # remove coloring from tag
         tag = self.text_coloring_remove(tag)
-        print('tag: ', tag)
+        # print('tag: ', tag)
 
         # print('current state: ', st.session_state.gallery_state)
 
@@ -237,7 +239,7 @@ class GalleryApp:
 
                     # remove coloring from prompt
                     selected_prompt = self.text_coloring_remove(selected_prompt)
-                    print('selected_prompt: ', selected_prompt)
+                    # print('selected_prompt: ', selected_prompt)
                     st.session_state.prompt_idx_last_time = prompts.index(selected_prompt) if selected_prompt else 0
 
             if selected_prompt is None:
@@ -282,15 +284,15 @@ class GalleryApp:
                         pass
 
                     if has_selection:
-                        checkout = st.button('Check out selections ➡️', use_container_width=True, type='primary')
-                        if 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 = 'check out'
-                            # print(st.session_state.gallery_state)
-                            st.experimental_rerun()
+                        checkout = st.button('Check out selections ➡️', use_container_width=True, type='primary', on_click=self.switch_to_checkout, args=(tag, selected_prompt))
+                        # if 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 = 'check out'
+                        #     # print(st.session_state.gallery_state)
+                        #     st.rerun()
                     else:
                         st.button(':orange[👇 **Select images you like below**]', disabled=True, use_container_width=True)
                 try:
@@ -308,6 +310,12 @@ class GalleryApp:
 
             self.checkout_mode(tag, items)
 
+    def switch_to_checkout(self, tag, selected_prompt):
+        # 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 = 'check out'
 
     def random_gallery_focus(self, tags):
         st.session_state.gallery_focus['tag'] = random.choice(tags)
@@ -315,7 +323,6 @@ class GalleryApp:
         prompts = self.promptBook[self.promptBook['tag'] == st.session_state.gallery_focus['tag']]['prompt'].unique()
         st.session_state.gallery_focus['prompt'] = random.choice(prompts)
 
-
     def graph_mode(self, prompt_id, items):
         graph_cols = st.columns([3, 1])
         # prompt = st.chat_input(f"Selected model version ids: {str(st.session_state.selected_dict.get(prompt_id, []))}",
@@ -353,24 +360,11 @@ class GalleryApp:
 
                     if checked:
                         # deselect = st.button('Deselect', key=f'select_{item["prompt_id"]}_{item["modelVersion_id"]}', use_container_width=True)
-                        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()
+                        deselect = st.form_submit_button('Deselect', use_container_width=True, on_click=self.image_selection_control, args=(item['tag'], item['prompt'], item['prompt_id'], item['modelVersion_id'], False))
 
                     else:
                         # select = st.button('Select', key=f'select_{item["prompt_id"]}_{item["modelVersion_id"]}', use_container_width=True, type='primary')
-                        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'])
-
-                            # add focus to session state
-                            st.session_state.gallery_focus['tag'] = item['tag']
-                            st.session_state.gallery_focus['prompt'] = item['prompt']
-
-                            st.experimental_rerun()
+                        select = st.form_submit_button('Select', use_container_width=True, type='primary', on_click=self.image_selection_control, args=(item['tag'], item['prompt'], item['prompt_id'], item['modelVersion_id'], True))
 
                     # st.write(item)
                     infos = ['model_name', 'modelVersion_name', 'model_download_count', 'clip_score', 'mcos_score',
@@ -384,6 +378,17 @@ class GalleryApp:
             else:
                 st.info('Please click on an image to show')
 
+    def image_selection_control(self, tag, prompt, prompt_id, modelVersion_id, to_select):
+        self.remove_ranking_states(prompt_id)
+        if to_select:
+            st.session_state.selected_dict[prompt_id].append(modelVersion_id)
+            # add focus to session state
+            st.session_state.gallery_focus['tag'] = tag
+            st.session_state.gallery_focus['prompt'] = prompt
+
+        else:
+            st.session_state.selected_dict[prompt_id].remove(modelVersion_id)
+
     def checkout_mode(self, tag, items):
         # st.write(items)
         if len(items) > 0:
@@ -407,7 +412,7 @@ class GalleryApp:
                                               ['model_name', 'model_id', 'modelVersion_name', 'modelVersion_id',
                                                'total_score', 'model_download_count', 'clip_score', 'mcos_score',
                                                'norm_nsfw'],
-                                              label_visibility='collapsed', key=f'info_{prompt_id}', placeholder='Select what infos to show')
+                                              label_visibility='collapsed', key=f'info_{prompt_id}', placeholder='Select what info to show')
 
                     with checkout_panel[-1]:
                         checkout_buttons = st.columns([1, 1, 1])
@@ -418,7 +423,7 @@ class GalleryApp:
                                 st.session_state.gallery_focus['prompt'] = prompt
                                 print(st.session_state.gallery_focus)
                                 st.session_state.gallery_state = 'graph'
-                                st.experimental_rerun()
+                                st.rerun()
 
                         with checkout_buttons[1]:
                             # init edit state
@@ -429,7 +434,7 @@ class GalleryApp:
                                 edit = st.button('Edit', key=f'checkout_edit_{prompt_id}', use_container_width=True)
                                 if edit:
                                     st.session_state.edit_state = True
-                                    st.experimental_rerun()
+                                    st.rerun()
                             else:
                                 done = st.button('Done', key=f'checkout_done_{prompt_id}', use_container_width=True)
                                 if done:
@@ -443,7 +448,7 @@ class GalleryApp:
                                                 st.session_state.selected_dict[prompt_id].append(int(keys[2]))
 
                                     st.session_state.edit_state = False
-                                    st.experimental_rerun()
+                                    st.rerun()
 
                         with checkout_buttons[-1]:
                             proceed = st.button('Proceed ➡️', key=f'checkout_proceed_{prompt_id}', use_container_width=True,
@@ -452,6 +457,40 @@ class GalleryApp:
                                 st.session_state.gallery_focus['tag'] = tag
                                 st.session_state.gallery_focus['prompt'] = prompt
                                 st.session_state.gallery_state = 'graph'
+
+                                print('selected_dict: ', st.session_state.selected_dict)
+
+                                # # save the user selection to database
+                                cursor = GALLERY_CONN.cursor()
+                                st.session_state.epoch['gallery'] += 1
+                                checkouttime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+                                # for modelVersion_id in st.session_state.selected_dict[prompt_id]:
+                                for key, values in st.session_state.selected_dict.items():
+                                    # print('key: ', key, 'values: ', values)
+                                    key_tag = self.promptBook[self.promptBook['prompt_id'] == key]['tag'].unique()[0]
+                                    for value in values:
+                                        query = "INSERT INTO gallery_selections (username, timestamp, tag, prompt_id, modelVersion_id, checkouttime, epoch) VALUES ('{}', '{}', '{}', '{}', {}, '{}', {})".format(st.session_state.user_id[0], st.session_state.user_id[1], key_tag, key, value, checkouttime, st.session_state.epoch['gallery'])
+                                        print(query)
+                                        cursor.execute(query)
+                                GALLERY_CONN.commit()
+                                cursor.close()
+
+                                # get the largest epoch number of this user and prompt
+                                cursor = GALLERY_CONN.cursor()
+                                db_table = 'battle_results' if st.session_state.assigned_rank_mode=='Battle' else 'sort_results'
+                                query = "SELECT MAX(epoch) FROM {} WHERE username = '{}' AND timestamp = '{}' AND prompt_id = {}".format(db_table, st.session_state.user_id[0], st.session_state.user_id[1], prompt_id)
+                                cursor.execute(query)
+                                max_epoch = cursor.fetchone()['MAX(epoch)'],
+                                # print('max epoch: ', max_epoch, type(max_epoch))
+                                cursor.close()
+
+                                try:
+                                    st.session_state.epoch['ranking'][prompt_id] = max_epoch[0] + 1
+                                except TypeError:
+                                    st.session_state.epoch['ranking'][prompt_id] = 1
+                                # st.session_state.epoch['summary'][tag] = st.session_state.epoch['summary'].get(tag, 0) + 1
+                                # st.session_state.epoch['summary']['overall'] += 1
+                                print('epoch: ', st.session_state.epoch)
                                 switch_page('ranking')
 
                     self.gallery_standard(items[items['prompt_id'] == prompt_id].reset_index(drop=True), 4, info, show_checkbox=st.session_state.edit_state)
@@ -463,7 +502,7 @@ class GalleryApp:
                 st.session_state.gallery_focus['tag'] = tag
                 st.session_state.gallery_focus['prompt'] = None
                 st.session_state.gallery_state = 'graph'
-                st.experimental_rerun()
+                st.rerun()
 
     def remove_ranking_states(self, prompt_id):
         # for drag sort
@@ -547,6 +586,7 @@ if __name__ == "__main__":
         if home_btn:
             switch_page("home")
     else:
+        GALLERY_CONN = connect_to_db()
         roster, promptBook, images_ds = load_hf_dataset(st.session_state.show_NSFW)
 
         app = GalleryApp(promptBook=promptBook, images_ds=images_ds)

pages/Ranking.py

@@ -11,6 +11,7 @@ from streamlit_elements import elements, mui, html, dashboard, nivo
 from streamlit_extras.switch_page_button import switch_page
 
 from pages.Gallery import load_hf_dataset
+from Home import connect_to_db
 
 
 class RankingApp:
@@ -122,7 +123,7 @@ 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):
+    def dragsort_mode(self, tag, 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]
 
@@ -149,16 +150,16 @@ 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, type='primary')
+                          kwargs={'tag': tag, '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, type='primary')
+                          kwargs={'tag': tag, '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',
                           kwargs={'prompt_id': prompt_id}, use_container_width=True)
 
-    def next_batch(self, prompt_id, progress=None):
+    def next_batch(self, tag, prompt_id, progress=None):
 
         curser = RANKING_CONN.cursor()
 
@@ -168,27 +169,36 @@ class RankingApp:
             modelVersion_id = self.promptBook[self.promptBook['image_id'] == image_id]['modelVersion_id'].values[0]
             position_version_dict[position] = modelVersion_id
 
-        # get all records of this user and prompt
-        query = "SELECT * FROM sort_results WHERE username = %s AND timestamp = %s AND prompt_id = %s"
-        curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id))
-        results = curser.fetchall()
-        print(results)
-
-        # remove the old ranking with the same modelVersion_id if exists
-        for result in results:
-            prev_ids = [result['position1'], result['position2'], result['position3'], result['position4']]
-            curr_ids = [position_version_dict[0], position_version_dict[1], position_version_dict[2], position_version_dict[3]]
-            if len(set(prev_ids).intersection(set(curr_ids))) == 4:
-                query = "DELETE FROM sort_results WHERE username = %s AND timestamp = %s AND prompt_id = %s AND position1 = %s AND position2 = %s AND position3 = %s AND position4 = %s"
-                curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id, result['position1'], result['position2'], result['position3'], result['position4']))
-
-        query = "INSERT INTO sort_results (username, timestamp, tag, prompt_id, position1, position2, position3, position4) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)"
-        curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], self.promptBook[self.promptBook['prompt_id'] == prompt_id]['tag'].values[0], prompt_id, position_version_dict[0], position_version_dict[1], position_version_dict[2], position_version_dict[3]))
+        # # get all records of this user and prompt
+        # query = "SELECT * FROM sort_results WHERE username = %s AND timestamp = %s AND prompt_id = %s"
+        # curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id))
+        # results = curser.fetchall()
+        # print(results)
+        #
+        # # remove the old ranking with the same modelVersion_id if exists
+        # for result in results:
+        #     prev_ids = [result['position1'], result['position2'], result['position3'], result['position4']]
+        #     curr_ids = [position_version_dict[0], position_version_dict[1], position_version_dict[2], position_version_dict[3]]
+        #     if len(set(prev_ids).intersection(set(curr_ids))) == 4:
+        #         query = "DELETE FROM sort_results WHERE username = %s AND timestamp = %s AND prompt_id = %s AND position1 = %s AND position2 = %s AND position3 = %s AND position4 = %s"
+        #         curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id, result['position1'], result['position2'], result['position3'], result['position4']))
+
+        sorttime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        # handle the case where user press the 'prev' button
+        query = "DELETE FROM sort_results WHERE username = %s AND timestamp = %s AND prompt_id = %s AND epoch = %s"
+        curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id, st.session_state.epoch['ranking'][prompt_id]))
+        query = "INSERT INTO sort_results (username, timestamp, tag, prompt_id, position1, position2, position3, position4, sorttime, epoch) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)"
+        curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], self.promptBook[self.promptBook['prompt_id'] == prompt_id]['tag'].values[0], prompt_id, position_version_dict[0], position_version_dict[1], position_version_dict[2], position_version_dict[3], sorttime, st.session_state.epoch['ranking'][prompt_id]))
 
         curser.close()
         RANKING_CONN.commit()
 
         if progress == 'finished':
+            st.session_state.epoch['ranking'][prompt_id] += 1
+
+            st.session_state.epoch['summary'][tag] = st.session_state.epoch['summary'].get(tag, 0) + 1
+            st.session_state.epoch['summary']['overall'] += 1
+
             st.session_state.progress[prompt_id] = 'finished'
             # drop 'modelVersion_standings' from session state if exists
             st.session_state.pop('modelVersion_standings', None)
@@ -198,7 +208,7 @@ class RankingApp:
     def prev_batch(self, prompt_id):
         st.session_state.counter[prompt_id] -= 1
 
-    def battle_images(self, items, prompt_id):
+    def battle_images(self, tag, items, prompt_id):
         if 'pointer' not in st.session_state:
             st.session_state.pointer = {}
 
@@ -221,7 +231,7 @@ class RankingApp:
             # 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)
+            btn_left = st.button('Left is better', key='left', on_click=self.next_battle, kwargs={'tag': tag, '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:
@@ -232,10 +242,10 @@ class RankingApp:
             # 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)
+            btn_right = st.button('Right is better', key='right', on_click=self.next_battle, kwargs={'tag': tag, '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):
+    def next_battle(self, tag, prompt_id, image_ids, winner, curr_position, total_num):
         loser = 'left' if winner == 'right' else 'right'
         battletime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
 
@@ -244,19 +254,24 @@ class RankingApp:
         winner_modelVersion_id = self.promptBook[self.promptBook['image_id'] == image_ids[st.session_state.pointer[prompt_id][winner]]]['modelVersion_id'].values[0]
         loser_modelVersion_id = self.promptBook[self.promptBook['image_id'] == image_ids[st.session_state.pointer[prompt_id][loser]]]['modelVersion_id'].values[0]
 
-        # remove the old battle result if exists
-        query = "DELETE FROM battle_results WHERE username = %s AND timestamp = %s AND prompt_id = %s AND winner = %s AND loser = %s"
-        curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id, winner_modelVersion_id, loser_modelVersion_id))
-        curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id, loser_modelVersion_id, winner_modelVersion_id))
+        # # remove the old battle result if exists
+        # query = "DELETE FROM battle_results WHERE username = %s AND timestamp = %s AND prompt_id = %s AND winner = %s AND loser = %s"
+        # curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id, winner_modelVersion_id, loser_modelVersion_id))
+        # curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id, loser_modelVersion_id, winner_modelVersion_id))
 
         # insert the battle result into the database
-        query = "INSERT INTO battle_results (username, timestamp, tag, prompt_id, winner, loser, battletime) VALUES (%s, %s, %s, %s, %s, %s, %s)"
-        curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], self.promptBook[self.promptBook['prompt_id'] == prompt_id]['tag'].values[0], prompt_id, winner_modelVersion_id, loser_modelVersion_id, battletime))
+        query = "INSERT INTO battle_results (username, timestamp, tag, prompt_id, winner, loser, battletime, epoch) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)"
+        curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], self.promptBook[self.promptBook['prompt_id'] == prompt_id]['tag'].values[0], prompt_id, winner_modelVersion_id, loser_modelVersion_id, battletime, st.session_state.epoch['ranking'][prompt_id]))
 
         curser.close()
         RANKING_CONN.commit()
 
         if curr_position == total_num - 1:
+            st.session_state.epoch['ranking'][prompt_id] += 1
+
+            st.session_state.epoch['summary'][tag] = st.session_state.epoch['summary'].get(tag, 0) + 1
+            st.session_state.epoch['summary']['overall'] += 1
+
             st.session_state.progress[prompt_id] = 'finished'
 
             # drop 'modelVersion_standings' from session state if exists
@@ -266,8 +281,8 @@ class RankingApp:
         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 battle_mode(self, tag, items, prompt_id):
+        self.battle_images(tag, items, prompt_id)
 
     def app(self):
         st.write('### Generative Model Ranking')
@@ -288,15 +303,43 @@ class RankingApp:
                                                                                   'tag'] in prompt_tags else 0
         print(tag_idx)
 
-        tag = st.radio('Select a tag', prompt_tags, index=tag_idx, horizontal=True, label_visibility='collapsed')
+        # color the finished tags
+        finished_tags = []
+        for tag in prompt_tags:
+            append_tag = True
+            for prompt_id in self.promptBook[self.promptBook['tag'] == tag]['prompt_id'].unique():
+                if prompt_id not in st.session_state.progress or st.session_state.progress[prompt_id] != 'finished':
+                    append_tag = False
+                    break
+            if append_tag:
+                finished_tags.append(tag)
+        tag_tobe_colored = self.promptBook[self.promptBook['tag'].isin(finished_tags)]['tag'].unique().tolist()
+        colored_tags = self.text_coloring_add(tag_tobe_colored, prompt_tags, color_name='orange')
+
+        tag = st.radio('Select a tag', colored_tags, index=tag_idx, horizontal=True, label_visibility='collapsed')
+        tag = self.text_coloring_remove(tag)
+        print(tag)
         items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)
+        # pick out prompts such that st.session_state.progress[prompt_id] == 'finished'
+
         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, index=prompt_idx, label_visibility='collapsed')
+        # color the finished prompts
+        finished_prompts = []
+        for prompt_id in items['prompt_id'].unique():
+            if prompt_id in st.session_state.progress and st.session_state.progress[prompt_id] == 'finished':
+                finished_prompts.append(prompt_id)
+        prompt_tobe_colored = items[items['prompt_id'].isin(finished_prompts)]['prompt'].unique().tolist()
+        colored_prompts = self.text_coloring_add(prompt_tobe_colored, prompts, color_name='✅')
+
+        selected_prompt = st.selectbox('Select a prompt', colored_prompts, index=prompt_idx, label_visibility='collapsed')
+        selected_prompt = self.text_coloring_remove(selected_prompt)
+
+        st.session_state.gallery_focus = {'tag': tag, 'prompt': selected_prompt}
 
         # mode = st.radio('Select a mode', ['Drag and Sort', 'Battle'], index=1)
         mode = st.session_state.assigned_rank_mode
@@ -317,13 +360,15 @@ class RankingApp:
             st.session_state.progress[prompt_id] = 'ranking'
 
         if st.session_state.progress[prompt_id] == 'ranking':
+            st.session_state.epoch['ranking'][prompt_id] = st.session_state.epoch['ranking'].get(prompt_id, 1)
             st.caption("We might pair some other images that you haven't selected based on our evaluation matrix.")
             if mode == 'Drag and Sort':
-                self.dragsort_mode(items, prompt_id, batch_num)
+                self.dragsort_mode(tag, items, prompt_id, batch_num)
             elif mode == 'Battle':
-                self.battle_mode(items, prompt_id)
+                self.battle_mode(tag, items, prompt_id)
 
         elif st.session_state.progress[prompt_id] == 'finished':
+            print(st.session_state.gallery_focus)
             # st.toast('**Summary is available now!**')
             # st.write('---')
             with st.form(key='ranking_finished'):
@@ -348,25 +393,45 @@ class RankingApp:
                     st.form_submit_button('👆 Rank other prompts', use_container_width=True)
 
                 with options_panel[3]:
-                    restart_btn = st.form_submit_button('🎖️ Re-rank this prompt', use_container_width=True)
-                    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 connect_to_db():
-    conn = pymysql.connect(
-        host=os.environ.get('RANKING_DB_HOST'),
-        port=3306,
-        database='myRanking',
-        user=os.environ.get('RANKING_DB_USER'),
-        password=os.environ.get('RANKING_DB_PASSWORD'),
-        charset='utf8mb4',
-        cursorclass=pymysql.cursors.DictCursor
-    )
-
-    return conn
+                    restart_btn = st.form_submit_button('🎖️ Re-rank this prompt', use_container_width=True, on_click=self.rerank, kwargs={'prompt_id': prompt_id})
+
+        # with st.sidebar:
+        #     st.write('epoch: ', st.session_state.epoch['ranking'][prompt_id])
+    def text_coloring_add(self, tobe_colored:list, total_items, color_name='orange'):
+        if color_name in ['orange', 'red', 'green', 'blue', 'violet', 'yellow']:
+            colored = [f':{color_name}[{item}]' if item in tobe_colored else item for item in total_items]
+        else:
+            colored = [f'[{color_name}] {item}' if item in tobe_colored else item for item in total_items]
+        return colored
+
+    def text_coloring_remove(self, tobe_removed):
+        if isinstance(tobe_removed, str):
+            if tobe_removed.startswith(':'):
+                tobe_removed = tobe_removed.split('[')[-1][:-1]
+
+            elif tobe_removed.startswith('['):
+                tobe_removed = tobe_removed.split(']')[-1][1:]
+        return tobe_removed
+
+    def rerank(self, prompt_id):
+        st.session_state.progress[prompt_id] = 'ranking'
+        if st.session_state.assigned_rank_mode == 'Drag and Sort':
+            st.session_state.counter[prompt_id] = 0
+        elif st.session_state.assigned_rank_mode == 'Battle':
+            st.session_state.pointer[prompt_id] = {'left': 0, 'right': 1}
+
+# def connect_to_db():
+#     conn = pymysql.connect(
+#         host=os.environ.get('RANKING_DB_HOST'),
+#         port=3306,
+#         database='myRanking',
+#         user=os.environ.get('RANKING_DB_USER'),
+#         password=os.environ.get('RANKING_DB_PASSWORD'),
+#         charset='utf8mb4',
+#         cursorclass=pymysql.cursors.DictCursor
+#     )
+#
+#     return conn
 
 
 if __name__ == "__main__":

pages/Summary.py

@@ -15,7 +15,7 @@ from streamlit_extras.stylable_container import stylable_container
 from st_clickable_images import clickable_images
 
 from pages.Gallery import load_hf_dataset
-from pages.Ranking import connect_to_db
+from Home import connect_to_db
 
 
 class DashboardApp:
@@ -42,35 +42,52 @@ class DashboardApp:
                 if back_to_ranking:
                     switch_page('ranking')
 
-            # with st.form('overall_feedback'):
-            #     feedback = st.text_area('Please leave your comments here.', key='comment')
-            #     submit_feedback = st.form_submit_button('Submit Feedback')
-            #     if submit_feedback:
-            #         print(feedback)
+            with st.form('overall_feedback'):
+                comment = st.text_area('Please leave your comments here.', key='comment')
+                submit_feedback = st.form_submit_button('Submit Feedback')
+                if submit_feedback:
+                    commenttime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+                    curser = RANKING_CONN.cursor()
+                    # parse the comment to at most 300 to avoid SQL injection
+                    for i in range(0, len(comment), 300):
+                        curser.execute(f"INSERT INTO comments (username, timestamp, comment, commenttime) VALUES ('{st.session_state.user_id[0]}', '{st.session_state.user_id[1]}', '{comment[i:i+300]}', '{commenttime}')")
+                    RANKING_CONN.commit()
+                    curser.close()
 
-        # return tag
+                    st.sidebar.info('🙏 **Thanks for your feedback! We will take it into consideration in our future work.**')
 
     def leaderboard(self, tag, db_table):
         tag = '%' if tag == 'overview' else tag
 
-        # get the ranking results of the current user
+        # print('tag', tag)
+
+        # get the ranking results of the current user with the lastest epoch
         curser = RANKING_CONN.cursor()
-        curser.execute(f"SELECT * FROM {db_table} WHERE username = '{st.session_state.user_id[0]}' AND timestamp = '{st.session_state.user_id[1]}' AND tag LIKE '{tag}'")
+        # curser.execute(f"SELECT * FROM {db_table} WHERE username = '{st.session_state.user_id[0]}' AND timestamp = '{st.session_state.user_id[1]}' AND tag LIKE '{tag}'")
+        # curser.execute(f"SELECT * FROM {db_table} WHERE username = '{st.session_state.user_id[0]}' AND timestamp = '{st.session_state.user_id[1]}' AND tag LIKE '{tag}' ORDER BY epoch DESC LIMIT 1")
+        curser.execute(
+            f"SELECT * FROM {db_table}\
+            WHERE username = '{st.session_state.user_id[0]}'\
+            AND timestamp = '{st.session_state.user_id[1]}'\
+            AND tag LIKE '{tag}'\
+            AND epoch =\
+                (SELECT MAX(epoch) FROM {db_table}\
+                WHERE username = '{st.session_state.user_id[0]}'\
+                AND timestamp = '{st.session_state.user_id[1]}'\
+                AND tag LIKE '{tag}')")
+
+
         results = curser.fetchall()
         curser.close()
 
+        # print('results', results, len(results))
+
         if tag not in st.session_state.modelVersion_standings:
             st.session_state.modelVersion_standings[tag] = self.score_calculator(results, db_table)
 
             # sort the modelVersion_standings by value into a list of tuples in descending order
             st.session_state.modelVersion_standings[tag] = sorted(st.session_state.modelVersion_standings[tag].items(), key=lambda x: x[1], reverse=True)
-
-        # tab1, tab2 = st.tabs(['Top Picks', 'Detailed Info'])
-
-        # with tab1:
-            # self.podium(modelVersion_standings)
-        # switch_stage = st.toggle('Manual Reorder', key='switch_stage')
-
+        print(st.session_state.modelVersion_standings[tag])
         example_prompts = []
         # get example images
         for key, value in st.session_state.selected_dict.items():
@@ -78,19 +95,8 @@ class DashboardApp:
                 if model[0] in value:
                     example_prompts.append(key)
 
-        # if switch_stage:
-        #     self.podium_expander(tag, n=len(st.session_state.modelVersion_standings[tag]), summary_mode='edit', example_prompts=example_prompts)
-        # else:
         self.podium_expander(tag, n=len(st.session_state.modelVersion_standings[tag]), summary_mode='display', example_prompts=example_prompts)
-        # if st.session_state.summary_mode == 'display':
-        #     switch_stage = st.button('Manual Reorder', key='switch_stage_edit', on_click=lambda: st.session_state.__setitem__('summary_mode', 'edit'))
-        #     self.podium_expander(tag, n=3, summary_mode='display')
-        #
-        # elif st.session_state.summary_mode == 'edit':
-        #     switch_stage = st.button('Done', key='switch_stage_done', type='primary', on_click=lambda: st.session_state.__setitem__('summary_mode', 'display'))
-        #     self.podium_expander(tag, n=len(st.session_state.modelVersion_standings[tag]), summary_mode='edit')
 
-        # with tab2:
         st.write('---')
         st.write('**Detailed information of all selected models**')
         detailed_info = pd.merge(pd.DataFrame(st.session_state.modelVersion_standings[tag], columns=['modelVersion_id', 'ranking_score']), self.roster, on='modelVersion_id')
@@ -101,6 +107,7 @@ class DashboardApp:
         st.caption('You can click the header to sort the table by that column.')
 
     def podium_expander(self, tag, example_prompts, n=3, summary_mode: ['display', 'edit'] = 'display'):
+        self.save_summary(tag)
 
         for i in range(n):
             modelVersion_id = st.session_state.modelVersion_standings[tag][i][0]
@@ -137,7 +144,7 @@ class DashboardApp:
                     example_images = [f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{image}.png" for image in example_images]
                     clickable_images(
                         example_images,
-                        img_style={"margin": "5px", "height": "120px"}
+                        img_style={"margin": "5px", "height": "120px"},
                     )
 
                 else:
@@ -173,19 +180,33 @@ class DashboardApp:
             if i != n - 1:
                 st.write('---')
 
+    def save_summary(self, tag):
+        # get the lastest summary_results epoch of the current user
+        tag_name = 'overview' if tag == '%' else tag
+        curser = RANKING_CONN.cursor()
+        curser.execute(f"SELECT epoch FROM summary_results WHERE username = '{st.session_state.user_id[0]}' AND timestamp = '{st.session_state.user_id[1]}' AND tag = '{tag_name}' ORDER BY epoch DESC LIMIT 1")
+        latest_epoch = curser.fetchone()
+        curser.close()
+        # print('latest_epoch',latest_epoch)
+        if latest_epoch is None or latest_epoch['epoch'] < st.session_state.epoch['summary'][tag_name]:
+            # save the current ranking results to the database
+            summarytime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            curser = RANKING_CONN.cursor()
+            for i in range(len(st.session_state.modelVersion_standings[tag])):
+                curser.execute(f"INSERT INTO summary_results (username, timestamp, tag, modelVersion_id, position, ranking_score, summarytime, epoch, customized) VALUES ('{st.session_state.user_id[0]}', '{st.session_state.user_id[1]}', '{tag_name}', '{st.session_state.modelVersion_standings[tag][i][0]}', {i+1}, {st.session_state.modelVersion_standings[tag][i][1]}, '{summarytime}', {st.session_state.epoch['summary'][tag_name]}, 0)")
+            RANKING_CONN.commit()
+            curser.close()
+
     def switch_order(self, tag, current, target):
-        # st.session_state.modelVersion_standings[tag][current], st.session_state.modelVersion_standings[tag][target] = st.session_state.modelVersion_standings[tag][target], st.session_state.modelVersion_standings[tag][current]
         # insert the current before the target
         st.session_state.modelVersion_standings[tag].insert(target, st.session_state.modelVersion_standings[tag].pop(current))
-        # print(st.session_state.modelVersion_standings[tag])
-        # curser = RANKING_CONN.cursor()
-        # RANKING_CONN.commit()
-        # curser.close()
+        tag_name = 'overview' if tag == '%' else tag
+        summarytime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         curser = RANKING_CONN.cursor()
         # clear the current user's ranking results
-        curser.execute(f"DELETE FROM summary_results WHERE username = '{st.session_state.user_id[0]}' AND timestamp = '{st.session_state.user_id[1]}' AND tag = '{tag}'")
+        curser.execute(f"DELETE FROM summary_results WHERE username = '{st.session_state.user_id[0]}' AND timestamp = '{st.session_state.user_id[1]}' AND tag = '{tag_name}' AND epoch = {st.session_state.epoch['summary'][tag_name]}")
         for i in range(len(st.session_state.modelVersion_standings[tag])):
-            curser.execute(f"INSERT INTO summary_results (username, timestamp, tag, modelVersion_id, position, ranking_score) VALUES ('{st.session_state.user_id[0]}', '{st.session_state.user_id[1]}', '{tag}', '{st.session_state.modelVersion_standings[tag][i][0]}', {i+1}, {st.session_state.modelVersion_standings[tag][i][1]})")
+            curser.execute(f"INSERT INTO summary_results (username, timestamp, tag, modelVersion_id, position, ranking_score, summarytime, epoch, customized) VALUES ('{st.session_state.user_id[0]}', '{st.session_state.user_id[1]}', '{tag_name}', '{st.session_state.modelVersion_standings[tag][i][0]}', {i+1}, {st.session_state.modelVersion_standings[tag][i][1]}, '{summarytime}', {st.session_state.epoch['summary'][tag_name]}, 1)")
         RANKING_CONN.commit()
         curser.close()
 
@@ -197,7 +218,6 @@ class DashboardApp:
 
             for record in results:
                 modelVersion_standings[record['winner']] = modelVersion_standings.get(record['winner'], 0) + 1
-
                 # add the loser who never wins
                 if record['loser'] not in modelVersion_standings:
                     modelVersion_standings[record['loser']] = 0
@@ -221,7 +241,7 @@ class DashboardApp:
         # get tags from database of the current user
         db_table = 'sort_results' if mode == 'Drag and Sort' else 'battle_results'
 
-        tags = ['overview']
+        tags = []
         curser = RANKING_CONN.cursor()
         curser.execute(
             f"SELECT DISTINCT tag FROM {db_table} WHERE username = '{st.session_state.user_id[0]}' AND timestamp = '{st.session_state.user_id[1]}'")
@@ -229,29 +249,30 @@ class DashboardApp:
             tags.append(row['tag'])
         curser.close()
 
-        if tags == ['overview']:
+        if len(tags) == 0:
             st.info(f'No rankings are finished with {mode} mode yet.')
 
         else:
-            tags = tags[1:2] if len(tags) == 2 else tags
+            # tags = tags[1:2] if len(tags) == 2 else tags
+            tag = ['overview'] + tags if len(tags) > 1 else tags
             tag = st.radio('Select a tag', tags, index=0, horizontal=True, label_visibility='collapsed')
             self.sidebar(tags, mode)
             self.leaderboard(tag, db_table)
 
-        with st.sidebar:
-            with st.form('overall_feedback'):
-                comment = st.text_area('Please leave your comments here.', key='comment')
-                submit_feedback = st.form_submit_button('Submit Feedback')
-                if submit_feedback:
-                    commenttime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-                    curser = RANKING_CONN.cursor()
-                    # parse the comment to at most 300 to avoid SQL injection
-                    for i in range(0, len(comment), 300):
-                        curser.execute(f"INSERT INTO comments (username, timestamp, comment, commenttime) VALUES ('{st.session_state.user_id[0]}', '{st.session_state.user_id[1]}', '{comment[i:i+300]}', '{commenttime}')")
-                    RANKING_CONN.commit()
-                    curser.close()
-
-                    st.toast('🙏 **Thanks for your feedback! We will take it into consideration in our future work.**')
+        # with st.sidebar:
+        #     with st.form('overall_feedback'):
+        #         comment = st.text_area('Please leave your comments here.', key='comment')
+        #         submit_feedback = st.form_submit_button('Submit Feedback')
+        #         if submit_feedback:
+        #             commenttime = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        #             curser = RANKING_CONN.cursor()
+        #             # parse the comment to at most 300 to avoid SQL injection
+        #             for i in range(0, len(comment), 300):
+        #                 curser.execute(f"INSERT INTO comments (username, timestamp, comment, commenttime) VALUES ('{st.session_state.user_id[0]}', '{st.session_state.user_id[1]}', '{comment[i:i+300]}', '{commenttime}')")
+        #             RANKING_CONN.commit()
+        #             curser.close()
+        #
+        #             st.sidebar.info('🙏 **Thanks for your feedback! We will take it into consideration in our future work.**')
 
 
 if __name__ == "__main__":