test_point_mesh_distance.py 34.9 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.
Georgia Gkioxari's avatar
Georgia Gkioxari committed
6
7
8
9
10

import unittest

import numpy as np
import torch
11
from common_testing import get_random_cuda_device, TestCaseMixin
Georgia Gkioxari's avatar
Georgia Gkioxari committed
12
13
from pytorch3d import _C
from pytorch3d.loss import point_mesh_edge_distance, point_mesh_face_distance
14
from pytorch3d.structures import Meshes, packed_to_list, Pointclouds
Georgia Gkioxari's avatar
Georgia Gkioxari committed
15
16
17
18
19
20
21
22
23
24
25


class TestPointMeshDistance(TestCaseMixin, unittest.TestCase):
    def setUp(self) -> None:
        np.random.seed(42)
        torch.manual_seed(42)

    @staticmethod
    def eps():
        return 1e-8

26
27
28
29
    @staticmethod
    def min_triangle_area():
        return 5e-3

Georgia Gkioxari's avatar
Georgia Gkioxari committed
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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
    @staticmethod
    def init_meshes_clouds(
        batch_size: int = 10,
        num_verts: int = 1000,
        num_faces: int = 3000,
        num_points: int = 3000,
        device: str = "cuda:0",
    ):
        device = torch.device(device)
        nump = torch.randint(low=1, high=num_points, size=(batch_size,))
        numv = torch.randint(low=3, high=num_verts, size=(batch_size,))
        numf = torch.randint(low=1, high=num_faces, size=(batch_size,))
        verts_list = []
        faces_list = []
        points_list = []
        for i in range(batch_size):
            # Randomly choose vertices
            verts = torch.rand((numv[i], 3), dtype=torch.float32, device=device)
            verts.requires_grad_(True)

            # Randomly choose faces. Our tests below compare argmin indices
            # over faces and edges. Argmin is sensitive even to small numeral variations
            # thus we make sure that faces are valid
            # i.e. a face f = (i0, i1, i2) s.t. i0 != i1 != i2,
            # otherwise argmin due to numeral sensitivities cannot be resolved
            faces, allf = [], 0
            validf = numv[i].item() - numv[i].item() % 3
            while allf < numf[i]:
                ff = torch.randperm(numv[i], device=device)[:validf].view(-1, 3)
                faces.append(ff)
                allf += ff.shape[0]
            faces = torch.cat(faces, 0)
            if faces.shape[0] > numf[i]:
                faces = faces[: numf[i]]

            verts_list.append(verts)
            faces_list.append(faces)

            # Randomly choose points
            points = torch.rand((nump[i], 3), dtype=torch.float32, device=device)
            points.requires_grad_(True)

            points_list.append(points)

        meshes = Meshes(verts_list, faces_list)
        pcls = Pointclouds(points_list)

        return meshes, pcls

    @staticmethod
    def _point_to_bary(point: torch.Tensor, tri: torch.Tensor) -> torch.Tensor:
        """
        Computes the barycentric coordinates of point wrt triangle (tri)
        Note that point needs to live in the space spanned by tri = (a, b, c),
        i.e. by taking the projection of an arbitrary point on the space spanned by tri

        Args:
            point: FloatTensor of shape (3)
            tri: FloatTensor of shape (3, 3)
        Returns:
            bary: FloatTensor of shape (3)
        """
        assert point.dim() == 1 and point.shape[0] == 3
        assert tri.dim() == 2 and tri.shape[0] == 3 and tri.shape[1] == 3

        a, b, c = tri.unbind(0)

        v0 = b - a
        v1 = c - a
        v2 = point - a

        d00 = v0.dot(v0)
        d01 = v0.dot(v1)
        d11 = v1.dot(v1)
        d20 = v2.dot(v0)
        d21 = v2.dot(v1)

107
        denom = d00 * d11 - d01 * d01 + TestPointMeshDistance.eps()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
        s2 = (d11 * d20 - d01 * d21) / denom
        s3 = (d00 * d21 - d01 * d20) / denom
        s1 = 1.0 - s2 - s3

        bary = torch.tensor([s1, s2, s3])
        return bary

    @staticmethod
    def _is_inside_triangle(point: torch.Tensor, tri: torch.Tensor) -> torch.Tensor:
        """
        Computes whether point is inside triangle tri
        Note that point needs to live in the space spanned by tri = (a, b, c)
        i.e. by taking the projection of an arbitrary point on the space spanned by tri

        Args:
            point: FloatTensor of shape (3)
            tri: FloatTensor of shape (3, 3)
        Returns:
            inside: BoolTensor of shape (1)
        """
128
129
130
131
132
        v0 = tri[1] - tri[0]
        v1 = tri[2] - tri[0]
        area = torch.cross(v0, v1).norm() / 2.0

        # check if triangle is a line or a point. In that case, return False
133
        if area < 5e-3:
134
            return False
