test_render_meshes.py 60.3 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
9
10


"""
Sanity checks for output images from the renderer.
"""
Georgia Gkioxari's avatar
Georgia Gkioxari committed
11
import os
facebook-github-bot's avatar
facebook-github-bot committed
12
import unittest
Nikhila Ravi's avatar
Nikhila Ravi committed
13
from collections import namedtuple
14
15

import numpy as np
facebook-github-bot's avatar
facebook-github-bot committed
16
17
import torch
from PIL import Image
Nikhila Ravi's avatar
Nikhila Ravi committed
18
from pytorch3d.io import load_obj
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
19
20
from pytorch3d.renderer import (
    AmbientLights,
Georgia Gkioxari's avatar
Georgia Gkioxari committed
21
22
    FoVOrthographicCameras,
    FoVPerspectiveCameras,
23
    look_at_view_transform,
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
24
25
26
27
28
    Materials,
    MeshRasterizer,
    MeshRasterizerOpenGL,
    MeshRenderer,
    MeshRendererWithFragments,
Georgia Gkioxari's avatar
Georgia Gkioxari committed
29
30
    OrthographicCameras,
    PerspectiveCameras,
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
31
32
33
34
35
    PointLights,
    RasterizationSettings,
    TexturesAtlas,
    TexturesUV,
    TexturesVertex,
Georgia Gkioxari's avatar
Georgia Gkioxari committed
36
)
facebook-github-bot's avatar
facebook-github-bot committed
37
38
from pytorch3d.renderer.mesh.shader import (
    BlendParams,
39
    HardFlatShader,
Patrick Labatut's avatar
Patrick Labatut committed
40
    HardGouraudShader,
41
    HardPhongShader,
42
    SoftPhongShader,
43
    SoftSilhouetteShader,
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
44
    SplatterPhongShader,
45
    TexturedSoftPhongShader,
facebook-github-bot's avatar
facebook-github-bot committed
46
)
47
48
49
from pytorch3d.structures.meshes import (
    join_meshes_as_batch,
    join_meshes_as_scene,
50
    Meshes,
51
)
facebook-github-bot's avatar
facebook-github-bot committed
52
from pytorch3d.utils.ico_sphere import ico_sphere
53
from pytorch3d.utils.torus import torus
facebook-github-bot's avatar
facebook-github-bot committed
54

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
55
56
57
58
59
60
61
from .common_testing import (
    get_pytorch3d_dir,
    get_tests_dir,
    load_rgb_image,
    TestCaseMixin,
)

62

Nikhila Ravi's avatar
Nikhila Ravi committed
63
# If DEBUG=True, save out images generated in the tests for debugging.
facebook-github-bot's avatar
facebook-github-bot committed
64
65
# All saved images have prefix DEBUG_
DEBUG = False
66
DATA_DIR = get_tests_dir() / "data"
67
TUTORIAL_DATA_DIR = get_pytorch3d_dir() / "docs/tutorials/data"
facebook-github-bot's avatar
facebook-github-bot committed
68

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
69
70
71
RasterizerTest = namedtuple(
    "RasterizerTest", ["rasterizer", "shader", "reference_name", "debug_name"]
)
Nikhila Ravi's avatar
Nikhila Ravi committed
72

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

Nikhila Ravi's avatar
Nikhila Ravi committed
74
class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
75
    def test_simple_sphere(self, elevated_camera=False, check_depth=False):
facebook-github-bot's avatar
facebook-github-bot committed
76
        """
Patrick Labatut's avatar
Patrick Labatut committed
77
        Test output of phong and gouraud shading matches a reference image using
facebook-github-bot's avatar
facebook-github-bot committed
78
79
80
81
82
83
84
85
86
87
88
89
        the default values for the light sources.

        Args:
            elevated_camera: Defines whether the camera observing the scene should
                           have an elevation of 45 degrees.
        """
        device = torch.device("cuda:0")

        # Init mesh
        sphere_mesh = ico_sphere(5, device)
        verts_padded = sphere_mesh.verts_padded()
        faces_padded = sphere_mesh.faces_padded()
Nikhila Ravi's avatar
Nikhila Ravi committed
90
91
        feats = torch.ones_like(verts_padded, device=device)
        textures = TexturesVertex(verts_features=feats)
92
        sphere_mesh = Meshes(verts=verts_padded, faces=faces_padded, textures=textures)
facebook-github-bot's avatar
facebook-github-bot committed
93
94
95

        # Init rasterizer settings
        if elevated_camera:
96
97
            # Elevated and rotated camera
            R, T = look_at_view_transform(dist=2.7, elev=45.0, azim=45.0)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
98
            postfix = "_elevated_"
99
100
            # If y axis is up, the spot of light should
            # be on the bottom left of the sphere.
facebook-github-bot's avatar
facebook-github-bot committed
101
        else:
102
            # No elevation or azimuth rotation
facebook-github-bot's avatar
facebook-github-bot committed
103
            R, T = look_at_view_transform(2.7, 0.0, 0.0)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
            postfix = "_"
        for cam_type in (
            FoVPerspectiveCameras,
            FoVOrthographicCameras,
            PerspectiveCameras,
            OrthographicCameras,
        ):
            cameras = cam_type(device=device, R=R, T=T)

            # Init shader settings
            materials = Materials(device=device)
            lights = PointLights(device=device)
            lights.location = torch.tensor([0.0, 0.0, +2.0], device=device)[None]

            raster_settings = RasterizationSettings(
                image_size=512, blur_radius=0.0, faces_per_pixel=1
            )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
121
            blend_params = BlendParams(0.5, 1e-4, (0, 0, 0))
Georgia Gkioxari's avatar
Georgia Gkioxari committed
122
123

            # Test several shaders
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
124
125
126
127
128
129
130
131
132
133
134
135
            rasterizer_tests = [
                RasterizerTest(MeshRasterizer, HardPhongShader, "phong", "hard_phong"),
                RasterizerTest(
                    MeshRasterizer, HardGouraudShader, "gouraud", "hard_gouraud"
                ),
                RasterizerTest(MeshRasterizer, HardFlatShader, "flat", "hard_flat"),
                RasterizerTest(
                    MeshRasterizerOpenGL,
                    SplatterPhongShader,
                    "splatter",
                    "splatter_phong",
                ),
Nikhila Ravi's avatar
Nikhila Ravi committed
136
            ]
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
137
            for test in rasterizer_tests:
Nikhila Ravi's avatar
Nikhila Ravi committed
138
                shader = test.shader(
Georgia Gkioxari's avatar
Georgia Gkioxari committed
139
140
141
142
143
                    lights=lights,
                    cameras=cameras,
                    materials=materials,
                    blend_params=blend_params,
                )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
144
145
146
147
148
149
150
151
152
153
154
155
156
                if test.rasterizer == MeshRasterizer:
                    rasterizer = test.rasterizer(
                        cameras=cameras, raster_settings=raster_settings
                    )
                elif test.rasterizer == MeshRasterizerOpenGL:
                    if type(cameras) in [PerspectiveCameras, OrthographicCameras]:
                        # MeshRasterizerOpenGL is only compatible with FoV cameras.
                        continue
                    rasterizer = test.rasterizer(
                        cameras=cameras,
                        raster_settings=raster_settings,
                    )

157
158
159
160
161
162
                if check_depth:
                    renderer = MeshRendererWithFragments(
                        rasterizer=rasterizer, shader=shader
                    )
                    images, fragments = renderer(sphere_mesh)
                    self.assertClose(fragments.zbuf, rasterizer(sphere_mesh).zbuf)
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
163
164
165
166
167
168
169
170
                    # Check the alpha channel is the mask. For soft rasterizers, the
                    # boundary will not match exactly so we use quantiles to compare.
                    self.assertLess(
                        (
                            images[..., -1]
                            - (fragments.pix_to_face[..., 0] >= 0).float()
                        ).quantile(0.99),
                        0.005,
171
                    )
172
173
174
175
                else:
                    renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
                    images = renderer(sphere_mesh)

Georgia Gkioxari's avatar
Georgia Gkioxari committed
176
177
                rgb = images[0, ..., :3].squeeze().cpu()
                filename = "simple_sphere_light_%s%s%s.png" % (
Nikhila Ravi's avatar
Nikhila Ravi committed
178
                    test.reference_name,
Georgia Gkioxari's avatar
Georgia Gkioxari committed
179
180
181
                    postfix,
                    cam_type.__name__,
                )
facebook-github-bot's avatar
facebook-github-bot committed
182

Georgia Gkioxari's avatar
Georgia Gkioxari committed
183
184
                image_ref = load_rgb_image("test_%s" % filename, DATA_DIR)
                self.assertClose(rgb, image_ref, atol=0.05)
185

Georgia Gkioxari's avatar
Georgia Gkioxari committed
186
                if DEBUG:
