Spaces:
Paused
Paused
File size: 5,564 Bytes
8c02843 |
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 |
import os
import argparse
import torch
import trimesh
import numpy as np
import pytorch_lightning as pl
import gradio as gr
from omegaconf import OmegaConf
import sys
sys.path.append('./src')
from StructDiffusion.data.semantic_arrangement_demo import SemanticArrangementDataset
from StructDiffusion.language.tokenizer import Tokenizer
from StructDiffusion.models.pl_models import ConditionalPoseDiffusionModel
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.batch_inference import move_pc_and_create_scene_simple, visualize_batch_pcs
from StructDiffusion.utils.rearrangement import show_pcs_with_trimesh
class Infer_Wrapper:
def __init__(self, args, cfg):
# load
pl.seed_everything(args.eval_random_seed)
self.device = (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
checkpoint_dir = os.path.join(cfg.WANDB.save_dir, cfg.WANDB.project, args.checkpoint_id, "checkpoints")
checkpoint_path = get_checkpoint_path_from_dir(checkpoint_dir)
self.tokenizer = Tokenizer(cfg.DATASET.vocab_dir)
# override ignore_rgb for visualization
cfg.DATASET.ignore_rgb = False
self.dataset = SemanticArrangementDataset(tokenizer=self.tokenizer, **cfg.DATASET)
self.sampler = Sampler(ConditionalPoseDiffusionModel, checkpoint_path, self.device)
def run(self, di):
# di = np.random.choice(len(self.dataset))
raw_datum = self.dataset.get_raw_data(di)
print(self.tokenizer.convert_structure_params_to_natural_language(raw_datum["sentence"]))
datum = self.dataset.convert_to_tensors(raw_datum, self.tokenizer)
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)
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)
# vis
vis_obj_xyzs = new_obj_xyzs[:3]
if torch.is_tensor(vis_obj_xyzs):
if vis_obj_xyzs.is_cuda:
vis_obj_xyzs = vis_obj_xyzs.detach().cpu()
vis_obj_xyzs = vis_obj_xyzs.numpy()
# for bi, vis_obj_xyz in enumerate(vis_obj_xyzs):
# if verbose:
# print("example {}".format(bi))
# print(vis_obj_xyz.shape)
#
# if trimesh:
# show_pcs_with_trimesh([xyz[:, :3] for xyz in vis_obj_xyz], [xyz[:, 3:] for xyz in vis_obj_xyz])
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.export(scene_filename)
# pc_filename = "/home/weiyu/Research/StructDiffusion/StructDiffusion/interactive_demo/tmp_data/pc.glb"
# scene_filename = "/home/weiyu/Research/StructDiffusion/StructDiffusion/interactive_demo/tmp_data/scene.glb"
#
# vis_obj_xyz = vis_obj_xyz.reshape(-1, 6)
# vis_pc = trimesh.PointCloud(vis_obj_xyz[:, :3], colors=np.concatenate([vis_obj_xyz[:, 3:] * 255, np.ones([vis_obj_xyz.shape[0], 1]) * 255], axis=-1))
# vis_pc.export(pc_filename)
#
# scene = trimesh.Scene()
# # add the coordinate frame first
# # geom = trimesh.creation.axis(0.01)
# # scene.add_geometry(geom)
# table = trimesh.creation.box(extents=[1.0, 1.0, 0.02])
# table.apply_translation([0.5, 0, -0.01])
# table.visual.vertex_colors = [150, 111, 87, 125]
# scene.add_geometry(table)
# # bounds = trimesh.creation.box(extents=[4.0, 4.0, 4.0])
# # bounds = trimesh.creation.icosphere(subdivisions=3, radius=3.1)
# # bounds.apply_translation([0, 0, 0])
# # bounds.visual.vertex_colors = [30, 30, 30, 30]
# # scene.add_geometry(bounds)
# # RT_4x4 = np.array([[-0.39560353822208355, -0.9183993826406329, 0.006357240869497738, 0.2651463080169481],
# # [-0.797630370081598, 0.3401340617616391, -0.4980909683511864, 0.2225696480721997],
# # [0.45528412367406523, -0.2021172778236285, -0.8671014777611122, 0.9449050652025951],
# # [0.0, 0.0, 0.0, 1.0]])
# # RT_4x4 = np.linalg.inv(RT_4x4)
# # RT_4x4 = RT_4x4 @ np.diag([1, -1, -1, 1])
# # scene.camera_transform = RT_4x4
#
# mesh_list = trimesh.util.concatenate(scene.dump())
# print(mesh_list)
# trimesh.io.export.export_mesh(mesh_list, scene_filename, file_type='obj')
return scene_filename
args = OmegaConf.create()
args.base_config_file = "./configs/base.yaml"
args.config_file = "./configs/conditional_pose_diffusion.yaml"
args.checkpoint_id = "ConditionalPoseDiffusion"
args.eval_random_seed = 42
args.num_samples = 1
base_cfg = OmegaConf.load(args.base_config_file)
cfg = OmegaConf.load(args.config_file)
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")
)
demo.launch() |