new version! multiple pages!

app.py → Archive/app.py

@@ -100,10 +100,10 @@ class GalleryApp:
 
                             st.image(image, use_column_width=True)
 
-                            # # show checkbox
-                            # self.promptBook.loc[items.iloc[idx + j]['row_idx'].item(), 'checked'] = st.checkbox(
-                            #     'Select', value=self.promptBook.loc[items.iloc[idx + j]['row_idx'].item(), 'checked'],
-                            #     key=f'select_{idx + j}')
+                            # show checkbox
+                            self.promptBook.loc[items.iloc[idx + j]['row_idx'].item(), 'checked'] = st.checkbox(
+                                'Select', value=self.promptBook.loc[items.iloc[idx + j]['row_idx'].item(), 'checked'],
+                                key=f'select_{idx + j}')
 
                             # st.write(idx+j)
                             # show selected info
@@ -294,8 +294,8 @@ class GalleryApp:
         return items, info, col_num
 
     def app(self):
-        st.title('Model Coffer Gallery')
-        st.write('This is a gallery of images generated by the models in the Model Coffer')
+        st.title('Model Visualization and Retrieval')
+        st.write('This is a gallery of images generated by the models')
 
         with st.sidebar:
             prompt_tags = self.promptBook['tag'].unique()
@@ -367,7 +367,7 @@ class GalleryApp:
 
             with st.form(key=f'{prompt_id}', clear_on_submit=True):
                 # buttons = st.columns([1, 1, 1])
-                buttons_space = st.container()
+                buttons_space = st.columns([1, 1, 1, 1])
                 gallery_space = st.empty()
                 # with buttons[0]:
                 #     submit = st.form_submit_button('Save selections', on_click=self.save_checked, use_container_width=True, type='primary')
@@ -379,8 +379,17 @@ class GalleryApp:
                 with gallery_space.container():
                     self.gallery_standard(items, col_num, info)
 
-                with buttons_space:
-                    st.form_submit_button('Refresh', on_click=gallery_space.empty, use_container_width=True, type='primary')
+                with buttons_space[0]:
+                    st.form_submit_button('Confirm and Continue', use_container_width=True, type='primary')
+
+                with buttons_space[1]:
+                    st.form_submit_button('Select All', use_container_width=True)
+
+                with buttons_space[2]:
+                    st.form_submit_button('Deselect All', use_container_width=True)
+
+                with buttons_space[3]:
+                    st.form_submit_button('Refresh', on_click=gallery_space.empty, use_container_width=True)
 
 
     def reset_current_prompt(self, prompt_id):
@@ -416,7 +425,7 @@ 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'))
+    images_ds = load_from_disk(os.path.join(os.getcwd(), '../data', 'promptbook'))
 
     # process dataset
     roster = roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name',

Archive/test.py

@@ -0,0 +1,26 @@
+import streamlit as st
+
+if __name__ == "__main__":
+    if 'check_dict' not in st.session_state:
+        st.session_state.check_dict = {'check1': False, 'check2': False, 'check3': False}
+
+    with st.form('my_form'):
+        st.session_state.check_dict['check1'] = st.checkbox('Check 1 out')
+        st.session_state.check_dict['check2'] = st.checkbox('Check 2 out')
+        st.session_state.check_dict['check3'] = st.checkbox('Check 3 out')
+
+        check21 = st.checkbox('Check 21 out')
+        if check21:
+            st.write('check21 is checked')
+        check22 = st.checkbox('Check 22 out')
+        if check22:
+            st.write('check22 is checked')
+        check23 = st.checkbox('Check 23 out')
+        if check23:
+            st.write('check23 is checked')
+
+        # Every form must have a submit button.
+        submitted = st.form_submit_button('Submit')
+
+    for key, value in st.session_state.check_dict.items():
+        st.write(key, value)

README.md

