test_meshes.py 53.2 KB
Newer Older
1
# Copyright (c) Meta Platforms, Inc. and affiliates.
Patrick Labatut's avatar
Patrick Labatut committed
2
3
4
5
# 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.
facebook-github-bot's avatar
facebook-github-bot committed
6

7
import itertools
8
import random
facebook-github-bot's avatar
facebook-github-bot committed
9
10
import unittest

11
12
import numpy as np
import torch
facebook-github-bot's avatar
facebook-github-bot committed
13
from common_testing import TestCaseMixin
14
from pytorch3d.structures.meshes import Meshes
facebook-github-bot's avatar
facebook-github-bot committed
15
16


Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def init_mesh(
    num_meshes: int = 10,
    max_v: int = 100,
    max_f: int = 300,
    lists_to_tensors: bool = False,
    device: str = "cpu",
    requires_grad: bool = False,
):
    """
    Function to generate a Meshes object of N meshes with
    random numbers of vertices and faces.

    Args:
        num_meshes: Number of meshes to generate.
        max_v: Max number of vertices per mesh.
        max_f: Max number of faces per mesh.
        lists_to_tensors: Determines whether the generated meshes should be
                            constructed from lists (=False) or
                            a tensor (=True) of faces/verts.

    Returns:
        Meshes object.
    """
    device = torch.device(device)

    verts_list = []
    faces_list = []

    # Randomly generate numbers of faces and vertices in each mesh.
    if lists_to_tensors:
        # If we define faces/verts with tensors, f/v has to be the
        # same for each mesh in the batch.
        f = torch.randint(1, max_f, size=(1,), dtype=torch.int32)
        v = torch.randint(3, high=max_v, size=(1,), dtype=torch.int32)
        f = f.repeat(num_meshes)
        v = v.repeat(num_meshes)
    else:
        # For lists of faces and vertices, we can sample different v/f
        # per mesh.
        f = torch.randint(max_f, size=(num_meshes,), dtype=torch.int32)
        v = torch.randint(3, high=max_v, size=(num_meshes,), dtype=torch.int32)

    # Generate the actual vertices and faces.
    for i in range(num_meshes):
        verts = torch.rand(
            (v[i], 3),
            dtype=torch.float32,
            device=device,
            requires_grad=requires_grad,
        )
        faces = torch.randint(v[i], size=(f[i], 3), dtype=torch.int64, device=device)
        verts_list.append(verts)
        faces_list.append(faces)
facebook-github-bot's avatar
facebook-github-bot committed
70

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
71
72
73
    if lists_to_tensors:
        verts_list = torch.stack(verts_list)
        faces_list = torch.stack(faces_list)
facebook-github-bot's avatar
facebook-github-bot committed
74

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
75
    return Meshes(verts=verts_list, faces=faces_list)
facebook-github-bot's avatar
facebook-github-bot committed
76
77


Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
78
79
80
def init_simple_mesh(device: str = "cpu"):
    """
    Returns a Meshes data structure of simple mesh examples.
facebook-github-bot's avatar
facebook-github-bot committed
81

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
82
83
84
85
    Returns:
        Meshes object.
    """
    device = torch.device(device)
facebook-github-bot's avatar
facebook-github-bot committed
86

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
    verts = [
        torch.tensor(
            [[0.1, 0.3, 0.5], [0.5, 0.2, 0.1], [0.6, 0.8, 0.7]],
            dtype=torch.float32,
            device=device,
        ),
        torch.tensor(
            [[0.1, 0.3, 0.3], [0.6, 0.7, 0.8], [0.2, 0.3, 0.4], [0.1, 0.5, 0.3]],
            dtype=torch.float32,
            device=device,
        ),
        torch.tensor(
            [
                [0.7, 0.3, 0.6],
                [0.2, 0.4, 0.8],
                [0.9, 0.5, 0.2],
                [0.2, 0.3, 0.4],
                [0.9, 0.3, 0.8],
            ],
            dtype=torch.float32,
            device=device,
        ),
    ]
    faces = [
        torch.tensor([[0, 1, 2]], dtype=torch.int64, device=device),
        torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.int64, device=device),
        torch.tensor(
            [
                [1, 2, 0],
                [0, 1, 3],
                [2, 3, 1],
                [4, 3, 2],
                [4, 0, 1],
                [4, 3, 1],
                [4, 2, 1],
            ],
            dtype=torch.int64,
            device=device,
        ),
    ]
    return Meshes(verts=verts, faces=faces)


130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
def mesh_structures_equal(mesh1, mesh2) -> bool:
    """
    Two meshes are equal if they have identical verts_list and faces_list.

    Use to_sorted() before passing into this function to obtain meshes invariant to
    vertex permutations. Note that this operator treats two geometrically identical
    meshes as different if their vertices are in different coordinate frames.
    """
    if mesh1.__class__ != mesh1.__class__:
        return False

    if mesh1.textures is not None or mesh2.textures is not None:
        raise NotImplementedError(
            "mesh equality is not implemented for textured meshes."
        )

    if len(mesh1.verts_list()) != len(mesh2.verts_list()) or not all(
        torch.equal(verts_mesh1, verts_mesh2)
        for (verts_mesh1, verts_mesh2) in zip(mesh1.verts_list(), mesh2.verts_list())
    ):
        return False

    if len(mesh1.faces_list()) != len(mesh2.faces_list()) or not all(
        torch.equal(faces_mesh1, faces_mesh2)
        for (faces_mesh1, faces_mesh2) in zip(mesh1.faces_list(), mesh2.faces_list())
    ):
        return False

    if len(mesh1.verts_normals_list()) != len(mesh2.verts_normals_list()) or not all(
        torch.equal(normals_mesh1, normals_mesh2)
        for (normals_mesh1, normals_mesh2) in zip(
            mesh1.verts_normals_list(), mesh2.verts_normals_list()
        )
    ):
        return False

    return True


Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
def to_sorted(mesh: Meshes) -> "Meshes":
    """
    Create a new Meshes object, where each sub-mesh's vertices are sorted
    alphabetically.

    Returns:
        A Meshes object with the same topology as this mesh, with vertices sorted
        alphabetically.

    Example:

    For a mesh with verts [[2.3, .2, .4], [.0, .1, .2], [.0, .0, .1]] and a single
    face [[0, 1, 2]], to_sorted will create a new mesh with verts [[.0, .0, .1],
    [.0, .1, .2], [2.3, .2, .4]] and a single face [[2, 1, 0]]. This is useful to
    create a semi-canonical representation of the mesh that is invariant to vertex
    permutations, but not invariant to coordinate frame changes.
    """
    if mesh.textures is not None:
        raise NotImplementedError(
            "to_sorted is not implemented for meshes with "
            f"{type(mesh.textures).__name__} textures."
        )

    verts_list = mesh.verts_list()
    faces_list = mesh.faces_list()
    verts_sorted_list = []
    faces_sorted_list = []

    for verts, faces in zip(verts_list, faces_list):
        # Argsort the vertices alphabetically: sort_ids[k] corresponds to the id of
        # the vertex in the non-sorted mesh that should sit at index k in the sorted mesh.
        sort_ids = torch.tensor(
            [
                idx_and_val[0]
                for idx_and_val in sorted(
                    enumerate(verts.tolist()),
                    key=lambda idx_and_val: idx_and_val[1],
                )
            ],
            device=mesh.device,
        )

        # Resort the vertices. index_select allocates new memory.
        verts_sorted = verts[sort_ids]
        verts_sorted_list.append(verts_sorted)

        # The `faces` tensor contains vertex ids. Substitute old vertex ids for the
        # new ones. new_vertex_ids is the inverse of sort_ids: new_vertex_ids[k]
        # corresponds to the id of the vertex in the sorted mesh that is the same as
        # vertex k in the non-sorted mesh.
        new_vertex_ids = torch.argsort(sort_ids)
        faces_sorted = (
            torch.gather(new_vertex_ids, 0, faces.flatten())
            .reshape(faces.shape)
            .clone()
        )
        faces_sorted_list.append(faces_sorted)

    other = mesh.__class__(verts=verts_sorted_list, faces=faces_sorted_list)
    for k in mesh._INTERNAL_TENSORS:
        v = getattr(mesh, k)
        if torch.is_tensor(v):
            setattr(other, k, v.clone())

    return other


Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
236
237
238
239
class TestMeshes(TestCaseMixin, unittest.TestCase):
    def setUp(self) -> None:
        np.random.seed(42)
        torch.manual_seed(42)
facebook-github-bot's avatar
facebook-github-bot committed
240
241

    def test_simple(self):
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
242
        mesh = init_simple_mesh("cuda:0")
facebook-github-bot's avatar
facebook-github-bot committed
243

Nikhila Ravi's avatar
Nikhila Ravi committed
244
        # Check that faces/verts per mesh are set in init:
245
246
        self.assertClose(mesh._num_faces_per_mesh.cpu(), torch.tensor([1, 2, 7]))
        self.assertClose(mesh._num_verts_per_mesh.cpu(), torch.tensor([3, 4, 5]))
Nikhila Ravi's avatar
Nikhila Ravi committed
247
248

        # Check computed tensors
facebook-github-bot's avatar
facebook-github-bot committed
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
        self.assertClose(
            mesh.verts_packed_to_mesh_idx().cpu(),
            torch.tensor([0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2]),
        )
        self.assertClose(
            mesh.mesh_to_verts_packed_first_idx().cpu(), torch.tensor([0, 3, 7])
        )
        self.assertClose(
            mesh.verts_padded_to_packed_idx().cpu(),
            torch.tensor([0, 1, 2, 5, 6, 7, 8, 10, 11, 12, 13, 14]),
        )
        self.assertClose(
            mesh.faces_packed_to_mesh_idx().cpu(),
            torch.tensor([0, 1, 1, 2, 2, 2, 2, 2, 2, 2]),
        )
        self.assertClose(
            mesh.mesh_to_faces_packed_first_idx().cpu(), torch.tensor([0, 1, 3])
        )
        self.assertClose(
268
            mesh.num_edges_per_mesh().cpu(), torch.tensor([3, 5, 10], dtype=torch.int32)
facebook-github-bot's avatar
facebook-github-bot committed
269
        )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
270
271
272
273
        self.assertClose(
            mesh.mesh_to_edges_packed_first_idx().cpu(),
            torch.tensor([0, 3, 8], dtype=torch.int64),
        )
facebook-github-bot's avatar
facebook-github-bot committed
274

275
276
277
278
    def test_init_error(self):
        # Check if correct errors are raised when verts/faces are on
        # different devices

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
279
        mesh = init_mesh(10, 10, 100)
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
        verts_list = mesh.verts_list()  # all tensors on cpu
        verts_list = [
            v.to("cuda:0") if random.uniform(0, 1) > 0.5 else v for v in verts_list
        ]
        faces_list = mesh.faces_list()

        with self.assertRaises(ValueError) as cm:
            Meshes(verts=verts_list, faces=faces_list)
            self.assertTrue("same device" in cm.msg)

        verts_padded = mesh.verts_padded()  # on cpu
        verts_padded = verts_padded.to("cuda:0")
        faces_padded = mesh.faces_padded()

        with self.assertRaises(ValueError) as cm:
            Meshes(verts=verts_padded, faces=faces_padded)
            self.assertTrue("same device" in cm.msg)

facebook-github-bot's avatar
facebook-github-bot committed
298
299
300
301
302
    def test_simple_random_meshes(self):

        # Define the test mesh object either as a list or tensor of faces/verts.
        for lists_to_tensors in (False, True):
            N = 10
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
303
            mesh = init_mesh(N, 100, 300, lists_to_tensors=lists_to_tensors)
facebook-github-bot's avatar
facebook-github-bot committed
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
            verts_list = mesh.verts_list()
            faces_list = mesh.faces_list()

            # Check batch calculations.
            verts_padded = mesh.verts_padded()
            faces_padded = mesh.faces_padded()
            verts_per_mesh = mesh.num_verts_per_mesh()
            faces_per_mesh = mesh.num_faces_per_mesh()
            for n in range(N):
                v = verts_list[n].shape[0]
                f = faces_list[n].shape[0]
                self.assertClose(verts_padded[n, :v, :], verts_list[n])
                if verts_padded.shape[1] > v:
                    self.assertTrue(verts_padded[n, v:, :].eq(0).all())
                self.assertClose(faces_padded[n, :f, :], faces_list[n])
                if faces_padded.shape[1] > f:
                    self.assertTrue(faces_padded[n, f:, :].eq(-1).all())
                self.assertEqual(verts_per_mesh[n], v)
                self.assertEqual(faces_per_mesh[n], f)

            # Check compute packed.
            verts_packed = mesh.verts_packed()
            vert_to_mesh = mesh.verts_packed_to_mesh_idx()
            mesh_to_vert = mesh.mesh_to_verts_packed_first_idx()
            faces_packed = mesh.faces_packed()
            face_to_mesh = mesh.faces_packed_to_mesh_idx()
            mesh_to_face = mesh.mesh_to_faces_packed_first_idx()

            curv, curf = 0, 0
            for n in range(N):
                v = verts_list[n].shape[0]
                f = faces_list[n].shape[0]
336
337
                self.assertClose(verts_packed[curv : curv + v, :], verts_list[n])
                self.assertClose(faces_packed[curf : curf + f, :] - curv, faces_list[n])
facebook-github-bot's avatar
facebook-github-bot committed
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
                self.assertTrue(vert_to_mesh[curv : curv + v].eq(n).all())
                self.assertTrue(face_to_mesh[curf : curf + f].eq(n).all())
                self.assertTrue(mesh_to_vert[n] == curv)
                self.assertTrue(mesh_to_face[n] == curf)
                curv += v
                curf += f

            # Check compute edges and compare with numpy unique.
            edges = mesh.edges_packed().cpu().numpy()
            edge_to_mesh_idx = mesh.edges_packed_to_mesh_idx().cpu().numpy()
            num_edges_per_mesh = mesh.num_edges_per_mesh().cpu().numpy()

            npfaces_packed = mesh.faces_packed().cpu().numpy()
            e01 = npfaces_packed[:, [0, 1]]
            e12 = npfaces_packed[:, [1, 2]]
            e20 = npfaces_packed[:, [2, 0]]
            npedges = np.concatenate((e12, e20, e01), axis=0)
            npedges = np.sort(npedges, axis=1)

357
            unique_edges, unique_idx = np.unique(npedges, return_index=True, axis=0)
facebook-github-bot's avatar
facebook-github-bot committed
358
359
360
361
362
363
364
            self.assertTrue(np.allclose(edges, unique_edges))
            temp = face_to_mesh.cpu().numpy()
            temp = np.concatenate((temp, temp, temp), axis=0)
            edge_to_mesh = temp[unique_idx]
            self.assertTrue(np.allclose(edge_to_mesh_idx, edge_to_mesh))
            num_edges = np.bincount(edge_to_mesh, minlength=N)
            self.assertTrue(np.allclose(num_edges_per_mesh, num_edges))
Georgia Gkioxari's avatar
Georgia Gkioxari committed
365
366
367
368
369
370
371
            mesh_to_edges_packed_first_idx = (
                mesh.mesh_to_edges_packed_first_idx().cpu().numpy()
            )
            self.assertTrue(
                np.allclose(mesh_to_edges_packed_first_idx[1:], num_edges.cumsum()[:-1])
            )
            self.assertTrue(mesh_to_edges_packed_first_idx[0] == 0)
facebook-github-bot's avatar
facebook-github-bot committed
372
373

    def test_allempty(self):
374
        mesh = Meshes(verts=[], faces=[])
facebook-github-bot's avatar
facebook-github-bot committed
375
376
377
378
379
        self.assertEqual(len(mesh), 0)
        self.assertEqual(mesh.verts_padded().shape[0], 0)
        self.assertEqual(mesh.faces_padded().shape[0], 0)
        self.assertEqual(mesh.verts_packed().shape[0], 0)
        self.assertEqual(mesh.faces_packed().shape[0], 0)
Nikhila Ravi's avatar
Nikhila Ravi committed
380
381
        self.assertEqual(mesh.num_faces_per_mesh().shape[0], 0)
        self.assertEqual(mesh.num_verts_per_mesh().shape[0], 0)
facebook-github-bot's avatar
facebook-github-bot committed
382
383
384
385
386
387
388
389
390
391
392
393

    def test_empty(self):
        N, V, F = 10, 100, 300
        device = torch.device("cuda:0")
        verts_list = []
        faces_list = []
        valid = torch.randint(2, size=(N,), dtype=torch.uint8, device=device)
        for n in range(N):
            if valid[n]:
                v = torch.randint(
                    3, high=V, size=(1,), dtype=torch.int32, device=device
                )[0]
394
                f = torch.randint(F, size=(1,), dtype=torch.int32, device=device)[0]
facebook-github-bot's avatar
facebook-github-bot committed
395
                verts = torch.rand((v, 3), dtype=torch.float32, device=device)
396
                faces = torch.randint(v, size=(f, 3), dtype=torch.int64, device=device)
facebook-github-bot's avatar
facebook-github-bot committed
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
            else:
                verts = torch.tensor([], dtype=torch.float32, device=device)
                faces = torch.tensor([], dtype=torch.int64, device=device)
            verts_list.append(verts)
            faces_list.append(faces)

        mesh = Meshes(verts=verts_list, faces=faces_list)
        verts_padded = mesh.verts_padded()
        faces_padded = mesh.faces_padded()
        verts_per_mesh = mesh.num_verts_per_mesh()
        faces_per_mesh = mesh.num_faces_per_mesh()
        for n in range(N):
            v = len(verts_list[n])
            f = len(faces_list[n])
            if v > 0:
                self.assertClose(verts_padded[n, :v, :], verts_list[n])
                if verts_padded.shape[1] > v:
                    self.assertTrue(verts_padded[n, v:, :].eq(0).all())
            if f > 0:
                self.assertClose(faces_padded[n, :f, :], faces_list[n])
                if faces_padded.shape[1] > f:
                    self.assertTrue(faces_padded[n, f:, :].eq(-1).all())
            self.assertTrue(verts_per_mesh[n] == v)
            self.assertTrue(faces_per_mesh[n] == f)

    def test_padding(self):
        N, V, F = 10, 100, 300
        device = torch.device("cuda:0")
        verts, faces = [], []
        valid = torch.randint(2, size=(N,), dtype=torch.uint8, device=device)
        num_verts, num_faces = (
            torch.zeros(N, dtype=torch.int32),
            torch.zeros(N, dtype=torch.int32),
        )
        for n in range(N):
            verts.append(torch.rand((V, 3), dtype=torch.float32, device=device))
433
            this_faces = torch.full((F, 3), -1, dtype=torch.int64, device=device)
facebook-github-bot's avatar
facebook-github-bot committed
434
435
436
437
            if valid[n]:
                v = torch.randint(
                    3, high=V, size=(1,), dtype=torch.int32, device=device
                )[0]
438
                f = torch.randint(F, size=(1,), dtype=torch.int32, device=device)[0]
facebook-github-bot's avatar
facebook-github-bot committed
439
440
441
442
443
444
445
446
447
                this_faces[:f, :] = torch.randint(
                    v, size=(f, 3), dtype=torch.int64, device=device
                )
                num_verts[n] = v
                num_faces[n] = f
            faces.append(this_faces)

        mesh = Meshes(verts=torch.stack(verts), faces=torch.stack(faces))

Nikhila Ravi's avatar
Nikhila Ravi committed
448
        # Check verts/faces per mesh are set correctly in init.
449
        self.assertListEqual(mesh._num_faces_per_mesh.tolist(), num_faces.tolist())
Nikhila Ravi's avatar
Nikhila Ravi committed
450
        self.assertListEqual(mesh._num_verts_per_mesh.tolist(), [V] * N)
facebook-github-bot's avatar
facebook-github-bot committed
451
452
453
454
455
456

        for n, (vv, ff) in enumerate(zip(mesh.verts_list(), mesh.faces_list())):
            self.assertClose(ff, faces[n][: num_faces[n]])
            self.assertClose(vv, verts[n])

        new_faces = [ff.clone() for ff in faces]
457
458
        v = torch.randint(3, high=V, size=(1,), dtype=torch.int32, device=device)[0]
        f = torch.randint(F - 10, size=(1,), dtype=torch.int32, device=device)[0]
facebook-github-bot's avatar
facebook-github-bot committed
459
460
461
462
463
464
465
466
467
468
469
        this_faces = torch.full((F, 3), -1, dtype=torch.int64, device=device)
        this_faces[10 : f + 10, :] = torch.randint(
            v, size=(f, 3), dtype=torch.int64, device=device
        )
        new_faces[3] = this_faces

        with self.assertRaisesRegex(ValueError, "Padding of faces"):
            Meshes(verts=torch.stack(verts), faces=torch.stack(new_faces))

    def test_clone(self):
        N = 5
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
470
        mesh = init_mesh(N, 10, 100)
facebook-github-bot's avatar
facebook-github-bot committed
471
472
473
474
475
476
477
478
479
480
481
        for force in [0, 1]:
            if force:
                # force mesh to have computed attributes
                mesh.verts_packed()
                mesh.edges_packed()
                mesh.verts_padded()

            new_mesh = mesh.clone()

            # Modify tensors in both meshes.
            new_mesh._verts_list[0] = new_mesh._verts_list[0] * 5
Georgia Gkioxari's avatar
Georgia Gkioxari committed
482

facebook-github-bot's avatar
facebook-github-bot committed
483
484
485
486
487
488
489
490
491
492
            # Check cloned and original Meshes objects do not share tensors.
            self.assertFalse(
                torch.allclose(new_mesh._verts_list[0], mesh._verts_list[0])
            )
            self.assertSeparate(new_mesh.verts_packed(), mesh.verts_packed())
            self.assertSeparate(new_mesh.verts_padded(), mesh.verts_padded())
            self.assertSeparate(new_mesh.faces_packed(), mesh.faces_packed())
            self.assertSeparate(new_mesh.faces_padded(), mesh.faces_padded())
            self.assertSeparate(new_mesh.edges_packed(), mesh.edges_packed())

493
494
    def test_detach(self):
        N = 5
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
495
        mesh = init_mesh(N, 10, 100, requires_grad=True)
496
497
498
499
500
501
502
503
504
505
506
        for force in [0, 1]:
            if force:
                # force mesh to have computed attributes
                mesh.verts_packed()
                mesh.edges_packed()
                mesh.verts_padded()

            new_mesh = mesh.detach()

            self.assertFalse(new_mesh.verts_packed().requires_grad)
            self.assertClose(new_mesh.verts_packed(), mesh.verts_packed())
507
            self.assertFalse(new_mesh.verts_padded().requires_grad)
508
509
            self.assertClose(new_mesh.verts_padded(), mesh.verts_padded())
            for v, newv in zip(mesh.verts_list(), new_mesh.verts_list()):
510
                self.assertFalse(newv.requires_grad)
511
512
                self.assertClose(newv, v)

