# Copyright (c) 2020-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
import os
import torch
from . import texture
from . import mesh
from . import material
######################################################################################
# Utility functions
######################################################################################
def _find_mat(materials, name):
for mat in materials:
if mat['name'] == name:
return mat
return materials[0] # Materials 0 is the default
def normalize_mesh(vertices, scale_factor=1.0):
# 计算边界框
min_vals, _ = torch.min(vertices, dim=0)
max_vals, _ = torch.max(vertices, dim=0)
# 计算中心点
center = (max_vals + min_vals) / 2
# 平移顶点
vertices = vertices - center
# 计算缩放因子
max_extent = torch.max(max_vals - min_vals)
scale = 2.0 * scale_factor / max_extent
# 缩放顶点
vertices = vertices * scale
return vertices
######################################################################################
# Create mesh object from objfile
######################################################################################
def rotate_y_90(v_pos):
# 定义绕X轴旋转90度的旋转矩阵
rotate_y = torch.tensor([[0, 0, 1, 0],
[0, 1, 0, 0],
[-1, 0, 0, 0],
[0, 0, 0, 1]], dtype=torch.float32, device=v_pos.device)
return rotate_y
def load_obj(filename, clear_ks=True, mtl_override=None, return_attributes=False, path_is_attributrs=False, scale_factor=1.0):
obj_path = os.path.dirname(filename)
# Read entire file
with open(filename, 'r') as f:
lines = f.readlines()
# Load materials
all_materials = [
{
'name' : '_default_mat',
'bsdf' : 'pbr',
'kd' : texture.Texture2D(torch.tensor([0.5, 0.5, 0.5], dtype=torch.float32, device='cuda')),
'ks' : texture.Texture2D(torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32, device='cuda'))
}
]
if mtl_override is None:
for line in lines:
if len(line.split()) == 0:
continue
if line.split()[0] == 'mtllib':
all_materials += material.load_mtl(os.path.join(obj_path, line.split()[1]), clear_ks) # Read in entire material library
else:
all_materials += material.load_mtl(mtl_override)
# load vertices
vertices, texcoords, normals = [], [], []
for line in lines:
if len(line.split()) == 0:
continue
prefix = line.split()[0].lower()
if prefix == 'v':
vertices.append([float(v) for v in line.split()[1:]])
elif prefix == 'vt':
val = [float(v) for v in line.split()[1:]]
texcoords.append([val[0], 1.0 - val[1]])
elif prefix == 'vn':
normals.append([float(v) for v in line.split()[1:]])
# load faces
activeMatIdx = None
used_materials = []
faces, tfaces, nfaces, mfaces = [], [], [], []
for line in lines:
if len(line.split()) == 0:
continue
prefix = line.split()[0].lower()
if prefix == 'usemtl': # Track used materials
mat = _find_mat(all_materials, line.split()[1])
if not mat in used_materials:
used_materials.append(mat)
activeMatIdx = used_materials.index(mat)
elif prefix == 'f': # Parse face
vs = line.split()[1:]
nv = len(vs)
vv = vs[0].split('/')
v0 = int(vv[0]) - 1
t0 = int(vv[1]) - 1 if vv[1] != "" else -1
n0 = int(vv[2]) - 1 if vv[2] != "" else -1
for i in range(nv - 2): # Triangulate polygons
vv = vs[i + 1].split('/')
v1 = int(vv[0]) - 1
t1 = int(vv[1]) - 1 if vv[1] != "" else -1
n1 = int(vv[2]) - 1 if vv[2] != "" else -1
vv = vs[i + 2].split('/')
v2 = int(vv[0]) - 1
t2 = int(vv[1]) - 1 if vv[1] != "" else -1
n2 = int(vv[2]) - 1 if vv[2] != "" else -1
mfaces.append(activeMatIdx)
faces.append([v0, v1, v2])
tfaces.append([t0, t1, t2])
nfaces.append([n0, n1, n2])
assert len(tfaces) == len(faces) and len(nfaces) == len (faces)
# Create an "uber" material by combining all textures into a larger texture
if len(used_materials) > 1:
uber_material, texcoords, tfaces = material.merge_materials(used_materials, texcoords, tfaces, mfaces)
else:
uber_material = used_materials[0]
vertices = torch.tensor(vertices, dtype=torch.float32, device='cuda')
texcoords = torch.tensor(texcoords, dtype=torch.float32, device='cuda') if len(texcoords) > 0 else None
normals = torch.tensor(normals, dtype=torch.float32, device='cuda') if len(normals) > 0 else None
faces = torch.tensor(faces, dtype=torch.int64, device='cuda')
tfaces = torch.tensor(tfaces, dtype=torch.int64, device='cuda') if texcoords is not None else None
nfaces = torch.tensor(nfaces, dtype=torch.int64, device='cuda') if normals is not None else None
vertices = normalize_mesh(vertices, scale_factor=scale_factor)
# vertices = vertices @ rotate_y_90(vertices)[:3,:3]
if return_attributes:
return mesh.Mesh(vertices, faces, normals, nfaces, texcoords, tfaces, material=uber_material), vertices, faces, normals, nfaces, texcoords, tfaces, uber_material
return mesh.Mesh(vertices, faces, normals, nfaces, texcoords, tfaces, material=uber_material)
######################################################################################
# Save mesh object to objfile
######################################################################################
def write_obj(folder, mesh, save_material=True):
obj_file = os.path.join(folder, 'mesh.obj')
print("Writing mesh: ", obj_file)
with open(obj_file, "w") as f:
f.write("mtllib mesh.mtl\n")
f.write("g default\n")
v_pos = mesh.v_pos.detach().cpu().numpy() if mesh.v_pos is not None else None
v_nrm = mesh.v_nrm.detach().cpu().numpy() if mesh.v_nrm is not None else None
v_tex = mesh.v_tex.detach().cpu().numpy() if mesh.v_tex is not None else None
t_pos_idx = mesh.t_pos_idx.detach().cpu().numpy() if mesh.t_pos_idx is not None else None
t_nrm_idx = mesh.t_nrm_idx.detach().cpu().numpy() if mesh.t_nrm_idx is not None else None
t_tex_idx = mesh.t_tex_idx.detach().cpu().numpy() if mesh.t_tex_idx is not None else None
print(" writing %d vertices" % len(v_pos))
for v in v_pos:
f.write('v {} {} {} \n'.format(v[0], v[1], v[2]))
if v_tex is not None:
print(" writing %d texcoords" % len(v_tex))
assert(len(t_pos_idx) == len(t_tex_idx))
for v in v_tex:
f.write('vt {} {} \n'.format(v[0], 1.0 - v[1]))
if v_nrm is not None:
print(" writing %d normals" % len(v_nrm))
assert(len(t_pos_idx) == len(t_nrm_idx))
for v in v_nrm:
f.write('vn {} {} {}\n'.format(v[0], v[1], v[2]))
# faces
f.write("s 1 \n")
f.write("g pMesh1\n")
f.write("usemtl defaultMat\n")
# Write faces
print(" writing %d faces" % len(t_pos_idx))
for i in range(len(t_pos_idx)):
f.write("f ")
for j in range(3):
f.write(' %s/%s/%s' % (str(t_pos_idx[i][j]+1), '' if v_tex is None else str(t_tex_idx[i][j]+1), '' if v_nrm is None else str(t_nrm_idx[i][j]+1)))
f.write("\n")
if save_material:
mtl_file = os.path.join(folder, 'mesh.mtl')
print("Writing material: ", mtl_file)
material.save_mtl(mtl_file, mesh.material)
print("Done exporting mesh")