import torch
import numpy as np
import pymeshlab as ml
from pytorch3d.renderer import TexturesVertex
from pytorch3d.structures import Meshes
import torch
import torch.nn.functional as F
import trimesh
from pymeshlab import PercentageValue
import open3d as o3d
def tensor2variable(tensor, device):
# [1,23,3,3]
return torch.tensor(tensor, device=device, requires_grad=True)
def rot6d_to_rotmat(x):
"""Convert 6D rotation representation to 3x3 rotation matrix.
Based on Zhou et al., "On the Continuity of Rotation Representations in Neural Networks", CVPR 2019
Input:
(B,6) Batch of 6-D rotation representations
Output:
(B,3,3) Batch of corresponding rotation matrices
"""
x = x.view(-1, 3, 2)
a1 = x[:, :, 0]
a2 = x[:, :, 1]
b1 = F.normalize(a1)
b2 = F.normalize(a2 - torch.einsum("bi,bi->b", b1, a2).unsqueeze(-1) * b1)
b3 = torch.cross(b1, b2)
return torch.stack((b1, b2, b3), dim=-1)
def fix_vert_color_glb(mesh_path):
from pygltflib import GLTF2, Material, PbrMetallicRoughness
obj1 = GLTF2().load(mesh_path)
obj1.meshes[0].primitives[0].material = 0
obj1.materials.append(Material(
pbrMetallicRoughness = PbrMetallicRoughness(
baseColorFactor = [1.0, 1.0, 1.0, 1.0],
metallicFactor = 0.,
roughnessFactor = 1.0,
),
emissiveFactor = [0.0, 0.0, 0.0],
doubleSided = True,
))
obj1.save(mesh_path)
def srgb_to_linear(c_srgb):
c_linear = np.where(c_srgb <= 0.04045, c_srgb / 12.92, ((c_srgb + 0.055) / 1.055) ** 2.4)
return c_linear.clip(0, 1.)
def save_py3dmesh_with_trimesh_fast(meshes: Meshes, save_glb_path, apply_sRGB_to_LinearRGB=True):
# convert from pytorch3d meshes to trimesh mesh
vertices = meshes.verts_packed().cpu().float().numpy()
triangles = meshes.faces_packed().cpu().long().numpy()
np_color = meshes.textures.verts_features_packed().cpu().float().numpy()
if save_glb_path.endswith(".glb"):
# rotate 180 along +Y
vertices[:, [0, 2]] = -vertices[:, [0, 2]]
if apply_sRGB_to_LinearRGB:
np_color = srgb_to_linear(np_color)
assert vertices.shape[0] == np_color.shape[0]
assert np_color.shape[1] == 3
assert 0 <= np_color.min() and np_color.max() <= 1, f"min={np_color.min()}, max={np_color.max()}"
mesh = trimesh.Trimesh(vertices=vertices, faces=triangles, vertex_colors=np_color)
mesh.remove_unreferenced_vertices()
# save mesh
mesh.export(save_glb_path)
if save_glb_path.endswith(".glb"):
fix_vert_color_glb(save_glb_path)
print(f"saving to {save_glb_path}")
def load_mesh_with_trimesh(file_name, file_type=None):
mesh: trimesh.Trimesh = trimesh.load(file_name, file_type=file_type)
if isinstance(mesh, trimesh.Scene):
assert len(mesh.geometry) > 0
# save to obj first and load again to avoid offset issue
from io import BytesIO
with BytesIO() as f:
mesh.export(f, file_type="obj")
f.seek(0)
mesh = trimesh.load(f, file_type="obj")
if isinstance(mesh, trimesh.Scene):
# we lose texture information here
mesh = trimesh.util.concatenate(
tuple(trimesh.Trimesh(vertices=g.vertices, faces=g.faces)
for g in mesh.geometry.values()))
assert isinstance(mesh, trimesh.Trimesh)
vertices = torch.from_numpy(mesh.vertices).T
faces = torch.from_numpy(mesh.faces).T
colors = None
if mesh.visual is not None:
if hasattr(mesh.visual, 'vertex_colors'):
colors = torch.from_numpy(mesh.visual.vertex_colors)[..., :3].T / 255.
if colors is None:
# print("Warning: no vertex color found in mesh! Filling it with gray.")
colors = torch.ones_like(vertices) * 0.5
return vertices, faces, colors
def meshlab_mesh_to_py3dmesh(mesh: ml.Mesh) -> Meshes:
verts = torch.from_numpy(mesh.vertex_matrix()).float()
faces = torch.from_numpy(mesh.face_matrix()).long()
colors = torch.from_numpy(mesh.vertex_color_matrix()[..., :3]).float()
textures = TexturesVertex(verts_features=[colors])
return Meshes(verts=[verts], faces=[faces], textures=textures)
def py3dmesh_to_meshlab_mesh(meshes: Meshes) -> ml.Mesh:
colors_in = F.pad(meshes.textures.verts_features_packed().cpu().float(), [0,1], value=1).numpy().astype(np.float64)
m1 = ml.Mesh(
vertex_matrix=meshes.verts_packed().cpu().float().numpy().astype(np.float64),
face_matrix=meshes.faces_packed().cpu().long().numpy().astype(np.int32),
v_normals_matrix=meshes.verts_normals_packed().cpu().float().numpy().astype(np.float64),
v_color_matrix=colors_in)
return m1
def to_pyml_mesh(vertices,faces):
m1 = ml.Mesh(
vertex_matrix=vertices.cpu().float().numpy().astype(np.float64),
face_matrix=faces.cpu().long().numpy().astype(np.int32),
)
return m1
def to_py3d_mesh(vertices, faces, normals=None):
from pytorch3d.structures import Meshes
from pytorch3d.renderer.mesh.textures import TexturesVertex
mesh = Meshes(verts=[vertices], faces=[faces], textures=None)
if normals is None:
normals = mesh.verts_normals_packed()
# set normals as vertext colors
mesh.textures = TexturesVertex(verts_features=[normals / 2 + 0.5])
return mesh
def from_py3d_mesh(mesh):
return mesh.verts_list()[0], mesh.faces_list()[0], mesh.textures.verts_features_packed()
def simple_clean_mesh(pyml_mesh: ml.Mesh, apply_smooth=True, stepsmoothnum=1, apply_sub_divide=False, sub_divide_threshold=0.25):
ms = ml.MeshSet()
ms.add_mesh(pyml_mesh, "cube_mesh")
if apply_smooth:
ms.apply_filter("apply_coord_laplacian_smoothing", stepsmoothnum=stepsmoothnum, cotangentweight=False)
if apply_sub_divide: # 5s, slow
ms.apply_filter("meshing_repair_non_manifold_vertices")
ms.apply_filter("meshing_repair_non_manifold_edges", method='Remove Faces')
ms.apply_filter("meshing_surface_subdivision_loop", iterations=2, threshold=PercentageValue(sub_divide_threshold))
return meshlab_mesh_to_py3dmesh(ms.current_mesh())
def post_process_mesh(mesh, cluster_to_keep=1000):
"""
Post-process a mesh to filter out floaters and disconnected parts
"""
import copy
print("post processing the mesh to have {} clusterscluster_to_kep".format(cluster_to_keep))
mesh_0 = copy.deepcopy(mesh)
with o3d.utility.VerbosityContextManager(o3d.utility.VerbosityLevel.Debug) as cm:
triangle_clusters, cluster_n_triangles, cluster_area = (mesh_0.cluster_connected_triangles())
triangle_clusters = np.asarray(triangle_clusters)
cluster_n_triangles = np.asarray(cluster_n_triangles)
cluster_area = np.asarray(cluster_area)
n_cluster = np.sort(cluster_n_triangles.copy())[-cluster_to_keep]
n_cluster = max(n_cluster, 50) # filter meshes smaller than 50
triangles_to_remove = cluster_n_triangles[triangle_clusters] < n_cluster
mesh_0.remove_triangles_by_mask(triangles_to_remove)
mesh_0.remove_unreferenced_vertices()
mesh_0.remove_degenerate_triangles()
print("num vertices raw {}".format(len(mesh.vertices)))
print("num vertices post {}".format(len(mesh_0.vertices)))
return mesh_0