mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-07-31 10:52:50 +08:00
9eeb456e82
Summary: Update all FB license strings to the new format. Reviewed By: patricklabatut Differential Revision: D33403538 fbshipit-source-id: 97a4596c5c888f3c54f44456dc07e718a387a02c
267 lines
9.1 KiB
Python
267 lines
9.1 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the BSD-style license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
|
|
import unittest
|
|
|
|
import torch
|
|
from pytorch3d.loss.mesh_normal_consistency import mesh_normal_consistency
|
|
from pytorch3d.structures.meshes import Meshes
|
|
from pytorch3d.utils.ico_sphere import ico_sphere
|
|
|
|
|
|
IS_TORCH_1_8 = torch.__version__.startswith("1.8.")
|
|
PROBLEMATIC_CUDA = torch.version.cuda in ("11.0", "11.1")
|
|
# TODO: There are problems with cuda 11.0 and 11.1 here.
|
|
# The symptom can be
|
|
# RuntimeError: radix_sort: failed on 1st step: cudaErrorInvalidDevice: invalid device ordinal
|
|
# or something like
|
|
# operator(): block: [0,0,0], thread: [96,0,0]
|
|
# Assertion `index >= -sizes[i] && index < sizes[i] && "index out of bounds"` failed.
|
|
AVOID_LARGE_MESH_CUDA = PROBLEMATIC_CUDA and IS_TORCH_1_8
|
|
|
|
|
|
class TestMeshNormalConsistency(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
torch.manual_seed(42)
|
|
|
|
@staticmethod
|
|
def init_faces(num_verts: int = 1000):
|
|
faces = []
|
|
for f0 in range(num_verts):
|
|
for f1 in range(f0 + 1, num_verts):
|
|
f2 = torch.arange(f1 + 1, num_verts)
|
|
n = f2.shape[0]
|
|
if n == 0:
|
|
continue
|
|
faces.append(
|
|
torch.stack(
|
|
[
|
|
torch.full((n,), f0, dtype=torch.int64),
|
|
torch.full((n,), f1, dtype=torch.int64),
|
|
f2,
|
|
],
|
|
dim=1,
|
|
)
|
|
)
|
|
faces = torch.cat(faces, 0)
|
|
return faces
|
|
|
|
@staticmethod
|
|
def init_meshes(num_meshes: int = 10, num_verts: int = 1000, num_faces: int = 3000):
|
|
if AVOID_LARGE_MESH_CUDA:
|
|
device = torch.device("cpu")
|
|
else:
|
|
device = torch.device("cuda:0")
|
|
valid_faces = TestMeshNormalConsistency.init_faces(num_verts).to(device)
|
|
verts_list = []
|
|
faces_list = []
|
|
for _ in range(num_meshes):
|
|
verts = (
|
|
torch.rand((num_verts, 3), dtype=torch.float32, device=device) * 2.0
|
|
- 1.0
|
|
) # verts in the space of [-1, 1]
|
|
"""
|
|
faces = torch.stack(
|
|
[
|
|
torch.randperm(num_verts, device=device)[:3]
|
|
for _ in range(num_faces)
|
|
],
|
|
dim=0,
|
|
)
|
|
# avoids duplicate vertices in a face
|
|
"""
|
|
idx = torch.randperm(valid_faces.shape[0], device=device)[
|
|
: min(valid_faces.shape[0], num_faces)
|
|
]
|
|
faces = valid_faces[idx]
|
|
verts_list.append(verts)
|
|
faces_list.append(faces)
|
|
meshes = Meshes(verts_list, faces_list)
|
|
return meshes
|
|
|
|
@staticmethod
|
|
def mesh_normal_consistency_naive(meshes):
|
|
"""
|
|
Naive iterative implementation of mesh normal consistency.
|
|
"""
|
|
N = len(meshes)
|
|
verts_packed = meshes.verts_packed()
|
|
faces_packed = meshes.faces_packed()
|
|
edges_packed = meshes.edges_packed()
|
|
face_to_edge = meshes.faces_packed_to_edges_packed()
|
|
edges_packed_to_mesh_idx = meshes.edges_packed_to_mesh_idx()
|
|
|
|
E = edges_packed.shape[0]
|
|
loss = []
|
|
mesh_idx = []
|
|
|
|
for e in range(E):
|
|
face_idx = face_to_edge.eq(e).any(1).nonzero() # indexed to faces
|
|
v0 = verts_packed[edges_packed[e, 0]]
|
|
v1 = verts_packed[edges_packed[e, 1]]
|
|
normals = []
|
|
for f in face_idx:
|
|
v2 = -1
|
|
for j in range(3):
|
|
if (
|
|
faces_packed[f, j] != edges_packed[e, 0]
|
|
and faces_packed[f, j] != edges_packed[e, 1]
|
|
):
|
|
v2 = faces_packed[f, j]
|
|
assert v2 > -1
|
|
v2 = verts_packed[v2]
|
|
normals.append((v1 - v0).view(-1).cross((v2 - v0).view(-1)))
|
|
for i in range(len(normals) - 1):
|
|
for j in range(i + 1, len(normals)):
|
|
mesh_idx.append(edges_packed_to_mesh_idx[e])
|
|
loss.append(
|
|
(
|
|
1
|
|
- torch.cosine_similarity(
|
|
normals[i].view(1, 3), -normals[j].view(1, 3)
|
|
)
|
|
)
|
|
)
|
|
|
|
mesh_idx = torch.tensor(mesh_idx, device=meshes.device)
|
|
num = mesh_idx.bincount(minlength=N)
|
|
weights = 1.0 / num[mesh_idx].float()
|
|
|
|
loss = torch.cat(loss) * weights
|
|
return loss.sum() / N
|
|
|
|
def test_mesh_normal_consistency_simple(self):
|
|
r"""
|
|
Mesh 1:
|
|
v3
|
|
/\
|
|
/ \
|
|
e4 / f1 \ e3
|
|
/ \
|
|
v2 /___e2___\ v1
|
|
\ /
|
|
\ /
|
|
e1 \ f0 / e0
|
|
\ /
|
|
\/
|
|
v0
|
|
"""
|
|
device = torch.device("cuda:0")
|
|
# mesh1 shown above
|
|
verts1 = torch.rand((4, 3), dtype=torch.float32, device=device)
|
|
faces1 = torch.tensor([[0, 1, 2], [2, 1, 3]], dtype=torch.int64, device=device)
|
|
|
|
# mesh2 is a cuboid with 8 verts, 12 faces and 18 edges
|
|
verts2 = torch.tensor(
|
|
[
|
|
[0, 0, 0],
|
|
[0, 0, 1],
|
|
[0, 1, 0],
|
|
[0, 1, 1],
|
|
[1, 0, 0],
|
|
[1, 0, 1],
|
|
[1, 1, 0],
|
|
[1, 1, 1],
|
|
],
|
|
dtype=torch.float32,
|
|
device=device,
|
|
)
|
|
faces2 = torch.tensor(
|
|
[
|
|
[0, 1, 2],
|
|
[1, 3, 2], # left face: 0, 1
|
|
[2, 3, 6],
|
|
[3, 7, 6], # bottom face: 2, 3
|
|
[0, 2, 6],
|
|
[0, 6, 4], # front face: 4, 5
|
|
[0, 5, 1],
|
|
[0, 4, 5], # up face: 6, 7
|
|
[6, 7, 5],
|
|
[6, 5, 4], # right face: 8, 9
|
|
[1, 7, 3],
|
|
[1, 5, 7], # back face: 10, 11
|
|
],
|
|
dtype=torch.int64,
|
|
device=device,
|
|
)
|
|
|
|
# mesh3 is like mesh1 but with another face added to e2
|
|
verts3 = torch.rand((5, 3), dtype=torch.float32, device=device)
|
|
faces3 = torch.tensor(
|
|
[[0, 1, 2], [2, 1, 3], [2, 1, 4]], dtype=torch.int64, device=device
|
|
)
|
|
|
|
meshes = Meshes(verts=[verts1, verts2, verts3], faces=[faces1, faces2, faces3])
|
|
|
|
# mesh1: normal consistency computation
|
|
n0 = (verts1[1] - verts1[2]).cross(verts1[3] - verts1[2])
|
|
n1 = (verts1[1] - verts1[2]).cross(verts1[0] - verts1[2])
|
|
loss1 = 1.0 - torch.cosine_similarity(n0.view(1, 3), -(n1.view(1, 3)))
|
|
|
|
# mesh2: normal consistency computation
|
|
# In the cube mesh, 6 edges are shared with coplanar faces (loss=0),
|
|
# 12 edges are shared by perpendicular faces (loss=1)
|
|
loss2 = 12.0 / 18
|
|
|
|
# mesh3
|
|
n0 = (verts3[1] - verts3[2]).cross(verts3[3] - verts3[2])
|
|
n1 = (verts3[1] - verts3[2]).cross(verts3[0] - verts3[2])
|
|
n2 = (verts3[1] - verts3[2]).cross(verts3[4] - verts3[2])
|
|
loss3 = (
|
|
3.0
|
|
- torch.cosine_similarity(n0.view(1, 3), -(n1.view(1, 3)))
|
|
- torch.cosine_similarity(n0.view(1, 3), -(n2.view(1, 3)))
|
|
- torch.cosine_similarity(n1.view(1, 3), -(n2.view(1, 3)))
|
|
)
|
|
loss3 /= 3.0
|
|
|
|
loss = (loss1 + loss2 + loss3) / 3.0
|
|
|
|
out = mesh_normal_consistency(meshes)
|
|
|
|
self.assertTrue(torch.allclose(out, loss))
|
|
|
|
def test_mesh_normal_consistency(self):
|
|
"""
|
|
Test Mesh Normal Consistency for random meshes.
|
|
"""
|
|
meshes = TestMeshNormalConsistency.init_meshes(5, 100, 300)
|
|
|
|
out1 = mesh_normal_consistency(meshes)
|
|
out2 = TestMeshNormalConsistency.mesh_normal_consistency_naive(meshes)
|
|
|
|
self.assertTrue(torch.allclose(out1, out2))
|
|
|
|
def test_no_intersection(self):
|
|
"""
|
|
Test Mesh Normal Consistency for a mesh known to have no
|
|
intersecting faces.
|
|
"""
|
|
verts = torch.rand(1, 6, 3)
|
|
faces = torch.arange(6).reshape(1, 2, 3)
|
|
meshes = Meshes(verts=verts, faces=faces)
|
|
out = mesh_normal_consistency(meshes)
|
|
self.assertEqual(out.item(), 0)
|
|
|
|
@staticmethod
|
|
def mesh_normal_consistency_with_ico(
|
|
num_meshes: int, level: int = 3, device: str = "cpu"
|
|
):
|
|
device = torch.device(device)
|
|
mesh = ico_sphere(level, device)
|
|
verts, faces = mesh.get_mesh_verts_faces(0)
|
|
verts_list = [verts.clone() for _ in range(num_meshes)]
|
|
faces_list = [faces.clone() for _ in range(num_meshes)]
|
|
meshes = Meshes(verts_list, faces_list)
|
|
torch.cuda.synchronize()
|
|
|
|
def loss():
|
|
mesh_normal_consistency(meshes)
|
|
torch.cuda.synchronize()
|
|
|
|
return loss
|