Spaces:
Paused
Paused
| 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) |