| 345 | |
| 346 | |
| 347 | def mesh_normals_areas(vertices, triangles=None, scale=[1.0], batch=None, normals=None): |
| 348 | # Single- or Multi-scale mode: |
| 349 | if hasattr(scale, "__len__"): |
| 350 | scales, single_scale = scale, False |
| 351 | else: |
| 352 | scales, single_scale = [scale], True |
| 353 | scales = torch.Tensor(scales).type_as(vertices) # (S,) |
| 354 | |
| 355 | # Compute the "raw" field of normals: |
| 356 | if triangles is not None: |
| 357 | # Vertices of all triangles in the mesh: |
| 358 | A = vertices[triangles[0, :]] # (N, 3) |
| 359 | B = vertices[triangles[1, :]] # (N, 3) |
| 360 | C = vertices[triangles[2, :]] # (N, 3) |
| 361 | |
| 362 | # Triangle centers and normals (length = surface area): |
| 363 | centers = (A + B + C) / 3 # (N, 3) |
| 364 | V = (B - A).cross(C - A) # (N, 3) |
| 365 | |
| 366 | # Vertice areas: |
| 367 | S = (V ** 2).sum(-1).sqrt() / 6 # (N,) 1/3 of a triangle area |
| 368 | areas = torch.zeros(len(vertices)).type_as(vertices) # (N,) |
| 369 | areas.scatter_add_(0, triangles[0, :], S) # Aggregate from "A's" |
| 370 | areas.scatter_add_(0, triangles[1, :], S) # Aggregate from "B's" |
| 371 | areas.scatter_add_(0, triangles[2, :], S) # Aggregate from "C's" |
| 372 | |
| 373 | else: # Use "normals" instead |
| 374 | areas = None |
| 375 | V = normals |
| 376 | centers = vertices |
| 377 | |
| 378 | # Normal of a vertex = average of all normals in a ball of size "scale": |
| 379 | x_i = LazyTensor(vertices[:, None, :]) # (N, 1, 3) |
| 380 | y_j = LazyTensor(centers[None, :, :]) # (1, M, 3) |
| 381 | v_j = LazyTensor(V[None, :, :]) # (1, M, 3) |
| 382 | s = LazyTensor(scales[None, None, :]) # (1, 1, S) |
| 383 | |
| 384 | D_ij = ((x_i - y_j) ** 2).sum(-1) # (N, M, 1) |
| 385 | K_ij = (-D_ij / (2 * s ** 2)).exp() # (N, M, S) |
| 386 | |
| 387 | # Support for heterogeneous batch processing: |
| 388 | if batch is not None: |
| 389 | batch_vertices = batch |
| 390 | batch_centers = batch[triangles[0, :]] if triangles is not None else batch |
| 391 | #K_ij.ranges = diagonal_ranges(batch_vertices, batch_centers) |
| 392 | |
| 393 | if single_scale: |
| 394 | U = (K_ij * v_j).sum(dim=1) # (N, 3) |
| 395 | else: |
| 396 | U = (K_ij.tensorprod(v_j)).sum(dim=1) # (N, S*3) |
| 397 | U = U.view(-1, len(scales), 3) # (N, S, 3) |
| 398 | |
| 399 | normals = F.normalize(U, p=2, dim=-1) # (N, 3) or (N, S, 3) |
| 400 | |
| 401 | return normals, areas |
| 402 | |
| 403 | |
| 404 | def tangent_vectors(normals): |