Georgia Gkioxari's avatar
Georgia Gkioxari committed
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
        bary = TestPointMeshDistance._point_to_bary(point, tri)
        inside = ((bary >= 0.0) * (bary <= 1.0)).all()
        return inside

    @staticmethod
    def _point_to_edge_distance(
        point: torch.Tensor, edge: torch.Tensor
    ) -> torch.Tensor:
        """
        Computes the squared euclidean distance of points to edges
        Args:
            point: FloatTensor of shape (3)
            edge: FloatTensor of shape (2, 3)
        Returns:
            dist: FloatTensor of shape (1)
Nikhila Ravi's avatar
Nikhila Ravi committed
150

Georgia Gkioxari's avatar
Georgia Gkioxari committed
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
        If a, b are the start and end points of the segments, we
        parametrize a point p as
            x(t) = a + t * (b - a)
        To find t which describes p we minimize (x(t) - p) ^ 2
        Note that p does not need to live in the space spanned by (a, b)
        """
        s0, s1 = edge.unbind(0)

        s01 = s1 - s0
        norm_s01 = s01.dot(s01)

        same_edge = norm_s01 < TestPointMeshDistance.eps()
        if same_edge:
            dist = 0.5 * (point - s0).dot(point - s0) + 0.5 * (point - s1).dot(
                point - s1
            )
            return dist

        t = s01.dot(point - s0) / norm_s01
        t = torch.clamp(t, min=0.0, max=1.0)
        x = s0 + t * s01
        dist = (x - point).dot(x - point)
        return dist

    @staticmethod
    def _point_to_tri_distance(point: torch.Tensor, tri: torch.Tensor) -> torch.Tensor:
        """
        Computes the squared euclidean distance of points to edges
        Args:
            point: FloatTensor of shape (3)
            tri: FloatTensor of shape (3, 3)
        Returns:
Nikhila Ravi's avatar
Nikhila Ravi committed
183
            dist: FloatTensor of shape (1)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
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
        """
        a, b, c = tri.unbind(0)
        cross = torch.cross(b - a, c - a)
        norm = cross.norm()
        normal = torch.nn.functional.normalize(cross, dim=0)

        # p0 is the projection of p onto the plane spanned by (a, b, c)
        # p0 = p + tt * normal, s.t. (p0 - a) is orthogonal to normal
        # => tt = dot(a - p, n)
        tt = normal.dot(a) - normal.dot(point)
        p0 = point + tt * normal
        dist_p = tt * tt

        # Compute the distance of p to all edge segments
        e01_dist = TestPointMeshDistance._point_to_edge_distance(point, tri[[0, 1]])
        e02_dist = TestPointMeshDistance._point_to_edge_distance(point, tri[[0, 2]])
        e12_dist = TestPointMeshDistance._point_to_edge_distance(point, tri[[1, 2]])

        with torch.no_grad():
            inside_tri = TestPointMeshDistance._is_inside_triangle(p0, tri)

        if inside_tri and (norm > TestPointMeshDistance.eps()):
            return dist_p
        else:
            if e01_dist.le(e02_dist) and e01_dist.le(e12_dist):
                return e01_dist
            elif e02_dist.le(e01_dist) and e02_dist.le(e12_dist):
                return e02_dist
            else:
                return e12_dist

    def test_point_edge_array_distance(self):
        """
        Test CUDA implementation for PointEdgeArrayDistanceForward
            &  PointEdgeArrayDistanceBackward
        """
        P, E = 16, 32
Nikhila Ravi's avatar
Nikhila Ravi committed
221
        device = get_random_cuda_device()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
222
223
224
225
226
227
228
        points = torch.rand((P, 3), dtype=torch.float32, device=device)
        edges = torch.rand((E, 2, 3), dtype=torch.float32, device=device)

        # randomly make some edge points equal
        same = torch.rand((E,), dtype=torch.float32, device=device) > 0.5
        edges[same, 1] = edges[same, 0].clone().detach()

229
230
231
        points_cpu = points.clone().cpu()
        edges_cpu = edges.clone().cpu()

Georgia Gkioxari's avatar
Georgia Gkioxari committed
232
233
234
235
236
237
238
239
240
241
242
243
244
        points.requires_grad = True
        edges.requires_grad = True
        grad_dists = torch.rand((P, E), dtype=torch.float32, device=device)

        # Naive python implementation
        dists_naive = torch.zeros((P, E), dtype=torch.float32, device=device)
        for p in range(P):
            for e in range(E):
                dist = self._point_to_edge_distance(points[p], edges[e])
                dists_naive[p, e] = dist

        # Cuda Forward Implementation
        dists_cuda = _C.point_edge_array_dist_forward(points, edges)
245
        dists_cpu = _C.point_edge_array_dist_forward(points_cpu, edges_cpu)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
246
247
248

        # Compare
        self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
249
        self.assertClose(dists_naive.cpu(), dists_cpu)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
250
251
252
253
254

        # CUDA Bacwkard Implementation
        grad_points_cuda, grad_edges_cuda = _C.point_edge_array_dist_backward(
            points, edges, grad_dists
        )
255
256
257
        grad_points_cpu, grad_edges_cpu = _C.point_edge_array_dist_backward(
            points_cpu, edges_cpu, grad_dists.cpu()
        )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
258
259

        dists_naive.backward(grad_dists)
260
261
        grad_points_naive = points.grad.cpu()
        grad_edges_naive = edges.grad.cpu()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
262
263

        # Compare
264
265
266
267
        self.assertClose(grad_points_naive, grad_points_cuda.cpu())
        self.assertClose(grad_edges_naive, grad_edges_cuda.cpu())
        self.assertClose(grad_points_naive, grad_points_cpu)
        self.assertClose(grad_edges_naive, grad_edges_cpu)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
268
269
270
271
272
273

    def test_point_edge_distance(self):
        """
        Test CUDA implementation for PointEdgeDistanceForward
            &  PointEdgeDistanceBackward
        """
Nikhila Ravi's avatar
Nikhila Ravi committed
274
        device = get_random_cuda_device()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
275
        N, V, F, P = 4, 32, 16, 24
Nikhila Ravi's avatar
Nikhila Ravi committed
276
        meshes, pcls = self.init_meshes_clouds(N, V, F, P, device=device)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297

        # make points packed a leaf node
        points_packed = pcls.points_packed().detach().clone()  # (P, 3)

        points_first_idx = pcls.cloud_to_packed_first_idx()
        max_p = pcls.num_points_per_cloud().max().item()

        # make edges packed a leaf node
        verts_packed = meshes.verts_packed()
        edges_packed = verts_packed[meshes.edges_packed()]  # (E, 2, 3)
        edges_packed = edges_packed.clone().detach()

        edges_first_idx = meshes.mesh_to_edges_packed_first_idx()

        # leaf nodes
        points_packed.requires_grad = True
        edges_packed.requires_grad = True
        grad_dists = torch.rand(
            (points_packed.shape[0],), dtype=torch.float32, device=device
        )

298
        # Cuda Implementation: forward
Georgia Gkioxari's avatar
Georgia Gkioxari committed
299
300
301
302
303
304
305
        dists_cuda, idx_cuda = _C.point_edge_dist_forward(
            points_packed, points_first_idx, edges_packed, edges_first_idx, max_p
        )
        # Cuda Implementation: backward
        grad_points_cuda, grad_edges_cuda = _C.point_edge_dist_backward(
            points_packed, edges_packed, idx_cuda, grad_dists
        )
306
307
308
309
310
311
312
313
314
315
316
317
318
319
        # Cpu Implementation: forward
        dists_cpu, idx_cpu = _C.point_edge_dist_forward(
            points_packed.cpu(),
            points_first_idx.cpu(),
            edges_packed.cpu(),
            edges_first_idx.cpu(),
            max_p,
        )

        # Cpu Implementation: backward
        # Note that using idx_cpu doesn't pass - there seems to be a problem with tied results.
        grad_points_cpu, grad_edges_cpu = _C.point_edge_dist_backward(
            points_packed.cpu(), edges_packed.cpu(), idx_cuda.cpu(), grad_dists.cpu()
        )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353

        # Naive Implementation: forward
        edges_list = packed_to_list(edges_packed, meshes.num_edges_per_mesh().tolist())
        dists_naive = []
        for i in range(N):
            points = pcls.points_list()[i]
            edges = edges_list[i]
            dists_temp = torch.zeros(
                (points.shape[0], edges.shape[0]), dtype=torch.float32, device=device
            )
            for p in range(points.shape[0]):
                for e in range(edges.shape[0]):
                    dist = self._point_to_edge_distance(points[p], edges[e])
                    dists_temp[p, e] = dist
            # torch.min() doesn't necessarily return the first index of the
            # smallest value, our warp_reduce does. So it's not straightforward
            # to directly compare indices, nor the gradients of grad_edges which
            # also depend on the indices of the minimum value.
            # To be able to compare, we will compare dists_temp.min(1) and
            # then feed the cuda indices to the naive output

            start = points_first_idx[i]
            end = points_first_idx[i + 1] if i < N - 1 else points_packed.shape[0]

            min_idx = idx_cuda[start:end] - edges_first_idx[i]
            iidx = torch.arange(points.shape[0], device=device)
            min_dist = dists_temp[iidx, min_idx]

            dists_naive.append(min_dist)

        dists_naive = torch.cat(dists_naive)

        # Compare
        self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
354
        self.assertClose(dists_naive.cpu(), dists_cpu)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
355
356
357
358

        # Naive Implementation: backward
        dists_naive.backward(grad_dists)
        grad_points_naive = torch.cat([cloud.grad for cloud in pcls.points_list()])
359
        grad_edges_naive = edges_packed.grad.cpu()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
360
361
362

        # Compare
        self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7)
363
364
365
        self.assertClose(grad_edges_naive, grad_edges_cuda.cpu(), atol=5e-7)
        self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7)
        self.assertClose(grad_edges_naive, grad_edges_cpu, atol=5e-7)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
366
367
368
369
370
371

    def test_edge_point_distance(self):
        """
        Test CUDA implementation for EdgePointDistanceForward
            &  EdgePointDistanceBackward
        """
Nikhila Ravi's avatar
Nikhila Ravi committed
372
        device = get_random_cuda_device()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
373
        N, V, F, P = 4, 32, 16, 24
Nikhila Ravi's avatar
Nikhila Ravi committed
374
        meshes, pcls = self.init_meshes_clouds(N, V, F, P, device=device)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405

        # make points packed a leaf node
        points_packed = pcls.points_packed().detach().clone()  # (P, 3)

        points_first_idx = pcls.cloud_to_packed_first_idx()

        # make edges packed a leaf node
        verts_packed = meshes.verts_packed()
        edges_packed = verts_packed[meshes.edges_packed()]  # (E, 2, 3)
        edges_packed = edges_packed.clone().detach()

        edges_first_idx = meshes.mesh_to_edges_packed_first_idx()
        max_e = meshes.num_edges_per_mesh().max().item()

        # leaf nodes
        points_packed.requires_grad = True
        edges_packed.requires_grad = True
        grad_dists = torch.rand(
            (edges_packed.shape[0],), dtype=torch.float32, device=device
        )

        # Cuda Implementation: forward
        dists_cuda, idx_cuda = _C.edge_point_dist_forward(
            points_packed, points_first_idx, edges_packed, edges_first_idx, max_e
        )

        # Cuda Implementation: backward
        grad_points_cuda, grad_edges_cuda = _C.edge_point_dist_backward(
            points_packed, edges_packed, idx_cuda, grad_dists
        )

406
407
408
409
410
411
412
413
414
415
416
417
418
419
        # Cpu Implementation: forward
        dists_cpu, idx_cpu = _C.edge_point_dist_forward(
            points_packed.cpu(),
            points_first_idx.cpu(),
            edges_packed.cpu(),
            edges_first_idx.cpu(),
            max_e,
        )

        # Cpu Implementation: backward
        grad_points_cpu, grad_edges_cpu = _C.edge_point_dist_backward(
            points_packed.cpu(), edges_packed.cpu(), idx_cpu, grad_dists.cpu()
        )

Georgia Gkioxari's avatar
Georgia Gkioxari committed
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
        # Naive Implementation: forward
        edges_list = packed_to_list(edges_packed, meshes.num_edges_per_mesh().tolist())
        dists_naive = []
        for i in range(N):
            points = pcls.points_list()[i]
            edges = edges_list[i]
            dists_temp = torch.zeros(
                (edges.shape[0], points.shape[0]), dtype=torch.float32, device=device
            )
            for e in range(edges.shape[0]):
                for p in range(points.shape[0]):
                    dist = self._point_to_edge_distance(points[p], edges[e])
                    dists_temp[e, p] = dist

            # torch.min() doesn't necessarily return the first index of the
            # smallest value, our warp_reduce does. So it's not straightforward
            # to directly compare indices, nor the gradients of grad_edges which
            # also depend on the indices of the minimum value.
            # To be able to compare, we will compare dists_temp.min(1) and
            # then feed the cuda indices to the naive output

            start = edges_first_idx[i]
            end = edges_first_idx[i + 1] if i < N - 1 else edges_packed.shape[0]

444
            min_idx = idx_cuda.cpu()[start:end] - points_first_idx[i].cpu()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
445
446
447
448
449
450
451
452
453
            iidx = torch.arange(edges.shape[0], device=device)
            min_dist = dists_temp[iidx, min_idx]

            dists_naive.append(min_dist)

        dists_naive = torch.cat(dists_naive)

        # Compare
        self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
454
        self.assertClose(dists_naive.cpu(), dists_cpu)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
455
456
457
458

        # Naive Implementation: backward
        dists_naive.backward(grad_dists)
        grad_points_naive = torch.cat([cloud.grad for cloud in pcls.points_list()])
459
        grad_edges_naive = edges_packed.grad.cpu()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
460
461
462

        # Compare
        self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7)
463
464
465
        self.assertClose(grad_edges_naive, grad_edges_cuda.cpu(), atol=5e-7)
        self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7)
        self.assertClose(grad_edges_naive, grad_edges_cpu, atol=5e-7)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
466
467
468
469
470

    def test_point_mesh_edge_distance(self):
        """
        Test point_mesh_edge_distance from pytorch3d.loss
        """
Nikhila Ravi's avatar
Nikhila Ravi committed
471
        device = get_random_cuda_device()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
472
        N, V, F, P = 4, 32, 16, 24
Nikhila Ravi's avatar
Nikhila Ravi committed
473
        meshes, pcls = self.init_meshes_clouds(N, V, F, P, device=device)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492

        # clone and detach for another backward pass through the op
        verts_op = [verts.clone().detach() for verts in meshes.verts_list()]
        for i in range(N):
            verts_op[i].requires_grad = True

        faces_op = [faces.clone().detach() for faces in meshes.faces_list()]
        meshes_op = Meshes(verts=verts_op, faces=faces_op)
        points_op = [points.clone().detach() for points in pcls.points_list()]
        for i in range(N):
            points_op[i].requires_grad = True
        pcls_op = Pointclouds(points_op)

        # Cuda implementation: forward & backward
        loss_op = point_mesh_edge_distance(meshes_op, pcls_op)

        # Naive implementation: forward & backward
        edges_packed = meshes.edges_packed()
        edges_list = packed_to_list(edges_packed, meshes.num_edges_per_mesh().tolist())
Nikhila Ravi's avatar
Nikhila Ravi committed
493
        loss_naive = torch.zeros(N, dtype=torch.float32, device=device)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
        for i in range(N):
            points = pcls.points_list()[i]
            verts = meshes.verts_list()[i]
            v_first_idx = meshes.mesh_to_verts_packed_first_idx()[i]
            edges = verts[edges_list[i] - v_first_idx]

            num_p = points.shape[0]
            num_e = edges.shape[0]
            dists = torch.zeros((num_p, num_e), dtype=torch.float32, device=device)
            for p in range(num_p):
                for e in range(num_e):
                    dist = self._point_to_edge_distance(points[p], edges[e])
                    dists[p, e] = dist

            min_dist_p, min_idx_p = dists.min(1)
            min_dist_e, min_idx_e = dists.min(0)

            loss_naive[i] = min_dist_p.mean() + min_dist_e.mean()
        loss_naive = loss_naive.mean()

        # NOTE that hear the comparison holds despite the discrepancy
        # due to the argmin indices returned by min(). This is because
        # we don't will compare gradients on the verts and not on the
        # edges or faces.

        # Compare forward pass
        self.assertClose(loss_op, loss_naive)

        # Compare backward pass
Nikhila Ravi's avatar
Nikhila Ravi committed
523
        rand_val = torch.rand(1).item()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
        grad_dist = torch.tensor(rand_val, dtype=torch.float32, device=device)

        loss_naive.backward(grad_dist)
        loss_op.backward(grad_dist)

        # check verts grad
        for i in range(N):
            self.assertClose(
                meshes.verts_list()[i].grad, meshes_op.verts_list()[i].grad
            )
            self.assertClose(pcls.points_list()[i].grad, pcls_op.points_list()[i].grad)

    def test_point_face_array_distance(self):
        """
        Test CUDA implementation for PointFaceArrayDistanceForward
            &  PointFaceArrayDistanceBackward
        """
        P, T = 16, 32
Nikhila Ravi's avatar
Nikhila Ravi committed
542
        device = get_random_cuda_device()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
543
544
        points = torch.rand((P, 3), dtype=torch.float32, device=device)
        tris = torch.rand((T, 3, 3), dtype=torch.float32, device=device)
545
546
        points_cpu = points.clone().cpu()
        tris_cpu = tris.clone().cpu()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565

        points.requires_grad = True
        tris.requires_grad = True
        grad_dists = torch.rand((P, T), dtype=torch.float32, device=device)

        points_temp = points.clone().detach()
        points_temp.requires_grad = True
        tris_temp = tris.clone().detach()
        tris_temp.requires_grad = True

        # Naive python implementation
        dists_naive = torch.zeros((P, T), dtype=torch.float32, device=device)
        for p in range(P):
            for t in range(T):
                dist = self._point_to_tri_distance(points[p], tris[t])
                dists_naive[p, t] = dist

        # Naive Backward
        dists_naive.backward(grad_dists)
566
567
        grad_points_naive = points.grad.cpu()
        grad_tris_naive = tris.grad.cpu()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
568
569

        # Cuda Forward Implementation
570
571
572
573
574
575
        dists_cuda = _C.point_face_array_dist_forward(
            points, tris, TestPointMeshDistance.min_triangle_area()
        )
        dists_cpu = _C.point_face_array_dist_forward(
            points_cpu, tris_cpu, TestPointMeshDistance.min_triangle_area()
        )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
576
577
578

        # Compare
        self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
579
        self.assertClose(dists_naive.cpu(), dists_cpu)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
580
581
582

        # CUDA Backward Implementation
        grad_points_cuda, grad_tris_cuda = _C.point_face_array_dist_backward(
583
            points, tris, grad_dists, TestPointMeshDistance.min_triangle_area()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
584
        )
585
        grad_points_cpu, grad_tris_cpu = _C.point_face_array_dist_backward(
586
587
588
589
            points_cpu,
            tris_cpu,
            grad_dists.cpu(),
            TestPointMeshDistance.min_triangle_area(),
590
        )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
591
592

        # Compare
593
594
595
596
        self.assertClose(grad_points_naive, grad_points_cuda.cpu())
        self.assertClose(grad_tris_naive, grad_tris_cuda.cpu(), atol=5e-6)
        self.assertClose(grad_points_naive, grad_points_cpu)
        self.assertClose(grad_tris_naive, grad_tris_cpu, atol=5e-6)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
597
598
599
600
601
602

    def test_point_face_distance(self):
        """
        Test CUDA implementation for PointFaceDistanceForward
            &  PointFaceDistanceBackward
        """
Nikhila Ravi's avatar
Nikhila Ravi committed
603
        device = get_random_cuda_device()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
604
        N, V, F, P = 4, 32, 16, 24
Nikhila Ravi's avatar
Nikhila Ravi committed
605
        meshes, pcls = self.init_meshes_clouds(N, V, F, P, device=device)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628

        # make points packed a leaf node
        points_packed = pcls.points_packed().detach().clone()  # (P, 3)

        points_first_idx = pcls.cloud_to_packed_first_idx()
        max_p = pcls.num_points_per_cloud().max().item()

        # make edges packed a leaf node
        verts_packed = meshes.verts_packed()
        faces_packed = verts_packed[meshes.faces_packed()]  # (T, 3, 3)
        faces_packed = faces_packed.clone().detach()

        faces_first_idx = meshes.mesh_to_faces_packed_first_idx()

        # leaf nodes
        points_packed.requires_grad = True
        faces_packed.requires_grad = True
        grad_dists = torch.rand(
            (points_packed.shape[0],), dtype=torch.float32, device=device
        )

        # Cuda Implementation: forward
        dists_cuda, idx_cuda = _C.point_face_dist_forward(
629
630
631
632
633
634
            points_packed,
            points_first_idx,
            faces_packed,
            faces_first_idx,
            max_p,
            TestPointMeshDistance.min_triangle_area(),
Georgia Gkioxari's avatar
Georgia Gkioxari committed
635
636
637
638
        )

        # Cuda Implementation: backward
        grad_points_cuda, grad_faces_cuda = _C.point_face_dist_backward(
639
640
641
642
643
            points_packed,
            faces_packed,
            idx_cuda,
            grad_dists,
            TestPointMeshDistance.min_triangle_area(),
Georgia Gkioxari's avatar
Georgia Gkioxari committed
644
645
        )

646
647
648
649
650
651
652
        # Cpu Implementation: forward
        dists_cpu, idx_cpu = _C.point_face_dist_forward(
            points_packed.cpu(),
            points_first_idx.cpu(),
            faces_packed.cpu(),
            faces_first_idx.cpu(),
            max_p,
653
            TestPointMeshDistance.min_triangle_area(),
654
655
656
657
658
        )

        # Cpu Implementation: backward
        # Note that using idx_cpu doesn't pass - there seems to be a problem with tied results.
        grad_points_cpu, grad_faces_cpu = _C.point_face_dist_backward(
659
660
661
662
663
            points_packed.cpu(),
            faces_packed.cpu(),
            idx_cuda.cpu(),
            grad_dists.cpu(),
            TestPointMeshDistance.min_triangle_area(),
664
665
        )

Georgia Gkioxari's avatar
Georgia Gkioxari committed
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
        # Naive Implementation: forward
        faces_list = packed_to_list(faces_packed, meshes.num_faces_per_mesh().tolist())
        dists_naive = []
        for i in range(N):
            points = pcls.points_list()[i]
            tris = faces_list[i]
            dists_temp = torch.zeros(
                (points.shape[0], tris.shape[0]), dtype=torch.float32, device=device
            )
            for p in range(points.shape[0]):
                for t in range(tris.shape[0]):
                    dist = self._point_to_tri_distance(points[p], tris[t])
                    dists_temp[p, t] = dist

            # torch.min() doesn't necessarily return the first index of the
            # smallest value, our warp_reduce does. So it's not straightforward
            # to directly compare indices, nor the gradients of grad_tris which
            # also depend on the indices of the minimum value.
            # To be able to compare, we will compare dists_temp.min(1) and
            # then feed the cuda indices to the naive output

            start = points_first_idx[i]
            end = points_first_idx[i + 1] if i < N - 1 else points_packed.shape[0]

690
            min_idx = idx_cuda.cpu()[start:end] - faces_first_idx[i].cpu()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
691
692
693
694
695
696
697
698
699
            iidx = torch.arange(points.shape[0], device=device)
            min_dist = dists_temp[iidx, min_idx]

            dists_naive.append(min_dist)

        dists_naive = torch.cat(dists_naive)

        # Compare
        self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
700
        self.assertClose(dists_naive.cpu(), dists_cpu)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
701
702
703
704

        #  Naive Implementation: backward
        dists_naive.backward(grad_dists)
        grad_points_naive = torch.cat([cloud.grad for cloud in pcls.points_list()])
705
        grad_faces_naive = faces_packed.grad.cpu()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
706
707
708

        # Compare
        self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7)
709
710
711
        self.assertClose(grad_faces_naive, grad_faces_cuda.cpu(), atol=5e-7)
        self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7)
        self.assertClose(grad_faces_naive, grad_faces_cpu, atol=5e-7)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
712
713
714
715
716
717

    def test_face_point_distance(self):
        """
        Test CUDA implementation for FacePointDistanceForward
            &  FacePointDistanceBackward
        """
Nikhila Ravi's avatar
Nikhila Ravi committed
718
        device = get_random_cuda_device()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
719
        N, V, F, P = 4, 32, 16, 24
Nikhila Ravi's avatar
Nikhila Ravi committed
720
        meshes, pcls = self.init_meshes_clouds(N, V, F, P, device=device)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743

        # make points packed a leaf node
        points_packed = pcls.points_packed().detach().clone()  # (P, 3)

        points_first_idx = pcls.cloud_to_packed_first_idx()

        # make edges packed a leaf node
        verts_packed = meshes.verts_packed()
        faces_packed = verts_packed[meshes.faces_packed()]  # (T, 3, 3)
        faces_packed = faces_packed.clone().detach()

        faces_first_idx = meshes.mesh_to_faces_packed_first_idx()
        max_f = meshes.num_faces_per_mesh().max().item()

        # leaf nodes
        points_packed.requires_grad = True
        faces_packed.requires_grad = True
        grad_dists = torch.rand(
            (faces_packed.shape[0],), dtype=torch.float32, device=device
        )

        # Cuda Implementation: forward
        dists_cuda, idx_cuda = _C.face_point_dist_forward(
744
745
746
747
748
749
            points_packed,
            points_first_idx,
            faces_packed,
            faces_first_idx,
            max_f,
            TestPointMeshDistance.min_triangle_area(),
Georgia Gkioxari's avatar
Georgia Gkioxari committed
750
751
752
753
        )

        # Cuda Implementation: backward
        grad_points_cuda, grad_faces_cuda = _C.face_point_dist_backward(
754
755
756
757
758
            points_packed,
            faces_packed,
            idx_cuda,
            grad_dists,
            TestPointMeshDistance.min_triangle_area(),
Georgia Gkioxari's avatar
Georgia Gkioxari committed
759
760
        )

761
762
763
764
765
766
767
        # Cpu Implementation: forward
        dists_cpu, idx_cpu = _C.face_point_dist_forward(
            points_packed.cpu(),
            points_first_idx.cpu(),
            faces_packed.cpu(),
            faces_first_idx.cpu(),
            max_f,
768
            TestPointMeshDistance.min_triangle_area(),
769
770
771
772
        )

        # Cpu Implementation: backward
        grad_points_cpu, grad_faces_cpu = _C.face_point_dist_backward(
773
774
775
776
777
            points_packed.cpu(),
            faces_packed.cpu(),
            idx_cpu,
            grad_dists.cpu(),
            TestPointMeshDistance.min_triangle_area(),
778
779
        )

Georgia Gkioxari's avatar
Georgia Gkioxari committed
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
        # Naive Implementation: forward
        faces_list = packed_to_list(faces_packed, meshes.num_faces_per_mesh().tolist())
        dists_naive = []
        for i in range(N):
            points = pcls.points_list()[i]
            tris = faces_list[i]
            dists_temp = torch.zeros(
                (tris.shape[0], points.shape[0]), dtype=torch.float32, device=device
            )
            for t in range(tris.shape[0]):
                for p in range(points.shape[0]):
                    dist = self._point_to_tri_distance(points[p], tris[t])
                    dists_temp[t, p] = dist

            # torch.min() doesn't necessarily return the first index of the
            # smallest value, our warp_reduce does. So it's not straightforward
            # to directly compare indices, nor the gradients of grad_tris which
            # also depend on the indices of the minimum value.
            # To be able to compare, we will compare dists_temp.min(1) and
            # then feed the cuda indices to the naive output

            start = faces_first_idx[i]
            end = faces_first_idx[i + 1] if i < N - 1 else faces_packed.shape[0]

804
            min_idx = idx_cuda.cpu()[start:end] - points_first_idx[i].cpu()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
805
806
807
808
809
810
811
812
813
            iidx = torch.arange(tris.shape[0], device=device)
            min_dist = dists_temp[iidx, min_idx]

            dists_naive.append(min_dist)

        dists_naive = torch.cat(dists_naive)

        # Compare
        self.assertClose(dists_naive.cpu(), dists_cuda.cpu())
814
        self.assertClose(dists_naive.cpu(), dists_cpu)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
815
816
817
818
819
820
821
822
823

        # Naive Implementation: backward
        dists_naive.backward(grad_dists)
        grad_points_naive = torch.cat([cloud.grad for cloud in pcls.points_list()])
        grad_faces_naive = faces_packed.grad

        # Compare
        self.assertClose(grad_points_naive.cpu(), grad_points_cuda.cpu(), atol=1e-7)
        self.assertClose(grad_faces_naive.cpu(), grad_faces_cuda.cpu(), atol=5e-7)
824
825
        self.assertClose(grad_points_naive.cpu(), grad_points_cpu, atol=1e-7)
        self.assertClose(grad_faces_naive.cpu(), grad_faces_cpu, atol=5e-7)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
