# 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")