LTT/PRM
24
1# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.2#3# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property4# and proprietary rights in and to this software, related documentation5# and any modifications thereto. Any use, reproduction, disclosure or6# distribution of this software and related documentation without an express7# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.8 9import torch10import xatlas11import trimesh12import cv213import numpy as np14import nvdiffrast.torch as dr15from PIL import Image16 17 18def save_obj(pointnp_px3, facenp_fx3, colornp_px3, fpath):19 20 pointnp_px3 = pointnp_px3 @ np.array([[1, 0, 0], [0, 1, 0], [0, 0, -1]])21 facenp_fx3 = facenp_fx3[:, [2, 1, 0]]22 23 mesh = trimesh.Trimesh(24 vertices=pointnp_px3, 25 faces=facenp_fx3, 26 vertex_colors=colornp_px3,27 )28 mesh.export(fpath, 'obj')29 30 31def save_glb(pointnp_px3, facenp_fx3, colornp_px3, fpath):32 33 pointnp_px3 = pointnp_px3 @ np.array([[-1, 0, 0], [0, 1, 0], [0, 0, -1]])34 35 mesh = trimesh.Trimesh(36 vertices=pointnp_px3, 37 faces=facenp_fx3, 38 vertex_colors=colornp_px3,39 )40 mesh.export(fpath, 'glb')41 42def save_ply(pointnp_px3, facenp_fx3, colornp_px3, fpath):43 pointnp_px3 = pointnp_px3 @ np.array([[1, 0, 0], [0, 1, 0], [0, 0, -1]])44 facenp_fx3 = facenp_fx3[:, [2, 1, 0]]45 46 mesh = trimesh.Trimesh(47 vertices=pointnp_px3, 48 faces=facenp_fx349 )50 mesh.export(fpath, 'ply')51 52 53def save_obj_with_mtl(pointnp_px3, tcoords_px2, facenp_fx3, facetex_fx3, texmap_hxwx3, fname):54 import os55 fol, na = os.path.split(fname)56 na, _ = os.path.splitext(na)57 58 matname = '%s/%s.mtl' % (fol, na)59 fid = open(matname, 'w')60 fid.write('newmtl material_0\n')61 fid.write('Kd 1 1 1\n')62 fid.write('Ka 0 0 0\n')63 fid.write('Ks 0.4 0.4 0.4\n')64 fid.write('Ns 10\n')65 fid.write('illum 2\n')66 fid.write('map_Kd %s.png\n' % na)67 fid.close()68 ####69 70 fid = open(fname, 'w')71 fid.write('mtllib %s.mtl\n' % na)72 73 for pidx, p in enumerate(pointnp_px3):74 pp = p75 fid.write('v %f %f %f\n' % (pp[0], pp[1], pp[2]))76 77 for pidx, p in enumerate(tcoords_px2):78 pp = p79 fid.write('vt %f %f\n' % (pp[0], pp[1]))80 81 fid.write('usemtl material_0\n')82 for i, f in enumerate(facenp_fx3):83 f1 = f + 184 f2 = facetex_fx3[i] + 185 fid.write('f %d/%d %d/%d %d/%d\n' % (f1[0], f2[0], f1[1], f2[1], f1[2], f2[2]))86 fid.close()87 88 # save texture map89 lo, hi = 0, 190 img = np.asarray(texmap_hxwx3, dtype=np.float32)91 img = (img - lo) * (255 / (hi - lo))92 img = img.clip(0, 255)93 mask = np.sum(img.astype(np.float32), axis=-1, keepdims=True)94 mask = (mask <= 3.0).astype(np.float32)95 kernel = np.ones((3, 3), 'uint8')96 dilate_img = cv2.dilate(img, kernel, iterations=1)97 img = img * (1 - mask) + dilate_img * mask98 img = img.clip(0, 255).astype(np.uint8)99 Image.fromarray(np.ascontiguousarray(img[::-1, :, :]), 'RGB').save(f'{fol}/{na}.png')100 101 102def loadobj(meshfile):103 v = []104 f = []105 meshfp = open(meshfile, 'r')106 for line in meshfp.readlines():107 data = line.strip().split(' ')108 data = [da for da in data if len(da) > 0]109 if len(data) != 4:110 continue111 if data[0] == 'v':112 v.append([float(d) for d in data[1:]])113 if data[0] == 'f':114 data = [da.split('/')[0] for da in data]115 f.append([int(d) for d in data[1:]])116 meshfp.close()117 118 # torch need int64119 facenp_fx3 = np.array(f, dtype=np.int64) - 1120 pointnp_px3 = np.array(v, dtype=np.float32)121 return pointnp_px3, facenp_fx3122 123 124def loadobjtex(meshfile):125 v = []126 vt = []127 f = []128 ft = []129 meshfp = open(meshfile, 'r')130 for line in meshfp.readlines():131 data = line.strip().split(' ')132 data = [da for da in data if len(da) > 0]133 if not ((len(data) == 3) or (len(data) == 4) or (len(data) == 5)):134 continue135 if data[0] == 'v':136 assert len(data) == 4137 138 v.append([float(d) for d in data[1:]])139 if data[0] == 'vt':140 if len(data) == 3 or len(data) == 4:141 vt.append([float(d) for d in data[1:3]])142 if data[0] == 'f':143 data = [da.split('/') for da in data]144 if len(data) == 4:145 f.append([int(d[0]) for d in data[1:]])146 ft.append([int(d[1]) for d in data[1:]])147 elif len(data) == 5:148 idx1 = [1, 2, 3]149 data1 = [data[i] for i in idx1]150 f.append([int(d[0]) for d in data1])151 ft.append([int(d[1]) for d in data1])152 idx2 = [1, 3, 4]153 data2 = [data[i] for i in idx2]154 f.append([int(d[0]) for d in data2])155 ft.append([int(d[1]) for d in data2])156 meshfp.close()157 158 # torch need int64159 facenp_fx3 = np.array(f, dtype=np.int64) - 1160 ftnp_fx3 = np.array(ft, dtype=np.int64) - 1161 pointnp_px3 = np.array(v, dtype=np.float32)162 uvs = np.array(vt, dtype=np.float32)163 return pointnp_px3, facenp_fx3, uvs, ftnp_fx3164 165 166# ==============================================================================================167def interpolate(attr, rast, attr_idx, rast_db=None):168 return dr.interpolate(attr.contiguous(), rast, attr_idx, rast_db=rast_db, diff_attrs=None if rast_db is None else 'all')169 170 171def xatlas_uvmap(ctx, mesh_v, mesh_pos_idx, resolution):172 vmapping, indices, uvs = xatlas.parametrize(mesh_v.detach().cpu().numpy(), mesh_pos_idx.detach().cpu().numpy())173 174 # Convert to tensors175 indices_int64 = indices.astype(np.uint64, casting='same_kind').view(np.int64)176 177 uvs = torch.tensor(uvs, dtype=torch.float32, device=mesh_v.device)178 mesh_tex_idx = torch.tensor(indices_int64, dtype=torch.int64, device=mesh_v.device)179 # mesh_v_tex. ture180 uv_clip = uvs[None, ...] * 2.0 - 1.0181 182 # pad to four component coordinate183 uv_clip4 = torch.cat((uv_clip, torch.zeros_like(uv_clip[..., 0:1]), torch.ones_like(uv_clip[..., 0:1])), dim=-1)184 185 # rasterize186 rast, _ = dr.rasterize(ctx, uv_clip4, mesh_tex_idx.int(), (resolution, resolution))187 188 # Interpolate world space position189 gb_pos, _ = interpolate(mesh_v[None, ...], rast, mesh_pos_idx.int())190 mask = rast[..., 3:4] > 0191 return uvs, mesh_tex_idx, gb_pos, mask192