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() |