new summary page

Home.py

@@ -18,6 +18,8 @@ def login():
         )
         st.write('You can leave it blank to be anonymous.')
 
+        st.session_state.show_NSFW = st.checkbox(':orange[show potentially mature content]', help='Inevitably, a few images might be NSFW, even if we tried to elimiate NFSW content in our prompts. We calculate a NSFW score to filter them out. Please check only if you are 18+ and want to take a look at the whole GEMRec-18k dataset', value=False, key='mature_content')
+
         # Every form must have a submit button.
         submitted = st.form_submit_button("Start")
         if submitted:
@@ -40,6 +42,8 @@ def logout():
     st.session_state.pop('gallery_state', None)
     st.session_state.pop('progress', None)
     st.session_state.pop('gallery_focus', None)
+    st.session_state.pop('assigned_rank_mode', None)
+    st.session_state.pop('show_NSFW', None)
 
 
 def info():

data/unsafe_prompts.json

@@ -0,0 +1,4 @@
+{
+  "world knowledge": [83],
+  "abstract": [1, 3]
+}

pages/Gallery.py

@@ -1,3 +1,4 @@
+import json
 import os
 import requests
 
@@ -326,13 +327,15 @@ class GalleryApp:
 
             # add safety check for some prompts
             safety_check = True
-            unsafe_prompts = {}
-            # initialize unsafe prompts
+
+            # load unsafe prompts
+            unsafe_prompts = json.load(open('./data/unsafe_prompts.json', 'r'))
             for prompt_tag in prompt_tags:
-                unsafe_prompts[prompt_tag] = []
-            # manually add unsafe prompts
-            unsafe_prompts['world knowledge'] = [83]
-            unsafe_prompts['abstract'] = [1, 3]
+                if prompt_tag not in unsafe_prompts:
+                    unsafe_prompts[prompt_tag] = []
+            # # manually add unsafe prompts
+            # unsafe_prompts['world knowledge'] = [83]
+            # unsafe_prompts['abstract'] = [1, 3]
 
             if int(prompt_id.item()) in unsafe_prompts[tag]:
                 st.warning('This prompt may contain unsafe content. They might be offensive, depressing, or sexual.')
@@ -643,7 +646,7 @@ def altair_histogram(hist_data, sort_by, mini, maxi):
 
 
 @st.cache_data
-def load_hf_dataset():
+def load_hf_dataset(show_NSFW=False):
     # login to huggingface
     login(token=os.environ.get("HF_TOKEN"))
 
@@ -669,7 +672,9 @@ def load_hf_dataset():
     promptBook.loc[:, 'row_idx'] = promptBook.index
 
     # apply a nsfw filter
-    promptBook = promptBook[promptBook['nsfw_score'] <= 0.84].reset_index(drop=True)
+    if not show_NSFW:
+        promptBook = promptBook[promptBook['norm_nsfw'] <= 0.8].reset_index(drop=True)
+        print('nsfw filter applied', len(promptBook))
 
     # add a column that adds up 'norm_clip', 'norm_mcos', and 'norm_pop'
     score_weights = [1.0, 0.8, 0.2]
@@ -693,9 +698,6 @@ def load_tsne_coordinates(items):
 if __name__ == "__main__":
     st.set_page_config(page_title="Model Coffer Gallery", page_icon="🖼️", layout="wide")
 
-    with open('./pages/style.css') as f:
-        st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
-
     if 'user_id' not in st.session_state:
         st.warning('Please log in first.')
         home_btn = st.button('Go to Home Page')
@@ -703,7 +705,7 @@ if __name__ == "__main__":
             switch_page("home")
     else:
         # st.write('You have already logged in as ' + st.session_state.user_id[0])
-        roster, promptBook, images_ds = load_hf_dataset()
+        roster, promptBook, images_ds = load_hf_dataset(st.session_state.show_NSFW)
         # print(promptBook.columns)
 
         # # initialize selected_dict
@@ -713,15 +715,6 @@ if __name__ == "__main__":
         app = GalleryApp(promptBook=promptBook, images_ds=images_ds)
         app.app()
 
-        # components.html(
-        #     """
-        #     <script>
-        #       var iframe = window.parent.document.querySelector('[title="streamlit_agraph.agraph"]');
-        #       console.log(iframe);
-        #       var targetElement = iframe.contentDocument.querySelector('div.vis-network div.vis-navigation div.vis-button.vis-zoomExtends');
-        #       console.log(targetElement);
-        #       targetElement.style.background-image = "url(https://www.flaticon.com/free-icon-font/menu-burger_3917215?related_id=3917215#)";
-        #     </script>
-        #     """,
-        #     # unsafe_allow_html=True,
-        # )
+    with open('./css/style.css') as f:
+        st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
+

pages/Ranking.py

@@ -366,9 +366,6 @@ def connect_to_db():
 if __name__ == "__main__":
     st.set_page_config(page_title="Personal Image Ranking", page_icon="🎖️️", layout="wide")
 
-    with open('./pages/style.css') as f:
-        st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
-
     if 'user_id' not in st.session_state:
         st.warning('Please log in first.')
         home_btn = st.button('Go to Home Page')
@@ -391,7 +388,7 @@ if __name__ == "__main__":
                 switch_page('gallery')
         else:
             # st.write('You have checked ' + str(len(selected_modelVersions)) + ' images.')
