Spaces:
Sleeping
Sleeping
| import numpy as np | |
| import pandas as pd | |
| import streamlit as st | |
| from streamlit_elements import elements, mui, html, dashboard, nivo | |
| from streamlit_extras.switch_page_button import switch_page | |
| from pages.Gallery import load_hf_dataset | |
| class RankingApp: | |
| def __init__(self, promptBook, images_endpoint, batch_size=4): | |
| self.promptBook = promptBook | |
| self.images_endpoint = images_endpoint | |
| self.batch_size = batch_size | |
| # self.batch_num = len(self.promptBook) // self.batch_size | |
| # 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 | |
| def sidebar(self): | |
| with st.sidebar: | |
| prompt_tags = self.promptBook['tag'].unique() | |
| prompt_tags = np.sort(prompt_tags) | |
| tag = st.selectbox('Select a prompt tag', prompt_tags) | |
| items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True) | |
| prompts = np.sort(items['prompt'].unique())[::-1] | |
| selected_prompt = st.selectbox('Select a prompt', prompts) | |
| items = items[items['prompt'] == selected_prompt].reset_index(drop=True) | |
| prompt_id = items['prompt_id'].unique()[0] | |
| with st.form(key='prompt_form'): | |
| # 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) | |
| return prompt_tags, tag, prompt_id, items | |
| def draggable_images(self, items, layout='portrait'): | |
| # init ranking by the order of items | |
| if 'ranking' not in st.session_state: | |
| st.session_state.ranking = {} | |
| for i in range(len(items)): | |
| st.session_state.ranking[str(items['image_id'][i])] = i | |
| print(items) | |
| with elements('dashboard'): | |
| if layout == 'portrait': | |
| col_num = 4 | |
| layout = [dashboard.Item(str(items['image_id'][i]), i % col_num, i//col_num, 1, 2, isResizable=False) for i in range(len(items))] | |
| 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 | |
| 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 | |
| mui.Chip(label=rank, | |
| # variant="outlined" if rank!=1 else "default", | |
| color="primary" if rank == 1 else "warning" if rank == 2 else "info", | |
| size="small", | |
| sx={"position": "absolute", "left": "-0.3rem", "top": "-0.3rem"}) | |
| img_url = self.images_endpoint + str(items['image_id'][i]) + '.png' | |
| mui.CardMedia( | |
| component="img", | |
| # image={"data:image/png;base64", img_str}, | |
| image=img_url, | |
| alt="There should be an image", | |
| sx={"height": "100%", "object-fit": "fit", 'bgcolor': 'black'}, | |
| ) | |
| def handle_layout_change(self, updated_layout): | |
| # print(updated_layout) | |
| 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) | |
| 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) | |
| prompt_tags, tag, prompt_id, items = self.sidebar() | |
| 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:") | |
| 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: | |
| selected_modelVersions = [] | |
| 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 len(selected_modelVersions) == 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(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() | |