import os import datasets import numpy as np import pandas as pd import pymysql.cursors import streamlit as st from datetime import datetime 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 from Home import connect_to_db 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 = {} def sidebar(self, selected_prompt, items): with st.sidebar: # st.title('Personal Image Ranking') # st.write('Here you can test out your selected images with any prompt you like. ') # prompt_tags = self.promptBook['tag'].unique() # prompt_tags = np.sort(prompt_tags).tolist() # # print(st.session_state.gallery_focus) # tag_idx = prompt_tags.index(st.session_state.gallery_focus['tag']) if st.session_state.gallery_focus['tag'] in prompt_tags else 0 # print(tag_idx) # # tag = st.selectbox('Select a prompt tag', prompt_tags, index=tag_idx) # items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True) # prompts = np.sort(items['prompt'].unique())[::-1].tolist() # # prompt_idx = prompts.index(st.session_state.gallery_focus['prompt']) if st.session_state.gallery_focus['prompt'] in prompts else 0 # print(prompt_idx) # # selected_prompt = st.selectbox('Select a prompt', prompts, index=prompt_idx) # # # mode = st.radio('Select a mode', ['Drag and Sort', 'Battle'], index=1) # mode = st.session_state.assigned_rank_mode # # 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 [Coming Soon]', type='primary', use_container_width=True, disabled=True) return None def draggable_images(self, items, prompt_id, layout='portrait'): # init ranking by the order of items if 'ranking' not in st.session_state: st.session_state.ranking = {} if prompt_id not in st.session_state.ranking: st.session_state.ranking[prompt_id] = {} if st.session_state.counter[prompt_id] not in st.session_state.ranking[prompt_id]: st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]] = {} for i in range(len(items)): st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]][str(items['image_id'][i])] = i else: # set the index of items to the corresponding ranking value of the image_id items.index = items['image_id'].apply(lambda x: st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]][str(x)]) 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.6, 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"): prompt_id = st.session_state.prompt_id_tmp batch_idx = st.session_state.counter[prompt_id] rank = st.session_state.ranking[prompt_id][batch_idx][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": "contain", '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] prompt_id = st.session_state.prompt_id_tmp batch_idx = st.session_state.counter[prompt_id] for k in st.session_state.ranking[prompt_id][batch_idx].keys(): st.session_state.ranking[prompt_id][batch_idx][k] = sorted_list.index(k) def dragsort_mode(self, tag, items, prompt_id, batch_num): st.session_state.counter[prompt_id] = 0 if prompt_id not in st.session_state.counter else \ st.session_state.counter[prompt_id] 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[prompt_id] + 1) / batch_num, text=f"Batch {st.session_state.counter[prompt_id] + 1} / {batch_num}") # st.write(items.iloc[self.batch_size*st.session_state.counter[prompt_id]: self.batch_size*(st.session_state.counter[prompt_id]+1)]) width, height = items.loc[0, 'size'].split('x') if int(height) >= int(width): self.draggable_images(items.iloc[ self.batch_size * st.session_state.counter[prompt_id]: self.batch_size * ( st.session_state.counter[prompt_id] + 1)].reset_index(drop=True), prompt_id=prompt_id, layout='portrait') else: self.draggable_images(items.iloc[ self.batch_size * st.session_state.counter[prompt_id]: self.batch_size * ( st.session_state.counter[prompt_id] + 1)].reset_index(drop=True), prompt_id=prompt_id, layout='landscape') # st.write(str(st.session_state.ranking)) with control: if st.session_state.counter[prompt_id] < batch_num - 1: st.button(":arrow_right:", key='next', on_click=self.next_batch, help='Next Batch', kwargs={'tag': tag, 'prompt_id': prompt_id}, use_container_width=True, type='primary') else: st.button(":heavy_check_mark:", key='finished', on_click=self.next_batch, help='Finished', kwargs={'tag': tag, 'prompt_id': prompt_id, 'progress': 'finished'}, use_container_width=True, type='primary') if st.session_state.counter[prompt_id] > 0: st.button(":arrow_left:", key='prev', on_click=self.prev_batch, help='Previous Batch', kwargs={'prompt_id': prompt_id}, use_container_width=True) def next_batch(self, tag, prompt_id, progress=None): curser = RANKING_CONN.cursor() # a not so elegant way to get the modelVersion_id of each image, but it works position_version_dict = {} for image_id, position in st.session_state.ranking[prompt_id][st.session_state.counter[prompt_id]].items(): modelVersion_id = self.promptBook[self.promptBook['image_id'] == image_id]['modelVersion_id'].values[0] position_version_dict[position] = modelVersion_id # # get all records of this user and prompt # query = "SELECT * FROM sort_results WHERE username = %s AND timestamp = %s AND prompt_id = %s" # curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id)) # results = curser.fetchall() # print(results) # # # remove the old ranking with the same modelVersion_id if exists # for result in results: # prev_ids = [result['position1'], result['position2'], result['position3'], result['position4']] # curr_ids = [position_version_dict[0], position_version_dict[1], position_version_dict[2], position_version_dict[3]] # if len(set(prev_ids).intersection(set(curr_ids))) == 4: # query = "DELETE FROM sort_results WHERE username = %s AND timestamp = %s AND prompt_id = %s AND position1 = %s AND position2 = %s AND position3 = %s AND position4 = %s" # curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id, result['position1'], result['position2'], result['position3'], result['position4'])) sorttime = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S") # handle the case where user press the 'prev' button query = "DELETE FROM sort_results WHERE username = %s AND timestamp = %s AND prompt_id = %s AND epoch = %s" curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id, st.session_state.epoch['ranking'][prompt_id])) query = "INSERT INTO sort_results (username, timestamp, tag, prompt_id, position1, position2, position3, position4, sorttime, epoch) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)" curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], self.promptBook[self.promptBook['prompt_id'] == prompt_id]['tag'].values[0], prompt_id, position_version_dict[0], position_version_dict[1], position_version_dict[2], position_version_dict[3], sorttime, st.session_state.epoch['ranking'][prompt_id])) curser.close() RANKING_CONN.commit() if progress == 'finished': st.session_state.epoch['ranking'][prompt_id] += 1 st.session_state.epoch['summary'][tag] = st.session_state.epoch['summary'].get(tag, 0) + 1 st.session_state.epoch['summary']['overall'] += 1 st.session_state.progress[prompt_id] = 'finished' # drop 'modelVersion_standings' from session state if exists st.session_state.pop('modelVersion_standings', None) else: st.session_state.counter[prompt_id] += 1 def prev_batch(self, prompt_id): st.session_state.counter[prompt_id] -= 1 def battle_images(self, tag, items, prompt_id): if 'pointer' not in st.session_state: st.session_state.pointer = {} if prompt_id not in st.session_state.pointer: st.session_state.pointer[prompt_id] = {'left': 0, 'right': 1} curr_position = max(st.session_state.pointer[prompt_id]['left'], st.session_state.pointer[prompt_id]['right']) progress = st.progress(curr_position / (len(items)-1), text=f"Progress {curr_position} / {len(items)-1}") # if curr_position == len(items) - 1: # st.session_state.progress[prompt_id] = 'finished' # # else: left, right = st.columns(2) with left: image_id = items['image_id'][st.session_state.pointer[prompt_id]['left']] img_url = self.images_endpoint + str(image_id) + '.png' # # write the total score of this image # total_score = items['total_score'][st.session_state.pointer[prompt_id]['left']] # st.write(f'Total Score: {total_score}') btn_left = st.button('Left is better', key='left', on_click=self.next_battle, kwargs={'tag': tag, 'prompt_id': prompt_id, 'image_ids': items['image_id'], 'winner': 'left', 'curr_position': curr_position, 'total_num': len(items)}, use_container_width=True) st.image(img_url, use_column_width=True) with right: image_id = items['image_id'][st.session_state.pointer[prompt_id]['right']] img_url = self.images_endpoint + str(image_id) + '.png' # # write the total score of this image # total_score = items['total_score'][st.session_state.pointer[prompt_id]['right']] # st.write(f'Total Score: {total_score}') btn_right = st.button('Right is better', key='right', on_click=self.next_battle, kwargs={'tag': tag, 'prompt_id': prompt_id, 'image_ids': items['image_id'], 'winner': 'right', 'curr_position': curr_position, 'total_num': len(items)}, use_container_width=True) st.image(img_url, use_column_width=True) def next_battle(self, tag, prompt_id, image_ids, winner, curr_position, total_num): loser = 'left' if winner == 'right' else 'right' battletime = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S") curser = RANKING_CONN.cursor() winner_modelVersion_id = self.promptBook[self.promptBook['image_id'] == image_ids[st.session_state.pointer[prompt_id][winner]]]['modelVersion_id'].values[0] loser_modelVersion_id = self.promptBook[self.promptBook['image_id'] == image_ids[st.session_state.pointer[prompt_id][loser]]]['modelVersion_id'].values[0] # # remove the old battle result if exists # query = "DELETE FROM battle_results WHERE username = %s AND timestamp = %s AND prompt_id = %s AND winner = %s AND loser = %s" # curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id, winner_modelVersion_id, loser_modelVersion_id)) # curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], prompt_id, loser_modelVersion_id, winner_modelVersion_id)) # insert the battle result into the database query = "INSERT INTO battle_results (username, timestamp, tag, prompt_id, winner, loser, battletime, epoch) VALUES (%s, %s, %s, %s, %s, %s, %s, %s)" curser.execute(query, (st.session_state.user_id[0], st.session_state.user_id[1], self.promptBook[self.promptBook['prompt_id'] == prompt_id]['tag'].values[0], prompt_id, winner_modelVersion_id, loser_modelVersion_id, battletime, st.session_state.epoch['ranking'][prompt_id])) curser.close() RANKING_CONN.commit() if curr_position == total_num - 1: st.session_state.epoch['ranking'][prompt_id] += 1 st.session_state.epoch['summary'][tag] = st.session_state.epoch['summary'].get(tag, 0) + 1 st.session_state.epoch['summary']['overall'] += 1 st.session_state.progress[prompt_id] = 'finished' # drop 'modelVersion_standings' from session state if exists st.session_state.pop('modelVersion_standings', None) # st.experimental_rerun() else: st.session_state.pointer[prompt_id][loser] = curr_position + 1 def battle_mode(self, tag, items, prompt_id): self.battle_images(tag, items, prompt_id) def app(self): st.write('### Generative Model Ranking') # st.write('Here you can test out your selected images with any prompt you like. ') # st.write(self.promptBook) # save the current progress to session state if 'progress' not in st.session_state: st.session_state.progress = {} print('current progress: ', st.session_state.progress) # select tag and prompt prompt_tags = self.promptBook['tag'].unique() prompt_tags = np.sort(prompt_tags).tolist() print(st.session_state.gallery_focus) tag_idx = prompt_tags.index(st.session_state.gallery_focus['tag']) if st.session_state.gallery_focus[ 'tag'] in prompt_tags else 0 print(tag_idx) # color the finished tags finished_tags = [] for tag in prompt_tags: append_tag = True for prompt_id in self.promptBook[self.promptBook['tag'] == tag]['prompt_id'].unique(): if prompt_id not in st.session_state.progress or st.session_state.progress[prompt_id] != 'finished': append_tag = False break if append_tag: finished_tags.append(tag) tag_tobe_colored = self.promptBook[self.promptBook['tag'].isin(finished_tags)]['tag'].unique().tolist() colored_tags = self.text_coloring_add(tag_tobe_colored, prompt_tags, color_name='orange') tag = st.radio('Select a tag', colored_tags, index=tag_idx, horizontal=True, label_visibility='collapsed') tag = self.text_coloring_remove(tag) print(tag) items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True) # pick out prompts such that st.session_state.progress[prompt_id] == 'finished' prompts = np.sort(items['prompt'].unique())[::-1].tolist() prompt_idx = prompts.index(st.session_state.gallery_focus['prompt']) if st.session_state.gallery_focus[ 'prompt'] in prompts else 0 print(prompt_idx) # color the finished prompts finished_prompts = [] for prompt_id in items['prompt_id'].unique(): if prompt_id in st.session_state.progress and st.session_state.progress[prompt_id] == 'finished': finished_prompts.append(prompt_id) prompt_tobe_colored = items[items['prompt_id'].isin(finished_prompts)]['prompt'].unique().tolist() colored_prompts = self.text_coloring_add(prompt_tobe_colored, prompts, color_name='✅') selected_prompt = st.selectbox('Select a prompt', colored_prompts, index=prompt_idx, label_visibility='collapsed') selected_prompt = self.text_coloring_remove(selected_prompt) st.session_state.gallery_focus = {'tag': tag, 'prompt': selected_prompt} # mode = st.radio('Select a mode', ['Drag and Sort', 'Battle'], index=1) mode = st.session_state.assigned_rank_mode items = items[items['prompt'] == selected_prompt].reset_index(drop=True) prompt_id = items['prompt_id'].unique()[0] self.sidebar(selected_prompt, items) batch_num = len(items) // self.batch_size batch_num += 1 if len(items) % self.batch_size != 0 else 0 # st.session_state.counter[prompt_id] = 0 if prompt_id not in st.session_state.counter else st.session_state.counter[prompt_id] # save prompt_id in session state st.session_state.prompt_id_tmp = prompt_id if prompt_id not in st.session_state.progress: st.session_state.progress[prompt_id] = 'ranking' if st.session_state.progress[prompt_id] == 'ranking': st.session_state.epoch['ranking'][prompt_id] = st.session_state.epoch['ranking'].get(prompt_id, 1) st.caption("We might pair some other images that you haven't selected based on our evaluation matrix.") if mode == 'Drag and Sort': self.dragsort_mode(tag, items, prompt_id, batch_num) elif mode == 'Battle': self.battle_mode(tag, items, prompt_id) elif st.session_state.progress[prompt_id] == 'finished': print(st.session_state.gallery_focus) # st.toast('**Summary is available now!**') # st.write('---') with st.form(key='ranking_finished'): st.info('**🎉 You have ranked all models for this prompt!**') st.write('Feel free to do the following things:') # st.write('* Back to the gallery page to see more images.') # st.write('* Rank again for this tag and prompt.') options_panel = st.columns(4) with options_panel[0]: summary_btn = st.form_submit_button('📊 See Summary', use_container_width=True, type='primary') if summary_btn: switch_page('summary') with options_panel[1]: gallery_btn = st.form_submit_button('🖼️ Back to Gallery', use_container_width=True) if gallery_btn: switch_page('gallery') with options_panel[2]: st.form_submit_button('👆 Rank other prompts', use_container_width=True) with options_panel[3]: restart_btn = st.form_submit_button('🎖️ Re-rank this prompt', use_container_width=True, on_click=self.rerank, kwargs={'prompt_id': prompt_id}) # with st.sidebar: # st.write('epoch: ', st.session_state.epoch['ranking'][prompt_id]) def text_coloring_add(self, tobe_colored:list, total_items, color_name='orange'): if color_name in ['orange', 'red', 'green', 'blue', 'violet', 'yellow']: colored = [f':{color_name}[{item}]' if item in tobe_colored else item for item in total_items] else: colored = [f'[{color_name}] {item}' if item in tobe_colored else item for item in total_items] return colored def text_coloring_remove(self, tobe_removed): if isinstance(tobe_removed, str): if tobe_removed.startswith(':'): tobe_removed = tobe_removed.split('[')[-1][:-1] elif tobe_removed.startswith('['): tobe_removed = tobe_removed.split(']')[-1][1:] return tobe_removed def rerank(self, prompt_id): st.session_state.progress[prompt_id] = 'ranking' if st.session_state.assigned_rank_mode == 'Drag and Sort': st.session_state.counter[prompt_id] = 0 elif st.session_state.assigned_rank_mode == 'Battle': st.session_state.pointer[prompt_id] = {'left': 0, 'right': 1} 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: has_selection = False for key, value in st.session_state.selected_dict.items(): for v in value: if v: has_selection = True break if not has_selection: 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(st.session_state.show_NSFW) print(st.session_state.selected_dict) # st.write("# Full function is coming soon.") RANKING_CONN = connect_to_db() # 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))]) # replace append with pd.concat user_selections = promptBook[(promptBook['prompt_id'] == key) & (promptBook['modelVersion_id'].isin(value))] # auto complete the selection with random images residual = len(user_selections) % 4 if residual != 0: # select 4-residual items from the promptbook outside the user_selections npc = promptBook[(promptBook['prompt_id'] == key) & (~promptBook['modelVersion_id'].isin(value))].sort_values(by=['total_score'], ascending=False).reset_index(drop=True).iloc[:4-residual] user_selections = pd.concat([user_selections, npc]) promptBook_selected = pd.concat([promptBook_selected, user_selections]) promptBook_selected = promptBook_selected.reset_index(drop=True) # sort promptBook by total_score promptBook_selected = promptBook_selected.sort_values(by=['total_score'], ascending=True).reset_index(drop=True) # st.write(promptBook_selected) images_endpoint = "https://modelcofferbucket.s3-accelerate.amazonaws.com/" app = RankingApp(promptBook_selected, images_endpoint, batch_size=4) app.app() with open('./css/style.css') as f: st.markdown(f'', unsafe_allow_html=True)