-            roster, promptBook, images_ds = load_hf_dataset()
+            roster, promptBook, images_ds = load_hf_dataset(st.session_state.show_NSFW)
             print(st.session_state.selected_dict)
 
             # st.write("# Full function is coming soon.")
@@ -422,3 +419,5 @@ if __name__ == "__main__":
             app = RankingApp(promptBook_selected, images_endpoint, batch_size=4)
             app.app()
 
+    with open('./css/style.css') as f:
+        st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)

pages/Summary.py

@@ -24,15 +24,16 @@ class DashboardApp:
     def sidebar(self, tags, mode):
         with st.sidebar:
             tag = st.selectbox('Select a tag', tags, key='tag')
-            st.write('---')
-            st.write('## Want a more comprehensive summary?')
-            st.write('Jump back to gallery and select more images to rank!')
-            back_to_gallery = st.button('🖼️ Go to Gallery', key='summary_sidebar_gallery')
-            if back_to_gallery:
-                switch_page('gallery')
-            back_to_ranking = st.button('🎖️ Go to Ranking', key='summary_sidebar_ranking')
-            if back_to_ranking:
-                switch_page('ranking')
+            # st.write('---')
+            with st.form('summary_sidebar_form'):
+                st.write('## Want a more comprehensive summary?')
+                st.write('Jump back to gallery and select more images to rank!')
+                back_to_gallery = st.form_submit_button('🖼️ Go to Gallery')
+                if back_to_gallery:
+                    switch_page('gallery')
+                back_to_ranking = st.form_submit_button('🎖️ Go to Ranking')
+                if back_to_ranking:
+                    switch_page('ranking')
 
         return tag
 
@@ -108,15 +109,8 @@ class DashboardApp:
                                 image = f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{images[i]}.png"
                                 st.image(image, use_column_width=True)
 
-
-            # # st.write('---')
-            # expander = st.expander(f'# {icon} {model_name}, [{modelVersion_name}](https://civitai.com/models/{modelVersion_id})')
-            # with expander:
-            #     images = self.promptBook[self.promptBook['modelVersion_id'] == modelVersion_id]['image_id'].values
-            #     st.write(images)
-
     def podium_expander(self, modelVersion_standings, n=3):
-        st.write('## Top picks')
+        # st.write('## Top picks')
         # metric_cols = st.columns(n)
         for i in range(n):
             # with metric_cols[i]:
@@ -126,15 +120,24 @@ class DashboardApp:
             model_id, model_name, modelVersion_name, url = self.roster[self.roster['modelVersion_id'] == modelVersion_id][['model_id', 'model_name', 'modelVersion_name', 'modelVersion_url']].values[0]
 
             icon = '🥇'if i == 0 else '🥈' if i == 1 else '🥉'
-            with st.expander(f'**{icon} {model_name}, [{modelVersion_name}](https://civitai.com/models/{model_id}?modelVersionId={modelVersion_id})**, Ranking Score: {winning_times}'):
-                images = self.promptBook[self.promptBook['modelVersion_id'] == modelVersion_id]['image_id'].values
-                st.write(f'### Images generated with {icon} {model_name}, {modelVersion_name}')
-                col_num = 4
-                image_cols = st.columns(col_num)
-                for i in range(len(images)):
-                    with image_cols[i % col_num]:
-                        image = f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{images[i]}.png"
-                        st.image(image, use_column_width=True)
+            podium_display = st.columns([1, 14])
+            with podium_display[0]:
+                st.title(f'{icon}')
+            with podium_display[1]:
+                st.write(f'##### {model_name}, {modelVersion_name}')
+                st.write(f'[Civitai Page](https://civitai.com/models/{model_id}?modelVersionId={modelVersion_id}), [Model Download Link]({url}), Ranking Score: {winning_times}')
+                # with st.expander(f'**{icon} {model_name}, [{modelVersion_name}](https://civitai.com/models/{model_id}?modelVersionId={modelVersion_id})**, Ranking Score: {winning_times}'):
+                with st.expander(f'Show Images'):
+                    images = self.promptBook[self.promptBook['modelVersion_id'] == modelVersion_id]['image_id'].values
+                    # st.write(f'### Images generated with {icon} {model_name}, {modelVersion_name}')
+                    col_num = 4
+                    image_cols = st.columns(col_num)
+                    for j in range(len(images)):
+                        with image_cols[j % col_num]:
+                            image = f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{images[j]}.png"
+                            st.image(image, use_column_width=True)
+            if i != n - 1:
+                st.write('---')
 
     def score_calculator(self, results, db_table):
         modelVersion_standings = {}
@@ -160,9 +163,8 @@ class DashboardApp:
 
         return modelVersion_standings
 
-
     def app(self):
-        st.title('Your Preferred Models', help="Scores are calculated based on your ranking results.")
+        st.write('### Your Preferred Models')
 
         # mode = st.sidebar.radio('Ranking mode', ['Drag and Sort', 'Battle'], horizontal=True, index=1)
         mode = st.session_state.assigned_rank_mode
@@ -219,7 +221,7 @@ if __name__ == "__main__":
                 switch_page('gallery')
 
         else:
-            roster, promptBook, images_ds = load_hf_dataset()
+            roster, promptBook, images_ds = load_hf_dataset(st.session_state.show_NSFW)
             RANKING_CONN = connect_to_db()
             app = DashboardApp(roster, promptBook, session_finished)
             app.app()