Spaces:
Sleeping
Sleeping
File size: 1,941 Bytes
b2a27a7 b3806ac b2a27a7 |
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 |
import argparse
import numpy as np
import torch
import trimesh
from triposg.inference_utils import hierarchical_extract_geometry
from triposg.models.autoencoders import TripoSGVAEModel
from huggingface_hub import snapshot_download
def load_surface(data_path, num_pc=204800):
data = np.load(data_path, allow_pickle=True).tolist()
surface = data["surface_points"] # Nx3
normal = data["surface_normals"] # Nx3
rng = np.random.default_rng()
ind = rng.choice(surface.shape[0], num_pc, replace=False)
surface = torch.FloatTensor(surface[ind])
normal = torch.FloatTensor(normal[ind])
surface = torch.cat([surface, normal], dim=-1).unsqueeze(0).cuda()
return surface
if __name__ == "__main__":
device = "cuda"
dtype = torch.float16
parser = argparse.ArgumentParser()
parser.add_argument("--surface-input", type=str, required=True)
args = parser.parse_args()
# download pretrained weights
triposg_weights_dir = "pretrained_weights/TripoSG"
snapshot_download(repo_id="VAST-AI/TripoSG", local_dir=triposg_weights_dir)
vae: TripoSGVAEModel = TripoSGVAEModel.from_pretrained(
triposg_weights_dir,
subfolder="vae",
).to(device, dtype=dtype)
# load surface from sdf and encode
surface = load_surface(
args.surface_input, num_pc=204800
).to(device, dtype=dtype)
sample = vae.encode(surface).latent_dist.sample()
# vae infer
with torch.no_grad():
geometric_func = lambda x: vae.decode(sample, sampled_points=x).sample
output = hierarchical_extract_geometry(
geometric_func,
device,
bounds=(-1.005, -1.005, -1.005, 1.005, 1.005, 1.005),
dense_octree_depth=9,
hierarchical_octree_depth=9,
)
meshes = [trimesh.Trimesh(mesh_v_f[0].astype(np.float32), mesh_v_f[1]) for mesh_v_f in output]
meshes[0].export("test_vae.glb")
|