File size: 7,212 Bytes
4e60da3
 
 
 
 
ec1c2f8
e9722ce
4e60da3
 
 
 
 
 
 
 
 
 
 
 
e9722ce
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7150512
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ec1c2f8
 
 
 
 
 
 
 
 
 
 
4e60da3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e9722ce
 
 
 
 
 
 
 
 
4e60da3
 
7150512
 
 
 
 
 
 
 
 
4e60da3
 
 
ec1c2f8
 
 
 
 
 
 
 
 
 
 
4e60da3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
97fb5a2
 
4e60da3
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
import os
import gradio as gr
import random
import pickle
import numpy as np
import zipfile
import trimesh
from PIL import Image
from huggingface_hub import hf_hub_download

def pose_generation(scene, count):
    assert isinstance(scene, str)
    results_path = hf_hub_download('SceneDiffuser/SceneDiffuser', 'results/pose_generation/results.pkl')
    with open(results_path, 'rb') as f:
        results = pickle.load(f)
    
    images = [Image.fromarray(results[scene][random.randint(0, 19)]) for i in range(count)]
    return images

def pose_generation_mesh(scene, count):
    assert isinstance(scene, str)
    scene_path = f"./results/pose_generation/mesh_results/{scene}/scene_downsample.ply"
    if not os.path.exists(scene_path):
        results_path = hf_hub_download('SceneDiffuser/SceneDiffuser', 'results/pose_generation/mesh_results.zip')
        os.makedirs('./results/pose_generation/', exist_ok=True)
        with zipfile.ZipFile(results_path, 'r') as zip_ref:
            zip_ref.extractall('./results/pose_generation/')
    
    res = './results/pose_generation/tmp.glb'
    S = trimesh.Scene()
    S.add_geometry(trimesh.load(scene_path))
    for i in range(count):
        rid = random.randint(0, 19)
        S.add_geometry(trimesh.load(
            f"./results/pose_generation/mesh_results/{scene}/body{rid:0>3d}.ply"
        ))
    S.export(res)
    
    return res

def motion_generation(scene):
    assert isinstance(scene, str)
    cnt = {
        'MPH1Library': 3,
        'MPH16': 6,
        'N0SittingBooth': 7,
        'N3OpenArea': 5
    }[scene]

    res = f"./results/motion_generation/results/{scene}/{random.randint(0, cnt-1)}.gif"
    if not os.path.exists(res):
        results_path = hf_hub_download('SceneDiffuser/SceneDiffuser', 'results/motion_generation/results.zip')
        os.makedirs('./results/motion_generation/', exist_ok=True)
        with zipfile.ZipFile(results_path, 'r') as zip_ref:
            zip_ref.extractall('./results/motion_generation/')
    
    return res

def grasp_generation(case_id):
    assert isinstance(case_id, str)
    res = f"./results/grasp_generation/results/{case_id}/{random.randint(0, 19)}.glb"
    if not os.path.exists(res):
        results_path = hf_hub_download('SceneDiffuser/SceneDiffuser', 'results/grasp_generation/results.zip')
        os.makedirs('./results/grasp_generation/', exist_ok=True)
        with zipfile.ZipFile(results_path, 'r') as zip_ref:
            zip_ref.extractall('./results/grasp_generation/')
    
    return res

def path_planning(case_id):
    assert isinstance(case_id, str)
    results_path = hf_hub_download('SceneDiffuser/SceneDiffuser', 'results/path_planning/results.pkl')
    with open(results_path, 'rb') as f:
        results = pickle.load(f)

    case = results[case_id]
    steps = case['step']
    image = Image.fromarray(case['image'])
    return image, steps

