update

app.py

@@ -17,6 +17,7 @@ from StructDiffusion.diffusion.sampler import Sampler
 from StructDiffusion.diffusion.pose_conversion import get_struct_objs_poses
 from StructDiffusion.utils.files import get_checkpoint_path_from_dir
 from StructDiffusion.utils.rearrangement import show_pcs_with_trimesh
+import StructDiffusion.utils.transformations as tra
 
 
 def move_pc_and_create_scene_simple(obj_xyzs, struct_pose, pc_poses_in_struct):
@@ -78,7 +79,21 @@ class Infer_Wrapper:
 
         self.sampler = Sampler(ConditionalPoseDiffusionModel, checkpoint_path, self.device)
 
-    def run(self, di):
+    def visualize_scene(self, di, session_id):
+        raw_datum = self.dataset.get_raw_data(di)
+        language_command = self.tokenizer.convert_structure_params_to_natural_language(raw_datum["sentence"])
+
+        obj_xyz = raw_datum["pcs"]
+        scene = show_pcs_with_trimesh([xyz[:, :3] for xyz in obj_xyz], [xyz[:, 3:] for xyz in obj_xyz], return_scene=True)
+
+        scene.apply_transform(tra.euler_matrix(np.pi, 0, np.pi/2))
+
+        scene_filename = "./tmp_data/input_scene_{}.glb".format(session_id)
+        scene.export(scene_filename)
+
+        return language_command, scene_filename
+
+    def infer(self, di, session_id, progress=gr.Progress()):
 
         # di = np.random.choice(len(self.dataset))
 
@@ -88,7 +103,7 @@ class Infer_Wrapper:
         batch = self.dataset.single_datum_to_batch(datum, args.num_samples, self.device, inference_mode=True)
 
         num_poses = datum["goal_poses"].shape[0]
-        xs = self.sampler.sample(batch, num_poses)
+        xs = self.sampler.sample(batch, num_poses, progress)
 
         struct_pose, pc_poses_in_struct = get_struct_objs_poses(xs[0])
         new_obj_xyzs = move_pc_and_create_scene_simple(batch["pcs"], struct_pose, pc_poses_in_struct)
@@ -110,7 +125,9 @@ class Infer_Wrapper:
         vis_obj_xyz = vis_obj_xyzs[0]
         scene = show_pcs_with_trimesh([xyz[:, :3] for xyz in vis_obj_xyz], [xyz[:, 3:] for xyz in vis_obj_xyz], return_scene=True)
 
-        scene_filename = "./tmp_data/scene.glb"
+        scene.apply_transform(tra.euler_matrix(np.pi, 0, np.pi/2))
+
+        scene_filename = "./tmp_data/output_scene_{}.glb".format(session_id)
         scene.export(scene_filename)
 
         # pc_filename = "/home/weiyu/Research/StructDiffusion/StructDiffusion/interactive_demo/tmp_data/pc.glb"
@@ -161,11 +178,34 @@ cfg = OmegaConf.merge(base_cfg, cfg)
 
 infer_wrapper = Infer_Wrapper(args, cfg)
 
-demo = gr.Interface(
-    fn=infer_wrapper.run,
-    inputs=gr.Slider(0, len(infer_wrapper.dataset)),
-    # clear color range [0-1.0]
-    outputs=gr.Model3D(clear_color=[0, 0, 0, 0],  label="3D Model")
-)
-
+# version 0
+# demo = gr.Interface(
+#     fn=infer_wrapper.run,
+#     inputs=gr.Slider(0, len(infer_wrapper.dataset)),
+#     # clear color range [0-1.0]
+#     outputs=gr.Model3D(clear_color=[0, 0, 0, 0],  label="3D Model")
+# )
+#
+# demo.launch()
+
+# version 1
+demo = gr.Blocks(theme=gr.themes.Soft())
+with demo:
+    gr.Markdown("<p style='text-align:center;font-size:18px'><b>StructDiffusion Demo</b></p>")
+    # font-size:18px
+    gr.Markdown("<p style='text-align:center'>StructDiffusion combines a diffusion model and an object-centric transformer to construct structures given partial-view point clouds and high-level language goals.<br><a href='https://structdiffusion.github.io/'>Website</a> | <a href='https://github.com/StructDiffusion/StructDiffusion'>Code</a></p>")
+
+    session_id = gr.State(value=np.random.randint(0, 1000))
+    data_selection = gr.Number(label="Example No.", minimum=0, maximum=len(infer_wrapper.dataset) - 1, precision=0)
+    input_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Input 3D Scene")
+    language_command = gr.Textbox(label="Input Language Command")
+    output_scene = gr.Model3D(clear_color=[0, 0, 0, 0],  label="Generated 3D Structure")
+
+    b1 = gr.Button("Show Input Language and Scene")
+    b2 = gr.Button("Generate 3D Structure")
+
+    b1.click(infer_wrapper.visualize_scene, inputs=[data_selection, session_id], outputs=[language_command, input_scene])
+    b2.click(infer_wrapper.infer, inputs=[data_selection, session_id], outputs=output_scene)
+
+demo.queue(concurrency_count=10)
 demo.launch()

src/StructDiffusion/diffusion/sampler.py

@@ -14,7 +14,7 @@ class Sampler:
         self.backbone.to(device)
         self.backbone.eval()
 
-    def sample(self, batch, num_poses):
+    def sample(self, batch, num_poses, progress):
 
         noise_schedule = self.model.noise_schedule
 
@@ -23,7 +23,7 @@ class Sampler:
         x_noisy = torch.randn((B, num_poses, 9), device=self.device)
 
         xs = []
-        for t_index in tqdm(reversed(range(0, noise_schedule.timesteps)),
+        for t_index in progress.tqdm(reversed(range(0, noise_schedule.timesteps)),
                             desc='sampling loop time step', total=noise_schedule.timesteps):
 
             t = torch.full((B,), t_index, device=self.device, dtype=torch.long)

src/StructDiffusion/utils/rearrangement.py

@@ -582,6 +582,7 @@ def show_pcs_with_trimesh(obj_xyzs, obj_rgbs, return_scene=False):
     RT_4x4 = np.linalg.inv(RT_4x4)
     RT_4x4 = RT_4x4 @ np.diag([1, -1, -1, 1])
     scene.camera_transform = RT_4x4
+
     if return_scene:
         return scene
     else: