test_blending.py 16.4 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
8

import unittest

9
import torch
facebook-github-bot's avatar
facebook-github-bot committed
10
11
12
13
14
15
from pytorch3d.renderer.blending import (
    BlendParams,
    hard_rgb_blend,
    sigmoid_alpha_blend,
    softmax_rgb_blend,
)
16
from pytorch3d.renderer.cameras import FoVPerspectiveCameras
facebook-github-bot's avatar
facebook-github-bot committed
17
from pytorch3d.renderer.mesh.rasterizer import Fragments
18
from pytorch3d.renderer.splatter_blend import SplatterBlender
facebook-github-bot's avatar
facebook-github-bot committed
19

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
20
21
from .common_testing import TestCaseMixin

facebook-github-bot's avatar
facebook-github-bot committed
22

Nikhila Ravi's avatar
Nikhila Ravi committed
23
def sigmoid_blend_naive_loop(colors, fragments, blend_params):
facebook-github-bot's avatar
facebook-github-bot committed
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
    """
    Naive for loop based implementation of distance based alpha calculation.
    Only for test purposes.
    """
    pix_to_face = fragments.pix_to_face
    dists = fragments.dists
    sigma = blend_params.sigma

    N, H, W, K = pix_to_face.shape
    device = pix_to_face.device
    pixel_colors = torch.ones((N, H, W, 4), dtype=colors.dtype, device=device)

    for n in range(N):
        for h in range(H):
            for w in range(W):
                alpha = 1.0

                # Loop over k faces and calculate 2D distance based probability
                # map.
                for k in range(K):
                    if pix_to_face[n, h, w, k] >= 0:
                        prob = torch.sigmoid(-dists[n, h, w, k] / sigma)
                        alpha *= 1.0 - prob  # cumulative product
                pixel_colors[n, h, w, :3] = colors[n, h, w, 0, :]
                pixel_colors[n, h, w, 3] = 1.0 - alpha

50
    return pixel_colors
facebook-github-bot's avatar
facebook-github-bot committed
51
52


53
54
55
56
57
58
59
60
61
62
def sigmoid_alpha_blend_vectorized(colors, fragments, blend_params) -> torch.Tensor:
    N, H, W, K = fragments.pix_to_face.shape
    pixel_colors = torch.ones((N, H, W, 4), dtype=colors.dtype, device=colors.device)
    mask = fragments.pix_to_face >= 0
    prob = torch.sigmoid(-fragments.dists / blend_params.sigma) * mask
    pixel_colors[..., :3] = colors[..., 0, :]
    pixel_colors[..., 3] = 1.0 - torch.prod((1.0 - prob), dim=-1)
    return pixel_colors


63
def sigmoid_blend_naive_loop_backward(grad_images, images, fragments, blend_params):
Nikhila Ravi's avatar
Nikhila Ravi committed
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
    pix_to_face = fragments.pix_to_face
    dists = fragments.dists
    sigma = blend_params.sigma

    N, H, W, K = pix_to_face.shape
    device = pix_to_face.device
    grad_distances = torch.zeros((N, H, W, K), dtype=dists.dtype, device=device)

    for n in range(N):
        for h in range(H):
            for w in range(W):
                alpha = 1.0 - images[n, h, w, 3]
                grad_alpha = grad_images[n, h, w, 3]
                # Loop over k faces and calculate 2D distance based probability
                # map.
                for k in range(K):
                    if pix_to_face[n, h, w, k] >= 0:
                        prob = torch.sigmoid(-dists[n, h, w, k] / sigma)
                        grad_distances[n, h, w, k] = (
                            grad_alpha * (-1.0 / sigma) * prob * alpha
                        )
    return grad_distances


facebook-github-bot's avatar
facebook-github-bot committed
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
def softmax_blend_naive(colors, fragments, blend_params):
    """
    Naive for loop based implementation of softmax blending.
    Only for test purposes.
    """
    pix_to_face = fragments.pix_to_face
    dists = fragments.dists
    zbuf = fragments.zbuf
    sigma = blend_params.sigma
    gamma = blend_params.gamma

    N, H, W, K = pix_to_face.shape
    device = pix_to_face.device
    pixel_colors = torch.ones((N, H, W, 4), dtype=colors.dtype, device=device)

    # Near and far clipping planes
    zfar = 100.0
    znear = 1.0
Nikhila Ravi's avatar
Nikhila Ravi committed
106
    eps = 1e-10
facebook-github-bot's avatar
facebook-github-bot committed
107
108
109
110
111
112
113
114
115

    bk_color = blend_params.background_color
    if not torch.is_tensor(bk_color):
        bk_color = torch.tensor(bk_color, dtype=colors.dtype, device=device)

    for n in range(N):
        for h in range(H):
            for w in range(W):
                alpha = 1.0
Nikhila Ravi's avatar
Nikhila Ravi committed
116
                weights_k = torch.zeros(K, device=device)
Nikhila Ravi's avatar
Nikhila Ravi committed
117
                zmax = torch.tensor(0.0, device=device)
facebook-github-bot's avatar
facebook-github-bot committed
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

                # Loop over K to find max z.
                for k in range(K):
                    if pix_to_face[n, h, w, k] >= 0:
                        zinv = (zfar - zbuf[n, h, w, k]) / (zfar - znear)
                        if zinv > zmax:
                            zmax = zinv

                # Loop over K faces to calculate 2D distance based probability
                # map and zbuf based weights for colors.
                for k in range(K):
                    if pix_to_face[n, h, w, k] >= 0:
                        zinv = (zfar - zbuf[n, h, w, k]) / (zfar - znear)
                        prob = torch.sigmoid(-dists[n, h, w, k] / sigma)
                        alpha *= 1.0 - prob  # cumulative product
                        weights_k[k] = prob * torch.exp((zinv - zmax) / gamma)

Nikhila Ravi's avatar
Nikhila Ravi committed
135
136
137
                # Clamp to ensure delta is never 0
                delta = torch.exp((eps - zmax) / blend_params.gamma).clamp(min=eps)
                delta = delta.to(device)
facebook-github-bot's avatar
facebook-github-bot committed
138
                denom = weights_k.sum() + delta
Nikhila Ravi's avatar
Nikhila Ravi committed
139
140
141
                cols = (weights_k[..., None] * colors[n, h, w, :, :]).sum(dim=0)
                pixel_colors[n, h, w, :3] = cols + delta * bk_color
                pixel_colors[n, h, w, :3] /= denom
facebook-github-bot's avatar
facebook-github-bot committed
142
143
                pixel_colors[n, h, w, 3] = 1.0 - alpha

144
    return pixel_colors
facebook-github-bot's avatar
facebook-github-bot committed
145
146


147
class TestBlending(TestCaseMixin, unittest.TestCase):
facebook-github-bot's avatar
facebook-github-bot committed
148
149
150
    def setUp(self) -> None:
        torch.manual_seed(42)

Nikhila Ravi's avatar
Nikhila Ravi committed
151
    def _compare_impls(
152
        self, fn1, fn2, args1, args2, grad_var1=None, grad_var2=None, compare_grads=True
Nikhila Ravi's avatar
Nikhila Ravi committed
153
154
155
    ):
        out1 = fn1(*args1)
        out2 = fn2(*args2)
156
        self.assertClose(out1.cpu()[..., 3], out2.cpu()[..., 3], atol=1e-7)
Nikhila Ravi's avatar
Nikhila Ravi committed
157
158
159
160
161
162
163
164
165
166
167

        # Check gradients
        if not compare_grads:
            return

        grad_out = torch.randn_like(out1)
        (out1 * grad_out).sum().backward()
        self.assertTrue(hasattr(grad_var1, "grad"))

        (out2 * grad_out).sum().backward()
        self.assertTrue(hasattr(grad_var2, "grad"))
Nikhila Ravi's avatar
Nikhila Ravi committed
168

169
        self.assertClose(grad_var1.grad.cpu(), grad_var2.grad.cpu(), atol=2e-5)
Nikhila Ravi's avatar
Nikhila Ravi committed
170

facebook-github-bot's avatar
facebook-github-bot committed
171
172
    def test_hard_rgb_blend(self):
        N, H, W, K = 5, 10, 10, 20
173
        pix_to_face = torch.randint(low=-1, high=100, size=(N, H, W, K))
facebook-github-bot's avatar
facebook-github-bot committed
174
175
176
177
178
179
180
        bary_coords = torch.ones((N, H, W, K, 3))
        fragments = Fragments(
            pix_to_face=pix_to_face,
            bary_coords=bary_coords,
            zbuf=pix_to_face,  # dummy
            dists=pix_to_face,  # dummy
        )
181
182
183
184
185
186
187
188
189
190
191
192
193
        colors = torch.randn((N, H, W, K, 3))
        blend_params = BlendParams(1e-4, 1e-4, (0.5, 0.5, 1))
        images = hard_rgb_blend(colors, fragments, blend_params)

        # Examine if the foreground colors are correct.
        is_foreground = pix_to_face[..., 0] >= 0
        self.assertClose(images[is_foreground][:, :3], colors[is_foreground][..., 0, :])

        # Examine if the background colors are correct.
        for i in range(3):  # i.e. RGB
            channel_color = blend_params.background_color[i]
            self.assertTrue(images[~is_foreground][..., i].eq(channel_color).all())

194
195
        # Examine the alpha channel
        self.assertClose(images[..., 3], (pix_to_face[..., 0] >= 0).float())
facebook-github-bot's avatar
facebook-github-bot committed
196

Nikhila Ravi's avatar
Nikhila Ravi committed
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
    def test_sigmoid_alpha_blend_manual_gradients(self):
        # Create dummy outputs of rasterization
        torch.manual_seed(231)
        F = 32  # number of faces in the mesh
        # The python loop version is really slow so only using small input sizes.
        N, S, K = 2, 3, 2
        device = torch.device("cuda")
        pix_to_face = torch.randint(F + 1, size=(N, S, S, K), device=device) - 1
        colors = torch.randn((N, S, S, K, 3), device=device)
        empty = torch.tensor([], device=device)

        # # randomly flip the sign of the distance
        # # (-) means inside triangle, (+) means outside triangle.
        random_sign_flip = torch.rand((N, S, S, K))
        random_sign_flip[random_sign_flip > 0.5] *= -1.0
212
        dists = torch.randn(size=(N, S, S, K), requires_grad=True, device=device)
Nikhila Ravi's avatar
Nikhila Ravi committed
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
        fragments = Fragments(
            pix_to_face=pix_to_face,
            bary_coords=empty,  # dummy
            zbuf=empty,  # dummy
            dists=dists,
        )
        blend_params = BlendParams(sigma=1e-3)
        pix_cols = sigmoid_blend_naive_loop(colors, fragments, blend_params)
        grad_out = torch.randn_like(pix_cols)

        # Backward pass
        pix_cols.backward(grad_out)
        grad_dists = sigmoid_blend_naive_loop_backward(
            grad_out, pix_cols, fragments, blend_params
        )
        self.assertTrue(torch.allclose(dists.grad, grad_dists, atol=1e-7))

    def test_sigmoid_alpha_blend_python(self):
facebook-github-bot's avatar
facebook-github-bot committed
231
        """
Nikhila Ravi's avatar
Nikhila Ravi committed
232
        Test outputs of python tensorised function and python loop
facebook-github-bot's avatar
facebook-github-bot committed
233
234
        """

Nikhila Ravi's avatar
Nikhila Ravi committed
235
236
237
238
        # Create dummy outputs of rasterization
        torch.manual_seed(231)
        F = 32  # number of faces in the mesh
        # The python loop version is really slow so only using small input sizes.
239
        N, S, K = 1, 4, 1
Nikhila Ravi's avatar
Nikhila Ravi committed
240
        device = torch.device("cuda")
241
        pix_to_face = torch.randint(low=-1, high=F, size=(N, S, S, K), device=device)
Nikhila Ravi's avatar
Nikhila Ravi committed
242
243
244
        colors = torch.randn((N, S, S, K, 3), device=device)
        empty = torch.tensor([], device=device)

245
246
247
        dists1 = torch.randn(size=(N, S, S, K), device=device)
        dists2 = dists1.clone()
        dists1.requires_grad = True
facebook-github-bot's avatar
facebook-github-bot committed
248
        dists2.requires_grad = True
Nikhila Ravi's avatar
Nikhila Ravi committed
249

facebook-github-bot's avatar
facebook-github-bot committed
250
251
        fragments1 = Fragments(
            pix_to_face=pix_to_face,
Nikhila Ravi's avatar
Nikhila Ravi committed
252
253
            bary_coords=empty,  # dummy
            zbuf=empty,  # dummy
facebook-github-bot's avatar
facebook-github-bot committed
254
255
256
257
            dists=dists1,
        )
        fragments2 = Fragments(
            pix_to_face=pix_to_face,
Nikhila Ravi's avatar
Nikhila Ravi committed
258
259
            bary_coords=empty,  # dummy
            zbuf=empty,  # dummy
facebook-github-bot's avatar
facebook-github-bot committed
260
261
262
            dists=dists2,
        )

Nikhila Ravi's avatar
Nikhila Ravi committed
263
264
265
266
267
268
        blend_params = BlendParams(sigma=1e-2)
        args1 = (colors, fragments1, blend_params)
        args2 = (colors, fragments2, blend_params)

        self._compare_impls(
            sigmoid_alpha_blend,
269
            sigmoid_alpha_blend_vectorized,
Nikhila Ravi's avatar
Nikhila Ravi committed
270
271
272
273
274
275
            args1,
            args2,
            dists1,
            dists2,
            compare_grads=True,
        )
facebook-github-bot's avatar
facebook-github-bot committed
276
277

    def test_softmax_rgb_blend(self):
278
        # Create dummy outputs of rasterization simulating a cube in the center
facebook-github-bot's avatar
facebook-github-bot committed
279
280
        # of the image with surrounding padded values.
        N, S, K = 1, 8, 2
Nikhila Ravi's avatar
Nikhila Ravi committed
281
        device = torch.device("cuda")
Nikhila Ravi's avatar
Nikhila Ravi committed
282
283
284
        pix_to_face = torch.full(
            (N, S, S, K), fill_value=-1, dtype=torch.int64, device=device
        )
facebook-github-bot's avatar
facebook-github-bot committed
285
        h = int(S / 2)
Nikhila Ravi's avatar
Nikhila Ravi committed
286
287
288
        pix_to_face_full = torch.randint(
            size=(N, h, h, K), low=0, high=100, device=device
        )
facebook-github-bot's avatar
facebook-github-bot committed
289
290
291
        s = int(S / 4)
        e = int(0.75 * S)
        pix_to_face[:, s:e, s:e, :] = pix_to_face_full
Nikhila Ravi's avatar
Nikhila Ravi committed
292
        empty = torch.tensor([], device=device)
facebook-github-bot's avatar
facebook-github-bot committed
293

Nikhila Ravi's avatar
Nikhila Ravi committed
294
        random_sign_flip = torch.rand((N, S, S, K), device=device)
facebook-github-bot's avatar
facebook-github-bot committed
295
        random_sign_flip[random_sign_flip > 0.5] *= -1.0
Nikhila Ravi's avatar
Nikhila Ravi committed
296
        zbuf1 = torch.randn(size=(N, S, S, K), device=device)
facebook-github-bot's avatar
facebook-github-bot committed
297
298
299

        # randomly flip the sign of the distance
        # (-) means inside triangle, (+) means outside triangle.
300
        dists1 = torch.randn(size=(N, S, S, K), device=device) * random_sign_flip
facebook-github-bot's avatar
facebook-github-bot committed
301
302
303
304
        dists2 = dists1.clone()
        zbuf2 = zbuf1.clone()
        dists1.requires_grad = True
        dists2.requires_grad = True
Nikhila Ravi's avatar
Nikhila Ravi committed
305
        colors = torch.randn((N, S, S, K, 3), device=device)
facebook-github-bot's avatar
facebook-github-bot committed
306
307
        fragments1 = Fragments(
            pix_to_face=pix_to_face,
Nikhila Ravi's avatar
Nikhila Ravi committed
308
            bary_coords=empty,  # dummy
facebook-github-bot's avatar
facebook-github-bot committed
309
310
311
312
313
            zbuf=zbuf1,
            dists=dists1,
        )
        fragments2 = Fragments(
            pix_to_face=pix_to_face,
Nikhila Ravi's avatar
Nikhila Ravi committed
314
            bary_coords=empty,  # dummy
facebook-github-bot's avatar
facebook-github-bot committed
315
316
317
318
            zbuf=zbuf2,
            dists=dists2,
        )

Nikhila Ravi's avatar
Nikhila Ravi committed
319
320
321
322
323
324
325
326
327
328
329
330
        blend_params = BlendParams(sigma=1e-3)
        args1 = (colors, fragments1, blend_params)
        args2 = (colors, fragments2, blend_params)
        self._compare_impls(
            softmax_rgb_blend,
            softmax_blend_naive,
            args1,
            args2,
            dists1,
            dists2,
            compare_grads=True,
        )
facebook-github-bot's avatar
facebook-github-bot committed
331

Nikhila Ravi's avatar
Nikhila Ravi committed
332
333
334
335
336
    @staticmethod
    def bm_sigmoid_alpha_blending(
        num_meshes: int = 16,
        image_size: int = 128,
        faces_per_pixel: int = 100,
337
338
        device="cuda",
        backend: str = "pytorch",
Nikhila Ravi's avatar
Nikhila Ravi committed
339
340
341
342
343
344
345
    ):
        device = torch.device(device)
        torch.manual_seed(231)

        # Create dummy outputs of rasterization
        N, S, K = num_meshes, image_size, faces_per_pixel
        F = 32  # num faces in the mesh
346
347
348
        pix_to_face = torch.randint(
            low=-1, high=F + 1, size=(N, S, S, K), device=device
        )
Nikhila Ravi's avatar
Nikhila Ravi committed
349
350
351
        colors = torch.randn((N, S, S, K, 3), device=device)
        empty = torch.tensor([], device=device)

352
        dists1 = torch.randn(size=(N, S, S, K), requires_grad=True, device=device)
Nikhila Ravi's avatar
Nikhila Ravi committed
353
354
355
356
357
358
359
        fragments = Fragments(
            pix_to_face=pix_to_face,
            bary_coords=empty,  # dummy
            zbuf=empty,  # dummy
            dists=dists1,
        )
        blend_params = BlendParams(sigma=1e-3)
360
361
362
363
364
365
366

        blend_fn = (
            sigmoid_alpha_blend_vectorized
            if backend == "pytorch"
            else sigmoid_alpha_blend
        )

Nikhila Ravi's avatar
Nikhila Ravi committed
367
368
369
370
        torch.cuda.synchronize()

        def fn():
            # test forward and backward pass
371
            images = blend_fn(colors, fragments, blend_params)
