new data cache method!

app.py

@@ -3,7 +3,6 @@ import numpy as np
 import random
 import pandas as pd
 import glob
-import csv
 from PIL import Image
 import datasets
 from datasets import load_dataset, Dataset, load_from_disk
@@ -13,13 +12,28 @@ import requests
 from bs4 import BeautifulSoup
 import re
 
+import altair as alt
+from streamlit_vega_lite import vega_lite_component, altair_component, _component_func
+
 SCORE_NAME_MAPPING = {'clip': 'clip_score', 'rank': 'avg_rank', 'pop': 'model_download_count'}
 
 
+# hist_data = pd.DataFrame(np.random.normal(42, 10, (200, 1)), columns=["x"])
+@st.cache_resource
+def altair_histogram(hist_data, sort_by):
+    brushed = alt.selection_interval(encodings=['x'], name="brushed")
+    return (
+        alt.Chart(hist_data)
+        .mark_bar()
+        .encode(alt.X(f"{sort_by}:Q", bin=True), y="count()")
+        .add_selection(brushed)
+        .properties(width=600, height=300)
+    )
+
 class GalleryApp:
-    def __init__(self, promptBook):
+    def __init__(self, promptBook, images_ds):
         self.promptBook = promptBook
-        st.set_page_config(layout="wide")
+        self.images_ds = images_ds
 
     def gallery_masonry(self, items, col_num, info):
         cols = st.columns(col_num)
@@ -27,7 +41,7 @@ class GalleryApp:
         # items = items.sort_values(by=['brisque'], ascending=True).reset_index(drop=True)
         for idx in range(len(items)):
             with cols[idx % col_num]:
-                image = st.session_state.images[items.iloc[idx]['row_idx'].item()]['image']
+                image = self.images_ds[items.iloc[idx]['row_idx'].item()]['image']
                 st.image(image,
                          use_column_width=True,
                 )
@@ -58,7 +72,7 @@ class GalleryApp:
                     if idx + j < len(items):
                         with cols[j]:
                             # show image