Nikhila Ravi's avatar
Nikhila Ravi committed
187
188
189
190
191
192
                    debug_filename = "simple_sphere_light_%s%s%s.png" % (
                        test.debug_name,
                        postfix,
                        cam_type.__name__,
                    )
                    filename = "DEBUG_%s" % debug_filename
Georgia Gkioxari's avatar
Georgia Gkioxari committed
193
194
195
                    Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                        DATA_DIR / filename
                    )
Nikhila Ravi's avatar
Nikhila Ravi committed
196

Georgia Gkioxari's avatar
Georgia Gkioxari committed
197
198
199
200
201
202
203
            ########################################################
            # Move the light to the +z axis in world space so it is
            # behind the sphere. Note that +Z is in, +Y up,
            # +X left for both world and camera space.
            ########################################################
            lights.location[..., 2] = -2.0
            phong_shader = HardPhongShader(
204
205
206
207
208
                lights=lights,
                cameras=cameras,
                materials=materials,
                blend_params=blend_params,
            )
209
210
211
212
213
214
215
216
            if check_depth:
                phong_renderer = MeshRendererWithFragments(
                    rasterizer=rasterizer, shader=phong_shader
                )
                images, fragments = phong_renderer(sphere_mesh, lights=lights)
                self.assertClose(
                    fragments.zbuf, rasterizer(sphere_mesh, lights=lights).zbuf
                )
217
                # Check the alpha channel is the mask
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
218
219
220
221
222
                self.assertLess(
                    (
                        images[..., -1] - (fragments.pix_to_face[..., 0] >= 0).float()
                    ).quantile(0.99),
                    0.005,
223
                )
224
225
226
227
228
            else:
                phong_renderer = MeshRenderer(
                    rasterizer=rasterizer, shader=phong_shader
                )
                images = phong_renderer(sphere_mesh, lights=lights)
Nikhila Ravi's avatar
Nikhila Ravi committed
229
230
            rgb = images[0, ..., :3].squeeze().cpu()
            if DEBUG:
Georgia Gkioxari's avatar
Georgia Gkioxari committed
231
232
233
234
                filename = "DEBUG_simple_sphere_dark%s%s.png" % (
                    postfix,
                    cam_type.__name__,
                )
Nikhila Ravi's avatar
Nikhila Ravi committed
235
236
237
                Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / filename
                )
facebook-github-bot's avatar
facebook-github-bot committed
238

Georgia Gkioxari's avatar
Georgia Gkioxari committed
239
240
241
            image_ref_phong_dark = load_rgb_image(
                "test_simple_sphere_dark%s%s.png" % (postfix, cam_type.__name__),
                DATA_DIR,
facebook-github-bot's avatar
facebook-github-bot committed
242
            )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
243
244
245
            # Soft shaders (SplatterPhong) will have a different boundary than hard
            # ones, but should be identical otherwise.
            self.assertLess((rgb - image_ref_phong_dark).quantile(0.99), 0.005)
facebook-github-bot's avatar
facebook-github-bot committed
246
247
248

    def test_simple_sphere_elevated_camera(self):
        """
Patrick Labatut's avatar
Patrick Labatut committed
249
        Test output of phong and gouraud shading matches a reference image using
facebook-github-bot's avatar
facebook-github-bot committed
250
251
252
253
254
255
        the default values for the light sources.

        The rendering is performed with a camera that has non-zero elevation.
        """
        self.test_simple_sphere(elevated_camera=True)

256
257
258
259
260
261
262
263
264
    def test_simple_sphere_depth(self):
        """
        Test output of phong and gouraud shading matches a reference image using
        the default values for the light sources.

        The rendering is performed with a camera that has non-zero elevation.
        """
        self.test_simple_sphere(check_depth=True)

Georgia Gkioxari's avatar
Georgia Gkioxari committed
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
    def test_simple_sphere_screen(self):

        """
        Test output when rendering with PerspectiveCameras & OrthographicCameras
        in NDC vs screen space.
        """
        device = torch.device("cuda:0")

        # Init mesh
        sphere_mesh = ico_sphere(5, device)
        verts_padded = sphere_mesh.verts_padded()
        faces_padded = sphere_mesh.faces_padded()
        feats = torch.ones_like(verts_padded, device=device)
        textures = TexturesVertex(verts_features=feats)
        sphere_mesh = Meshes(verts=verts_padded, faces=faces_padded, textures=textures)

        R, T = look_at_view_transform(2.7, 0.0, 0.0)

        # Init shader settings
        materials = Materials(device=device)
        lights = PointLights(device=device)
        lights.location = torch.tensor([0.0, 0.0, +2.0], device=device)[None]

        raster_settings = RasterizationSettings(
            image_size=512, blur_radius=0.0, faces_per_pixel=1
        )
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
291
        half_half = (512.0 / 2.0, 512.0 / 2.0)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
292
293
294
295
296
        for cam_type in (PerspectiveCameras, OrthographicCameras):
            cameras = cam_type(
                device=device,
                R=R,
                T=T,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
297
298
                principal_point=(half_half,),
                focal_length=(half_half,),
Georgia Gkioxari's avatar
Georgia Gkioxari committed
299
                image_size=((512, 512),),
300
                in_ndc=False,
Georgia Gkioxari's avatar
Georgia Gkioxari committed
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
            )
            rasterizer = MeshRasterizer(
                cameras=cameras, raster_settings=raster_settings
            )
            blend_params = BlendParams(1e-4, 1e-4, (0, 0, 0))

            shader = HardPhongShader(
                lights=lights,
                cameras=cameras,
                materials=materials,
                blend_params=blend_params,
            )
            renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
            images = renderer(sphere_mesh)
            rgb = images[0, ..., :3].squeeze().cpu()
            filename = "test_simple_sphere_light_phong_%s.png" % cam_type.__name__
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
317
318
319
320
            if DEBUG:
                Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / f"{filename}_.png"
                )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
321
322
323
324

            image_ref = load_rgb_image(filename, DATA_DIR)
            self.assertClose(rgb, image_ref, atol=0.05)

facebook-github-bot's avatar
facebook-github-bot committed
325
326
    def test_simple_sphere_batched(self):
        """
Nikhila Ravi's avatar
Nikhila Ravi committed
327
        Test a mesh with vertex textures can be extended to form a batch, and
Nikhila Ravi's avatar
Nikhila Ravi committed
328
329
        is rendered correctly with Phong, Gouraud and Flat Shaders with batched
        lighting and hard and soft blending.
facebook-github-bot's avatar
facebook-github-bot committed
330
        """
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
331
        batch_size = 3
facebook-github-bot's avatar
facebook-github-bot committed
332
333
        device = torch.device("cuda:0")

Nikhila Ravi's avatar
Nikhila Ravi committed
334
        # Init mesh with vertex textures.
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
335
        sphere_meshes = ico_sphere(3, device).extend(batch_size)
facebook-github-bot's avatar
facebook-github-bot committed
336
337
        verts_padded = sphere_meshes.verts_padded()
        faces_padded = sphere_meshes.faces_padded()
Nikhila Ravi's avatar
Nikhila Ravi committed
338
339
        feats = torch.ones_like(verts_padded, device=device)
        textures = TexturesVertex(verts_features=feats)
facebook-github-bot's avatar
facebook-github-bot committed
340
341
342
343
344
        sphere_meshes = Meshes(
            verts=verts_padded, faces=faces_padded, textures=textures
        )

        # Init rasterizer settings
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
345
        dist = torch.tensor([2, 4, 6]).to(device)
facebook-github-bot's avatar
facebook-github-bot committed
346
347
348
        elev = torch.zeros_like(dist)
        azim = torch.zeros_like(dist)
        R, T = look_at_view_transform(dist, elev, azim)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
349
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
facebook-github-bot's avatar
facebook-github-bot committed
350
        raster_settings = RasterizationSettings(
Nikhila Ravi's avatar
Nikhila Ravi committed
351
            image_size=512, blur_radius=0.0, faces_per_pixel=4
facebook-github-bot's avatar
facebook-github-bot committed
352
353
354
355
        )

        # Init shader settings
        materials = Materials(device=device)
Nikhila Ravi's avatar
Nikhila Ravi committed
356
357
358
        lights_location = torch.tensor([0.0, 0.0, +2.0], device=device)
        lights_location = lights_location[None].expand(batch_size, -1)
        lights = PointLights(device=device, location=lights_location)
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
359
        blend_params = BlendParams(0.5, 1e-4, (0, 0, 0))