facebook-github-bot's avatar
facebook-github-bot committed
513
514
515
516
517
518
519
520
521
522
523
    def test_offset_verts(self):
        def naive_offset_verts(mesh, vert_offsets_packed):
            # new Meshes class
            new_verts_packed = mesh.verts_packed() + vert_offsets_packed
            new_verts_list = list(
                new_verts_packed.split(mesh.num_verts_per_mesh().tolist(), 0)
            )
            new_faces_list = [f.clone() for f in mesh.faces_list()]
            return Meshes(verts=new_verts_list, faces=new_faces_list)

        N = 5
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
524
        mesh = init_mesh(N, 30, 100, lists_to_tensors=True)
facebook-github-bot's avatar
facebook-github-bot committed
525
526
        all_v = mesh.verts_packed().size(0)
        verts_per_mesh = mesh.num_verts_per_mesh()
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
527
        for force, deform_shape in itertools.product([False, True], [(all_v, 3), 3]):
facebook-github-bot's avatar
facebook-github-bot committed
528
529
530
531
532
533
534
535
536
            if force:
                # force mesh to have computed attributes
                mesh._compute_packed(refresh=True)
                mesh._compute_padded()
                mesh._compute_edges_packed()
                mesh.verts_padded_to_packed_idx()
                mesh._compute_face_areas_normals(refresh=True)
                mesh._compute_vertex_normals(refresh=True)

537
            deform = torch.rand(deform_shape, dtype=torch.float32, device=mesh.device)
facebook-github-bot's avatar
facebook-github-bot committed
538
539
540
541
542
543
544
545
546
            # new meshes class to hold the deformed mesh
            new_mesh_naive = naive_offset_verts(mesh, deform)

            new_mesh = mesh.offset_verts(deform)

            # check verts_list & faces_list
            verts_cumsum = torch.cumsum(verts_per_mesh, 0).tolist()
            verts_cumsum.insert(0, 0)
            for i in range(N):
547
548
549
550
551
                item_offset = (
                    deform
                    if deform.ndim == 1
                    else deform[verts_cumsum[i] : verts_cumsum[i + 1]]
                )
facebook-github-bot's avatar
facebook-github-bot committed
552
553
                self.assertClose(
                    new_mesh.verts_list()[i],
554
                    mesh.verts_list()[i] + item_offset,
facebook-github-bot's avatar
facebook-github-bot committed
555
556
557
558
                )
                self.assertClose(
                    new_mesh.verts_list()[i], new_mesh_naive.verts_list()[i]
                )
559
                self.assertClose(mesh.faces_list()[i], new_mesh_naive.faces_list()[i])
facebook-github-bot's avatar
facebook-github-bot committed
560
561
562
                self.assertClose(
                    new_mesh.faces_list()[i], new_mesh_naive.faces_list()[i]
                )
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
563

facebook-github-bot's avatar
facebook-github-bot committed
564
565
566
567
                # check faces and vertex normals
                self.assertClose(
                    new_mesh.verts_normals_list()[i],
                    new_mesh_naive.verts_normals_list()[i],
568
                    atol=1e-6,
facebook-github-bot's avatar
facebook-github-bot committed
569
570
571
572
                )
                self.assertClose(
                    new_mesh.faces_normals_list()[i],
                    new_mesh_naive.faces_normals_list()[i],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
573
                    atol=1e-6,
facebook-github-bot's avatar
facebook-github-bot committed
574
575
576
                )

            # check padded & packed
577
578
579
580
581
            self.assertClose(new_mesh.faces_padded(), new_mesh_naive.faces_padded())
            self.assertClose(new_mesh.verts_padded(), new_mesh_naive.verts_padded())
            self.assertClose(new_mesh.faces_packed(), new_mesh_naive.faces_packed())
            self.assertClose(new_mesh.verts_packed(), new_mesh_naive.verts_packed())
            self.assertClose(new_mesh.edges_packed(), new_mesh_naive.edges_packed())
facebook-github-bot's avatar
facebook-github-bot committed
582
583
584
585
586
587
588
589
590
            self.assertClose(
                new_mesh.verts_packed_to_mesh_idx(),
                new_mesh_naive.verts_packed_to_mesh_idx(),
            )
            self.assertClose(
                new_mesh.mesh_to_verts_packed_first_idx(),
                new_mesh_naive.mesh_to_verts_packed_first_idx(),
            )
            self.assertClose(
591
                new_mesh.num_verts_per_mesh(), new_mesh_naive.num_verts_per_mesh()
facebook-github-bot's avatar
facebook-github-bot committed
592
593
594
595
596
597
598
599
600
601
            )
            self.assertClose(
                new_mesh.faces_packed_to_mesh_idx(),
                new_mesh_naive.faces_packed_to_mesh_idx(),
            )
            self.assertClose(
                new_mesh.mesh_to_faces_packed_first_idx(),
                new_mesh_naive.mesh_to_faces_packed_first_idx(),
            )
            self.assertClose(
602
                new_mesh.num_faces_per_mesh(), new_mesh_naive.num_faces_per_mesh()
facebook-github-bot's avatar
facebook-github-bot committed
603
604
605
606
607
608
609
610
611
612
613
614
615
616
            )
            self.assertClose(
                new_mesh.edges_packed_to_mesh_idx(),
                new_mesh_naive.edges_packed_to_mesh_idx(),
            )
            self.assertClose(
                new_mesh.verts_padded_to_packed_idx(),
                new_mesh_naive.verts_padded_to_packed_idx(),
            )
            self.assertTrue(all(new_mesh.valid == new_mesh_naive.valid))
            self.assertTrue(new_mesh.equisized == new_mesh_naive.equisized)

            # check face areas, normals and vertex normals
            self.assertClose(
617
618
619
                new_mesh.verts_normals_packed(),
                new_mesh_naive.verts_normals_packed(),
                atol=1e-6,
facebook-github-bot's avatar
facebook-github-bot committed
620
621
            )
            self.assertClose(
622
623
624
                new_mesh.verts_normals_padded(),
                new_mesh_naive.verts_normals_padded(),
                atol=1e-6,
facebook-github-bot's avatar
facebook-github-bot committed
625
626
            )
            self.assertClose(
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
627
628
629
                new_mesh.faces_normals_packed(),
                new_mesh_naive.faces_normals_packed(),
                atol=1e-6,
facebook-github-bot's avatar
facebook-github-bot committed
630
631
            )
            self.assertClose(
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
632
633
634
                new_mesh.faces_normals_padded(),
                new_mesh_naive.faces_normals_padded(),
                atol=1e-6,
facebook-github-bot's avatar
facebook-github-bot committed
635
636
            )
            self.assertClose(
637
                new_mesh.faces_areas_packed(), new_mesh_naive.faces_areas_packed()
facebook-github-bot's avatar
facebook-github-bot committed
638
            )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
639
640
641
642
            self.assertClose(
                new_mesh.mesh_to_edges_packed_first_idx(),
                new_mesh_naive.mesh_to_edges_packed_first_idx(),
            )
facebook-github-bot's avatar
facebook-github-bot committed
643
644
645
646
647
648
649
650
651
652
653
654
655
656

    def test_scale_verts(self):
        def naive_scale_verts(mesh, scale):
            if not torch.is_tensor(scale):
                scale = torch.ones(len(mesh)).mul_(scale)
            # new Meshes class
            new_verts_list = [
                scale[i] * v.clone() for (i, v) in enumerate(mesh.verts_list())
            ]
            new_faces_list = [f.clone() for f in mesh.faces_list()]
            return Meshes(verts=new_verts_list, faces=new_faces_list)

        N = 5
        for test in ["tensor", "scalar"]:
Georgia Gkioxari's avatar
Georgia Gkioxari committed
657
            for force in (False, True):
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
658
                mesh = init_mesh(N, 10, 100, lists_to_tensors=True)