Nikhila Ravi's avatar
Nikhila Ravi committed
372
373
374
375
376
377
378
379
380
381
382
            images.sum().backward()
            torch.cuda.synchronize()

        return fn

    @staticmethod
    def bm_softmax_blending(
        num_meshes: int = 16,
        image_size: int = 128,
        faces_per_pixel: int = 100,
        device: str = "cpu",
383
        backend: str = "pytorch",
Nikhila Ravi's avatar
Nikhila Ravi committed
384
385
386
387
388
389
390
391
392
393
394
    ):
        if torch.cuda.is_available() and "cuda:" in device:
            # If a device other than the default is used, set the device explicity.
            torch.cuda.set_device(device)

        device = torch.device(device)
        torch.manual_seed(231)

        # Create dummy outputs of rasterization
        N, S, K = num_meshes, image_size, faces_per_pixel
        F = 32  # num faces in the mesh
395
396
397
        pix_to_face = torch.randint(
            low=-1, high=F + 1, size=(N, S, S, K), device=device
        )
Nikhila Ravi's avatar
Nikhila Ravi committed
398
399
400
        colors = torch.randn((N, S, S, K, 3), device=device)
        empty = torch.tensor([], device=device)

401
        dists1 = torch.randn(size=(N, S, S, K), requires_grad=True, device=device)
Nikhila Ravi's avatar
Nikhila Ravi committed
402
403
        zbuf = torch.randn(size=(N, S, S, K), requires_grad=True, device=device)
        fragments = Fragments(
404
            pix_to_face=pix_to_face, bary_coords=empty, zbuf=zbuf, dists=dists1  # dummy
Nikhila Ravi's avatar
Nikhila Ravi committed
405
406
407
408
409
410
411
412
413
414
415
416
        )
        blend_params = BlendParams(sigma=1e-3)

        torch.cuda.synchronize()

        def fn():
            # test forward and backward pass
            images = softmax_rgb_blend(colors, fragments, blend_params)
            images.sum().backward()
            torch.cuda.synchronize()

        return fn
417

418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
    @staticmethod
    def bm_splatter_blending(
        num_meshes: int = 16,
        image_size: int = 128,
        faces_per_pixel: int = 2,
        use_jit: bool = False,
        device: str = "cpu",
        backend: str = "pytorch",
    ):
        if torch.cuda.is_available() and "cuda:" in device:
            # If a device other than the default is used, set the device explicity.
            torch.cuda.set_device(device)

        device = torch.device(device)
        torch.manual_seed(231)

        # Create dummy outputs of rasterization
        N, S, K = num_meshes, image_size, faces_per_pixel
        F = 32  # num faces in the mesh

        pixel_coords_camera = torch.randn(
            (N, S, S, K, 3), device=device, requires_grad=True
        )
        cameras = FoVPerspectiveCameras(device=device)
        colors = torch.randn((N, S, S, K, 3), device=device)
        background_mask = torch.randint(
            low=-1, high=F + 1, size=(N, S, S, K), device=device
        )
        background_mask = torch.full((N, S, S, K), False, dtype=bool, device=device)
        blend_params = BlendParams(sigma=0.5)

        torch.cuda.synchronize()
        splatter_blender = SplatterBlender((N, S, S, K), colors.device)

        def fn():
            # test forward and backward pass
            images = splatter_blender(
                colors,
                pixel_coords_camera,
                cameras,
                background_mask,
                blend_params,
            )
            images.sum().backward()
            torch.cuda.synchronize()

        return fn

466
    def test_blend_params(self):
467
        """Test color parameter of BlendParams().
468
469
        Assert passed value overrides default value.
        """
470
471
472
473
        bp_default = BlendParams()
        bp_new = BlendParams(background_color=(0.5, 0.5, 0.5))
        self.assertEqual(bp_new.background_color, (0.5, 0.5, 0.5))
        self.assertEqual(bp_default.background_color, (1.0, 1.0, 1.0))