87 lines
3.3 KiB
Python
87 lines
3.3 KiB
Python
import torch
|
|
import trimesh
|
|
import o_voxel
|
|
|
|
from PIL import Image
|
|
|
|
def make_texture_square_pow2(img: Image.Image, target_size=None):
|
|
w, h = img.size
|
|
max_side = max(w, h)
|
|
pow2 = 1
|
|
while pow2 < max_side:
|
|
pow2 *= 2
|
|
if target_size is not None:
|
|
pow2 = target_size
|
|
pow2 = min(pow2, 2048)
|
|
return img.resize((pow2, pow2), Image.BILINEAR)
|
|
|
|
def preprocess_scene_textures(asset):
|
|
if not isinstance(asset, trimesh.Scene):
|
|
return asset
|
|
TEX_KEYS = ["baseColorTexture", "normalTexture", "metallicRoughnessTexture", "emissiveTexture", "occlusionTexture"]
|
|
for geom in asset.geometry.values():
|
|
visual = getattr(geom, "visual", None)
|
|
mat = getattr(visual, "material", None)
|
|
if mat is None:
|
|
continue
|
|
for key in TEX_KEYS:
|
|
if not hasattr(mat, key):
|
|
continue
|
|
tex = getattr(mat, key)
|
|
if tex is None:
|
|
continue
|
|
if isinstance(tex, Image.Image):
|
|
setattr(mat, key, make_texture_square_pow2(tex))
|
|
elif hasattr(tex, "image") and tex.image is not None:
|
|
img = tex.image
|
|
if not isinstance(img, Image.Image):
|
|
img = Image.fromarray(img)
|
|
tex.image = make_texture_square_pow2(img)
|
|
if hasattr(mat, "image") and mat.image is not None:
|
|
img = mat.image
|
|
if not isinstance(img, Image.Image):
|
|
img = Image.fromarray(img)
|
|
mat.image = make_texture_square_pow2(img)
|
|
return asset
|
|
|
|
def glb_to_vxz(glb_path, vxz_path):
|
|
asset = trimesh.load(glb_path, force='scene')
|
|
asset = preprocess_scene_textures(asset)
|
|
aabb = asset.bounding_box.bounds
|
|
center = (aabb[0] + aabb[1]) / 2
|
|
scale = 0.99999 / (aabb[1] - aabb[0]).max()
|
|
asset.apply_translation(-center)
|
|
asset.apply_scale(scale)
|
|
mesh = asset.to_mesh()
|
|
vertices = torch.from_numpy(mesh.vertices).float()
|
|
faces = torch.from_numpy(mesh.faces).long()
|
|
|
|
voxel_indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid(
|
|
vertices, faces, grid_size=512, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
|
|
face_weight=1.0, boundary_weight=0.2, regularization_weight=1e-2, timing=False
|
|
)
|
|
vid = o_voxel.serialize.encode_seq(voxel_indices)
|
|
mapping = torch.argsort(vid)
|
|
voxel_indices = voxel_indices[mapping]
|
|
dual_vertices = dual_vertices[mapping]
|
|
intersected = intersected[mapping]
|
|
|
|
voxel_indices_mat, attributes = o_voxel.convert.textured_mesh_to_volumetric_attr(
|
|
asset, grid_size=512, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]], timing=False
|
|
)
|
|
vid_mat = o_voxel.serialize.encode_seq(voxel_indices_mat)
|
|
mapping_mat = torch.argsort(vid_mat)
|
|
attributes = {k: v[mapping_mat] for k, v in attributes.items()}
|
|
|
|
dual_vertices = dual_vertices * 512 - voxel_indices
|
|
dual_vertices = (torch.clamp(dual_vertices, 0, 1) * 255).type(torch.uint8)
|
|
intersected = (intersected[:, 0:1] + 2 * intersected[:, 1:2] + 4 * intersected[:, 2:3]).type(torch.uint8)
|
|
|
|
attributes['dual_vertices'] = dual_vertices
|
|
attributes['intersected'] = intersected
|
|
o_voxel.io.write(vxz_path, voxel_indices, attributes)
|
|
|
|
if __name__ == "__main__":
|
|
glb_path = "./assets/example.glb"
|
|
vxz_path = "./assets/input.vxz"
|
|
glb_to_vxz(glb_path, vxz_path) |