first version of ranking page!

data/ranking_script.py

@@ -0,0 +1,16 @@
+from datasets import Dataset
+
+
+def init_ranking_data():
+    ds = Dataset.from_dict({'image_id': [], 'modelVersion_id': [], 'ranking': [], "user_name": [], "timestamp": []})\
+
+    # add example data
+    # note that image_id is a string, other ids are int
+    ds = ds.add_item({'image_id': '0', 'modelVersion_id': 0, 'ranking': 0, "user_name": "example_data", "timestamp": 0.0})
+
+    ds.push_to_hub("MAPS-research/GEMRec-Ranking", split='train')
+
+
+if __name__ == '__main__':
+    init_ranking_data()
+

pages/Gallery.py

@@ -278,6 +278,8 @@ class GalleryApp:
                 switch_page("ranking")
 
     def submit_actions(self, status, prompt_id):
+        # remove counter from session state
+        st.session_state.pop('counter', None)
         if status == 'Select':
             modelVersions = self.promptBook[self.promptBook['prompt_id'] == prompt_id]['modelVersion_id'].unique()
             st.session_state.selected_dict[prompt_id] = modelVersions.tolist()
@@ -400,8 +402,8 @@ def load_hf_dataset():
     login(token=os.environ.get("HF_TOKEN"))
 
     # load from huggingface
-    roster = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferRoster', split='train'))
-    promptBook = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferMetadata', split='train'))
+    roster = pd.DataFrame(load_dataset('MAPS-research/GEMRec-Roster', split='train'))
+    promptBook = pd.DataFrame(load_dataset('MAPS-research/GEMRec-Metadata', split='train'))
     # images_ds = load_from_disk(os.path.join(os.getcwd(), 'data', 'promptbook'))
     images_ds = None  # set to None for now since we use s3 bucket to store images
 

pages/Ranking.py

@@ -1,3 +1,4 @@
+import datasets
 import numpy as np
 import pandas as pd
 import streamlit as st
@@ -17,7 +18,7 @@ class RankingApp:
         # self.batch_num += 1 if len(self.promptBook) % self.batch_size != 0 else 0
 
         if 'counter' not in st.session_state:
-            st.session_state.counter = 0
+            st.session_state.counter = {}
 
     def sidebar(self):
         with st.sidebar:
@@ -37,18 +38,27 @@ class RankingApp:
                 # 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)
+                st.form_submit_button('Generate Images [Coming Soon]', type='primary', use_container_width=True, disabled=True)
 
         return prompt_tags, tag, prompt_id, items
 
-    def draggable_images(self, items, layout='portrait'):
+    def draggable_images(self, items, prompt_id, layout='portrait'):
         # init ranking by the order of items
+
         if 'ranking' not in st.session_state:
             st.session_state.ranking = {}
+
+        if prompt_id not in st.session_state.ranking:
+            st.session_state.ranking[prompt_id] = {}
+
+        if st.session_state.counter[prompt_id] not in st.session_state.ranking[prompt_id]:
+            st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]] = {}
             for i in range(len(items)):
-                st.session_state.ranking[str(items['image_id'][i])] = i
+                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
+            items.index = items['image_id'].apply(lambda x: st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]][str(x)])
 
-        print(items)
         with elements('dashboard'):
             if layout == 'portrait':
                 col_num = 4
@@ -57,14 +67,17 @@ class RankingApp:
             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
+                    dashboard.Item(str(items['image_id'][i]), i % col_num * 2, i // col_num, 2, 1.6, 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
+                        prompt_id = st.session_state.prompt_id_tmp
+                        batch_idx = st.session_state.counter[prompt_id]
+
+                        rank = st.session_state.ranking[prompt_id][batch_idx][str(items['image_id'][i])] + 1
 
                         mui.Chip(label=rank,
                                  # variant="outlined" if rank!=1 else "default",
@@ -79,7 +92,7 @@ class RankingApp:
                             # image={"data:image/png;base64", img_str},
                             image=img_url,
                             alt="There should be an image",
-                            sx={"height": "100%", "object-fit": "fit", 'bgcolor': 'black'},
+                            sx={"height": "100%", "object-fit": "contain", 'bgcolor': 'black'},
                         )
 
     def handle_layout_change(self, updated_layout):
@@ -87,26 +100,95 @@ class RankingApp:
         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)
+        prompt_id = st.session_state.prompt_id_tmp
+        batch_idx = st.session_state.counter[prompt_id]
+
+        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 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)
 
+        # save the current progress to session state
+        if 'progress' not in st.session_state:
+            st.session_state.progress = {}
+        # print('current progress: ', st.session_state.progress)
+
         prompt_tags, tag, prompt_id, items = 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]
+
+        # save prompt_id in session state
+        st.session_state.prompt_id_tmp = prompt_id
+
+        if prompt_id not in st.session_state.progress:
+            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})
+                else:
+                    st.button(":ballot_box_with_check:", key='finished', on_click=self.next_batch, help='Finished', kwargs={'prompt_id': prompt_id, 'progress': 'finished'})
+
+                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})
+
+        elif st.session_state.progress[prompt_id] == 'finished':
+            st.write('## You have ranked all models for this tag!')
+            st.write('Thank you for your participation! Feel free to do the following things:')
+            st.write('* Rank for other tags and prompts.')
+            st.write('* Back to the gallery page to see more images.')
+            st.write('* Rank again for this tag and prompt.')
+            st.write('*More functions are coming soon... Please stay tuned*')
+
+            gallery_btn = st.button('🖼️ Back to Gallery')
+            if gallery_btn:
+                switch_page('gallery')
+
+            restart_btn = st.button('🎖️ Rank Again')
+            if restart_btn:
+                st.session_state.progress['prompt_id'] = 'ranking'
+                st.session_state.counter[prompt_id] = 0
+                st.experimental_rerun()
 
-        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:")
+    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')
+        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]})
+        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
 
 
 if __name__ == "__main__":
@@ -119,13 +201,15 @@ if __name__ == "__main__":
             switch_page("home")
 
     else:
-        selected_modelVersions = []
+        has_selection = False
+
         for key, value in st.session_state.selected_dict.items():
             for v in value:
-                if v not in selected_modelVersions:
-                    selected_modelVersions.append(v)
+                if v:
+                    has_selection = True
+                    break
 
-        if len(selected_modelVersions) == 0:
+        if not has_selection:
             st.info('You have not checked any image yet. Please go back to the gallery page and check some images.')
             gallery_btn = st.button('Go to Gallery')
             if gallery_btn:
@@ -134,21 +218,16 @@ if __name__ == "__main__":
             # st.write('You have checked ' + str(len(selected_modelVersions)) + ' images.')
             roster, promptBook, images_ds = load_hf_dataset()
             print(st.session_state.selected_dict)
-            st.write("# Full function is coming soon.")
-            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()
+            # st.write("# Full function is coming soon.")
+
+            # 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)
+            # st.write(promptBook_selected)
+            images_endpoint = "https://modelcofferbucket.s3-accelerate.amazonaws.com/"
+
+            app = RankingApp(promptBook_selected, images_endpoint, batch_size=4)
+            app.app()