LTT/PRM
24
1# Copyright (c) 2020-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2#3# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual4# property and proprietary rights in and to this material, related5# documentation and any modifications thereto. Any use, reproduction, 6# disclosure or distribution of this material and related documentation 7# without an express license agreement from NVIDIA CORPORATION or 8# its affiliates is strictly prohibited.9 10import os11import torch12 13from . import texture14from . import mesh15from . import material16 17######################################################################################18# Utility functions19######################################################################################20 21def _find_mat(materials, name):22 for mat in materials:23 if mat['name'] == name:24 return mat25 return materials[0] # Materials 0 is the default26 27 28def normalize_mesh(vertices, scale_factor=1.0):29 # 计算边界框30 min_vals, _ = torch.min(vertices, dim=0)31 max_vals, _ = torch.max(vertices, dim=0)32 33 # 计算中心点34 center = (max_vals + min_vals) / 235 36 # 平移顶点37 vertices = vertices - center38 39 # 计算缩放因子40 max_extent = torch.max(max_vals - min_vals)41 scale = 2.0 * scale_factor / max_extent42 43 # 缩放顶点44 vertices = vertices * scale45 46 return vertices47 48######################################################################################49# Create mesh object from objfile50######################################################################################51def rotate_y_90(v_pos):52 # 定义绕X轴旋转90度的旋转矩阵53 rotate_y = torch.tensor([[0, 0, 1, 0], 54 [0, 1, 0, 0], 55 [-1, 0, 0, 0], 56 [0, 0, 0, 1]], dtype=torch.float32, device=v_pos.device)57 return rotate_y58 59def load_obj(filename, clear_ks=True, mtl_override=None, return_attributes=False, path_is_attributrs=False, scale_factor=1.0):60 obj_path = os.path.dirname(filename)61 62 # Read entire file63 with open(filename, 'r') as f:64 lines = f.readlines()65 66 # Load materials67 all_materials = [68 {69 'name' : '_default_mat',70 'bsdf' : 'pbr',71 'kd' : texture.Texture2D(torch.tensor([0.5, 0.5, 0.5], dtype=torch.float32, device='cuda')),72 'ks' : texture.Texture2D(torch.tensor([0.0, 0.0, 0.0], dtype=torch.float32, device='cuda'))73 }74 ]75 if mtl_override is None: 76 for line in lines:77 if len(line.split()) == 0:78 continue79 if line.split()[0] == 'mtllib':80 all_materials += material.load_mtl(os.path.join(obj_path, line.split()[1]), clear_ks) # Read in entire material library81 else:82 all_materials += material.load_mtl(mtl_override)83 84 # load vertices85 vertices, texcoords, normals = [], [], []86 for line in lines:87 if len(line.split()) == 0:88 continue89 90 prefix = line.split()[0].lower()91 if prefix == 'v':92 vertices.append([float(v) for v in line.split()[1:]])93 elif prefix == 'vt':94 val = [float(v) for v in line.split()[1:]]95 texcoords.append([val[0], 1.0 - val[1]])96 elif prefix == 'vn':97 normals.append([float(v) for v in line.split()[1:]])98 99 # load faces100 activeMatIdx = None101 used_materials = []102 faces, tfaces, nfaces, mfaces = [], [], [], []103 for line in lines:104 if len(line.split()) == 0:105 continue106 107 prefix = line.split()[0].lower()108 if prefix == 'usemtl': # Track used materials109 mat = _find_mat(all_materials, line.split()[1])110 if not mat in used_materials:111 used_materials.append(mat)112 activeMatIdx = used_materials.index(mat)113 elif prefix == 'f': # Parse face114 vs = line.split()[1:]115 nv = len(vs)116 vv = vs[0].split('/')117 v0 = int(vv[0]) - 1118 t0 = int(vv[1]) - 1 if vv[1] != "" else -1119 n0 = int(vv[2]) - 1 if vv[2] != "" else -1120 for i in range(nv - 2): # Triangulate polygons121 vv = vs[i + 1].split('/')122 v1 = int(vv[0]) - 1123 t1 = int(vv[1]) - 1 if vv[1] != "" else -1124 n1 = int(vv[2]) - 1 if vv[2] != "" else -1125 vv = vs[i + 2].split('/')126 v2 = int(vv[0]) - 1127 t2 = int(vv[1]) - 1 if vv[1] != "" else -1128 n2 = int(vv[2]) - 1 if vv[2] != "" else -1129 mfaces.append(activeMatIdx)130 faces.append([v0, v1, v2])131 tfaces.append([t0, t1, t2])132 nfaces.append([n0, n1, n2])133 assert len(tfaces) == len(faces) and len(nfaces) == len (faces)134 135 # Create an "uber" material by combining all textures into a larger texture136 if len(used_materials) > 1:137 uber_material, texcoords, tfaces = material.merge_materials(used_materials, texcoords, tfaces, mfaces)138 else:139 uber_material = used_materials[0]140 141 vertices = torch.tensor(vertices, dtype=torch.float32, device='cuda')142 texcoords = torch.tensor(texcoords, dtype=torch.float32, device='cuda') if len(texcoords) > 0 else None143 normals = torch.tensor(normals, dtype=torch.float32, device='cuda') if len(normals) > 0 else None144 145 faces = torch.tensor(faces, dtype=torch.int64, device='cuda')146 tfaces = torch.tensor(tfaces, dtype=torch.int64, device='cuda') if texcoords is not None else None147 nfaces = torch.tensor(nfaces, dtype=torch.int64, device='cuda') if normals is not None else None148 149 vertices = normalize_mesh(vertices, scale_factor=scale_factor)150 # vertices = vertices @ rotate_y_90(vertices)[:3,:3]151 152 if return_attributes:153 return mesh.Mesh(vertices, faces, normals, nfaces, texcoords, tfaces, material=uber_material), vertices, faces, normals, nfaces, texcoords, tfaces, uber_material154 return mesh.Mesh(vertices, faces, normals, nfaces, texcoords, tfaces, material=uber_material)155 156######################################################################################157# Save mesh object to objfile158######################################################################################159 160def write_obj(folder, mesh, save_material=True):161 obj_file = os.path.join(folder, 'mesh.obj')162 print("Writing mesh: ", obj_file)163 with open(obj_file, "w") as f:164 f.write("mtllib mesh.mtl\n")165 f.write("g default\n")166 167 v_pos = mesh.v_pos.detach().cpu().numpy() if mesh.v_pos is not None else None168 v_nrm = mesh.v_nrm.detach().cpu().numpy() if mesh.v_nrm is not None else None169 v_tex = mesh.v_tex.detach().cpu().numpy() if mesh.v_tex is not None else None170 171 t_pos_idx = mesh.t_pos_idx.detach().cpu().numpy() if mesh.t_pos_idx is not None else None172 t_nrm_idx = mesh.t_nrm_idx.detach().cpu().numpy() if mesh.t_nrm_idx is not None else None173 t_tex_idx = mesh.t_tex_idx.detach().cpu().numpy() if mesh.t_tex_idx is not None else None174 175 print(" writing %d vertices" % len(v_pos))176 for v in v_pos:177 f.write('v {} {} {} \n'.format(v[0], v[1], v[2]))178 179 if v_tex is not None:180 print(" writing %d texcoords" % len(v_tex))181 assert(len(t_pos_idx) == len(t_tex_idx))182 for v in v_tex:183 f.write('vt {} {} \n'.format(v[0], 1.0 - v[1]))184 185 if v_nrm is not None:186 print(" writing %d normals" % len(v_nrm))187 assert(len(t_pos_idx) == len(t_nrm_idx))188 for v in v_nrm:189 f.write('vn {} {} {}\n'.format(v[0], v[1], v[2]))190 191 # faces192 f.write("s 1 \n")193 f.write("g pMesh1\n")194 f.write("usemtl defaultMat\n")195 196 # Write faces197 print(" writing %d faces" % len(t_pos_idx))198 for i in range(len(t_pos_idx)):199 f.write("f ")200 for j in range(3):201 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)))202 f.write("\n")203 204 if save_material:205 mtl_file = os.path.join(folder, 'mesh.mtl')206 print("Writing material: ", mtl_file)207 material.save_mtl(mtl_file, mesh.material)208 209 print("Done exporting mesh")210