with gr.Blocks() as demo:
    gr.Markdown("# **<p align='center'>Diffusion-based Generation, Optimization, and Planning in 3D Scenes</p>**")
    gr.HTML(value="<img src='file/figures/teaser.png' alt='Teaser' width='710px' height='284px' style='display: block; margin: auto;'>")
    gr.HTML(value="<p align='center' style='font-size: 1.25em; color: #485fc7;'><a href='' target='_blank'>Paper</a> | <a href='' target='_blank'>Project Page</a> | <a href='' target='_blank'>Github</a></p>")
    gr.Markdown("<p align='center'><i>\"SceneDiffuser provides a unified model for solving scene-conditioned generation, optimization, and planning.\"</i></p>")

    ## five task
    ## pose generation
    with gr.Tab("Pose Generation"):
        with gr.Row():
            with gr.Column():
                input1 = [
                    gr.Dropdown(choices=['MPH16', 'MPH1Library', 'N0SittingBooth', 'N3OpenArea'], label='Scenes'),
                    gr.Slider(minimum=1, maximum=4, step=1, label='Count', interactive=True)
                ]
                button1 = gr.Button("Generate")
            with gr.Column():
                output1 = [
                    gr.Gallery(label="Result").style(grid=[1], height="auto")
                ]
    button1.click(pose_generation, inputs=input1, outputs=output1)

    with gr.Tab("Pose Generation Mesh"):
        input11 = [
            gr.Dropdown(choices=['MPH16', 'MPH1Library', 'N0SittingBooth', 'N3OpenArea'], label='Scenes'),
            gr.Slider(minimum=1, maximum=4, step=1, label='Count', interactive=True)
        ]
        button11 = gr.Button("Generate")
        output11 = gr.Model3D(clear_color=[255, 255, 255, 255], label="Result")
    button11.click(pose_generation_mesh, inputs=input11, outputs=output11)

    ## motion generation
    with gr.Tab("Motion Generation"):
        with gr.Row():
            with gr.Column():
                input2 = [
                    gr.Dropdown(choices=['MPH16', 'MPH1Library', 'N0SittingBooth', 'N3OpenArea'], label='Scenes')
                ]
                button2 = gr.Button("Generate")
            with gr.Column():
                output2 = gr.Image(label="Result")
    button2.click(motion_generation, inputs=input2, outputs=output2)
    
    ## grasp generation
    with gr.Tab("Grasp Generation"):
        with gr.Row():
            with gr.Column():
                input3 = [
                    gr.Dropdown(choices=['contactdb+apple', 'contactdb+camera', 'contactdb+cylinder_medium', 'contactdb+door_knob', 'contactdb+rubber_duck', 'contactdb+water_bottle', 'ycb+baseball', 'ycb+pear', 'ycb+potted_meat_can', 'ycb+tomato_soup_can'], label='Objects')
                ]
                button3 = gr.Button("Run")
            with gr.Column():
                output3 = [
                    gr.Model3D(clear_color=[255, 255, 255, 255], label="Result")
                ]
    button3.click(grasp_generation, inputs=input3, outputs=output3)
    
    ## path planning
    with gr.Tab("Path Planing"):
        with gr.Row():
            with gr.Column():
                input4 = [
                    gr.Dropdown(choices=['scene0603_00_N0pT', 'scene0621_00_cJ4H', 'scene0634_00_48Y3', 'scene0634_00_gIRH', 'scene0637_00_YgjR', 'scene0640_00_BO94', 'scene0641_00_3K6J', 'scene0641_00_KBKx', 'scene0641_00_cb7l', 'scene0645_00_35Hy', 'scene0645_00_47D1', 'scene0645_00_XfLE', 'scene0667_00_DK4F', 'scene0667_00_o7XB', 'scene0667_00_rUMp', 'scene0672_00_U250', 'scene0673_00_Jyw8', 'scene0673_00_u1lJ', 'scene0678_00_QbNL', 'scene0678_00_RrY0', 'scene0678_00_aE1p', 'scene0678_00_hnXu', 'scene0694_00_DgAL', 'scene0694_00_etF5', 'scene0698_00_tT3Q'], label='Scenes'),
                ]
                button4 = gr.Button("Run")
            with gr.Column():
                # output4 = gr.Gallery(label="Result").style(grid=[1], height="auto")
                output4 = [
                    gr.Image(label="Result"),
                    gr.Number(label="Steps", precision=0)
                ]
    button4.click(path_planning, inputs=input4, outputs=output4)

    ## arm motion planning
    with gr.Tab("Arm Motion Planning"):
        gr.Markdown('Coming soon!')

    gr.Markdown("<p>Note: Currently, the output results are pre-sampled results. We will deploy a real-time model after we release the code.</p>")
    
demo.launch()