facebook-github-bot's avatar
facebook-github-bot committed
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
                if force:
                    # force mesh to have computed attributes
                    mesh.verts_packed()
                    mesh.edges_packed()
                    mesh.verts_padded()
                    mesh._compute_face_areas_normals(refresh=True)
                    mesh._compute_vertex_normals(refresh=True)

                if test == "tensor":
                    scales = torch.rand(N)
                elif test == "scalar":
                    scales = torch.rand(1)[0].item()
                new_mesh_naive = naive_scale_verts(mesh, scales)
                new_mesh = mesh.scale_verts(scales)
                for i in range(N):
                    if test == "tensor":
                        self.assertClose(
676
                            scales[i] * mesh.verts_list()[i], new_mesh.verts_list()[i]
facebook-github-bot's avatar
facebook-github-bot committed
677
678
679
                        )
                    else:
                        self.assertClose(
680
                            scales * mesh.verts_list()[i], new_mesh.verts_list()[i]
facebook-github-bot's avatar
facebook-github-bot committed
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
                        )
                    self.assertClose(
                        new_mesh.verts_list()[i], new_mesh_naive.verts_list()[i]
                    )
                    self.assertClose(
                        mesh.faces_list()[i], new_mesh_naive.faces_list()[i]
                    )
                    self.assertClose(
                        new_mesh.faces_list()[i], new_mesh_naive.faces_list()[i]
                    )
                    # check face and vertex normals
                    self.assertClose(
                        new_mesh.verts_normals_list()[i],
                        new_mesh_naive.verts_normals_list()[i],
                    )
                    self.assertClose(
                        new_mesh.faces_normals_list()[i],
                        new_mesh_naive.faces_normals_list()[i],
                    )

                # check padded & packed
702
703
704
705
706
                self.assertClose(new_mesh.faces_padded(), new_mesh_naive.faces_padded())
                self.assertClose(new_mesh.verts_padded(), new_mesh_naive.verts_padded())
                self.assertClose(new_mesh.faces_packed(), new_mesh_naive.faces_packed())
                self.assertClose(new_mesh.verts_packed(), new_mesh_naive.verts_packed())
                self.assertClose(new_mesh.edges_packed(), new_mesh_naive.edges_packed())
facebook-github-bot's avatar
facebook-github-bot committed
707
708
709
710
711
712
713
714
715
                self.assertClose(
                    new_mesh.verts_packed_to_mesh_idx(),
                    new_mesh_naive.verts_packed_to_mesh_idx(),
                )
                self.assertClose(
                    new_mesh.mesh_to_verts_packed_first_idx(),
                    new_mesh_naive.mesh_to_verts_packed_first_idx(),
                )
                self.assertClose(
716
                    new_mesh.num_verts_per_mesh(), new_mesh_naive.num_verts_per_mesh()
facebook-github-bot's avatar
facebook-github-bot committed
717
718
719
720
721
722
723
724
725
726
                )
                self.assertClose(
                    new_mesh.faces_packed_to_mesh_idx(),
                    new_mesh_naive.faces_packed_to_mesh_idx(),
                )
                self.assertClose(
                    new_mesh.mesh_to_faces_packed_first_idx(),
                    new_mesh_naive.mesh_to_faces_packed_first_idx(),
                )
                self.assertClose(
727
                    new_mesh.num_faces_per_mesh(), new_mesh_naive.num_faces_per_mesh()
facebook-github-bot's avatar
facebook-github-bot committed
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
                )
                self.assertClose(
                    new_mesh.edges_packed_to_mesh_idx(),
                    new_mesh_naive.edges_packed_to_mesh_idx(),
                )
                self.assertClose(
                    new_mesh.verts_padded_to_packed_idx(),
                    new_mesh_naive.verts_padded_to_packed_idx(),
                )
                self.assertTrue(all(new_mesh.valid == new_mesh_naive.valid))
                self.assertTrue(new_mesh.equisized == new_mesh_naive.equisized)

                # check face areas, normals and vertex normals
                self.assertClose(
                    new_mesh.verts_normals_packed(),
                    new_mesh_naive.verts_normals_packed(),
                )
                self.assertClose(
                    new_mesh.verts_normals_padded(),
                    new_mesh_naive.verts_normals_padded(),
                )
                self.assertClose(
                    new_mesh.faces_normals_packed(),
                    new_mesh_naive.faces_normals_packed(),
                )
                self.assertClose(
                    new_mesh.faces_normals_padded(),
                    new_mesh_naive.faces_normals_padded(),
                )
                self.assertClose(
758
                    new_mesh.faces_areas_packed(), new_mesh_naive.faces_areas_packed()
facebook-github-bot's avatar
facebook-github-bot committed
759
                )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
760
761
762
763
                self.assertClose(
                    new_mesh.mesh_to_edges_packed_first_idx(),
                    new_mesh_naive.mesh_to_edges_packed_first_idx(),
                )
facebook-github-bot's avatar
facebook-github-bot committed
764
765
766

    def test_extend_list(self):
        N = 10
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
767
        mesh = init_mesh(5, 10, 100)
facebook-github-bot's avatar
facebook-github-bot committed
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
        for force in [0, 1]:
            if force:
                # force some computes to happen
                mesh._compute_packed(refresh=True)
                mesh._compute_padded()
                mesh._compute_edges_packed()
                mesh.verts_padded_to_packed_idx()
            new_mesh = mesh.extend(N)
            self.assertEqual(len(mesh) * 10, len(new_mesh))
            for i in range(len(mesh)):
                for n in range(N):
                    self.assertClose(
                        mesh.verts_list()[i], new_mesh.verts_list()[i * N + n]
                    )
                    self.assertClose(
                        mesh.faces_list()[i], new_mesh.faces_list()[i * N + n]
                    )
                    self.assertTrue(mesh.valid[i] == new_mesh.valid[i * N + n])
            self.assertAllSeparate(
                mesh.verts_list()
                + new_mesh.verts_list()
                + mesh.faces_list()
                + new_mesh.faces_list()
            )
            self.assertTrue(new_mesh._verts_packed is None)
            self.assertTrue(new_mesh._faces_packed is None)
            self.assertTrue(new_mesh._verts_padded is None)
            self.assertTrue(new_mesh._faces_padded is None)
            self.assertTrue(new_mesh._edges_packed is None)

        with self.assertRaises(ValueError):
            mesh.extend(N=-1)

    def test_to(self):
802
803
804
805
806
807
808
809
810
811
812
813
814
815
        mesh = init_mesh(5, 10, 100)

        cpu_device = torch.device("cpu")

        converted_mesh = mesh.to("cpu")
        self.assertEqual(cpu_device, converted_mesh.device)
        self.assertEqual(cpu_device, mesh.device)
        self.assertIs(mesh, converted_mesh)

        converted_mesh = mesh.to(cpu_device)
        self.assertEqual(cpu_device, converted_mesh.device)
        self.assertEqual(cpu_device, mesh.device)
        self.assertIs(mesh, converted_mesh)

816
        cuda_device = torch.device("cuda:0")
817

818
        converted_mesh = mesh.to("cuda:0")
819
820
821
        self.assertEqual(cuda_device, converted_mesh.device)
        self.assertEqual(cpu_device, mesh.device)
        self.assertIsNot(mesh, converted_mesh)
facebook-github-bot's avatar
facebook-github-bot committed
822

823
824
825
826
        converted_mesh = mesh.to(cuda_device)
        self.assertEqual(cuda_device, converted_mesh.device)
        self.assertEqual(cpu_device, mesh.device)
        self.assertIsNot(mesh, converted_mesh)
facebook-github-bot's avatar
facebook-github-bot committed
827
828

    def test_split_mesh(self):
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
829
        mesh = init_mesh(5, 10, 100)
