import json import gradio as gr from pathlib import Path import os import pickle from constants import OBJECTS, RECEPTACLES import pandas as pd from generate_video import generate_video from utils import * def get_scene_info(scene_file_name: str) -> tuple[str, str]: """ Return: scene_image_path: str, the path of the scene image markdown_description: str, the markdown description of the scene """ scene_dir = get_scene_dir_path(scene_file_name) scene_image_path = get_scene_image_path(scene_file_name) scene_data = get_scene_data(scene_dir) all_object_type_list = [object_info['id'].split('|', 1)[0] for object_info in scene_data['objects']] room_counter = get_room_counter(scene_data) receptacle_type_list = [object_type for object_type in all_object_type_list if object_type in RECEPTACLES] object_type_list = [object_type for object_type in all_object_type_list if object_type in OBJECTS - RECEPTACLES] receptacle_counter = {receptacle_type: receptacle_type_list.count(receptacle_type) for receptacle_type in set(receptacle_type_list)} object_type_couter = {object_type: object_type_list.count(object_type) for object_type in set(object_type_list)} receptacle_type_items = list(receptacle_counter.items()) object_type_items = list(object_type_couter.items()) receptacle_type_items.sort(key=lambda x: x[1], reverse=True) object_type_items.sort(key=lambda x: x[1], reverse=True) receptacle_type_items = receptacle_type_items[:10] object_type_items = object_type_items[:10] receptacle_type_items += [('' , '')] * (10 - len(receptacle_type_items)) object_type_items += [('' , '')] * (10 - len(object_type_items)) object_type_couter_df = pd.DataFrame(object_type_items, columns=['Object Type', 'Count']) object_type_couter_df = object_type_couter_df.reset_index(drop=True) receptacle_counter_df = pd.DataFrame(receptacle_type_items, columns=['Receptacle Type', 'Count']) receptacle_counter_df = receptacle_counter_df.reset_index(drop=True) recetpace_number = len(receptacle_type_list) object_number = len(object_type_list) # room type can be like 4 Bedroom, 1 Living Room, 1 Kitchen, 1 Bathroom markdown_description = f"""Scene Information - Number of Rooms: {sum(room_counter.values())} - Room Types: {', '.join([f'{count} {room_type}' for room_type, count in room_counter.items()])} - Number of Objects: {object_number} - Number of Receptacles: {recetpace_number} """ return scene_image_path, markdown_description, receptacle_counter_df, object_type_couter_df def visualize_scene(): # 使用 Grid 布局组件来组织界面 default_scene_file_name = SCENE_FILE_NAME_LIST[0] default_scene_image_path, default_text, default_receptacle_table, default_object_table = get_scene_info(default_scene_file_name) with gr.Row(): dropdown = gr.Dropdown(choices=SCENE_FILE_NAME_LIST, label="Select Scene ID", value=default_scene_file_name) with gr.Row(equal_height=True): image = gr.Image(label="Scene Overhead View", show_label=False, value=default_scene_image_path) with gr.Column(): text = gr.Textbox(label="Scene Statistics", value=default_text, lines=6) with gr.Row(equal_height=True): receptacle_table = gr.Dataframe(label="Receptacle Type Count", height=520, value=default_receptacle_table) object_table = gr.Dataframe(label="Object Type Count", height=520, value=default_object_table) dropdown.change(fn=get_scene_info, inputs=dropdown, outputs=[image, text, receptacle_table, object_table]) return dropdown PERSON_TABLE_COLUMNS = ['name', 'age', 'gender', 'personality', 'routine', 'occupation', 'thoughts', 'lifestyle'] TEMPLATE_ROW_NUMBER = 10 PERSON_ELEMENT_NUM = 6 def get_person_info_list(person_file_list: list[str|Path|dict]): person_info_list = [] for person_file in person_file_list: if isinstance(person_file, str) or isinstance(person_file, Path): with open(person_file, 'rb') as f: person_info = pickle.load(f) person_info = person_info['persona'] if 'image' not in person_info: image_path = os.path.join(os.path.dirname(person_file), 'avatar.jpg') person_info['image'] = image_path else: person_info = person_file['persona'] person_info_list.append(person_info) return person_info_list def person_info_to_description(person_info): return f"{person_info['name']}, a {', '.join(person_info['personality'])} {person_info['age']} years old {'man' if person_info['gender'] == 'Male' else 'woman'}." def person_info_to_elements(person_info): return ( person_info['image'], f"**Name:** {person_info['name']}", f"**Age:** {person_info['age']}", f"**Gender:** {person_info['gender']}", f"**Personality:** {', '.join(person_info['personality'])}", f"**Routine:** {person_info['routine']}" ) def get_person_elements_from_row_elements(elements: list[str], row_index: int): return elements[row_index * PERSON_ELEMENT_NUM: (row_index + 1) * PERSON_ELEMENT_NUM] def get_person_name_from_row_elements(row_elements: list[str]): return row_elements[1].replace('**Name:** ', '') def get_person_dataframe(row_index_to_person_name: list[str], person_name_to_info: dict[str, dict]): if len(row_index_to_person_name) == 0: return pd.DataFrame(columns=PERSON_TABLE_COLUMNS) else: return pd.DataFrame([person_name_to_info[person_name] for person_name in row_index_to_person_name])[PERSON_TABLE_COLUMNS] def get_max_person_number(scene_dir_path: Path) -> int: scene_data = get_scene_data(scene_dir_path) room_counter = get_room_counter(scene_data) return room_counter.get('Bedroom', 0) def create_person_page(): with gr.Row(visible=False) as row: image = gr.Image(width=200, scale=0.25, show_label=False, interactive=False) with gr.Column(): with gr.Row(): name = gr.Markdown() delete_button = gr.Button("Delete",size='sm') age = gr.Markdown() gender = gr.Markdown() personality = gr.Markdown() routine = gr.Markdown() # gr.Markdown(f"**Occupation:** {person['occupation']}") # gr.Markdown(f"**Thoughts:** {person['thoughts']}") # gr.Markdown(f"**Lifestyle:** {person['lifestyle']}") return { 'row': row, 'delete_button': delete_button, 'elements': (image, name, age, gender, personality, routine), } def full_view(person_name_to_description: dict[str, str]): with gr.Blocks(): person_pages = [create_person_page() for _ in range(TEMPLATE_ROW_NUMBER)] add_person_dropdown = gr.Dropdown(choices=[person_name_to_description[key] for key in sorted(person_name_to_description.keys())], label='Add Person') add_button = gr.Button('Add') return (person_pages, add_person_dropdown, add_button) def compact_view(): return gr.Dataframe(value=pd.DataFrame(columns=PERSON_TABLE_COLUMNS), label='People Information') def add_person_to_scene( max_person_number: int, person_dropdown_description: str, person_name_to_description: dict[str, str], row_index_to_person_name: list[str], person_name_to_info: dict[str, dict] ): current_row_number = len(row_index_to_person_name) + 1 add_button_visible = current_row_number < max_person_number add_person_name = [key for key, value in person_name_to_description.items() if value == person_dropdown_description][0] add_person_elements = person_info_to_elements(person_name_to_info[add_person_name]) row_elements_update_list = [gr.update() for _ in range(TEMPLATE_ROW_NUMBER * len(add_person_elements))] row_visible_list = [gr.update() for _ in range(TEMPLATE_ROW_NUMBER)] row_elements_update_list[(current_row_number-1) * PERSON_ELEMENT_NUM: current_row_number * PERSON_ELEMENT_NUM] = add_person_elements row_visible_list[current_row_number-1] = gr.update(visible=True) row_index_to_person_name.append(add_person_name) to_be_chose_person_name_list = [key for key in person_name_to_description.keys() if key not in row_index_to_person_name] to_be_chose_person_description_list = [person_name_to_description[key] for key in sorted(to_be_chose_person_name_list)] return row_elements_update_list + row_visible_list + [ row_index_to_person_name, gr.update(choices=to_be_chose_person_description_list, visible=add_button_visible), gr.update(visible=add_button_visible), get_person_dataframe(row_index_to_person_name, person_name_to_info), gr.update(visible=True) ] def delete_person_from_scene( person_name_to_info: dict[str, dict], row_index_to_person_name: list[int], button_row_index: int, person_name_to_description: dict[str, str], *row_elements: tuple[str] ): current_row_number = len(row_index_to_person_name) row_elements = list(row_elements) row_elements[button_row_index * PERSON_ELEMENT_NUM: (current_row_number-1) * PERSON_ELEMENT_NUM] = row_elements[(button_row_index + 1) * PERSON_ELEMENT_NUM: current_row_number * PERSON_ELEMENT_NUM] person_row_visible = [gr.update() for _ in range(TEMPLATE_ROW_NUMBER)] person_row_visible[current_row_number-1] = gr.update(visible=False) row_index_to_person_name.pop(button_row_index) to_be_chose_person_name_list = [key for key in person_name_to_description.keys() if key not in row_index_to_person_name] to_be_chose_person_description_list = [person_name_to_description[key] for key in sorted(to_be_chose_person_name_list)] return row_elements + person_row_visible + [ gr.update(visible=True), gr.update(choices=to_be_chose_person_description_list, visible=True), row_index_to_person_name, get_person_dataframe(row_index_to_person_name, person_name_to_info), gr.update(visible=len(row_index_to_person_name) > 0) ] def reset_person_rows(person_name_to_description: dict[str, str]): rows_visible = [gr.update(visible=False) for _ in range(TEMPLATE_ROW_NUMBER)] row_index_to_person_name = [] add_button_visible = gr.update(visible=True) add_person_dropdown = gr.update(choices=[person_name_to_description[key] for key in sorted(person_name_to_description.keys())], visible=True) person_dataframe = pd.DataFrame(columns=PERSON_TABLE_COLUMNS) return rows_visible + [row_index_to_person_name, add_button_visible, add_person_dropdown, person_dataframe] def reset_person_page( scene_file_name: str, person_name_to_description: dict[str, str], ): # get the bedroom number of the scene scene_dir = SCENE_ROOT_DIR / scene_file_name max_added_person = get_max_person_number(scene_dir) return reset_person_rows(person_name_to_description) + [max_added_person, gr.update(visible=False), gr.update(visible=False)] def visualize_person(person_name_to_info: gr.State, person_name_to_description: gr.State, max_added_person: int): row_index_to_person_name = gr.State([]) max_added_person = gr.State(max_added_person) with gr.Blocks(): with gr.Tab(label='Full View'): person_pages, add_person_dropdown, add_button = full_view(person_name_to_description.value) with gr.Tab(label='Compact View'): person_dataframe = compact_view() generate_button = gr.Button("Generate Video", visible=False) row_elements = [element for page in person_pages for element in page['elements']] rows = [page['row'] for page in person_pages] delete_buttons: list[gr.Button] = [page['delete_button'] for page in person_pages] add_button.click( fn=add_person_to_scene, inputs=[max_added_person, add_person_dropdown, person_name_to_description, row_index_to_person_name, person_name_to_info], outputs=row_elements + rows + [row_index_to_person_name, add_person_dropdown, add_button, person_dataframe, generate_button] ) for i, delete_button in enumerate(delete_buttons): delete_button.click( fn=delete_person_from_scene, inputs=[person_name_to_info, row_index_to_person_name, gr.State(i), person_name_to_description] + row_elements, outputs=row_elements + rows + [add_button, add_person_dropdown, row_index_to_person_name, person_dataframe, generate_button] ) return row_index_to_person_name, max_added_person, row_elements, rows, delete_buttons, add_button, add_person_dropdown, person_dataframe, generate_button def generate_button_click_change_state(): return gr.update(visible=False), gr.update(visible=True) def generate_button_click_change_video( scene_file_name: str, row_index_to_person_name: list[str], person_name_to_file: dict[str, dict[str, any]] ): video_path = generate_video(scene_file_name, [person_name_to_file[person_name] for person_name in row_index_to_person_name]) if video_path: return str(video_path) return gr.update(visible=False) def visualize_dynamic_generate(person_name_to_file, person_name_to_info, person_name_to_description): gr.Markdown("## Scene Information") scene_dropdown = visualize_scene() max_person_number = get_max_person_number(SCENE_ROOT_DIR / scene_dropdown.value) gr.Markdown("## Person Information") row_index_to_person_name, max_person_number, _, rows, _, add_button, add_person_dropdown, person_dataframe, generate_button \ = visualize_person(person_name_to_info, person_name_to_description, max_person_number) video = gr.Video(visible=False) clear_button = gr.Button("Clear") generate_button.click( fn=generate_button_click_change_video, inputs=[scene_dropdown, row_index_to_person_name, person_name_to_file], outputs=video ) generate_button.click( fn=generate_button_click_change_state, outputs=[generate_button, video] ) reset_inputs = [scene_dropdown, person_name_to_description] reset_outputs = rows + [row_index_to_person_name, add_button, add_person_dropdown, person_dataframe, max_person_number, generate_button, video] scene_dropdown.change( fn=reset_person_page, inputs=reset_inputs, outputs=reset_outputs ) clear_button.click( fn=reset_person_page, inputs=reset_inputs, outputs=reset_outputs ) return reset_inputs, reset_outputs if __name__ == '__main__': with gr.Blocks() as demo: visualize_dynamic_generate() demo.launch()