@@ -6,7 +6,7 @@ colorTo: purple
 sdk: streamlit
 sdk_version: 1.19.0
 python_version: 3.9.13
-app_file: app.py
+app_file: 🏠_Home.py
 pinned: false
 ---
 

data/download_script.py

@@ -20,5 +20,11 @@ def test():
     print(promptbook[0]['image'])
 
 
+# def drop_metadata_checked_column():
+#     ModelCofferMetadata = load_dataset('NYUSHPRP/ModelCofferMetadata', split='train')
+#     ModelCofferMetadata = ModelCofferMetadata.remove_columns(['checked'])
+#     ModelCofferMetadata.push_to_hub('NYUSHPRP/ModelCofferMetadata', split='train')
+
+
 if __name__ == '__main__':
-    main()
+    main()

pages/1_🖼️_Gallery.py

@@ -0,0 +1,348 @@
+import streamlit as st
+import numpy as np
+import pandas as pd
+import glob
+from datasets import load_dataset, Dataset, load_from_disk
+from huggingface_hub import login
+import os
+import requests
+from bs4 import BeautifulSoup
+import altair as alt
+from streamlit_extras.switch_page_button import switch_page
+
+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, mini, maxi):
+    brushed = alt.selection_interval(encodings=['x'], name="brushed")
+
+    chart = (
+        alt.Chart(hist_data)
+        .mark_bar(opacity=0.7, cornerRadius=2)
+        .encode(alt.X(f"{sort_by}:Q", bin=alt.Bin(maxbins=25)), y="count()")
+        # .add_selection(brushed)
+        # .properties(width=800, height=300)
+    )
+
+    # Create a transparent rectangle for highlighting the range
+    highlight = (
+        alt.Chart(pd.DataFrame({'x1': [mini], 'x2': [maxi]}))
+        .mark_rect(opacity=0.3)
+        .encode(x='x1', x2='x2')
+        # .properties(width=800, height=300)
+    )
+
+    # Layer the chart and the highlight rectangle
+    layered_chart = alt.layer(chart, highlight)
+
+    return layered_chart
+
+
+class GalleryApp:
+    def __init__(self, promptBook, images_ds):
+        self.promptBook = promptBook
+        self.images_ds = images_ds
+
+    def gallery_standard(self, items, col_num, info):
+        rows = len(items) // col_num + 1
+        containers = [st.container() for _ in range(rows)]
+        for idx in range(0, len(items), col_num):
+            row_idx = idx // col_num
+            with containers[row_idx]:
+                cols = st.columns(col_num)
+                for j in range(col_num):
+                    if idx + j < len(items):
+                        with cols[j]:
+                            # show image
+                            image = self.images_ds[items.iloc[idx + j]['row_idx'].item()]['image']
+                            st.image(image, use_column_width=True)
+
+                            # handel checkbox information
+                            prompt_id = items.iloc[idx + j]['prompt_id']
+                            modelVersion_id = items.iloc[idx + j]['modelVersion_id']
+                            check_init = True if modelVersion_id in st.session_state.selected_dict.get(prompt_id, []) else False
+
+                            # show checkbox
+                            checked = st.checkbox('Select', key=f'select_{idx + j}', value=check_init)
+                            if checked:
+                                st.session_state.selected_dict[prompt_id] = st.session_state.selected_dict.get(prompt_id, []) + [modelVersion_id]
+                            else:
+                                try:
+                                    st.session_state.selected_dict[prompt_id].remove(modelVersion_id)
+                                except:
+                                    pass
+
+                            # show selected info
+                            for key in info:
+                                st.write(f"**{key}**: {items.iloc[idx + j][key]}")
+
+    def selection_panel(self, items):
+        selecters = st.columns([1, 4])
+
+        # select sort type
+        with selecters[0]:
+            sort_type = st.selectbox('Sort by', ['Scores', 'IDs and Names'])
+            if sort_type == 'Scores':
+                sort_by = 'weighted_score_sum'
+
+        # select other options
+        with selecters[1]:
+            if sort_type == 'IDs and Names':
+                sub_selecters = st.columns([3, 1])
+                # select sort by
+                with sub_selecters[0]:
+                    sort_by = st.selectbox('Sort by',
+                                           ['model_name', 'model_id', 'modelVersion_name', 'modelVersion_id'],
+                                           label_visibility='hidden')
+
+                continue_idx = 1
+
+            else:
+                # add custom weights
+                sub_selecters = st.columns([1, 1, 1, 1])
+
+                if 'score_weights' not in st.session_state:
+                    st.session_state.score_weights = [1.0, 0.8, 0.2, 0.84]
+
+                with sub_selecters[0]:
+                    clip_weight = st.number_input('Clip Score Weight', min_value=-100.0, max_value=100.0, value=st.session_state.score_weights[0], step=0.1, help='the weight for normalized clip score')
+                with sub_selecters[1]:
+                    rank_weight = st.number_input('Distinctiveness Weight', min_value=-100.0, max_value=100.0, value=st.session_state.score_weights[1], step=0.1, help='the weight for average rank')
+                with sub_selecters[2]:
+                    pop_weight = st.number_input('Popularity Weight', min_value=-100.0, max_value=100.0, value=st.session_state.score_weights[2], step=0.1, help='the weight for normalized popularity score')
+
+                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
+
+            # select threshold
+            with sub_selecters[continue_idx]:
+                dist_threshold = st.number_input('Distinctiveness Threshold', min_value=0.0, max_value=1.0, value=st.session_state.score_weights[3], step=0.01, help='Only show models with distinctiveness score lower than this threshold, set 1.0 to show all images')
+                items = items[items['avg_rank'] < dist_threshold].reset_index(drop=True)
+
+            # save latest weights
+            st.session_state.score_weights = [clip_weight, rank_weight, pop_weight, dist_threshold]
+
+        # draw a distribution histogram
+        if sort_type == 'Scores':
+            try:
+                with st.expander('Show score distribution histogram and select score range'):
+                    st.write('**Score distribution histogram**')
+                    chart_space = st.container()
+                    # st.write('Select the range of scores to show')
+                    hist_data = pd.DataFrame(items[sort_by])
+                    mini = hist_data[sort_by].min().item()
+                    mini = mini//0.1 * 0.1
+                    maxi = hist_data[sort_by].max().item()
+                    maxi = maxi//0.1 * 0.1 + 0.1
+                    st.write('**Select the range of scores to show**')
+                    r = st.slider('Select the range of scores to show', min_value=mini, max_value=maxi, value=(mini, maxi), step=0.05, label_visibility='collapsed')
+                    with chart_space:
+                        st.altair_chart(altair_histogram(hist_data, sort_by, r[0], r[1]), use_container_width=True)
+                    # 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)
+            except:
+                pass
+
+        display_options = st.columns([1, 4])
+
+        with display_options[0]:
+            # select order
+            order = st.selectbox('Order', ['Ascending', 'Descending'], index=1 if sort_type == 'Scores' else 0)
+            if order == 'Ascending':
+                order = True
+            else:
+                order = False
+
+        with display_options[1]:
+
+            # select info to show
+            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',
+                                   'clip+rank+pop', 'weighted_score_sum'],
+                                  default=sort_by)
+
+        # apply sorting to dataframe
+        items = items.sort_values(by=[sort_by], ascending=order).reset_index(drop=True)
+
+        # select number of columns
+        col_num = st.slider('Number of columns', min_value=1, max_value=9, value=4, step=1, key='col_num')
+
+        return items, info, col_num
+
+    def sidebar(self):
+        with st.sidebar:
+            prompt_tags = self.promptBook['tag'].unique()
+            # sort tags by alphabetical order
+            prompt_tags = np.sort(prompt_tags)[::-1]
+
+            tag = st.selectbox('Select a tag', prompt_tags)
+
+            items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)
+
+            original_prompts = np.sort(items['prompt'].unique())[::-1]
+
+            # remove the first four items in the prompt, which are mostly the same
+            if tag != 'abstract':
+                prompts = [', '.join(x.split(', ')[4:]) for x in original_prompts]
+                prompt = st.selectbox('Select prompt', prompts)
+
+                idx = prompts.index(prompt)
+                prompt_full = ', '.join(original_prompts[idx].split(', ')[:4]) + ', ' + prompt
+            else:
+                prompt_full = st.selectbox('Select prompt', original_prompts)
+
+            items = items[items['prompt'] == prompt_full].reset_index(drop=True)
+            prompt_id = items['prompt_id'].unique()[0]
+
+            # show image metadata
+            image_metadatas = ['prompt_id', 'prompt', 'negativePrompt', 'sampler', 'cfgScale', 'size', 'seed']
+            for key in image_metadatas:
+                label = ' '.join(key.split('_')).capitalize()
+                st.write(f"**{label}**")
+                if items[key][0] == ' ':
+                    st.write('`None`')
+                else:
+                    st.caption(f"{items[key][0]}")
+
+            # for tag as civitai, add civitai reference
+            if tag == 'civitai':
+                try:
+                    st.write('**Civitai Reference**')
+                    res = requests.get(f'https://civitai.com/images/{prompt_id.item()}')
+                    # st.write(res.text)
+                    soup = BeautifulSoup(res.text, 'html.parser')
+                    image_section = soup.find('div', {'class': 'mantine-12rlksp'})
+                    image_url = image_section.find('img')['src']
+                    st.image(image_url, use_column_width=True)
+                except:
+                    pass
+
+        return prompt_tags, tag, prompt_id, items
+
+    def app(self):
+        st.title('Model Visualization and Retrieval')
+        st.write('This is a gallery of images generated by the models')
+
+        prompt_tags, tag, prompt_id, items = self.sidebar()
+
+        # add safety check for some prompts
+        safety_check = True
+        unsafe_prompts = {}
+        # initialize unsafe prompts
+        for prompt_tag in prompt_tags:
+            unsafe_prompts[prompt_tag] = []
+        # manually add unsafe prompts
+        unsafe_prompts['civitai'] = [375790, 366222, 295008, 256477]
+        unsafe_prompts['people'] = [53]
+        unsafe_prompts['art'] = [23]
+        unsafe_prompts['abstract'] = [10, 12]
+        unsafe_prompts['food'] = [34]
+
+        if int(prompt_id.item()) in unsafe_prompts[tag]:
+            st.warning('This prompt may contain unsafe content. They might be offensive, depressing, or sexual.')
+            safety_check = st.checkbox('I understand that this prompt may contain unsafe content. Show these images anyway.', key=f'{prompt_id}')
+
+        if safety_check:
+            items, info, col_num = self.selection_panel(items)
+            # self.gallery_standard(items, col_num, info)
+
+            with st.form(key=f'{prompt_id}'):
+                # buttons = st.columns([1, 1, 1])
+                buttons_space = st.columns([1, 1, 1, 1])
+                gallery_space = st.empty()
+
+                with buttons_space[0]:
+                    continue_btn = st.form_submit_button('Confirm Selection', use_container_width=True, type='primary')
+                    if continue_btn:
+                        self.submit_actions('Continue', prompt_id)
+
+                with buttons_space[1]:
+                    select_btn = st.form_submit_button('Select All', use_container_width=True)
+                    if select_btn:
+                        self.submit_actions('Select', prompt_id)
+
+                with buttons_space[2]:
+                    deselect_btn = st.form_submit_button('Deselect All', use_container_width=True)
+                    if deselect_btn:
+                        self.submit_actions('Deselect', prompt_id)
+
+                with buttons_space[3]:
+                    refresh_btn = st.form_submit_button('Refresh', on_click=gallery_space.empty, use_container_width=True)
+
+                with gallery_space.container():
+                    with st.spinner('Loading images...'):
+                        self.gallery_standard(items, col_num, info)
+
+    def submit_actions(self, status, 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()
+            print(st.session_state.selected_dict, 'select')
+        elif status == 'Deselect':
+            st.session_state.selected_dict[prompt_id] = []
+            print(st.session_state.selected_dict, 'deselect')
+            # self.promptBook.loc[self.promptBook['prompt_id'] == prompt_id, 'checked'] = False
+            pass
+        elif status == 'Continue':
+            # switch_page("ranking")
+            pass
+
+
+@st.cache_data
+def load_hf_dataset():
+    # login to huggingface
+    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'))
+    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__":
+    st.set_page_config(page_title="Model Coffer Gallery", page_icon="🖼️", layout="wide")
+    if 'user_id' not in st.session_state:
+        st.warning('Please log in first.')
+        home_btn = st.button('Go to Home Page')
+        if home_btn:
+            switch_page("home")
+    else:
+        st.write('You have already logged in as ' + st.session_state.user_id[0])
+        roster, promptBook, st.session_state["images_ds"] = load_hf_dataset()
+        # print(promptBook.columns)
+
+        # initialize selected_dict
+        if 'selected_dict' not in st.session_state:
+            st.session_state['selected_dict'] = {}
+
+        app = GalleryApp(promptBook=promptBook, images_ds=st.session_state.images_ds)
+        app.app()

pages/2_🎖️_Ranking.py

@@ -0,0 +1,28 @@
+import streamlit as st
+import numpy as np
+import pandas as pd
+from streamlit_extras.switch_page_button import switch_page
+
+if __name__ == "__main__":
+    st.set_page_config(page_title="Personal Image Ranking", page_icon="🎖️️", layout="wide")
+
+    if 'user_id' not in st.session_state:
+        st.warning('Please log in first.')
+        home_btn = st.button('Go to Home Page')
+        if home_btn:
+            switch_page("home")
+
+    else:
+        all_checked = 0
+        for key, value in st.session_state.selected_dict.items():
+            for v in value:
+                all_checked += 1
+
+        if all_checked == 0:
+            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:
+                switch_page('gallery')
+        else:
+            st.write('You have checked ' + str(all_checked) + ' images.')
+

🏠_Home.py

@@ -0,0 +1,48 @@
+import streamlit as st
+import random
+import time
+from streamlit_extras.switch_page_button import switch_page
+
+
+def login():
+    # skip customize user name for debug mode
+
+    with st.form("user_login"):
+        st.write('## Enter Your Name')
+        user_id = st.text_input(
+            "Enter your name for personalization 👇",
+            label_visibility='visible',
+            disabled=False,
+            placeholder='anonymous',
+        )
+        st.write('You can leave it blank to be anonymous.')
+
+        # Every form must have a submit button.
+        submitted = st.form_submit_button("Start")
+        if submitted:
+            save_user_id(user_id)
+            switch_page("gallery")
+
+
+def save_user_id(user_id):
+    print(user_id)
+    if not user_id:
+        user_id = 'anonymous' + str(random.randint(0, 100000))
+    st.session_state.user_id = [user_id, time.time()]
+
+
+if __name__ == '__main__':
+    st.set_page_config(page_title="Login", page_icon="🏠")
+
+    st.title("Personalized Image Ranking")
+    st.write(
+        "This is an web application to collect personal preference to ai generated images. \
+        You can know which model you like most after you finish the survey."
+    )
+
+    if 'user_id' not in st.session_state:
+        login()
+    else:
+        st.write('You have already logged in as ' + st.session_state.user_id[0])
+        st.button('Log out', on_click=lambda: st.session_state.pop('user_id'))
+