pages/1_🖼️_Gallery.py

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()