Spaces:
Sleeping
Sleeping
| import os | |
| import threading | |
| import gradio as gr | |
| from transcrib3d_main import gen_prompt, get_gpt_response, get_openai_config, extract_answer_id_from_last_line | |
| from code_interpreter import CodeInterpreter | |
| from display_model import * | |
| scan_id = "scene0132_00" | |
| ply_file = os.path.join("scenes", f"{scan_id}_vh_clean_2_aligned.ply") | |
| glb_file = os.path.join("scenes", f"{scan_id}_vh_clean_2_aligned.glb") | |
| new_ply_file = os.path.join("scenes", f"{scan_id}_vh_clean_2_aligned_AddBox.ply") | |
| new_glb_file = os.path.join("scenes", f"{scan_id}_vh_clean_2_aligned_AddBox.glb") | |
| objects_info_file = os.path.join("objects_info", f"objects_info_{scan_id}.npy") | |
| def insert_user_none_between_assistant(messages): | |
| # 初始化结果列表 | |
| result = [] | |
| # 初始状态设置为"user",以确保列表第一个条目为"assistant"时能正确插入 | |
| last_role = "user" | |
| for msg in messages: | |
| # 检查当前信息的角色 | |
| current_role = msg["role"] | |
| # 如果上一个和当前信息均为"assistant",插入content为None的"user"信息 | |
| if last_role == "assistant" and current_role == "assistant": | |
| result.append({"role": "user", "content": None}) | |
| # 将当前信息添加到结果列表 | |
| result.append(msg) | |
| # 更新上一条信息的角色 | |
| last_role = current_role | |
| return result | |
| def generate_answer_glb(answer_content): | |
| last_line = answer_content.splitlines()[-1] if len(answer_content) > 0 else '' | |
| answer_id, _ = extract_answer_id_from_last_line(last_line) | |
| print("extracted answer id:", answer_id) | |
| # get the bounding box of the answer object | |
| box = np.load(objects_info_file, allow_pickle=True)[answer_id]['extension'] | |
| print("box extension:",box) | |
| # add the box to ply | |
| add_1box_to_ply(box, ply_file, new_ply_file) | |
| ply_to_glb(new_ply_file, new_glb_file) | |
| def run_inferring(instruction, model3d, dialogue): | |
| # generate prompt from user instruction | |
| # scan_id = "scene0132_00" | |
| prompt = gen_prompt(instruction, scan_id) | |
| # get oepnai config | |
| openai_config = get_openai_config() | |
| # get LLM response | |
| code_interpreter = CodeInterpreter(**openai_config) | |
| get_gpt_response(prompt, code_interpreter) | |
| messages = code_interpreter.pretext | |
| # draw the answer bounding box to the scene | |
| generate_answer_glb(messages[-1]['content']) | |
| # global model3d | |
| # print(model3d.value) | |
| # model3d.postprocess(new_glb_file) | |
| # print(model3d.value) | |
| # form gradio chat history | |
| messages = insert_user_none_between_assistant(messages[1:]) | |
| # print(len(messages)) | |
| # print(messages) | |
| gradio_messages = [] | |
| for idx in range(int(len(messages)/2)): | |
| gradio_message = [messages[idx*2]['content'], messages[idx*2+1]['content']] | |
| gradio_messages.append(gradio_message) | |
| # return new_glb_file, gradio_messages | |
| model3d.update(value=new_glb_file) | |
| dialogue.update(gradio_messages) | |
| def process_instruction_callback(user_instruction, model3d, dialogue): | |
| threading.Thread(target=run_inferring, args=(user_instruction, model3d, dialogue)).start() | |
| # return "Processing your instruction, please wait...", | |
| with gr.Blocks() as demo: | |
| gr.Markdown("## Transcrib3D-Demo") | |
| with gr.Row(): | |
| model3d = gr.Model3D( | |
| value="scenes/scene0132_00_vh_clean_2_aligned.glb", | |
| # value="scenes/scene0132_00_vh_clean_2_aligned_AddBox.glb", | |
| # value="scenes/scene0132_00_vh_clean_2_aligned.ply", | |
| # value="scenes/scene0132_00_vh_clean_2_aligned.obj", | |
| # value="scenes/scene0132_00_gt_bboxes_aligned.ply", | |
| # value="scenes/cube.ply", | |
| label="ScanNet-scene0132_00", | |
| camera_position=(90,120,8), | |
| zoom_speed=0.25, | |
| height=635 | |
| ) | |
| # print("Type1:",type(model3d)) | |
| with gr.Column(): | |
| # with gr.Row(): | |
| user_instruction_textbox = gr.Textbox( | |
| label="Instruction", | |
| placeholder="Describe an object in the scene with its attributes and its relation with other objects.", | |
| # scale=4 | |
| ) | |
| bt = gr.Button( | |
| value="Submit", | |
| # scale=1 | |
| ) | |
| dialogue = gr.Chatbot( | |
| height=470 | |
| # value = [["1","2"], [None, '3']] | |
| ) | |
| # print("Type2:",type(model3d)) | |
| # 直接在 inputs列表里写model3d,会导致实际传给callback函数的是str | |
| # bt.click(fn=process_instruction_callback, inputs=user_instruction_textbox, outputs=dialogue) | |
| bt.click(fn=process_instruction_callback, inputs=[user_instruction_textbox, gr.State(model3d), gr.State(dialogue)])#, outputs=[model3d,dialogue]) | |
| # 直接用lambda函数定义一个映射 | |
| # type(user_instruction_textbox.value) | |
| # user_instruction_textbox. | |
| # user_instruction_textbox.submit(fn=lambda: process_instruction_callback(user_instruction_textbox, model3d), inputs=[], outputs=dialogue) | |
| # user_instruction_textbox. | |
| # bt.click(fn=lambda: process_instruction_callback(user_instruction_textbox, model3d), inputs=[], outputs=dialogue) | |
| demo.launch() |