DynamicScene/DynamicGeneration

1import json2import gradio as gr3from pathlib import Path4import os5import pickle6from constants import OBJECTS, RECEPTACLES7import pandas as pd8from generate_video import generate_video9from utils import *10 11def get_scene_info(scene_file_name: str) -> tuple[str, str]:12    """13    Return:14        scene_image_path: str, the path of the scene image15        markdown_description: str, the markdown description of the scene16    """17    scene_dir = get_scene_dir_path(scene_file_name)18    scene_image_path = get_scene_image_path(scene_file_name)19    scene_data = get_scene_data(scene_dir)20    all_object_type_list = [object_info['id'].split('|', 1)[0] for object_info in scene_data['objects']]21    room_counter = get_room_counter(scene_data)22    23    receptacle_type_list = [object_type for object_type in all_object_type_list if object_type in RECEPTACLES]24    object_type_list = [object_type for object_type in all_object_type_list if object_type in OBJECTS - RECEPTACLES]25    26    receptacle_counter = {receptacle_type: receptacle_type_list.count(receptacle_type) for receptacle_type in set(receptacle_type_list)}27    object_type_couter = {object_type: object_type_list.count(object_type) for object_type in set(object_type_list)}28    receptacle_type_items = list(receptacle_counter.items())29    object_type_items = list(object_type_couter.items())30    receptacle_type_items.sort(key=lambda x: x[1], reverse=True)31    object_type_items.sort(key=lambda x: x[1], reverse=True)32    receptacle_type_items = receptacle_type_items[:10]33    object_type_items = object_type_items[:10]34    receptacle_type_items += [('' , '')] * (10 - len(receptacle_type_items))35    object_type_items += [('' , '')] * (10 - len(object_type_items))36 37    object_type_couter_df = pd.DataFrame(object_type_items, columns=['Object Type', 'Count'])38    object_type_couter_df = object_type_couter_df.reset_index(drop=True)39    40    receptacle_counter_df = pd.DataFrame(receptacle_type_items, columns=['Receptacle Type', 'Count'])41    receptacle_counter_df = receptacle_counter_df.reset_index(drop=True)42 43    recetpace_number = len(receptacle_type_list)44    object_number = len(object_type_list)45    46    # room type can be like 4 Bedroom, 1 Living Room, 1 Kitchen, 1 Bathroom47    markdown_description = f"""Scene Information48- Number of Rooms: {sum(room_counter.values())}49- Room Types: {', '.join([f'{count} {room_type}' for room_type, count in room_counter.items()])}50- Number of Objects: {object_number}51- Number of Receptacles: {recetpace_number}52"""53    return scene_image_path, markdown_description, receptacle_counter_df, object_type_couter_df54 55def visualize_scene():56    # 使用 Grid 布局组件来组织界面57    default_scene_file_name = SCENE_FILE_NAME_LIST[0]58    default_scene_image_path, default_text, default_receptacle_table, default_object_table = get_scene_info(default_scene_file_name)59    with gr.Row():60        dropdown = gr.Dropdown(choices=SCENE_FILE_NAME_LIST, label="Select Scene ID", value=default_scene_file_name)61    with gr.Row(equal_height=True):62        image = gr.Image(label="Scene Overhead View", show_label=False, value=default_scene_image_path)63        with gr.Column():64            text = gr.Textbox(label="Scene Statistics", value=default_text, lines=6)65            with gr.Row(equal_height=True):66                receptacle_table = gr.Dataframe(label="Receptacle Type Count", height=520, value=default_receptacle_table)67                object_table = gr.Dataframe(label="Object Type Count", height=520, value=default_object_table)68 69    dropdown.change(fn=get_scene_info, inputs=dropdown, outputs=[image, text, receptacle_table, object_table])70    return dropdown71 72 73PERSON_TABLE_COLUMNS = ['name', 'age', 'gender', 'personality', 'routine', 'occupation', 'thoughts', 'lifestyle']74TEMPLATE_ROW_NUMBER = 10    75PERSON_ELEMENT_NUM = 676 77def get_person_info_list(person_file_list: list[str|Path|dict]):78    person_info_list = []79    for person_file in person_file_list:80        if isinstance(person_file, str) or isinstance(person_file, Path):81            with open(person_file, 'rb') as f:82                person_info = pickle.load(f)83                person_info = person_info['persona']84                if 'image' not in person_info:85                    image_path = os.path.join(os.path.dirname(person_file), 'avatar.jpg')86                    person_info['image'] = image_path87        else:88            person_info = person_file['persona']89        90        person_info_list.append(person_info)91    return person_info_list92 93def person_info_to_description(person_info):94    return f"{person_info['name']}, a {', '.join(person_info['personality'])} {person_info['age']} years old {'man' if person_info['gender'] == 'Male' else 'woman'}."95    96def person_info_to_elements(person_info):97    return (98        person_info['image'],99        f"**Name:** {person_info['name']}",100        f"**Age:** {person_info['age']}",101        f"**Gender:** {person_info['gender']}",102        f"**Personality:** {', '.join(person_info['personality'])}",103        f"**Routine:** {person_info['routine']}"104    )105 106def get_person_elements_from_row_elements(elements: list[str], row_index: int):107    return elements[row_index * PERSON_ELEMENT_NUM: (row_index + 1) * PERSON_ELEMENT_NUM]108 109def get_person_name_from_row_elements(row_elements: list[str]):110    return row_elements[1].replace('**Name:** ', '')111 112def get_person_dataframe(row_index_to_person_name: list[str], person_name_to_info: dict[str, dict]):113    if len(row_index_to_person_name) == 0:114        return pd.DataFrame(columns=PERSON_TABLE_COLUMNS)115    else:116        return pd.DataFrame([person_name_to_info[person_name] for person_name in row_index_to_person_name])[PERSON_TABLE_COLUMNS]117 118def get_max_person_number(scene_dir_path: Path) -> int:119    scene_data = get_scene_data(scene_dir_path)120    room_counter = get_room_counter(scene_data)121    return room_counter.get('Bedroom', 0)122    123def create_person_page():124    with gr.Row(visible=False) as row:125        image = gr.Image(width=200, scale=0.25, show_label=False, interactive=False)126        with gr.Column():127            with gr.Row():128                name = gr.Markdown()129                delete_button = gr.Button("Delete",size='sm')130            age = gr.Markdown()131            gender = gr.Markdown()132            personality = gr.Markdown()133            routine = gr.Markdown()134            # gr.Markdown(f"**Occupation:** {person['occupation']}")135            # gr.Markdown(f"**Thoughts:** {person['thoughts']}")136            # gr.Markdown(f"**Lifestyle:** {person['lifestyle']}")137    return {138        'row': row,139        'delete_button': delete_button,140        'elements': (image, name, age, gender, personality, routine),141    }142 143def full_view(person_name_to_description: dict[str, str]):144    with gr.Blocks():145        person_pages = [create_person_page() for _ in range(TEMPLATE_ROW_NUMBER)]146        add_person_dropdown = gr.Dropdown(choices=[person_name_to_description[key] for key in sorted(person_name_to_description.keys())], label='Add Person')147        add_button = gr.Button('Add')148    return (person_pages, add_person_dropdown, add_button)149 150def compact_view():151    return gr.Dataframe(value=pd.DataFrame(columns=PERSON_TABLE_COLUMNS), label='People Information')152 153 154def add_person_to_scene(155    max_person_number: int, 156    person_dropdown_description: str, person_name_to_description: dict[str, str],157    row_index_to_person_name: list[str],158    person_name_to_info: dict[str, dict]159):160    current_row_number = len(row_index_to_person_name) + 1161    add_button_visible = current_row_number < max_person_number162    add_person_name = [key for key, value in person_name_to_description.items() if value == person_dropdown_description][0]163    164    add_person_elements = person_info_to_elements(person_name_to_info[add_person_name])165    row_elements_update_list = [gr.update() for _ in range(TEMPLATE_ROW_NUMBER * len(add_person_elements))]166    row_visible_list = [gr.update() for _ in range(TEMPLATE_ROW_NUMBER)]167    row_elements_update_list[(current_row_number-1) * PERSON_ELEMENT_NUM: current_row_number * PERSON_ELEMENT_NUM] = add_person_elements168    row_visible_list[current_row_number-1] = gr.update(visible=True)169    170    row_index_to_person_name.append(add_person_name)171    to_be_chose_person_name_list = [key for key in person_name_to_description.keys() if key not in row_index_to_person_name]172    to_be_chose_person_description_list = [person_name_to_description[key] for key in sorted(to_be_chose_person_name_list)]173    174    return row_elements_update_list + row_visible_list + [175        row_index_to_person_name,176        gr.update(choices=to_be_chose_person_description_list, visible=add_button_visible),177        gr.update(visible=add_button_visible),178        get_person_dataframe(row_index_to_person_name, person_name_to_info),179        gr.update(visible=True)180    ]181 182 183def delete_person_from_scene(184    person_name_to_info: dict[str, dict],185    row_index_to_person_name: list[int],186    button_row_index: int,187    person_name_to_description: dict[str, str], 188    *row_elements: tuple[str]189):190    current_row_number = len(row_index_to_person_name)191    row_elements = list(row_elements)192    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]193    person_row_visible = [gr.update() for _ in range(TEMPLATE_ROW_NUMBER)]194    person_row_visible[current_row_number-1] = gr.update(visible=False)195    row_index_to_person_name.pop(button_row_index)196    to_be_chose_person_name_list = [key for key in person_name_to_description.keys() if key not in row_index_to_person_name]197    to_be_chose_person_description_list = [person_name_to_description[key] for key in sorted(to_be_chose_person_name_list)]198    199    return row_elements + person_row_visible + [200        gr.update(visible=True),201        gr.update(choices=to_be_chose_person_description_list, visible=True),202        row_index_to_person_name,203        get_person_dataframe(row_index_to_person_name, person_name_to_info),204        gr.update(visible=len(row_index_to_person_name) > 0)205    ]206    207 208def reset_person_rows(person_name_to_description: dict[str, str]):209    rows_visible = [gr.update(visible=False) for _ in range(TEMPLATE_ROW_NUMBER)]210    row_index_to_person_name = []211    add_button_visible = gr.update(visible=True)212    add_person_dropdown = gr.update(choices=[person_name_to_description[key] for key in sorted(person_name_to_description.keys())], visible=True)213    person_dataframe = pd.DataFrame(columns=PERSON_TABLE_COLUMNS)214    return rows_visible + [row_index_to_person_name, add_button_visible, add_person_dropdown, person_dataframe]215 216def reset_person_page(217    scene_file_name: str,218    person_name_to_description: dict[str, str],219):220    # get the bedroom number of the scene221    scene_dir = SCENE_ROOT_DIR / scene_file_name222    max_added_person = get_max_person_number(scene_dir)223    return reset_person_rows(person_name_to_description) + [max_added_person, gr.update(visible=False), gr.update(visible=False)]224 225def visualize_person(person_name_to_info: gr.State, person_name_to_description: gr.State, max_added_person: int):226 227    row_index_to_person_name = gr.State([])228    max_added_person = gr.State(max_added_person)229 230    with gr.Blocks():231        with gr.Tab(label='Full View'):232            person_pages, add_person_dropdown, add_button = full_view(person_name_to_description.value)233        with gr.Tab(label='Compact View'):234            person_dataframe = compact_view()235 236    generate_button = gr.Button("Generate Video", visible=False)237    row_elements = [element for page in person_pages for element in page['elements']]238    rows = [page['row'] for page in person_pages]239    delete_buttons: list[gr.Button] = [page['delete_button'] for page in person_pages]240    241    add_button.click(242        fn=add_person_to_scene, 243        inputs=[max_added_person, add_person_dropdown, person_name_to_description, row_index_to_person_name, person_name_to_info],244        outputs=row_elements + rows + [row_index_to_person_name, add_person_dropdown, add_button, person_dataframe, generate_button]245    )246    for i, delete_button in enumerate(delete_buttons):247        delete_button.click(248            fn=delete_person_from_scene,249            inputs=[person_name_to_info, row_index_to_person_name, gr.State(i), person_name_to_description] + row_elements,250            outputs=row_elements + rows + [add_button, add_person_dropdown, row_index_to_person_name, person_dataframe, generate_button]251        )252        253    return row_index_to_person_name, max_added_person, row_elements, rows, delete_buttons, add_button, add_person_dropdown, person_dataframe, generate_button254    255    256def generate_button_click_change_state():257    return gr.update(visible=False), gr.update(visible=True)258     259 260def generate_button_click_change_video(261    scene_file_name: str,262    row_index_to_person_name: list[str],263    person_name_to_file: dict[str, dict[str, any]]264):265    video_path = generate_video(scene_file_name, [person_name_to_file[person_name] for person_name in row_index_to_person_name])266    if video_path:267        return str(video_path)268    return gr.update(visible=False)269 270 271def visualize_dynamic_generate(person_name_to_file, person_name_to_info, person_name_to_description):272    gr.Markdown("## Scene Information")273    scene_dropdown = visualize_scene()274    max_person_number = get_max_person_number(SCENE_ROOT_DIR / scene_dropdown.value)275    gr.Markdown("## Person Information")276    row_index_to_person_name, max_person_number, _, rows, _, add_button, add_person_dropdown, person_dataframe, generate_button \277        = visualize_person(person_name_to_info, person_name_to_description, max_person_number)278        279    video = gr.Video(visible=False)280    clear_button = gr.Button("Clear")281    generate_button.click(282        fn=generate_button_click_change_video,283        inputs=[scene_dropdown, row_index_to_person_name, person_name_to_file],284        outputs=video285    )286    generate_button.click(287        fn=generate_button_click_change_state,288        outputs=[generate_button, video]289    )290    reset_inputs = [scene_dropdown, person_name_to_description]291    reset_outputs = rows + [row_index_to_person_name, add_button, add_person_dropdown, person_dataframe, max_person_number, generate_button, video]292    scene_dropdown.change(293        fn=reset_person_page,294        inputs=reset_inputs,295        outputs=reset_outputs296    )297    clear_button.click(298        fn=reset_person_page,299        inputs=reset_inputs,300        outputs=reset_outputs301    )302    return reset_inputs, reset_outputs303 304if __name__ == '__main__':305    with gr.Blocks() as demo:306        visualize_dynamic_generate()307    demo.launch()308 309