facebook-github-bot's avatar
facebook-github-bot committed
830
831
832
833
834
        split_sizes = [2, 3]
        split_meshes = mesh.split(split_sizes)
        self.assertTrue(len(split_meshes[0]) == 2)
        self.assertTrue(
            split_meshes[0].verts_list()
835
            == [mesh.get_mesh_verts_faces(0)[0], mesh.get_mesh_verts_faces(1)[0]]
facebook-github-bot's avatar
facebook-github-bot committed
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
        )
        self.assertTrue(len(split_meshes[1]) == 3)
        self.assertTrue(
            split_meshes[1].verts_list()
            == [
                mesh.get_mesh_verts_faces(2)[0],
                mesh.get_mesh_verts_faces(3)[0],
                mesh.get_mesh_verts_faces(4)[0],
            ]
        )

        split_sizes = [2, 0.3]
        with self.assertRaises(ValueError):
            mesh.split(split_sizes)

Georgia Gkioxari's avatar
Georgia Gkioxari committed
851
852
853
854
855
    def test_update_padded(self):
        # Define the test mesh object either as a list or tensor of faces/verts.
        N = 10
        for lists_to_tensors in (False, True):
            for force in (True, False):
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
856
                mesh = init_mesh(N, 100, 300, lists_to_tensors=lists_to_tensors)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
                num_verts_per_mesh = mesh.num_verts_per_mesh()
                if force:
                    # force mesh to have computed attributes
                    mesh.verts_packed()
                    mesh.edges_packed()
                    mesh.laplacian_packed()
                    mesh.faces_areas_packed()

                new_verts = torch.rand((mesh._N, mesh._V, 3), device=mesh.device)
                new_verts_list = [
                    new_verts[i, : num_verts_per_mesh[i]] for i in range(N)
                ]
                new_mesh = mesh.update_padded(new_verts)

                # check the attributes assigned at construction time
                self.assertEqual(new_mesh._N, mesh._N)
                self.assertEqual(new_mesh._F, mesh._F)
                self.assertEqual(new_mesh._V, mesh._V)
                self.assertEqual(new_mesh.equisized, mesh.equisized)
                self.assertTrue(all(new_mesh.valid == mesh.valid))
                self.assertNotSeparate(
                    new_mesh.num_verts_per_mesh(), mesh.num_verts_per_mesh()
                )
                self.assertClose(
                    new_mesh.num_verts_per_mesh(), mesh.num_verts_per_mesh()
                )
                self.assertNotSeparate(
                    new_mesh.num_faces_per_mesh(), mesh.num_faces_per_mesh()
                )
                self.assertClose(
                    new_mesh.num_faces_per_mesh(), mesh.num_faces_per_mesh()
                )

                # check that the following attributes are not assigned
                self.assertIsNone(new_mesh._verts_list)
                self.assertIsNone(new_mesh._faces_areas_packed)
                self.assertIsNone(new_mesh._faces_normals_packed)
                self.assertIsNone(new_mesh._verts_normals_packed)

                check_tensors = [
                    "_faces_packed",
                    "_verts_packed_to_mesh_idx",
                    "_faces_packed_to_mesh_idx",
                    "_mesh_to_verts_packed_first_idx",
                    "_mesh_to_faces_packed_first_idx",
                    "_edges_packed",
                    "_edges_packed_to_mesh_idx",
                    "_mesh_to_edges_packed_first_idx",
                    "_faces_packed_to_edges_packed",
                    "_num_edges_per_mesh",
                ]
                for k in check_tensors:
                    v = getattr(new_mesh, k)
                    if not force:
                        self.assertIsNone(v)
                    else:
                        v_old = getattr(mesh, k)
                        self.assertNotSeparate(v, v_old)
                        self.assertClose(v, v_old)

                # check verts/faces padded
                self.assertClose(new_mesh.verts_padded(), new_verts)
                self.assertNotSeparate(new_mesh.verts_padded(), new_verts)
                self.assertClose(new_mesh.faces_padded(), mesh.faces_padded())
                self.assertNotSeparate(new_mesh.faces_padded(), mesh.faces_padded())
                # check verts/faces list
                for i in range(N):
                    self.assertNotSeparate(
                        new_mesh.faces_list()[i], mesh.faces_list()[i]
                    )
                    self.assertClose(new_mesh.faces_list()[i], mesh.faces_list()[i])
                    self.assertSeparate(new_mesh.verts_list()[i], mesh.verts_list()[i])
                    self.assertClose(new_mesh.verts_list()[i], new_verts_list[i])
                # check verts/faces packed
                self.assertClose(new_mesh.verts_packed(), torch.cat(new_verts_list))
                self.assertSeparate(new_mesh.verts_packed(), mesh.verts_packed())
                self.assertClose(new_mesh.faces_packed(), mesh.faces_packed())
                # check pad_to_packed
                self.assertClose(
                    new_mesh.verts_padded_to_packed_idx(),
                    mesh.verts_padded_to_packed_idx(),
                )
                # check edges
                self.assertClose(new_mesh.edges_packed(), mesh.edges_packed())

facebook-github-bot's avatar
facebook-github-bot committed
942
943
944
945
946
947
948
    def test_get_mesh_verts_faces(self):
        device = torch.device("cuda:0")
        verts_list = []
        faces_list = []
        verts_faces = [(10, 100), (20, 200)]
        for (V, F) in verts_faces:
            verts = torch.rand((V, 3), dtype=torch.float32, device=device)
949
            faces = torch.randint(V, size=(F, 3), dtype=torch.int64, device=device)
facebook-github-bot's avatar
facebook-github-bot committed
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
            verts_list.append(verts)
            faces_list.append(faces)

        mesh = Meshes(verts=verts_list, faces=faces_list)

        for i, (V, F) in enumerate(verts_faces):
            verts, faces = mesh.get_mesh_verts_faces(i)
            self.assertTrue(len(verts) == V)
            self.assertClose(verts, verts_list[i])
            self.assertTrue(len(faces) == F)
            self.assertClose(faces, faces_list[i])

        with self.assertRaises(ValueError):
            mesh.get_mesh_verts_faces(5)
        with self.assertRaises(ValueError):
            mesh.get_mesh_verts_faces(0.2)

    def test_get_bounding_boxes(self):
        device = torch.device("cuda:0")
        verts_list = []
        faces_list = []
        for (V, F) in [(10, 100)]:
            verts = torch.rand((V, 3), dtype=torch.float32, device=device)
973
            faces = torch.randint(V, size=(F, 3), dtype=torch.int64, device=device)
facebook-github-bot's avatar
facebook-github-bot committed
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
            verts_list.append(verts)
            faces_list.append(faces)

        mins = torch.min(verts, dim=0)[0]
        maxs = torch.max(verts, dim=0)[0]
        bboxes_gt = torch.stack([mins, maxs], dim=1).unsqueeze(0)
        mesh = Meshes(verts=verts_list, faces=faces_list)
        bboxes = mesh.get_bounding_boxes()
        self.assertClose(bboxes_gt, bboxes)

    def test_padded_to_packed_idx(self):
        device = torch.device("cuda:0")
        verts_list = []
        faces_list = []
        verts_faces = [(10, 100), (20, 200), (30, 300)]
        for (V, F) in verts_faces:
            verts = torch.rand((V, 3), dtype=torch.float32, device=device)
991
            faces = torch.randint(V, size=(F, 3), dtype=torch.int64, device=device)
facebook-github-bot's avatar
facebook-github-bot committed
992
993
994
995
996
997
998
999
1000
            verts_list.append(verts)
            faces_list.append(faces)

        mesh = Meshes(verts=verts_list, faces=faces_list)
        verts_padded_to_packed_idx = mesh.verts_padded_to_packed_idx()
        verts_packed = mesh.verts_packed()
        verts_padded = mesh.verts_padded()
        verts_padded_flat = verts_padded.view(-1, 3)