facebook-github-bot's avatar
facebook-github-bot committed
360
361

        # Init renderer
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
362
363
364
365
366
367
368
369
370
371
372
373
        rasterizer_tests = [
            RasterizerTest(MeshRasterizer, HardPhongShader, "phong", "hard_phong"),
            RasterizerTest(
                MeshRasterizer, HardGouraudShader, "gouraud", "hard_gouraud"
            ),
            RasterizerTest(MeshRasterizer, HardFlatShader, "flat", "hard_flat"),
            RasterizerTest(
                MeshRasterizerOpenGL,
                SplatterPhongShader,
                "splatter",
                "splatter_phong",
            ),
Nikhila Ravi's avatar
Nikhila Ravi committed
374
        ]
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
375
        for test in rasterizer_tests:
Nikhila Ravi's avatar
Nikhila Ravi committed
376
377
            reference_name = test.reference_name
            debug_name = test.debug_name
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
378
379
380
381
            rasterizer = test.rasterizer(
                cameras=cameras, raster_settings=raster_settings
            )

Nikhila Ravi's avatar
Nikhila Ravi committed
382
            shader = test.shader(
383
384
385
386
387
                lights=lights,
                cameras=cameras,
                materials=materials,
                blend_params=blend_params,
            )
Nikhila Ravi's avatar
Nikhila Ravi committed
388
389
390
            renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
            images = renderer(sphere_meshes)
            for i in range(batch_size):
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
391
392
393
394
395
                image_ref = load_rgb_image(
                    "test_simple_sphere_batched_%s_%s_%s.png"
                    % (reference_name, type(cameras).__name__, i),
                    DATA_DIR,
                )
Nikhila Ravi's avatar
Nikhila Ravi committed
396
                rgb = images[i, ..., :3].squeeze().cpu()
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
397
398
                if DEBUG:
                    filename = "DEBUG_simple_sphere_batched_%s_%s_%s.png" % (
Nikhila Ravi's avatar
Nikhila Ravi committed
399
                        debug_name,
Georgia Gkioxari's avatar
Georgia Gkioxari committed
400
                        type(cameras).__name__,
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
401
                        i,
Georgia Gkioxari's avatar
Georgia Gkioxari committed
402
                    )
Nikhila Ravi's avatar
Nikhila Ravi committed
403
404
405
                    Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                        DATA_DIR / filename
                    )
Nikhila Ravi's avatar
Nikhila Ravi committed
406
                self.assertClose(rgb, image_ref, atol=0.05)
facebook-github-bot's avatar
facebook-github-bot committed
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421

    def test_silhouette_with_grad(self):
        """
        Test silhouette blending. Also check that gradient calculation works.
        """
        device = torch.device("cuda:0")
        sphere_mesh = ico_sphere(5, device)
        verts, faces = sphere_mesh.get_mesh_verts_faces(0)
        sphere_mesh = Meshes(verts=[verts], faces=[faces])

        blend_params = BlendParams(sigma=1e-4, gamma=1e-4)
        raster_settings = RasterizationSettings(
            image_size=512,
            blur_radius=np.log(1.0 / 1e-4 - 1.0) * blend_params.sigma,
            faces_per_pixel=80,
422
            clip_barycentric_coords=True,
facebook-github-bot's avatar
facebook-github-bot committed
423
424
425
        )

        # Init rasterizer settings
426
        R, T = look_at_view_transform(2.7, 0, 0)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
427
428
429
430
431
432
433
434
435
436
437
438
439
440
        for cam_type in (
            FoVPerspectiveCameras,
            FoVOrthographicCameras,
            PerspectiveCameras,
            OrthographicCameras,
        ):
            cameras = cam_type(device=device, R=R, T=T)

            # Init renderer
            renderer = MeshRenderer(
                rasterizer=MeshRasterizer(
                    cameras=cameras, raster_settings=raster_settings
                ),
                shader=SoftSilhouetteShader(blend_params=blend_params),
facebook-github-bot's avatar
facebook-github-bot committed
441
            )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
442
443
444
445
446
447
448
449
450
            images = renderer(sphere_mesh)
            alpha = images[0, ..., 3].squeeze().cpu()
            if DEBUG:
                filename = os.path.join(
                    DATA_DIR, "DEBUG_%s_silhouette.png" % (cam_type.__name__)
                )
                Image.fromarray((alpha.detach().numpy() * 255).astype(np.uint8)).save(
                    filename
                )
facebook-github-bot's avatar
facebook-github-bot committed
451

Georgia Gkioxari's avatar
Georgia Gkioxari committed
452
453
454
455
            ref_filename = "test_%s_silhouette.png" % (cam_type.__name__)
            image_ref_filename = DATA_DIR / ref_filename
            with Image.open(image_ref_filename) as raw_image_ref:
                image_ref = torch.from_numpy(np.array(raw_image_ref))
Nikhila Ravi's avatar
Nikhila Ravi committed
456

Georgia Gkioxari's avatar
Georgia Gkioxari committed
457
458
            image_ref = image_ref.to(dtype=torch.float32) / 255.0
            self.assertClose(alpha, image_ref, atol=0.055)
facebook-github-bot's avatar
facebook-github-bot committed
459

Georgia Gkioxari's avatar
Georgia Gkioxari committed
460
461
462
463
464
465
            # Check grad exist
            verts.requires_grad = True
            sphere_mesh = Meshes(verts=[verts], faces=[faces])
            images = renderer(sphere_mesh)
            images[0, ...].sum().backward()
            self.assertIsNotNone(verts.grad)
facebook-github-bot's avatar
facebook-github-bot committed
466
467
468

    def test_texture_map(self):
        """
469
470
        Test a mesh with a texture map is loaded and rendered correctly.
        The pupils in the eyes of the cow should always be looking to the left.
facebook-github-bot's avatar
facebook-github-bot committed
471
        """
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
472
473
474
475
476
477
478
479
480
481
        self._texture_map_per_rasterizer(MeshRasterizer)

    def test_texture_map_opengl(self):
        """
        Test a mesh with a texture map is loaded and rendered correctly.
        The pupils in the eyes of the cow should always be looking to the left.
        """
        self._texture_map_per_rasterizer(MeshRasterizerOpenGL)

    def _texture_map_per_rasterizer(self, rasterizer_type):
facebook-github-bot's avatar
facebook-github-bot committed
482
        device = torch.device("cuda:0")
483
484

        obj_filename = TUTORIAL_DATA_DIR / "cow_mesh/cow.obj"
facebook-github-bot's avatar
facebook-github-bot committed
485
486

        # Load mesh + texture
Nikhila Ravi's avatar
Nikhila Ravi committed
487
488
489
490
491
492
493
494
495
        verts, faces, aux = load_obj(
            obj_filename, device=device, load_textures=True, texture_wrap=None
        )
        tex_map = list(aux.texture_images.values())[0]
        tex_map = tex_map[None, ...].to(faces.textures_idx.device)
        textures = TexturesUV(
            maps=tex_map, faces_uvs=[faces.textures_idx], verts_uvs=[aux.verts_uvs]
        )
        mesh = Meshes(verts=[verts], faces=[faces.verts_idx], textures=textures)
facebook-github-bot's avatar
facebook-github-bot committed
496
497

        # Init rasterizer settings
498
        R, T = look_at_view_transform(2.7, 0, 0)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
499
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
Nikhila Ravi's avatar
Nikhila Ravi committed
500

facebook-github-bot's avatar
facebook-github-bot committed
501
        raster_settings = RasterizationSettings(
Nikhila Ravi's avatar
Nikhila Ravi committed
502
            image_size=512, blur_radius=0.0, faces_per_pixel=1
facebook-github-bot's avatar
facebook-github-bot committed
503
504
505
506
507
        )

        # Init shader settings
        materials = Materials(device=device)
        lights = PointLights(device=device)
508
509
510
511

        # Place light behind the cow in world space. The front of
        # the cow is facing the -z direction.
        lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None]
facebook-github-bot's avatar
facebook-github-bot committed
512

513
        blend_params = BlendParams(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
514
            sigma=1e-1 if rasterizer_type == MeshRasterizer else 0.5,
515
516
517
            gamma=1e-4,
            background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
        )
facebook-github-bot's avatar
facebook-github-bot committed
518
        # Init renderer
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
519
520
521
        rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
        if rasterizer_type == MeshRasterizer:
            shader = TexturedSoftPhongShader(
522
523
524
525
                lights=lights,
                cameras=cameras,
                materials=materials,
                blend_params=blend_params,
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
526
527
528
529
530
531
532
533
534
            )
        elif rasterizer_type == MeshRasterizerOpenGL:
            shader = SplatterPhongShader(
                lights=lights,
                cameras=cameras,
                materials=materials,
                blend_params=blend_params,
            )
        renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
facebook-github-bot's avatar
facebook-github-bot committed
535
536

        # Load reference image
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
537
538
539
        image_ref = load_rgb_image(
            f"test_texture_map_back_{rasterizer_type.__name__}.png", DATA_DIR
        )
facebook-github-bot's avatar
facebook-github-bot committed
540