826
827
828
829
830

    def test_point_mesh_face_distance(self):
        """
        Test point_mesh_face_distance from pytorch3d.loss
        """
Nikhila Ravi's avatar
Nikhila Ravi committed
831
        device = get_random_cuda_device()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
832
        N, V, F, P = 4, 32, 16, 24
Nikhila Ravi's avatar
Nikhila Ravi committed
833
        meshes, pcls = self.init_meshes_clouds(N, V, F, P, device=device)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
834
835
836
837
838
839
840
841
842
843
844
845
846
847

        # clone and detach for another backward pass through the op
        verts_op = [verts.clone().detach() for verts in meshes.verts_list()]
        for i in range(N):
            verts_op[i].requires_grad = True

        faces_op = [faces.clone().detach() for faces in meshes.faces_list()]
        meshes_op = Meshes(verts=verts_op, faces=faces_op)
        points_op = [points.clone().detach() for points in pcls.points_list()]
        for i in range(N):
            points_op[i].requires_grad = True
        pcls_op = Pointclouds(points_op)

        # naive implementation
Nikhila Ravi's avatar
Nikhila Ravi committed
848
        loss_naive = torch.zeros(N, dtype=torch.float32, device=device)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
        for i in range(N):
            points = pcls.points_list()[i]
            verts = meshes.verts_list()[i]
            faces = meshes.faces_list()[i]
            tris = verts[faces]

            num_p = points.shape[0]
            num_t = tris.shape[0]
            dists = torch.zeros((num_p, num_t), dtype=torch.float32, device=device)
            for p in range(num_p):
                for t in range(num_t):
                    dist = self._point_to_tri_distance(points[p], tris[t])
                    dists[p, t] = dist

            min_dist_p, min_idx_p = dists.min(1)
            min_dist_t, min_idx_t = dists.min(0)

            loss_naive[i] = min_dist_p.mean() + min_dist_t.mean()
        loss_naive = loss_naive.mean()

        # Op
        loss_op = point_mesh_face_distance(meshes_op, pcls_op)

        # Compare forward pass
        self.assertClose(loss_op, loss_naive)

        # Compare backward pass
Nikhila Ravi's avatar
Nikhila Ravi committed
876
        rand_val = torch.rand(1).item()
Georgia Gkioxari's avatar
Georgia Gkioxari committed
877
878
879
880
881
882
883
884
885
886
887
888
        grad_dist = torch.tensor(rand_val, dtype=torch.float32, device=device)

        loss_naive.backward(grad_dist)
        loss_op.backward(grad_dist)

        # check verts grad
        for i in range(N):
            self.assertClose(
                meshes.verts_list()[i].grad, meshes_op.verts_list()[i].grad
            )
            self.assertClose(pcls.points_list()[i].grad, pcls_op.points_list()[i].grad)

889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
    def test_small_faces_case(self):
        for device in [torch.device("cpu"), torch.device("cuda:0")]:
            mesh_vertices = torch.tensor(
                [
                    [-0.0021, -0.3769, 0.7146],
                    [-0.0161, -0.3771, 0.7146],
                    [-0.0021, -0.3771, 0.7147],
                ],
                dtype=torch.float32,
                device=device,
            )
            mesh1_faces = torch.tensor([[0, 2, 1]], device=device)
            mesh2_faces = torch.tensor([[2, 0, 1]], device=device)
            pcd_points = torch.tensor([[-0.3623, -0.5340, 0.7727]], device=device)
            mesh1 = Meshes(verts=[mesh_vertices], faces=[mesh1_faces])
            mesh2 = Meshes(verts=[mesh_vertices], faces=[mesh2_faces])
            pcd = Pointclouds(points=[pcd_points])

            loss1 = point_mesh_face_distance(mesh1, pcd)
            loss2 = point_mesh_face_distance(mesh2, pcd)
            self.assertClose(loss1, loss2)

Georgia Gkioxari's avatar
Georgia Gkioxari committed
911
912
913
    @staticmethod
    def point_mesh_edge(N: int, V: int, F: int, P: int, device: str):
        device = torch.device(device)
Nikhila Ravi's avatar
Nikhila Ravi committed
914
915
916
        meshes, pcls = TestPointMeshDistance.init_meshes_clouds(
            N, V, F, P, device=device
        )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
917
918
919
920
921
922
923
924
925
926
927
        torch.cuda.synchronize()

        def loss():
            point_mesh_edge_distance(meshes, pcls)
            torch.cuda.synchronize()

        return loss

    @staticmethod
    def point_mesh_face(N: int, V: int, F: int, P: int, device: str):
        device = torch.device(device)
Nikhila Ravi's avatar
Nikhila Ravi committed
928
929
930
        meshes, pcls = TestPointMeshDistance.init_meshes_clouds(
            N, V, F, P, device=device
        )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
931
932
933
934
935
936
937
        torch.cuda.synchronize()

        def loss():
            point_mesh_face_distance(meshes, pcls)
            torch.cuda.synchronize()

        return loss