1001
        self.assertClose(verts_padded_flat[verts_padded_to_packed_idx], verts_packed)
facebook-github-bot's avatar
facebook-github-bot committed
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012

        idx = verts_padded_to_packed_idx.view(-1, 1).expand(-1, 3)
        self.assertClose(verts_padded_flat.gather(0, idx), verts_packed)

    def test_getitem(self):
        device = torch.device("cuda:0")
        verts_list = []
        faces_list = []
        verts_faces = [(10, 100), (20, 200), (30, 300)]
        for (V, F) in verts_faces:
            verts = torch.rand((V, 3), dtype=torch.float32, device=device)
1013
            faces = torch.randint(V, size=(F, 3), dtype=torch.int64, device=device)
facebook-github-bot's avatar
facebook-github-bot committed
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
            verts_list.append(verts)
            faces_list.append(faces)

        mesh = Meshes(verts=verts_list, faces=faces_list)

        def check_equal(selected, indices):
            for selectedIdx, index in enumerate(indices):
                self.assertClose(
                    selected.verts_list()[selectedIdx], mesh.verts_list()[index]
                )
                self.assertClose(
                    selected.faces_list()[selectedIdx], mesh.faces_list()[index]
                )

        # int index
        index = 1
        mesh_selected = mesh[index]
        self.assertTrue(len(mesh_selected) == 1)
        check_equal(mesh_selected, [index])

        # list index
        index = [1, 2]
        mesh_selected = mesh[index]
        self.assertTrue(len(mesh_selected) == len(index))
        check_equal(mesh_selected, index)

        # slice index
        index = slice(0, 2, 1)
        mesh_selected = mesh[index]
        check_equal(mesh_selected, [0, 1])

        # bool tensor
        index = torch.tensor([1, 0, 1], dtype=torch.bool, device=device)
        mesh_selected = mesh[index]
        self.assertTrue(len(mesh_selected) == index.sum())
        check_equal(mesh_selected, [0, 2])

        # int tensor
        index = torch.tensor([1, 2], dtype=torch.int64, device=device)
        mesh_selected = mesh[index]
        self.assertTrue(len(mesh_selected) == index.numel())
        check_equal(mesh_selected, index.tolist())

        # invalid index
        index = torch.tensor([1, 0, 1], dtype=torch.float32, device=device)
        with self.assertRaises(IndexError):
            mesh_selected = mesh[index]
        index = 1.2
        with self.assertRaises(IndexError):
            mesh_selected = mesh[index]

    def test_compute_faces_areas(self):
        verts = torch.tensor(
            [
                [0.0, 0.0, 0.0],
                [0.5, 0.0, 0.0],
                [0.5, 0.5, 0.0],
                [0.5, 0.0, 0.0],
                [0.25, 0.8, 0.0],
            ],
            dtype=torch.float32,
        )
        faces = torch.tensor([[0, 1, 2], [0, 3, 4]], dtype=torch.int64)
        mesh = Meshes(verts=[verts], faces=[faces])

        face_areas = mesh.faces_areas_packed()
        expected_areas = torch.tensor([0.125, 0.2])
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
1081
        self.assertClose(face_areas, expected_areas)
facebook-github-bot's avatar
facebook-github-bot committed
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106

    def test_compute_normals(self):

        # Simple case with one mesh where normals point in either +/- ijk
        verts = torch.tensor(
            [
                [0.1, 0.3, 0.0],
                [0.5, 0.2, 0.0],
                [0.6, 0.8, 0.0],
                [0.0, 0.3, 0.2],
                [0.0, 0.2, 0.5],
                [0.0, 0.8, 0.7],
                [0.5, 0.0, 0.2],
                [0.6, 0.0, 0.5],
                [0.8, 0.0, 0.7],
                [0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0],
            ],
            dtype=torch.float32,
        )
        faces = torch.tensor(
            [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10, 11]], dtype=torch.int64
        )
        mesh = Meshes(verts=[verts], faces=[faces])
1107
        self.assertFalse(mesh.has_verts_normals())
facebook-github-bot's avatar
facebook-github-bot committed
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
        verts_normals_expected = torch.tensor(
            [
                [0.0, 0.0, 1.0],
                [0.0, 0.0, 1.0],
                [0.0, 0.0, 1.0],
                [-1.0, 0.0, 0.0],
                [-1.0, 0.0, 0.0],
                [-1.0, 0.0, 0.0],
                [0.0, 1.0, 0.0],
                [0.0, 1.0, 0.0],
                [0.0, 1.0, 0.0],
                [0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0],
                [0.0, 0.0, 0.0],
            ]
        )
        faces_normals_expected = verts_normals_expected[[0, 3, 6, 9], :]

        self.assertTrue(
            torch.allclose(mesh.verts_normals_list()[0], verts_normals_expected)
        )
1129
        self.assertTrue(mesh.has_verts_normals())
facebook-github-bot's avatar
facebook-github-bot committed
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
        self.assertTrue(
            torch.allclose(mesh.faces_normals_list()[0], faces_normals_expected)
        )
        self.assertTrue(
            torch.allclose(mesh.verts_normals_packed(), verts_normals_expected)
        )
        self.assertTrue(
            torch.allclose(mesh.faces_normals_packed(), faces_normals_expected)
        )

        # Multiple meshes in the batch with equal sized meshes
        meshes_extended = mesh.extend(3)
        for m in meshes_extended.verts_normals_list():
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
1143
            self.assertClose(m, verts_normals_expected)
facebook-github-bot's avatar
facebook-github-bot committed
1144
        for f in meshes_extended.faces_normals_list():
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
1145
            self.assertClose(f, faces_normals_expected)
facebook-github-bot's avatar
facebook-github-bot committed
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189

        # Multiple meshes in the batch with different sized meshes
        # Check padded and packed normals are the correct sizes.
        verts2 = torch.tensor(
            [
                [0.1, 0.3, 0.0],
                [0.5, 0.2, 0.0],
                [0.6, 0.8, 0.0],
                [0.0, 0.3, 0.2],
                [0.0, 0.2, 0.5],
                [0.0, 0.8, 0.7],
            ],
            dtype=torch.float32,
        )
        faces2 = torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int64)
        verts_list = [verts, verts2]
        faces_list = [faces, faces2]
        meshes = Meshes(verts=verts_list, faces=faces_list)
        verts_normals_padded = meshes.verts_normals_padded()
        faces_normals_padded = meshes.faces_normals_padded()

        for n in range(len(meshes)):
            v = verts_list[n].shape[0]
            f = faces_list[n].shape[0]
            if verts_normals_padded.shape[1] > v:
                self.assertTrue(verts_normals_padded[n, v:, :].eq(0).all())
                self.assertTrue(
                    torch.allclose(
                        verts_normals_padded[n, :v, :].view(-1, 3),
                        verts_normals_expected[:v, :],
                    )
                )
            if faces_normals_padded.shape[1] > f:
                self.assertTrue(faces_normals_padded[n, f:, :].eq(0).all())
                self.assertTrue(
                    torch.allclose(
                        faces_normals_padded[n, :f, :].view(-1, 3),
                        faces_normals_expected[:f, :],
                    )
                )

        verts_normals_packed = meshes.verts_normals_packed()
        faces_normals_packed = meshes.faces_normals_packed()
        self.assertTrue(
1190
            list(verts_normals_packed.shape) == [verts.shape[0] + verts2.shape[0], 3]
facebook-github-bot's avatar
facebook-github-bot committed
1191
1192
        )
        self.assertTrue(
1193
            list(faces_normals_packed.shape) == [faces.shape[0] + faces2.shape[0], 3]
facebook-github-bot's avatar
facebook-github-bot committed
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
        )

        # Single mesh where two faces share one vertex so the normal is
        # the weighted sum of the two face normals.
        verts = torch.tensor(
            [
                [0.1, 0.3, 0.0],
                [0.5, 0.2, 0.0],
                [0.0, 0.3, 0.2],  # vertex is shared between two faces
                [0.0, 0.2, 0.5],
                [0.0, 0.8, 0.7],
            ],
            dtype=torch.float32,
        )
        faces = torch.tensor([[0, 1, 2], [2, 3, 4]], dtype=torch.int64)
        mesh = Meshes(verts=[verts], faces=[faces])

        verts_normals_expected = torch.tensor(
            [
                [-0.2408, -0.9631, -0.1204],
                [-0.2408, -0.9631, -0.1204],
                [-0.9389, -0.3414, -0.0427],
                [-1.0000, 0.0000, 0.0000],
                [-1.0000, 0.0000, 0.0000],
            ]
        )
        faces_normals_expected = torch.tensor(
            [[-0.2408, -0.9631, -0.1204], [-1.0000, 0.0000, 0.0000]]
        )
        self.assertTrue(
            torch.allclose(
                mesh.verts_normals_list()[0], verts_normals_expected, atol=4e-5
            )
        )
        self.assertTrue(
            torch.allclose(
                mesh.faces_normals_list()[0], faces_normals_expected, atol=4e-5
            )
        )

        # Check empty mesh has empty normals
        meshes = Meshes(verts=[], faces=[])
        self.assertEqual(meshes.verts_normals_packed().shape[0], 0)
        self.assertEqual(meshes.verts_normals_padded().shape[0], 0)
        self.assertEqual(meshes.verts_normals_list(), [])
        self.assertEqual(meshes.faces_normals_packed().shape[0], 0)
        self.assertEqual(meshes.faces_normals_padded().shape[0], 0)
        self.assertEqual(meshes.faces_normals_list(), [])