Nikhila Ravi's avatar
Nikhila Ravi committed
541
        for bin_size in [0, None]:
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
542
543
544
            if rasterizer_type == MeshRasterizerOpenGL and bin_size == 0:
                # MeshRasterizerOpenGL does not use this parameter.
                continue
Nikhila Ravi's avatar
Nikhila Ravi committed
545
546
547
548
            # Check both naive and coarse to fine produce the same output.
            renderer.rasterizer.raster_settings.bin_size = bin_size
            images = renderer(mesh)
            rgb = images[0, ..., :3].squeeze().cpu()
facebook-github-bot's avatar
facebook-github-bot committed
549

Nikhila Ravi's avatar
Nikhila Ravi committed
550
551
            if DEBUG:
                Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
552
                    DATA_DIR / f"DEBUG_texture_map_back_{rasterizer_type.__name__}.png"
Nikhila Ravi's avatar
Nikhila Ravi committed
553
554
555
556
557
558
                )

            # NOTE some pixels can be flaky and will not lead to
            # `cond1` being true. Add `cond2` and check `cond1 or cond2`
            cond1 = torch.allclose(rgb, image_ref, atol=0.05)
            cond2 = ((rgb - image_ref).abs() > 0.05).sum() < 5
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
559
            # self.assertTrue(cond1 or cond2)
facebook-github-bot's avatar
facebook-github-bot committed
560
561

        # Check grad exists
562
        [verts] = mesh.verts_list()
facebook-github-bot's avatar
facebook-github-bot committed
563
        verts.requires_grad = True
564
        mesh2 = Meshes(verts=[verts], faces=mesh.faces_list(), textures=mesh.textures)
565
        images = renderer(mesh2)
facebook-github-bot's avatar
facebook-github-bot committed
566
567
        images[0, ...].sum().backward()
        self.assertIsNotNone(verts.grad)
568

569
570
571
572
573
        ##########################################
        # Check rendering of the front of the cow
        ##########################################

        R, T = look_at_view_transform(2.7, 0, 180)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
574
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
575
576
577
578
579

        # Move light to the front of the cow in world space
        lights.location = torch.tensor([0.0, 0.0, -2.0], device=device)[None]

        # Load reference image
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
580
581
582
        image_ref = load_rgb_image(
            f"test_texture_map_front_{rasterizer_type.__name__}.png", DATA_DIR
        )
583

Nikhila Ravi's avatar
Nikhila Ravi committed
584
        for bin_size in [0, None]:
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
585
586
587
            if rasterizer == MeshRasterizerOpenGL and bin_size == 0:
                # MeshRasterizerOpenGL does not use this parameter.
                continue
Nikhila Ravi's avatar
Nikhila Ravi committed
588
589
590
591
592
593
594
595
            # Check both naive and coarse to fine produce the same output.
            renderer.rasterizer.raster_settings.bin_size = bin_size

            images = renderer(mesh, cameras=cameras, lights=lights)
            rgb = images[0, ..., :3].squeeze().cpu()

            if DEBUG:
                Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
596
                    DATA_DIR / f"DEBUG_texture_map_front_{rasterizer_type.__name__}.png"
Nikhila Ravi's avatar
Nikhila Ravi committed
597
598
599
600
601
602
603
                )

            # NOTE some pixels can be flaky and will not lead to
            # `cond1` being true. Add `cond2` and check `cond1 or cond2`
            cond1 = torch.allclose(rgb, image_ref, atol=0.05)
            cond2 = ((rgb - image_ref).abs() > 0.05).sum() < 5
            self.assertTrue(cond1 or cond2)
604

605
606
607
        #################################
        # Add blurring to rasterization
        #################################
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
        if rasterizer_type == MeshRasterizer:
            # Note that MeshRasterizer can blur the images arbitrarily, however
            # MeshRasterizerOpenGL is limited by its kernel size (currently 3 px^2),
            # so this test only makes sense for MeshRasterizer.
            R, T = look_at_view_transform(2.7, 0, 180)
            cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
            # For MeshRasterizer, blurring is controlled by blur_radius. For
            # MeshRasterizerOpenGL, by sigma.
            blend_params = BlendParams(sigma=5e-4, gamma=1e-4)
            raster_settings = RasterizationSettings(
                image_size=512,
                blur_radius=np.log(1.0 / 1e-4 - 1.0) * blend_params.sigma,
                faces_per_pixel=100,
                clip_barycentric_coords=True,
                perspective_correct=rasterizer_type.__name__ == "MeshRasterizerOpenGL",
            )
624

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
625
626
            # Load reference image
            image_ref = load_rgb_image("test_blurry_textured_rendering.png", DATA_DIR)
Nikhila Ravi's avatar
Nikhila Ravi committed
627

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
628
629
630
            for bin_size in [0, None]:
                # Check both naive and coarse to fine produce the same output.
                renderer.rasterizer.raster_settings.bin_size = bin_size
Nikhila Ravi's avatar
Nikhila Ravi committed
631

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
632
633
634
635
636
                images = renderer(
                    mesh.clone(),
                    cameras=cameras,
                    raster_settings=raster_settings,
                    blend_params=blend_params,
Nikhila Ravi's avatar
Nikhila Ravi committed
637
                )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
638
                rgb = images[0, ..., :3].squeeze().cpu()
639

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
640
641
642
643
644
645
                if DEBUG:
                    Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                        DATA_DIR / "DEBUG_blurry_textured_rendering.png"
                    )

                self.assertClose(rgb, image_ref, atol=0.05)
646

647
    def test_batch_uvs(self):
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
648
649
650
651
652
653
        self._batch_uvs(MeshRasterizer)

    def test_batch_uvs_opengl(self):
        self._batch_uvs(MeshRasterizer)

    def _batch_uvs(self, rasterizer_type):
654
655
656
        """Test that two random tori with TexturesUV render the same as each individually."""
        torch.manual_seed(1)
        device = torch.device("cuda:0")
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
657