-                            image = st.session_state.images[items.iloc[idx+j]['row_idx'].item()]['image']
+                            image = self.images_ds[items.iloc[idx+j]['row_idx'].item()]['image']
 
                             st.image(image,
                                      use_column_width=True,
@@ -184,11 +198,12 @@ class GalleryApp:
                 with sub_selecters[2]:
                     pop_weight = st.number_input('Popularity Weight', min_value=-100.0, max_value=100.0, value=1.0, step=0.1)
 
-                items.loc[:, 'weighted_score_sum'] = items['norm_clip'] * clip_weight + items['avg_rank'] * rank_weight + items[
-                    'norm_pop'] * pop_weight
+                items.loc[:, 'weighted_score_sum'] = round(items['norm_clip'] * clip_weight + items['avg_rank'] * rank_weight + items[
+                    'norm_pop'] * pop_weight, 4)
 
                 continue_idx = 3
 
+
             with sub_selecters[continue_idx]:
                 order = st.selectbox('Order', ['Ascending', 'Descending'], index=1 if sort_type == 'Scores' else 0)
                 if order == 'Ascending':
@@ -211,6 +226,15 @@ class GalleryApp:
                     items = items[items['checked'] == True].reset_index(drop=True)
                     print(items)
 
+        if sort_type == 'Scores':
+            st.write('Select the range of scores to show')
+            hist_data = pd.DataFrame(items[sort_by])
+            event_dict = altair_component(altair_chart=altair_histogram(hist_data, sort_by))
+            r = event_dict.get(sort_by)
+            if r:
+                items = items[(items[sort_by] >= r[0]) & (items[sort_by] <= r[1])].reset_index(drop=True)
+                st.write(r)
+
         info = st.multiselect('Show Info',
                               ['model_download_count', 'clip_score', 'avg_rank', 'model_name', 'model_id',
                                'modelVersion_name', 'modelVersion_id', 'clip+rank', 'clip+pop', 'rank+pop',
@@ -303,6 +327,7 @@ class GalleryApp:
 
         if safety_check:
             items, info, col_num = self.selection_panel_2(items)
+
             # self.gallery_standard(items, col_num, info)
 
             with st.form(key=f'{prompt_id}', clear_on_submit=False):
@@ -340,44 +365,77 @@ class GalleryApp:
         dataset.push_to_hub('NYUSHPRP/ModelCofferMetadata', split='train')
 
 
+@st.cache_data
+def load_hf_dataset():
+    # load from huggingface
+    roster = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferRoster', split='train'))
+    promptBook = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferMetadata', split='train'))
+    images_ds = load_from_disk(os.path.join(os.getcwd(), 'data', 'promptbook'))
+
+    # process dataset
+    roster = roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name',
+                                                       'model_download_count']].drop_duplicates().reset_index(drop=True)
+
+    # add 'checked' column to promptBook if not exist
+    if 'checked' not in promptBook.columns:
+        promptBook.loc[:, 'checked'] = False
+
+    # add 'custom_score_weights' column to promptBook if not exist
+    if 'weighted_score_sum' not in promptBook.columns:
+        promptBook.loc[:, 'weighted_score_sum'] = 0
+
+    # merge roster and promptbook
+    promptBook = promptBook.merge(roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name', 'model_download_count']],
+                                                                    on=['model_id', 'modelVersion_id'], how='left')
+
+    # add column to record current row index
+    promptBook.loc[:, 'row_idx'] = promptBook.index
+
+    return roster, promptBook, images_ds
+
+
 if __name__ == '__main__':
     login(token=os.environ.get("HF_TOKEN"))
-
-    if 'roster' not in st.session_state:
-        print('loading roster')
-        # st.session_state.roster = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferRoster', split='train'))
-        st.session_state.roster = pd.DataFrame(load_from_disk(os.path.join(os.getcwd(), 'data', 'roster')))
-        st.session_state.roster = st.session_state.roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name',
-                                          'model_download_count']].drop_duplicates().reset_index(drop=True)
-    # add model download count from roster to promptbook dataframe
-    if 'promptBook' not in st.session_state:
-        print('loading promptBook')
-
-        st.session_state.promptBook = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferMetadata', split='train'))
-        # add 'checked' column to promptBook if not exist
-        if 'checked' not in st.session_state.promptBook.columns:
-            st.session_state.promptBook.loc[:, 'checked'] = False
-
-        # add 'custom_score_weights' column to promptBook if not exist
-        if 'weighted_score_sum' not in st.session_state.promptBook.columns:
-            st.session_state.promptBook.loc[:, 'weighted_score_sum'] = 0
-
-        st.session_state.images = load_from_disk(os.path.join(os.getcwd(), 'data', 'promptbook'))
-        # st.session_state.images = load_dataset('NYUSHPRP/ModelCofferPromptBook', split='train', streaming=True)
-        print(st.session_state.images)
-        print('images loaded')
-        # st.session_state.promptBook = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferPromptBook', split='train'))
-        st.session_state.promptBook = st.session_state.promptBook.merge(st.session_state.roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name', 'model_download_count']], on=['model_id', 'modelVersion_id'], how='left')
-
-        # add column to record current row index
-        st.session_state.promptBook['row_idx'] = st.session_state.promptBook.index
-        print('promptBook loaded')
-    # print(st.session_state.promptBook)
-
-    check_roster_error = False
-    if check_roster_error:
-        # print all rows with the same model_id and modelVersion_id but different model_download_count in roster
-        print(st.session_state.roster[st.session_state.roster.duplicated(subset=['model_id', 'modelVersion_id'], keep=False)].sort_values(by=['model_id', 'modelVersion_id']))
-
-    app = GalleryApp(promptBook=st.session_state.promptBook)
+    st.set_page_config(layout="wide")
+
+    # if 'roster' not in st.session_state:
+    #     print('loading roster')
+    #     # st.session_state.roster = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferRoster', split='train'))
+    #     st.session_state.roster = pd.DataFrame(load_from_disk(os.path.join(os.getcwd(), 'data', 'roster')))
+    #     st.session_state.roster = st.session_state.roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name',
+    #                                       'model_download_count']].drop_duplicates().reset_index(drop=True)
+    # # add model download count from roster to promptbook dataframe
+    # if 'promptBook' not in st.session_state:
+    #     print('loading promptBook')
+    #
+    #     st.session_state.promptBook = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferMetadata', split='train'))
+    #     # add 'checked' column to promptBook if not exist
+    #     if 'checked' not in st.session_state.promptBook.columns:
+    #         st.session_state.promptBook.loc[:, 'checked'] = False
+    #
+    #     # add 'custom_score_weights' column to promptBook if not exist
+    #     if 'weighted_score_sum' not in st.session_state.promptBook.columns:
+    #         st.session_state.promptBook.loc[:, 'weighted_score_sum'] = 0
+    #
+    #     st.session_state.images = load_from_disk(os.path.join(os.getcwd(), 'data', 'promptbook'))
+    #     # st.session_state.images = load_dataset('NYUSHPRP/ModelCofferPromptBook', split='train', streaming=True)
+    #     print(st.session_state.images)
+    #     print('images loaded')
+    #     # st.session_state.promptBook = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferPromptBook', split='train'))
+    #     st.session_state.promptBook = st.session_state.promptBook.merge(st.session_state.roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name', 'model_download_count']], on=['model_id', 'modelVersion_id'], how='left')
+    #
+    #     # add column to record current row index
+    #     st.session_state.promptBook['row_idx'] = st.session_state.promptBook.index
+    #     print('promptBook loaded')
+    # # print(st.session_state.promptBook)
+    #
+    # check_roster_error = False
+    # if check_roster_error:
+    #     # print all rows with the same model_id and modelVersion_id but different model_download_count in roster
+    #     print(st.session_state.roster[st.session_state.roster.duplicated(subset=['model_id', 'modelVersion_id'], keep=False)].sort_values(by=['model_id', 'modelVersion_id']))
+    roster, promptBook, images_ds = load_hf_dataset()
+    # if 'images' not in st.session_state:
+    #     st.session_state.images = load_from_disk(os.path.join(os.getcwd(), 'data', 'promptbook'))
+
+    app = GalleryApp(promptBook=promptBook, images_ds=images_ds)
     app.app()

requirements.txt

@@ -2,4 +2,5 @@ huggingface_hub
 streamlit-elements==0.1.0
 streamlit-extras
 altair<5
-streamlit-plotly-events
+streamlit-plotly-events
+streamlit-vega-lite