1243
1244
1245
    def test_assigned_normals(self):
        verts = torch.rand(2, 6, 3)
        faces = torch.randint(6, size=(2, 4, 3))
1246
1247
        no_normals = Meshes(verts=verts, faces=faces)
        self.assertFalse(no_normals.has_verts_normals())
1248
1249
1250
1251
1252

        for verts_normals in [list(verts.unbind(0)), verts]:
            yes_normals = Meshes(
                verts=verts.clone(), faces=faces, verts_normals=verts_normals
            )
1253
            self.assertTrue(yes_normals.has_verts_normals())
1254
1255
1256
1257
1258
1259
            self.assertClose(yes_normals.verts_normals_padded(), verts)
            yes_normals.offset_verts_(torch.FloatTensor([1, 2, 3]))
            self.assertClose(yes_normals.verts_normals_padded(), verts)
            yes_normals.offset_verts_(torch.FloatTensor([1, 2, 3]).expand(12, 3))
            self.assertFalse(torch.allclose(yes_normals.verts_normals_padded(), verts))

facebook-github-bot's avatar
facebook-github-bot committed
1260
1261
1262
1263
    def test_compute_faces_areas_cpu_cuda(self):
        num_meshes = 10
        max_v = 100
        max_f = 300
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
1264
        mesh_cpu = init_mesh(num_meshes, max_v, max_f, device="cpu")
facebook-github-bot's avatar
facebook-github-bot committed
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
        device = torch.device("cuda:0")
        mesh_cuda = mesh_cpu.to(device)

        face_areas_cpu = mesh_cpu.faces_areas_packed()
        face_normals_cpu = mesh_cpu.faces_normals_packed()
        face_areas_cuda = mesh_cuda.faces_areas_packed()
        face_normals_cuda = mesh_cuda.faces_normals_packed()
        self.assertClose(face_areas_cpu, face_areas_cuda.cpu(), atol=1e-6)
        # because of the normalization of the normals with arbitrarily small values,
        # normals can become unstable. Thus only compare normals, for faces
        # with areas > eps=1e-6
        nonzero = face_areas_cpu > 1e-6
        self.assertClose(
1278
            face_normals_cpu[nonzero], face_normals_cuda.cpu()[nonzero], atol=1e-6
facebook-github-bot's avatar
facebook-github-bot committed
1279
1280
        )

1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
    def test_equality(self):
        meshes1 = init_mesh(num_meshes=2)
        meshes2 = init_mesh(num_meshes=2)
        meshes3 = init_mesh(num_meshes=3)
        empty_mesh = Meshes([], [])
        self.assertTrue(mesh_structures_equal(empty_mesh, Meshes([], [])))
        self.assertTrue(mesh_structures_equal(meshes1, meshes1))
        self.assertTrue(mesh_structures_equal(meshes1, meshes1.clone()))
        self.assertFalse(mesh_structures_equal(empty_mesh, meshes1))
        self.assertFalse(mesh_structures_equal(meshes1, meshes2))
        self.assertFalse(mesh_structures_equal(meshes1, meshes3))

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
    def test_to_sorted(self):
        mesh = init_simple_mesh()
        sorted_mesh = to_sorted(mesh)

        expected_verts = [
            torch.tensor(
                [[0.1, 0.3, 0.5], [0.5, 0.2, 0.1], [0.6, 0.8, 0.7]],
                dtype=torch.float32,
            ),
            torch.tensor(
                # Vertex permutation: 0->0, 1->3, 2->2, 3->1
                [[0.1, 0.3, 0.3], [0.1, 0.5, 0.3], [0.2, 0.3, 0.4], [0.6, 0.7, 0.8]],
                dtype=torch.float32,
            ),
            torch.tensor(
                # Vertex permutation: 0->2, 1->1, 2->4, 3->0, 4->3
                [
                    [0.2, 0.3, 0.4],
                    [0.2, 0.4, 0.8],
                    [0.7, 0.3, 0.6],
                    [0.9, 0.3, 0.8],
                    [0.9, 0.5, 0.2],
                ],
                dtype=torch.float32,
            ),
        ]

        expected_faces = [
            torch.tensor([[0, 1, 2]], dtype=torch.int64),
            torch.tensor([[0, 3, 2], [3, 2, 1]], dtype=torch.int64),
            torch.tensor(
                [
                    [1, 4, 2],
                    [2, 1, 0],
                    [4, 0, 1],
                    [3, 0, 4],
                    [3, 2, 1],
                    [3, 0, 1],
                    [3, 4, 1],
                ],
                dtype=torch.int64,
            ),
        ]

        self.assertFalse(mesh_structures_equal(mesh, sorted_mesh))
        self.assertTrue(
            mesh_structures_equal(
                Meshes(verts=expected_verts, faces=expected_faces), sorted_mesh
            )
        )

facebook-github-bot's avatar
facebook-github-bot committed
1344
1345
    @staticmethod
    def compute_packed_with_init(
1346
        num_meshes: int = 10, max_v: int = 100, max_f: int = 300, device: str = "cpu"
facebook-github-bot's avatar
facebook-github-bot committed
1347
    ):
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
1348
        mesh = init_mesh(num_meshes, max_v, max_f, device=device)
facebook-github-bot's avatar
facebook-github-bot committed
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
        torch.cuda.synchronize()

        def compute_packed():
            mesh._compute_packed(refresh=True)
            torch.cuda.synchronize()

        return compute_packed

    @staticmethod
    def compute_padded_with_init(
1359
        num_meshes: int = 10, max_v: int = 100, max_f: int = 300, device: str = "cpu"
facebook-github-bot's avatar
facebook-github-bot committed
1360
    ):
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
1361
        mesh = init_mesh(num_meshes, max_v, max_f, device=device)
facebook-github-bot's avatar
facebook-github-bot committed
1362
1363
1364
1365
1366
1367
1368
        torch.cuda.synchronize()

        def compute_padded():
            mesh._compute_padded(refresh=True)
            torch.cuda.synchronize()

        return compute_padded