658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
        plain_torus = torus(r=1, R=4, sides=10, rings=10, device=device)
        [verts] = plain_torus.verts_list()
        [faces] = plain_torus.faces_list()
        nocolor = torch.zeros((100, 100), device=device)
        color_gradient = torch.linspace(0, 1, steps=100, device=device)
        color_gradient1 = color_gradient[None].expand_as(nocolor)
        color_gradient2 = color_gradient[:, None].expand_as(nocolor)
        colors1 = torch.stack([nocolor, color_gradient1, color_gradient2], dim=2)
        colors2 = torch.stack([color_gradient1, color_gradient2, nocolor], dim=2)
        verts_uvs1 = torch.rand(size=(verts.shape[0], 2), device=device)
        verts_uvs2 = torch.rand(size=(verts.shape[0], 2), device=device)

        textures1 = TexturesUV(
            maps=[colors1], faces_uvs=[faces], verts_uvs=[verts_uvs1]
        )
        textures2 = TexturesUV(
            maps=[colors2], faces_uvs=[faces], verts_uvs=[verts_uvs2]
        )
        mesh1 = Meshes(verts=[verts], faces=[faces], textures=textures1)
        mesh2 = Meshes(verts=[verts], faces=[faces], textures=textures2)
        mesh_both = join_meshes_as_batch([mesh1, mesh2])

        R, T = look_at_view_transform(10, 10, 0)
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)

        raster_settings = RasterizationSettings(
            image_size=128, blur_radius=0.0, faces_per_pixel=1
        )

        # Init shader settings
        lights = PointLights(device=device)
        lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None]

        blend_params = BlendParams(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
692
            sigma=0.5,
693
694
695
696
            gamma=1e-4,
            background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
        )
        # Init renderer
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
697
698
699
        rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
        if rasterizer_type == MeshRasterizer:
            shader = HardPhongShader(
700
                device=device, lights=lights, cameras=cameras, blend_params=blend_params
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
701
702
703
704
705
706
707
            )
        else:
            shader = SplatterPhongShader(
                device=device, lights=lights, cameras=cameras, blend_params=blend_params
            )

        renderer = MeshRenderer(rasterizer, shader)
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738

        outputs = []
        for meshes in [mesh_both, mesh1, mesh2]:
            outputs.append(renderer(meshes))

        if DEBUG:
            Image.fromarray(
                (outputs[0][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
            ).save(DATA_DIR / "test_batch_uvs0.png")
            Image.fromarray(
                (outputs[1][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
            ).save(DATA_DIR / "test_batch_uvs1.png")
            Image.fromarray(
                (outputs[0][1, ..., :3].cpu().numpy() * 255).astype(np.uint8)
            ).save(DATA_DIR / "test_batch_uvs2.png")
            Image.fromarray(
                (outputs[2][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
            ).save(DATA_DIR / "test_batch_uvs3.png")

            diff = torch.abs(outputs[0][0, ..., :3] - outputs[1][0, ..., :3])
            Image.fromarray(((diff > 1e-5).cpu().numpy().astype(np.uint8) * 255)).save(
                DATA_DIR / "test_batch_uvs01.png"
            )
            diff = torch.abs(outputs[0][1, ..., :3] - outputs[2][0, ..., :3])
            Image.fromarray(((diff > 1e-5).cpu().numpy().astype(np.uint8) * 255)).save(
                DATA_DIR / "test_batch_uvs23.png"
            )

        self.assertClose(outputs[0][0, ..., :3], outputs[1][0, ..., :3], atol=1e-5)
        self.assertClose(outputs[0][1, ..., :3], outputs[2][0, ..., :3], atol=1e-5)

739
    def test_join_uvs(self):
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
740
741
742
743
744
745
        self._join_uvs(MeshRasterizer)

    def test_join_uvs_opengl(self):
        self._join_uvs(MeshRasterizerOpenGL)

    def _join_uvs(self, rasterizer_type):
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
        """Meshes with TexturesUV joined into a scene"""
        # Test the result of rendering three tori with separate textures.
        # The expected result is consistent with rendering them each alone.
        # This tests TexturesUV.join_scene with rectangle flipping,
        # and we check the form of the merged map as well.
        torch.manual_seed(1)
        device = torch.device("cuda:0")

        R, T = look_at_view_transform(18, 0, 0)
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)

        raster_settings = RasterizationSettings(
            image_size=256, blur_radius=0.0, faces_per_pixel=1
        )

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
761
        lights = AmbientLights(device=device)
762
        blend_params = BlendParams(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
763
            sigma=0.5,
764
765
766
            gamma=1e-4,
            background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
        )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
767
768
769
        rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
        if rasterizer_type == MeshRasterizer:
            shader = HardPhongShader(
770
                device=device, blend_params=blend_params, cameras=cameras, lights=lights
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
771
772
773
774
775
776
            )
        else:
            shader = SplatterPhongShader(
                device=device, blend_params=blend_params, cameras=cameras, lights=lights
            )
        renderer = MeshRenderer(rasterizer, shader)
777
778
779
780
781
782
783
784
785

        plain_torus = torus(r=1, R=4, sides=5, rings=6, device=device)
        [verts] = plain_torus.verts_list()
        verts_shifted1 = verts.clone()
        verts_shifted1 *= 0.5
        verts_shifted1[:, 1] += 7
        verts_shifted2 = verts.clone()
        verts_shifted2 *= 0.5
        verts_shifted2[:, 1] -= 7
786
787
788
        verts_shifted3 = verts.clone()
        verts_shifted3 *= 0.5
        verts_shifted3[:, 1] -= 700
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833

        [faces] = plain_torus.faces_list()
        nocolor = torch.zeros((100, 100), device=device)
        color_gradient = torch.linspace(0, 1, steps=100, device=device)
        color_gradient1 = color_gradient[None].expand_as(nocolor)
        color_gradient2 = color_gradient[:, None].expand_as(nocolor)
        colors1 = torch.stack([nocolor, color_gradient1, color_gradient2], dim=2)
        colors2 = torch.stack([color_gradient1, color_gradient2, nocolor], dim=2)
        verts_uvs1 = torch.rand(size=(verts.shape[0], 2), device=device)
        verts_uvs2 = torch.rand(size=(verts.shape[0], 2), device=device)

        for i, align_corners, padding_mode in [
            (0, True, "border"),
            (1, False, "border"),
            (2, False, "zeros"),
        ]:
            textures1 = TexturesUV(
                maps=[colors1],
                faces_uvs=[faces],
                verts_uvs=[verts_uvs1],
                align_corners=align_corners,
                padding_mode=padding_mode,
            )

            # These downsamplings of colors2 are chosen to ensure a flip and a non flip
            # when the maps are merged.
            # We have maps of size (100, 100), (50, 99) and (99, 50).
            textures2 = TexturesUV(
                maps=[colors2[::2, :-1]],
                faces_uvs=[faces],
                verts_uvs=[verts_uvs2],
                align_corners=align_corners,
                padding_mode=padding_mode,
            )
            offset = torch.tensor([0, 0, 0.5], device=device)
            textures3 = TexturesUV(
                maps=[colors2[:-1, ::2] + offset],
                faces_uvs=[faces],
                verts_uvs=[verts_uvs2],
                align_corners=align_corners,
                padding_mode=padding_mode,
            )
            mesh1 = Meshes(verts=[verts], faces=[faces], textures=textures1)
            mesh2 = Meshes(verts=[verts_shifted1], faces=[faces], textures=textures2)
            mesh3 = Meshes(verts=[verts_shifted2], faces=[faces], textures=textures3)
834
835
836
837
838
            # mesh4 is like mesh1 but outside the field of view. It is here to test
            # that having another texture with the same map doesn't produce
            # two copies in the joined map.
            mesh4 = Meshes(verts=[verts_shifted3], faces=[faces], textures=textures1)
            mesh = join_meshes_as_scene([mesh1, mesh2, mesh3, mesh4])
839
840
841
842
843
844
845
846
847

            output = renderer(mesh)[0, ..., :3].cpu()
            output1 = renderer(mesh1)[0, ..., :3].cpu()
            output2 = renderer(mesh2)[0, ..., :3].cpu()
            output3 = renderer(mesh3)[0, ..., :3].cpu()
            # The background color is white and the objects do not overlap, so we can
            # predict the merged image by taking the minimum over every channel
            merged = torch.min(torch.min(output1, output2), output3)

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
848
849
850
            image_ref = load_rgb_image(
                f"test_joinuvs{i}_{rasterizer_type.__name__}_final.png", DATA_DIR
            )
851
852
853
854
            map_ref = load_rgb_image(f"test_joinuvs{i}_map.png", DATA_DIR)

            if DEBUG:
                Image.fromarray((output.numpy() * 255).astype(np.uint8)).save(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
855
856
                    DATA_DIR
                    / f"DEBUG_test_joinuvs{i}_{rasterizer_type.__name__}_final.png"
857
                )
858
                Image.fromarray((merged.numpy() * 255).astype(np.uint8)).save(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
859
860
                    DATA_DIR
                    / f"DEBUG_test_joinuvs{i}_{rasterizer_type.__name__}_merged.png"
861
862
863
                )

                Image.fromarray((output1.numpy() * 255).astype(np.uint8)).save(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
864
                    DATA_DIR / f"DEBUG_test_joinuvs{i}_{rasterizer_type.__name__}_1.png"
865
866
                )
                Image.fromarray((output2.numpy() * 255).astype(np.uint8)).save(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
867
                    DATA_DIR / f"DEBUG_test_joinuvs{i}_{rasterizer_type.__name__}_2.png"
868
869
                )
                Image.fromarray((output3.numpy() * 255).astype(np.uint8)).save(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
870
                    DATA_DIR / f"DEBUG_test_joinuvs{i}_{rasterizer_type.__name__}_3.png"
871
872
873
874
875
                )
                Image.fromarray(
                    (mesh.textures.maps_padded()[0].cpu().numpy() * 255).astype(
                        np.uint8
                    )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
876
                ).save(DATA_DIR / f"DEBUG_test_joinuvs{i}_map.png")
877
878
879
880
                Image.fromarray(
                    (mesh2.textures.maps_padded()[0].cpu().numpy() * 255).astype(
                        np.uint8
                    )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
881
                ).save(DATA_DIR / f"DEBUG_test_joinuvs{i}_map2.png")
882
883
884
885
                Image.fromarray(
                    (mesh3.textures.maps_padded()[0].cpu().numpy() * 255).astype(
                        np.uint8
                    )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
886
                ).save(DATA_DIR / f"DEBUG_test_joinuvs{i}_map3.png")
887

888
889
            self.assertClose(output, merged)
            self.assertClose(output, image_ref, atol=0.005)
890
891
            self.assertClose(mesh.textures.maps_padded()[0].cpu(), map_ref, atol=0.05)

892
893
894
895
896
897
898
899
900
901
902
903
    def test_join_uvs_simple(self):
        # Example from issue #826
        a = TexturesUV(
            maps=torch.full((1, 4000, 4000, 3), 0.8),
            faces_uvs=torch.arange(300).reshape(1, 100, 3),
            verts_uvs=torch.rand(1, 300, 2) * 0.4 + 0.1,
        )
        b = TexturesUV(
            maps=torch.full((1, 2000, 2000, 3), 0.7),
            faces_uvs=torch.arange(150).reshape(1, 50, 3),
            verts_uvs=torch.rand(1, 150, 2) * 0.2 + 0.3,
        )
904
905
        self.assertEqual(a._num_faces_per_mesh, [100])
        self.assertEqual(b._num_faces_per_mesh, [50])
906
        c = a.join_batch([b]).join_scene()
907
908
909
        self.assertEqual(a._num_faces_per_mesh, [100])
        self.assertEqual(b._num_faces_per_mesh, [50])
        self.assertEqual(c._num_faces_per_mesh, [150])
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927

        color = c.faces_verts_textures_packed()
        color1 = color[:100, :, 0].flatten()
        color2 = color[100:, :, 0].flatten()
        expect1 = color1.new_tensor(0.8)
        expect2 = color2.new_tensor(0.7)
        self.assertClose(color1.min(), expect1)
        self.assertClose(color1.max(), expect1)
        self.assertClose(color2.min(), expect2)
        self.assertClose(color2.max(), expect2)

        if DEBUG:
            from pytorch3d.vis.texture_vis import texturesuv_image_PIL as PI

            PI(a, radius=5).save(DATA_DIR / "test_join_uvs_simple_a.png")
            PI(b, radius=5).save(DATA_DIR / "test_join_uvs_simple_b.png")
            PI(c, radius=5).save(DATA_DIR / "test_join_uvs_simple_c.png")

928
    def test_join_verts(self):
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
929
930
931
932
933
934
        self._join_verts(MeshRasterizer)

    def test_join_verts_opengl(self):
        self._join_verts(MeshRasterizerOpenGL)

    def _join_verts(self, rasterizer_type):
935
936
937
938
939
        """Meshes with TexturesVertex joined into a scene"""
        # Test the result of rendering two tori with separate textures.
        # The expected result is consistent with rendering them each alone.
        torch.manual_seed(1)
        device = torch.device("cuda:0")
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
940

941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
        plain_torus = torus(r=1, R=4, sides=5, rings=6, device=device)
        [verts] = plain_torus.verts_list()
        verts_shifted1 = verts.clone()
        verts_shifted1 *= 0.5
        verts_shifted1[:, 1] += 7

        faces = plain_torus.faces_list()
        textures1 = TexturesVertex(verts_features=[torch.rand_like(verts)])
        textures2 = TexturesVertex(verts_features=[torch.rand_like(verts)])
        mesh1 = Meshes(verts=[verts], faces=faces, textures=textures1)
        mesh2 = Meshes(verts=[verts_shifted1], faces=faces, textures=textures2)
        mesh = join_meshes_as_scene([mesh1, mesh2])

        R, T = look_at_view_transform(18, 0, 0)
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)

        raster_settings = RasterizationSettings(
            image_size=256, blur_radius=0.0, faces_per_pixel=1
        )

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
961
        lights = AmbientLights(device=device)
962
        blend_params = BlendParams(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
963
            sigma=0.5,
964
965
966
            gamma=1e-4,
            background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
        )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
967
968
969
        rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
        if rasterizer_type == MeshRasterizer:
            shader = HardPhongShader(
970
                device=device, blend_params=blend_params, cameras=cameras, lights=lights
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
971
972
973
974
975
976
977
            )
        else:
            shader = SplatterPhongShader(
                device=device, blend_params=blend_params, cameras=cameras, lights=lights
            )

        renderer = MeshRenderer(rasterizer, shader)
978
979
980

        output = renderer(mesh)

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
981
982
983
        image_ref = load_rgb_image(
            f"test_joinverts_final_{rasterizer_type.__name__}.png", DATA_DIR
        )
984
985
986
987
988
989
990

        if DEBUG:
            debugging_outputs = []
            for mesh_ in [mesh1, mesh2]:
                debugging_outputs.append(renderer(mesh_))
            Image.fromarray(
                (output[0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
991
992
993
            ).save(
                DATA_DIR / f"DEBUG_test_joinverts_final_{rasterizer_type.__name__}.png"
            )
994
995
            Image.fromarray(
                (debugging_outputs[0][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
996
            ).save(DATA_DIR / "DEBUG_test_joinverts_1.png")
997
998
            Image.fromarray(
                (debugging_outputs[1][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
999
            ).save(DATA_DIR / "DEBUG_test_joinverts_2.png")
1000
1001
1002
1003
1004

        result = output[0, ..., :3].cpu()
        self.assertClose(result, image_ref, atol=0.05)

    def test_join_atlas(self):
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1005
1006
1007
1008
1009
1010
        self._join_atlas(MeshRasterizer)

    def test_join_atlas_opengl(self):
        self._join_atlas(MeshRasterizerOpenGL)

    def _join_atlas(self, rasterizer_type):
1011
1012
1013
1014
1015
        """Meshes with TexturesAtlas joined into a scene"""
        # Test the result of rendering two tori with separate textures.
        # The expected result is consistent with rendering them each alone.
        torch.manual_seed(1)
        device = torch.device("cuda:0")
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
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
        plain_torus = torus(r=1, R=4, sides=5, rings=6, device=device)
        [verts] = plain_torus.verts_list()
        verts_shifted1 = verts.clone()
        verts_shifted1 *= 1.2
        verts_shifted1[:, 0] += 4
        verts_shifted1[:, 1] += 5
        verts[:, 0] -= 4
        verts[:, 1] -= 4

        [faces] = plain_torus.faces_list()
        map_size = 3
        # Two random atlases.
        # The averaging of the random numbers here is not consistent with the
        # meaning of the atlases, but makes each face a bit smoother than
        # if everything had a random color.
        atlas1 = torch.rand(size=(faces.shape[0], map_size, map_size, 3), device=device)
        atlas1[:, 1] = 0.5 * atlas1[:, 0] + 0.5 * atlas1[:, 2]
        atlas1[:, :, 1] = 0.5 * atlas1[:, :, 0] + 0.5 * atlas1[:, :, 2]
        atlas2 = torch.rand(size=(faces.shape[0], map_size, map_size, 3), device=device)
        atlas2[:, 1] = 0.5 * atlas2[:, 0] + 0.5 * atlas2[:, 2]
        atlas2[:, :, 1] = 0.5 * atlas2[:, :, 0] + 0.5 * atlas2[:, :, 2]

        textures1 = TexturesAtlas(atlas=[atlas1])
        textures2 = TexturesAtlas(atlas=[atlas2])
        mesh1 = Meshes(verts=[verts], faces=[faces], textures=textures1)
        mesh2 = Meshes(verts=[verts_shifted1], faces=[faces], textures=textures2)
1043
1044
        self.assertEqual(textures1._num_faces_per_mesh, [len(faces)])
        self.assertEqual(textures2._num_faces_per_mesh, [len(faces)])
1045
        mesh_joined = join_meshes_as_scene([mesh1, mesh2])
1046
1047
1048
        self.assertEqual(textures1._num_faces_per_mesh, [len(faces)])
        self.assertEqual(textures2._num_faces_per_mesh, [len(faces)])
        self.assertEqual(mesh_joined.textures._num_faces_per_mesh, [len(faces) * 2])
1049
1050
1051
1052
1053

        R, T = look_at_view_transform(18, 0, 0)
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)

        raster_settings = RasterizationSettings(
1054
1055
1056
            image_size=512,
            blur_radius=0.0,
            faces_per_pixel=1,
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1057
            perspective_correct=rasterizer_type.__name__ == "MeshRasterizerOpenGL",
1058
1059
        )

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
1060
        lights = AmbientLights(device=device)
1061
        blend_params = BlendParams(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1062
            sigma=0.5,
1063
1064
1065
            gamma=1e-4,
            background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
        )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1066
1067
1068
1069

        rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
        if rasterizer_type == MeshRasterizer:
            shader = HardPhongShader(
1070
                device=device, blend_params=blend_params, cameras=cameras, lights=lights
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1071
1072
1073
1074
1075
1076
1077
            )
        else:
            shader = SplatterPhongShader(
                device=device, blend_params=blend_params, cameras=cameras, lights=lights
            )

        renderer = MeshRenderer(rasterizer, shader)
1078
1079
1080

        output = renderer(mesh_joined)

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1081
1082
1083
        image_ref = load_rgb_image(
            f"test_joinatlas_final_{rasterizer_type.__name__}.png", DATA_DIR
        )
1084
1085
1086
1087
1088
1089
1090

        if DEBUG:
            debugging_outputs = []
            for mesh_ in [mesh1, mesh2]:
                debugging_outputs.append(renderer(mesh_))
            Image.fromarray(
                (output[0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1091
1092
1093
            ).save(
                DATA_DIR / f"DEBUG_test_joinatlas_final_{rasterizer_type.__name__}.png"
            )
1094
1095
            Image.fromarray(
                (debugging_outputs[0][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1096
            ).save(DATA_DIR / f"test_joinatlas_1_{rasterizer_type.__name__}.png")
1097
1098
            Image.fromarray(
                (debugging_outputs[1][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1099
            ).save(DATA_DIR / f"test_joinatlas_2_{rasterizer_type.__name__}.png")
1100
1101
1102
1103

        result = output[0, ..., :3].cpu()
        self.assertClose(result, image_ref, atol=0.05)

1104
    def test_joined_spheres(self):
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1105
1106
1107
1108
1109
1110
        self._joined_spheres(MeshRasterizer)

    def test_joined_spheres_opengl(self):
        self._joined_spheres(MeshRasterizerOpenGL)

    def _joined_spheres(self, rasterizer_type):
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
        """
        Test a list of Meshes can be joined as a single mesh and
        the single mesh is rendered correctly with Phong, Gouraud
        and Flat Shaders.
        """
        device = torch.device("cuda:0")

        # Init mesh with vertex textures.
        # Initialize a list containing two ico spheres of different sizes.
        sphere_list = [ico_sphere(3, device), ico_sphere(4, device)]
        # [(42 verts, 80 faces), (162 verts, 320 faces)]
        # The scale the vertices need to be set at to resize the spheres
        scales = [0.25, 1]
        # The distance the spheres ought to be offset horizontally to prevent overlap.
        offsets = [1.2, -0.3]
        # Initialize a list containing the adjusted sphere meshes.
        sphere_mesh_list = []
        for i in range(len(sphere_list)):
            verts = sphere_list[i].verts_padded() * scales[i]
            verts[0, :, 0] += offsets[i]
            sphere_mesh_list.append(
                Meshes(verts=verts, faces=sphere_list[i].faces_padded())
            )
1134
        joined_sphere_mesh = join_meshes_as_scene(sphere_mesh_list)
Nikhila Ravi's avatar
Nikhila Ravi committed
1135
1136
        joined_sphere_mesh.textures = TexturesVertex(
            verts_features=torch.ones_like(joined_sphere_mesh.verts_padded())
1137
1138
1139
1140
        )

        # Init rasterizer settings
        R, T = look_at_view_transform(2.7, 0.0, 0.0)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
1141
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
1142
        raster_settings = RasterizationSettings(
1143
1144
1145
            image_size=512,
            blur_radius=0.0,
            faces_per_pixel=1,
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1146
            perspective_correct=rasterizer_type.__name__ == "MeshRasterizerOpenGL",
1147
1148
1149
1150
1151
1152
        )

        # Init shader settings
        materials = Materials(device=device)
        lights = PointLights(device=device)
        lights.location = torch.tensor([0.0, 0.0, +2.0], device=device)[None]
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1153
        blend_params = BlendParams(0.5, 1e-4, (0, 0, 0))
1154
1155

        # Init renderer
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1156
        rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
1157
1158
1159
1160
        shaders = {
            "phong": HardPhongShader,
            "gouraud": HardGouraudShader,
            "flat": HardFlatShader,
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1161
            "splatter": SplatterPhongShader,
1162
1163
        }
        for (name, shader_init) in shaders.items():
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1164
1165
1166
1167
1168
            if rasterizer_type == MeshRasterizerOpenGL and name != "splatter":
                continue
            if rasterizer_type == MeshRasterizer and name == "splatter":
                continue

1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
            shader = shader_init(
                lights=lights,
                cameras=cameras,
                materials=materials,
                blend_params=blend_params,
            )
            renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
            image = renderer(joined_sphere_mesh)
            rgb = image[..., :3].squeeze().cpu()
            if DEBUG:
                file_name = "DEBUG_joined_spheres_%s.png" % name
                Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / file_name
                )
            image_ref = load_rgb_image("test_joined_spheres_%s.png" % name, DATA_DIR)
            self.assertClose(rgb, image_ref, atol=0.05)
Nikhila Ravi's avatar
Nikhila Ravi committed
1185
1186

    def test_texture_map_atlas(self):
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1187
1188
1189
1190
1191
1192
        self._texture_map_atlas(MeshRasterizer)

    def test_texture_map_atlas_opengl(self):
        self._texture_map_atlas(MeshRasterizerOpenGL)

    def _texture_map_atlas(self, rasterizer_type):
Nikhila Ravi's avatar
Nikhila Ravi committed
1193
1194
        """
        Test a mesh with a texture map as a per face atlas is loaded and rendered correctly.
Nikhila Ravi's avatar
Nikhila Ravi committed
1195
        Also check that the backward pass for texture atlas rendering is differentiable.
Nikhila Ravi's avatar
Nikhila Ravi committed
1196
1197
        """
        device = torch.device("cuda:0")
1198
1199

        obj_filename = TUTORIAL_DATA_DIR / "cow_mesh/cow.obj"
Nikhila Ravi's avatar
Nikhila Ravi committed
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209

        # Load mesh and texture as a per face texture atlas.
        verts, faces, aux = load_obj(
            obj_filename,
            device=device,
            load_textures=True,
            create_texture_atlas=True,
            texture_atlas_size=8,
            texture_wrap=None,
        )
Nikhila Ravi's avatar
Nikhila Ravi committed
1210
        atlas = aux.texture_atlas
Nikhila Ravi's avatar
Nikhila Ravi committed
1211
1212
1213
        mesh = Meshes(
            verts=[verts],
            faces=[faces.verts_idx],
Nikhila Ravi's avatar
Nikhila Ravi committed
1214
            textures=TexturesAtlas(atlas=[atlas]),
Nikhila Ravi's avatar
Nikhila Ravi committed
1215
1216
1217
1218
        )

        # Init rasterizer settings
        R, T = look_at_view_transform(2.7, 0, 0)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
1219
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
Nikhila Ravi's avatar
Nikhila Ravi committed
1220
1221

        raster_settings = RasterizationSettings(
Nikhila Ravi's avatar
Nikhila Ravi committed
1222
1223
1224
1225
            image_size=512,
            blur_radius=0.0,
            faces_per_pixel=1,
            cull_backfaces=True,
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1226
            perspective_correct=rasterizer_type.__name__ == "MeshRasterizerOpenGL",
Nikhila Ravi's avatar
Nikhila Ravi committed
1227
1228
1229
1230
        )

        # Init shader settings
        materials = Materials(device=device, specular_color=((0, 0, 0),), shininess=0.0)
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1231
        blend_params = BlendParams(0.5, 1e-4, (1.0, 1.0, 1.0))
Nikhila Ravi's avatar
Nikhila Ravi committed
1232
1233
1234
1235
1236
1237
1238
        lights = PointLights(device=device)

        # Place light behind the cow in world space. The front of
        # the cow is facing the -z direction.
        lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None]

        # The HardPhongShader can be used directly with atlas textures.
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
        rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
        if rasterizer_type == MeshRasterizer:
            shader = HardPhongShader(
                device=device,
                blend_params=blend_params,
                cameras=cameras,
                lights=lights,
                materials=materials,
            )
        else:
            shader = SplatterPhongShader(
                device=device,
                blend_params=blend_params,
                cameras=cameras,
                lights=lights,
                materials=materials,
            )

        renderer = MeshRenderer(rasterizer, shader)
Nikhila Ravi's avatar
Nikhila Ravi committed
1258
1259

        images = renderer(mesh)
Nikhila Ravi's avatar
Nikhila Ravi committed
1260
        rgb = images[0, ..., :3].squeeze()
Nikhila Ravi's avatar
Nikhila Ravi committed
1261
1262

        # Load reference image
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1263
1264
1265
        image_ref = load_rgb_image(
            f"test_texture_atlas_8x8_back_{rasterizer_type.__name__}.png", DATA_DIR
        )
Nikhila Ravi's avatar
Nikhila Ravi committed
1266
1267

        if DEBUG:
Nikhila Ravi's avatar
Nikhila Ravi committed
1268
            Image.fromarray((rgb.detach().cpu().numpy() * 255).astype(np.uint8)).save(
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1269
1270
                DATA_DIR
                / f"DEBUG_texture_atlas_8x8_back_{rasterizer_type.__name__}.png"
Nikhila Ravi's avatar
Nikhila Ravi committed
1271
1272
            )

Nikhila Ravi's avatar
Nikhila Ravi committed
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
        self.assertClose(rgb.cpu(), image_ref, atol=0.05)

        # Check gradients are propagated
        # correctly back to the texture atlas.
        # Because of how texture sampling is implemented
        # for the texture atlas it is not possible to get
        # gradients back to the vertices.
        atlas.requires_grad = True
        mesh = Meshes(
            verts=[verts],
            faces=[faces.verts_idx],
            textures=TexturesAtlas(atlas=[atlas]),
        )
        raster_settings = RasterizationSettings(
            image_size=512,
            blur_radius=0.0001,
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1289
1290
            faces_per_pixel=5 if rasterizer_type.__name__ == "MeshRasterizer" else 1,
            cull_backfaces=rasterizer_type.__name__ == "MeshRasterizer",
Nikhila Ravi's avatar
Nikhila Ravi committed
1291
1292
1293
1294
1295
1296
            clip_barycentric_coords=True,
        )
        images = renderer(mesh, raster_settings=raster_settings)
        images[0, ...].sum().backward()

        fragments = rasterizer(mesh, raster_settings=raster_settings)
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1297
1298
1299
1300
        if rasterizer_type == MeshRasterizer:
            # Some of the bary coordinates are outside the
            # [0, 1] range as expected because the blur is > 0.
            self.assertTrue(fragments.bary_coords.ge(1.0).any())
Nikhila Ravi's avatar
Nikhila Ravi committed
1301
1302
        self.assertIsNotNone(atlas.grad)
        self.assertTrue(atlas.grad.sum().abs() > 0.0)
1303
1304

    def test_simple_sphere_outside_zfar(self):
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1305
1306
1307
1308
1309
1310
        self._simple_sphere_outside_zfar(MeshRasterizer)

    def test_simple_sphere_outside_zfar_opengl(self):
        self._simple_sphere_outside_zfar(MeshRasterizerOpenGL)

    def _simple_sphere_outside_zfar(self, rasterizer_type):
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
        """
        Test output when rendering a sphere that is beyond zfar with a SoftPhongShader.
        This renders a sphere of radius 500, with the camera at x=1500 for different
        settings of zfar.  This is intended to check 1) setting cameras.zfar propagates
        to the blender and that the rendered sphere is (soft) clipped if it is beyond
        zfar, 2) make sure there are no numerical precision/overflow errors associated
        with larger world coordinates
        """
        device = torch.device("cuda:0")

        # Init mesh
        sphere_mesh = ico_sphere(5, device)
        verts_padded = sphere_mesh.verts_padded() * 500
        faces_padded = sphere_mesh.faces_padded()
        feats = torch.ones_like(verts_padded, device=device)
        textures = TexturesVertex(verts_features=feats)
        sphere_mesh = Meshes(verts=verts_padded, faces=faces_padded, textures=textures)

        R, T = look_at_view_transform(1500, 0.0, 0.0)

        # Init shader settings
        materials = Materials(device=device)
        lights = PointLights(device=device)
        lights.location = torch.tensor([0.0, 0.0, +1000.0], device=device)[None]

        raster_settings = RasterizationSettings(
            image_size=256, blur_radius=0.0, faces_per_pixel=1
        )
        for zfar in (10000.0, 100.0):
            cameras = FoVPerspectiveCameras(
                device=device, R=R, T=T, aspect_ratio=1.0, fov=60.0, zfar=zfar
            )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1343
1344
            blend_params = BlendParams(
                1e-4 if rasterizer_type == MeshRasterizer else 0.5, 1e-4, (0, 0, 1.0)
1345
            )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1346
1347
            rasterizer = rasterizer_type(
                cameras=cameras, raster_settings=raster_settings
1348
            )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
            if rasterizer_type == MeshRasterizer:
                shader = SoftPhongShader(
                    blend_params=blend_params,
                    cameras=cameras,
                    lights=lights,
                    materials=materials,
                )
            else:
                shader = SplatterPhongShader(
                    device=device,
                    blend_params=blend_params,
                    cameras=cameras,
                    lights=lights,
                    materials=materials,
                )
1364
1365
1366
1367
            renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
            images = renderer(sphere_mesh)
            rgb = images[0, ..., :3].squeeze().cpu()

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1368
            filename = f"test_simple_sphere_outside_zfar_{int(zfar)}_{rasterizer_type.__name__}.png"
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378

            # Load reference image
            image_ref = load_rgb_image(filename, DATA_DIR)

            if DEBUG:
                Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / ("DEBUG_" + filename)
                )

            self.assertClose(rgb, image_ref, atol=0.05)
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395

    def test_cameras_kwarg(self):
        """
        Test that when cameras are passed in as a kwarg the rendering
        works as expected
        """
        device = torch.device("cuda:0")

        # Init mesh
        sphere_mesh = ico_sphere(5, device)
        verts_padded = sphere_mesh.verts_padded()
        faces_padded = sphere_mesh.faces_padded()
        feats = torch.ones_like(verts_padded, device=device)
        textures = TexturesVertex(verts_features=feats)
        sphere_mesh = Meshes(verts=verts_padded, faces=faces_padded, textures=textures)

        # No elevation or azimuth rotation
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1396
1397
1398
1399
1400
1401
1402
1403
1404
        rasterizer_tests = [
            RasterizerTest(MeshRasterizer, HardPhongShader, "phong", "hard_phong"),
            RasterizerTest(
                MeshRasterizerOpenGL,
                SplatterPhongShader,
                "splatter",
                "splatter_phong",
            ),
        ]
1405
1406
1407
1408
1409
1410
1411
        R, T = look_at_view_transform(2.7, 0.0, 0.0)
        for cam_type in (
            FoVPerspectiveCameras,
            FoVOrthographicCameras,
            PerspectiveCameras,
            OrthographicCameras,
        ):
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
            for test in rasterizer_tests:
                if test.rasterizer == MeshRasterizerOpenGL and cam_type in [
                    PerspectiveCameras,
                    OrthographicCameras,
                ]:
                    # MeshRasterizerOpenGL only works with FoV cameras.
                    continue

                cameras = cam_type(device=device, R=R, T=T)

                # Init shader settings
                materials = Materials(device=device)
                lights = PointLights(device=device)
                lights.location = torch.tensor([0.0, 0.0, +2.0], device=device)[None]

                raster_settings = RasterizationSettings(
                    image_size=512, blur_radius=0.0, faces_per_pixel=1
                )
                rasterizer = test.rasterizer(raster_settings=raster_settings)
                blend_params = BlendParams(0.5, 1e-4, (0, 0, 0))
                shader = test.shader(
                    lights=lights, materials=materials, blend_params=blend_params
                )
                renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)

                # Cameras can be passed into the renderer in the forward pass
                images = renderer(sphere_mesh, cameras=cameras)
                rgb = images.squeeze()[..., :3].cpu().numpy()
                image_ref = load_rgb_image(
                    f"test_simple_sphere_light_{test.reference_name}_{cam_type.__name__}.png",
                    DATA_DIR,
                )
                self.assertClose(rgb, image_ref, atol=0.05)

        def test_nd_sphere(self):
            """
            Test that the render can handle textures with more than 3 channels and
            not just 3 channel RGB.
            """
            torch.manual_seed(1)
            device = torch.device("cuda:0")
            C = 5
            WHITE = ((1.0,) * C,)
            BLACK = ((0.0,) * C,)

            # Init mesh
            sphere_mesh = ico_sphere(5, device)
            verts_padded = sphere_mesh.verts_padded()
            faces_padded = sphere_mesh.faces_padded()
            feats = torch.ones(*verts_padded.shape[:-1], C, device=device)
            n_verts = feats.shape[1]
            # make some non-uniform pattern
            feats *= torch.arange(0, 10, step=10 / n_verts, device=device).unsqueeze(1)
            textures = TexturesVertex(verts_features=feats)
            sphere_mesh = Meshes(
                verts=verts_padded, faces=faces_padded, textures=textures
            )

            # No elevation or azimuth rotation
            R, T = look_at_view_transform(2.7, 0.0, 0.0)

            cameras = PerspectiveCameras(device=device, R=R, T=T)
1474
1475

            # Init shader settings
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
            materials = Materials(
                device=device,
                ambient_color=WHITE,
                diffuse_color=WHITE,
                specular_color=WHITE,
            )
            lights = AmbientLights(
                device=device,
                ambient_color=WHITE,
            )
1486
1487
1488
1489
1490
            lights.location = torch.tensor([0.0, 0.0, +2.0], device=device)[None]

            raster_settings = RasterizationSettings(
                image_size=512, blur_radius=0.0, faces_per_pixel=1
            )
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1491
1492
1493
1494
1495
1496
1497
1498
            rasterizer = MeshRasterizer(
                cameras=cameras, raster_settings=raster_settings
            )
            blend_params = BlendParams(
                1e-4,
                1e-4,
                background_color=BLACK[0],
            )
1499

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1500
1501
1502
            # only test HardFlatShader since that's the only one that makes
            # sense for classification
            shader = HardFlatShader(
1503
                lights=lights,
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1504
                cameras=cameras,
1505
1506
1507
1508
                materials=materials,
                blend_params=blend_params,
            )
            renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1509
            images = renderer(sphere_mesh)
1510

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1511
1512
1513
            self.assertEqual(images.shape[-1], C + 1)
            self.assertClose(images.amax(), torch.tensor(10.0), atol=0.01)
            self.assertClose(images.amin(), torch.tensor(0.0), atol=0.01)
1514

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1515
1516
1517
            # grab last 3 color channels
            rgb = (images[0, ..., C - 3 : C] / 10).squeeze().cpu()
            filename = "test_nd_sphere.png"
1518

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1519
1520
1521
1522
1523
            if DEBUG:
                debug_filename = "DEBUG_%s" % filename
                Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / debug_filename
                )
1524

Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
1525
1526
            image_ref = load_rgb_image(filename, DATA_DIR)
            self.assertClose(rgb, image_ref, atol=0.05)