Heartsync/TRELLIS2
0
1from typing import *
2import torch
3from easydict import EasyDict as edict
4from ..representations.mesh import Mesh, MeshWithVoxel, MeshWithPbrMaterial, TextureFilterMode, AlphaMode, TextureWrapMode
5import torch.nn.functional as F
6
7
8def intrinsics_to_projection(
9 intrinsics: torch.Tensor,
10 near: float,
11 far: float,
12 ) -> torch.Tensor:
13 """
14 OpenCV intrinsics to OpenGL perspective matrix
15
16 Args:
17 intrinsics (torch.Tensor): [3, 3] OpenCV intrinsics matrix
18 near (float): near plane to clip
19 far (float): far plane to clip
20 Returns:
21 (torch.Tensor): [4, 4] OpenGL perspective matrix
22 """
23 fx, fy = intrinsics[0, 0], intrinsics[1, 1]
24 cx, cy = intrinsics[0, 2], intrinsics[1, 2]
25 ret = torch.zeros((4, 4), dtype=intrinsics.dtype, device=intrinsics.device)
26 ret[0, 0] = 2 * fx
27 ret[1, 1] = 2 * fy
28 ret[0, 2] = 2 * cx - 1
29 ret[1, 2] = - 2 * cy + 1
30 ret[2, 2] = (far + near) / (far - near)
31 ret[2, 3] = 2 * near * far / (near - far)
32 ret[3, 2] = 1.
33 return ret
34
35
36class MeshRenderer:
37 """
38 Renderer for the Mesh representation.
39
40 Args:
41 rendering_options (dict): Rendering options.
42 """
43 def __init__(self, rendering_options={}, device='cuda'):
44 if 'dr' not in globals():
45 import nvdiffrast.torch as dr
46
47 self.rendering_options = edict({
48 "resolution": None,
49 "near": None,
50 "far": None,
51 "ssaa": 1,
52 "chunk_size": None,
53 "antialias": True,
54 "clamp_barycentric_coords": False,
55 })
56 self.rendering_options.update(rendering_options)
57 self.glctx = dr.RasterizeCudaContext(device=device)
58 self.device=device
59
60 def render(
61 self,
62 mesh : Mesh,
63 extrinsics: torch.Tensor,
64 intrinsics: torch.Tensor,
65 return_types = ["mask", "normal", "depth"],
66 transformation : Optional[torch.Tensor] = None
67 ) -> edict:
68 """
69 Render the mesh.
70
71 Args:
72 mesh : meshmodel
73 extrinsics (torch.Tensor): (4, 4) camera extrinsics
74 intrinsics (torch.Tensor): (3, 3) camera intrinsics
75 return_types (list): list of return types, can be "attr", "mask", "depth", "coord", "normal"
76
77 Returns:
78 edict based on return_types containing:
79 attr (torch.Tensor): [C, H, W] rendered attr image
80 depth (torch.Tensor): [H, W] rendered depth image
81 normal (torch.Tensor): [3, H, W] rendered normal image
82 mask (torch.Tensor): [H, W] rendered mask image
83 """
84 if 'dr' not in globals():
85 import nvdiffrast.torch as dr
86
87 resolution = self.rendering_options["resolution"]
88 near = self.rendering_options["near"]
89 far = self.rendering_options["far"]
90 ssaa = self.rendering_options["ssaa"]
91 chunk_size = self.rendering_options["chunk_size"]
92 antialias = self.rendering_options["antialias"]
93 clamp_barycentric_coords = self.rendering_options["clamp_barycentric_coords"]
94
95 if mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:
96 ret_dict = edict()
97 for type in return_types:
98 if type == "mask" :
99 ret_dict[type] = torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device)
100 elif type == "depth":
101 ret_dict[type] = torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device)
102 elif type == "normal":
103 ret_dict[type] = torch.full((3, resolution, resolution), 0.5, dtype=torch.float32, device=self.device)
104 elif type == "coord":
105 ret_dict[type] = torch.zeros((3, resolution, resolution), dtype=torch.float32, device=self.device)
106 elif type == "attr":
107 if isinstance(mesh, MeshWithVoxel):
108 ret_dict[type] = torch.zeros((mesh.attrs.shape[-1], resolution, resolution), dtype=torch.float32, device=self.device)
109 else:
110 ret_dict[type] = torch.zeros((mesh.vertex_attrs.shape[-1], resolution, resolution), dtype=torch.float32, device=self.device)
111 return ret_dict
112
113 perspective = intrinsics_to_projection(intrinsics, near, far)
114
115 full_proj = (perspective @ extrinsics).unsqueeze(0)
116 extrinsics = extrinsics.unsqueeze(0)
117
118 vertices = mesh.vertices.unsqueeze(0)
119 vertices_homo = torch.cat([vertices, torch.ones_like(vertices[..., :1])], dim=-1)
120 if transformation is not None:
121 vertices_homo = torch.bmm(vertices_homo, transformation.unsqueeze(0).transpose(-1, -2))
122 vertices = vertices_homo[..., :3].contiguous()
123 vertices_camera = torch.bmm(vertices_homo, extrinsics.transpose(-1, -2))
124 vertices_clip = torch.bmm(vertices_homo, full_proj.transpose(-1, -2))
125 faces = mesh.faces
126
127 if 'normal' in return_types:
128 v0 = vertices_camera[0, mesh.faces[:, 0], :3]
129 v1 = vertices_camera[0, mesh.faces[:, 1], :3]
130 v2 = vertices_camera[0, mesh.faces[:, 2], :3]
131 e0 = v1 - v0
132 e1 = v2 - v0
133 face_normal = torch.cross(e0, e1, dim=1)
134 face_normal = F.normalize(face_normal, dim=1)
135 face_normal = torch.where(torch.sum(face_normal * v0, dim=1, keepdim=True) > 0, face_normal, -face_normal)
136
137 out_dict = edict()
138 if chunk_size is None:
139 rast, rast_db = dr.rasterize(
140 self.glctx, vertices_clip, faces, (resolution * ssaa, resolution * ssaa)
141 )
142 if clamp_barycentric_coords:
143 rast[..., :2] = torch.clamp(rast[..., :2], 0, 1)
144 rast[..., :2] /= torch.where(rast[..., :2].sum(dim=-1, keepdim=True) > 1, rast[..., :2].sum(dim=-1, keepdim=True), torch.ones_like(rast[..., :2]))
145 for type in return_types:
146 img = None
147 if type == "mask" :
148 img = (rast[..., -1:] > 0).float()
149 if antialias: img = dr.antialias(img, rast, vertices_clip, faces)
150 elif type == "depth":
151 img = dr.interpolate(vertices_camera[..., 2:3].contiguous(), rast, faces)[0]
152 if antialias: img = dr.antialias(img, rast, vertices_clip, faces)
153 elif type == "normal" :
154 img = dr.interpolate(face_normal.unsqueeze(0), rast, torch.arange(face_normal.shape[0], dtype=torch.int, device=self.device).unsqueeze(1).repeat(1, 3).contiguous())[0]
155 if antialias: img = dr.antialias(img, rast, vertices_clip, faces)
156 img = (img + 1) / 2
157 elif type == "coord":
158 img = dr.interpolate(vertices, rast, faces)[0]
159 if antialias: img = dr.antialias(img, rast, vertices_clip, faces)
160 elif type == "attr":
161 if isinstance(mesh, MeshWithVoxel):
162 if 'grid_sample_3d' not in globals():
163 from flex_gemm.ops.grid_sample import grid_sample_3d
164 mask = rast[..., -1:] > 0
165 xyz = dr.interpolate(vertices, rast, faces)[0]
166 xyz = ((xyz - mesh.origin) / mesh.voxel_size).reshape(1, -1, 3)
167 img = grid_sample_3d(
168 mesh.attrs,
169 torch.cat([torch.zeros_like(mesh.coords[..., :1]), mesh.coords], dim=-1),
170 mesh.voxel_shape,
171 xyz,
172 mode='trilinear'
173 )
174 img = img.reshape(1, resolution * ssaa, resolution * ssaa, mesh.attrs.shape[-1]) * mask
175 elif isinstance(mesh, MeshWithPbrMaterial):
176 tri_id = rast[0, :, :, -1:]
177 mask = tri_id > 0
178 uv_coords = mesh.uv_coords.reshape(1, -1, 2)
179 texc, texd = dr.interpolate(
180 uv_coords,
181 rast,
182 torch.arange(mesh.uv_coords.shape[0] * 3, dtype=torch.int, device=self.device).reshape(-1, 3),
183 rast_db=rast_db,
184 diff_attrs='all'
185 )
186 # Fix problematic texture coordinates
187 texc = torch.nan_to_num(texc, nan=0.0, posinf=1e3, neginf=-1e3)
188 texc = torch.clamp(texc, min=-1e3, max=1e3)
189 texd = torch.nan_to_num(texd, nan=0.0, posinf=1e3, neginf=-1e3)
190 texd = torch.clamp(texd, min=-1e3, max=1e3)
191 mid = mesh.material_ids[(tri_id - 1).long()]
192 imgs = {
193 'base_color': torch.zeros((resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device),
194 'metallic': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device),
195 'roughness': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device),
196 'alpha': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device)
197 }
198 for id, mat in enumerate(mesh.materials):
199 mat_mask = (mid == id).float() * mask.float()
200 mat_texc = texc * mat_mask
201 mat_texd = texd * mat_mask
202
203 if mat.base_color_texture is not None:
204 base_color = dr.texture(
205 mat.base_color_texture.image.unsqueeze(0),
206 mat_texc,
207 mat_texd,
208 filter_mode='linear-mipmap-linear' if mat.base_color_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
209 boundary_mode='clamp' if mat.base_color_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
210 )[0]
211 imgs['base_color'] += base_color * mat.base_color_factor * mat_mask
212 else:
213 imgs['base_color'] += mat.base_color_factor * mat_mask
214
215 if mat.metallic_texture is not None:
216 metallic = dr.texture(
217 mat.metallic_texture.image.unsqueeze(0),
218 mat_texc,
219 mat_texd,
220 filter_mode='linear-mipmap-linear' if mat.metallic_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
221 boundary_mode='clamp' if mat.metallic_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
222 )[0]
223 imgs['metallic'] += metallic * mat.metallic_factor * mat_mask
224 else:
225 imgs['metallic'] += mat.metallic_factor * mat_mask
226
227 if mat.roughness_texture is not None:
228 roughness = dr.texture(
229 mat.roughness_texture.image.unsqueeze(0),
230 mat_texc,
231 mat_texd,
232 filter_mode='linear-mipmap-linear' if mat.roughness_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
233 boundary_mode='clamp' if mat.roughness_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
234 )[0]
235 imgs['roughness'] += roughness * mat.roughness_factor * mat_mask
236 else:
237 imgs['roughness'] += mat.roughness_factor * mat_mask
238
239 if mat.alpha_mode == AlphaMode.OPAQUE:
240 imgs['alpha'] += 1.0 * mat_mask
241 else:
242 if mat.alpha_texture is not None:
243 alpha = dr.texture(
244 mat.alpha_texture.image.unsqueeze(0),
245 mat_texc,
246 mat_texd,
247 filter_mode='linear-mipmap-linear' if mat.alpha_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
248 boundary_mode='clamp' if mat.alpha_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
249 )[0]
250 if mat.alpha_mode == AlphaMode.MASK:
251 imgs['alpha'] += (alpha * mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask
252 elif mat.alpha_mode == AlphaMode.BLEND:
253 imgs['alpha'] += alpha * mat.alpha_factor * mat_mask
254 else:
255 if mat.alpha_mode == AlphaMode.MASK:
256 imgs['alpha'] += (mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask
257 elif mat.alpha_mode == AlphaMode.BLEND:
258 imgs['alpha'] += mat.alpha_factor * mat_mask
259
260 img = torch.cat([imgs[name] for name in imgs.keys()], dim=-1).unsqueeze(0)
261 else:
262 img = dr.interpolate(mesh.vertex_attrs.unsqueeze(0), rast, faces)[0]
263 if antialias: img = dr.antialias(img, rast, vertices_clip, faces)
264
265 out_dict[type] = img
266 else:
267 z_buffer = torch.full((1, resolution * ssaa, resolution * ssaa), torch.inf, device=self.device, dtype=torch.float32)
268 for i in range(0, faces.shape[0], chunk_size):
269 faces_chunk = faces[i:i+chunk_size]
270 rast, rast_db = dr.rasterize(
271 self.glctx, vertices_clip, faces_chunk, (resolution * ssaa, resolution * ssaa)
272 )
273 z_filter = torch.logical_and(
274 rast[..., 3] != 0,
275 rast[..., 2] < z_buffer
276 )
277 z_buffer[z_filter] = rast[z_filter][..., 2]
278
279 for type in return_types:
280 img = None
281 if type == "mask" :
282 img = (rast[..., -1:] > 0).float()
283 elif type == "depth":
284 img = dr.interpolate(vertices_camera[..., 2:3].contiguous(), rast, faces_chunk)[0]
285 elif type == "normal" :
286 face_normal_chunk = face_normal[i:i+chunk_size]
287 img = dr.interpolate(face_normal_chunk.unsqueeze(0), rast, torch.arange(face_normal_chunk.shape[0], dtype=torch.int, device=self.device).unsqueeze(1).repeat(1, 3).contiguous())[0]
288 img = (img + 1) / 2
289 elif type == "coord":
290 img = dr.interpolate(vertices, rast, faces_chunk)[0]
291 elif type == "attr":
292 if isinstance(mesh, MeshWithVoxel):
293 if 'grid_sample_3d' not in globals():
294 from flex_gemm.ops.grid_sample import grid_sample_3d
295 mask = rast[..., -1:] > 0
296 xyz = dr.interpolate(vertices, rast, faces_chunk)[0]
297 xyz = ((xyz - mesh.origin) / mesh.voxel_size).reshape(1, -1, 3)
298 img = grid_sample_3d(
299 mesh.attrs,
300 torch.cat([torch.zeros_like(mesh.coords[..., :1]), mesh.coords], dim=-1),
301 mesh.voxel_shape,
302 xyz,
303 mode='trilinear'
304 )
305 img = img.reshape(1, resolution * ssaa, resolution * ssaa, mesh.attrs.shape[-1]) * mask
306 elif isinstance(mesh, MeshWithPbrMaterial):
307 tri_id = rast[0, :, :, -1:]
308 mask = tri_id > 0
309 uv_coords = mesh.uv_coords.reshape(1, -1, 2)
310 texc, texd = dr.interpolate(
311 uv_coords,
312 rast,
313 torch.arange(mesh.uv_coords.shape[0] * 3, dtype=torch.int, device=self.device).reshape(-1, 3),
314 rast_db=rast_db,
315 diff_attrs='all'
316 )
317 # Fix problematic texture coordinates
318 texc = torch.nan_to_num(texc, nan=0.0, posinf=1e3, neginf=-1e3)
319 texc = torch.clamp(texc, min=-1e3, max=1e3)
320 texd = torch.nan_to_num(texd, nan=0.0, posinf=1e3, neginf=-1e3)
321 texd = torch.clamp(texd, min=-1e3, max=1e3)
322 mid = mesh.material_ids[(tri_id - 1).long()]
323 imgs = {
324 'base_color': torch.zeros((resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device),
325 'metallic': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device),
326 'roughness': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device),
327 'alpha': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device)
328 }
329 for id, mat in enumerate(mesh.materials):
330 mat_mask = (mid == id).float() * mask.float()
331 mat_texc = texc * mat_mask
332 mat_texd = texd * mat_mask
333
334 if mat.base_color_texture is not None:
335 base_color = dr.texture(
336 mat.base_color_texture.image.unsqueeze(0),
337 mat_texc,
338 mat_texd,
339 filter_mode='linear-mipmap-linear' if mat.base_color_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
340 boundary_mode='clamp' if mat.base_color_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
341 )[0]
342 imgs['base_color'] += base_color * mat.base_color_factor * mat_mask
343 else:
344 imgs['base_color'] += mat.base_color_factor * mat_mask
345
346 if mat.metallic_texture is not None:
347 metallic = dr.texture(
348 mat.metallic_texture.image.unsqueeze(0),
349 mat_texc,
350 mat_texd,
351 filter_mode='linear-mipmap-linear' if mat.metallic_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
352 boundary_mode='clamp' if mat.metallic_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
353 )[0]
354 imgs['metallic'] += metallic * mat.metallic_factor * mat_mask
355 else:
356 imgs['metallic'] += mat.metallic_factor * mat_mask
357
358 if mat.roughness_texture is not None:
359 roughness = dr.texture(
360 mat.roughness_texture.image.unsqueeze(0),
361 mat_texc,
362 mat_texd,
363 filter_mode='linear-mipmap-linear' if mat.roughness_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
364 boundary_mode='clamp' if mat.roughness_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
365 )[0]
366 imgs['roughness'] += roughness * mat.roughness_factor * mat_mask
367 else:
368 imgs['roughness'] += mat.roughness_factor * mat_mask
369
370 if mat.alpha_mode == AlphaMode.OPAQUE:
371 imgs['alpha'] += 1.0 * mat_mask
372 else:
373 if mat.alpha_texture is not None:
374 alpha = dr.texture(
375 mat.alpha_texture.image.unsqueeze(0),
376 mat_texc,
377 mat_texd,
378 filter_mode='linear-mipmap-linear' if mat.alpha_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
379 boundary_mode='clamp' if mat.alpha_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
380 )[0]
381 if mat.alpha_mode == AlphaMode.MASK:
382 imgs['alpha'] += (alpha * mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask
383 elif mat.alpha_mode == AlphaMode.BLEND:
384 imgs['alpha'] += alpha * mat.alpha_factor * mat_mask
385 else:
386 if mat.alpha_mode == AlphaMode.MASK:
387 imgs['alpha'] += (mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask
388 elif mat.alpha_mode == AlphaMode.BLEND:
389 imgs['alpha'] += mat.alpha_factor * mat_mask
390
391 img = torch.cat([imgs[name] for name in imgs.keys()], dim=-1).unsqueeze(0)
392 else:
393 img = dr.interpolate(mesh.vertex_attrs.unsqueeze(0), rast, faces_chunk)[0]
394
395 if type not in out_dict:
396 out_dict[type] = img
397 else:
398 out_dict[type][z_filter] = img[z_filter]
399
400 for type in return_types:
401 img = out_dict[type]
402 if ssaa > 1:
403 img = F.interpolate(img.permute(0, 3, 1, 2), (resolution, resolution), mode='bilinear', align_corners=False, antialias=True)
404 img = img.squeeze()
405 else:
406 img = img.permute(0, 3, 1, 2).squeeze()
407 out_dict[type] = img
408
409 if isinstance(mesh, (MeshWithVoxel, MeshWithPbrMaterial)) and 'attr' in return_types:
410 for k, s in mesh.layout.items():
411 out_dict[k] = out_dict['attr'][s]
412 del out_dict['attr']
413
414 return out_dict
415