Spaces:
Running on Zero
Running on Zero
File size: 2,822 Bytes
43ca9e0 | 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 | import os
import torch
import numpy as np
import imageio
import trimesh
import third_party.TRELLIS.trellis.modules.sparse as sp
from third_party.TRELLIS.trellis.utils import render_utils, postprocessing_utils
def _white_geometry_mesh(vertices, faces):
mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
mesh.visual = trimesh.visual.TextureVisuals(
material=trimesh.visual.material.PBRMaterial(
name='geometry_white',
baseColorFactor=[1.0, 1.0, 1.0, 1.0],
metallicFactor=0.0,
roughnessFactor=0.85,
)
)
mesh.vertex_normals
mesh.face_normals
return mesh
def _postprocess_mesh_geometry(mesh, simplify=0.95):
vertices = mesh.vertices.cpu().numpy()
faces = mesh.faces.cpu().numpy()
vertices, faces = postprocessing_utils.postprocess_mesh(
vertices,
faces,
simplify=simplify > 0,
simplify_ratio=simplify,
fill_holes=True,
fill_holes_max_hole_size=0.04,
fill_holes_max_hole_nbe=int(250 * np.sqrt(1 - simplify)),
fill_holes_resolution=1024,
fill_holes_num_views=1000,
debug=False,
verbose=False,
)
return vertices, faces
def decode_slat(generation_pipeline, feats, coords, out_meshpath, out_gspath, texture=True):
# Decode Output SLAT
slat = sp.SparseTensor(
feats = feats.float(),
coords = coords.int(),
).cuda()
formats = ['mesh', 'gaussian'] if texture else ['mesh']
with torch.no_grad():
outputs = generation_pipeline.decode_slat(slat, formats)
out_geometry_path = f"{os.path.splitext(out_meshpath)[0]}_geometry.glb"
if texture:
mesh_textured = postprocessing_utils.to_glb(
outputs['gaussian'][0],
outputs['mesh'][0],
# Optional parameters
simplify=0.95, # Ratio of triangles to remove in the simplification process
texture_size=1024, # Size of the texture used for the GLB
verbose=False, # Print logs
)
mesh_textured.export(out_meshpath)
mesh_geometry = _white_geometry_mesh(
mesh_textured.vertices.copy(),
mesh_textured.faces.copy(),
)
mesh_geometry.export(out_geometry_path)
if out_gspath:
video = render_utils.render_video(outputs['gaussian'][0])['color']
imageio.mimsave(out_gspath, video, fps=30)
return
vertices, faces = _postprocess_mesh_geometry(outputs['mesh'][0], simplify=0.95)
mesh_geometry = _white_geometry_mesh(vertices, faces)
mesh_geometry.export(out_meshpath)
mesh_geometry.export(out_geometry_path)
|