test_render_meshes.py 49.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
Georgia Gkioxari's avatar
Georgia Gkioxari committed
19
20
21
from pytorch3d.renderer.cameras import (
    FoVOrthographicCameras,
    FoVPerspectiveCameras,
22
    look_at_view_transform,
Georgia Gkioxari's avatar
Georgia Gkioxari committed
23
24
25
    OrthographicCameras,
    PerspectiveCameras,
)
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
26
from pytorch3d.renderer.lighting import AmbientLights, PointLights
facebook-github-bot's avatar
facebook-github-bot committed
27
from pytorch3d.renderer.materials import Materials
Nikhila Ravi's avatar
Nikhila Ravi committed
28
from pytorch3d.renderer.mesh import TexturesAtlas, TexturesUV, TexturesVertex
29
from pytorch3d.renderer.mesh.rasterizer import MeshRasterizer, RasterizationSettings
30
from pytorch3d.renderer.mesh.renderer import MeshRenderer, MeshRendererWithFragments
facebook-github-bot's avatar
facebook-github-bot committed
31
32
from pytorch3d.renderer.mesh.shader import (
    BlendParams,
33
    HardFlatShader,
Patrick Labatut's avatar
Patrick Labatut committed
34
    HardGouraudShader,
35
    HardPhongShader,
36
    SoftPhongShader,
37
    SoftSilhouetteShader,
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
38
    SplatterPhongShader,
39
    TexturedSoftPhongShader,
facebook-github-bot's avatar
facebook-github-bot committed
40
)
41
42
43
from pytorch3d.structures.meshes import (
    join_meshes_as_batch,
    join_meshes_as_scene,
44
    Meshes,
45
)
facebook-github-bot's avatar
facebook-github-bot committed
46
from pytorch3d.utils.ico_sphere import ico_sphere
47
from pytorch3d.utils.torus import torus
facebook-github-bot's avatar
facebook-github-bot committed
48

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
49
50
51
52
53
54
55
from .common_testing import (
    get_pytorch3d_dir,
    get_tests_dir,
    load_rgb_image,
    TestCaseMixin,
)

56

Nikhila Ravi's avatar
Nikhila Ravi committed
57
# If DEBUG=True, save out images generated in the tests for debugging.
facebook-github-bot's avatar
facebook-github-bot committed
58
59
# All saved images have prefix DEBUG_
DEBUG = False
60
DATA_DIR = get_tests_dir() / "data"
61
TUTORIAL_DATA_DIR = get_pytorch3d_dir() / "docs/tutorials/data"
facebook-github-bot's avatar
facebook-github-bot committed
62

Nikhila Ravi's avatar
Nikhila Ravi committed
63
64
ShaderTest = namedtuple("ShaderTest", ["shader", "reference_name", "debug_name"])

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

Nikhila Ravi's avatar
Nikhila Ravi committed
66
class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
67
    def test_simple_sphere(self, elevated_camera=False, check_depth=False):
facebook-github-bot's avatar
facebook-github-bot committed
68
        """
Patrick Labatut's avatar
Patrick Labatut committed
69
        Test output of phong and gouraud shading matches a reference image using
facebook-github-bot's avatar
facebook-github-bot committed
70
71
72
73
74
75
76
77
78
79
80
81
        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
82
83
        feats = torch.ones_like(verts_padded, device=device)
        textures = TexturesVertex(verts_features=feats)
84
        sphere_mesh = Meshes(verts=verts_padded, faces=faces_padded, textures=textures)
facebook-github-bot's avatar
facebook-github-bot committed
85
86
87

        # Init rasterizer settings
        if elevated_camera:
88
89
            # 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
90
            postfix = "_elevated_"
91
92
            # 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
93
        else:
94
            # No elevation or azimuth rotation
facebook-github-bot's avatar
facebook-github-bot committed
95
            R, T = look_at_view_transform(2.7, 0.0, 0.0)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
            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
            )
            rasterizer = MeshRasterizer(
                cameras=cameras, raster_settings=raster_settings
            )
            blend_params = BlendParams(1e-4, 1e-4, (0, 0, 0))

            # Test several shaders
Nikhila Ravi's avatar
Nikhila Ravi committed
119
120
121
122
123
124
125
            shader_tests = [
                ShaderTest(HardPhongShader, "phong", "hard_phong"),
                ShaderTest(HardGouraudShader, "gouraud", "hard_gouraud"),
                ShaderTest(HardFlatShader, "flat", "hard_flat"),
            ]
            for test in shader_tests:
                shader = test.shader(
Georgia Gkioxari's avatar
Georgia Gkioxari committed
126
127
128
129
130
                    lights=lights,
                    cameras=cameras,
                    materials=materials,
                    blend_params=blend_params,
                )
131
132
133
134
135
136
                if check_depth:
                    renderer = MeshRendererWithFragments(
                        rasterizer=rasterizer, shader=shader
                    )
                    images, fragments = renderer(sphere_mesh)
                    self.assertClose(fragments.zbuf, rasterizer(sphere_mesh).zbuf)
137
138
139
140
                    # Check the alpha channel is the mask
                    self.assertClose(
                        images[..., -1], (fragments.pix_to_face[..., 0] >= 0).float()
                    )
141
142
143
144
                else:
                    renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
                    images = renderer(sphere_mesh)

Georgia Gkioxari's avatar
Georgia Gkioxari committed
145
146
                rgb = images[0, ..., :3].squeeze().cpu()
                filename = "simple_sphere_light_%s%s%s.png" % (
Nikhila Ravi's avatar
Nikhila Ravi committed
147
                    test.reference_name,
Georgia Gkioxari's avatar
Georgia Gkioxari committed
148
149
150
                    postfix,
                    cam_type.__name__,
                )
facebook-github-bot's avatar
facebook-github-bot committed
151

Georgia Gkioxari's avatar
Georgia Gkioxari committed
152
153
                image_ref = load_rgb_image("test_%s" % filename, DATA_DIR)
                self.assertClose(rgb, image_ref, atol=0.05)
154

Georgia Gkioxari's avatar
Georgia Gkioxari committed
155
                if DEBUG:
Nikhila Ravi's avatar
Nikhila Ravi committed
156
157
158
159
160
161
                    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
162
163
164
                    Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                        DATA_DIR / filename
                    )
Nikhila Ravi's avatar
Nikhila Ravi committed
165

Georgia Gkioxari's avatar
Georgia Gkioxari committed
166
167
168
169
170
171
172
            ########################################################
            # 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(
173
174
175
176
177
                lights=lights,
                cameras=cameras,
                materials=materials,
                blend_params=blend_params,
            )
178
179
180
181
182
183
184
185
            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
                )
186
187
188
189
                # Check the alpha channel is the mask
                self.assertClose(
                    images[..., -1], (fragments.pix_to_face[..., 0] >= 0).float()
                )
190
191
192
193
194
            else:
                phong_renderer = MeshRenderer(
                    rasterizer=rasterizer, shader=phong_shader
                )
                images = phong_renderer(sphere_mesh, lights=lights)
Nikhila Ravi's avatar
Nikhila Ravi committed
195
196
            rgb = images[0, ..., :3].squeeze().cpu()
            if DEBUG:
Georgia Gkioxari's avatar
Georgia Gkioxari committed
197
198
199
200
                filename = "DEBUG_simple_sphere_dark%s%s.png" % (
                    postfix,
                    cam_type.__name__,
                )
Nikhila Ravi's avatar
Nikhila Ravi committed
201
202
203
                Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / filename
                )
facebook-github-bot's avatar
facebook-github-bot committed
204

Georgia Gkioxari's avatar
Georgia Gkioxari committed
205
206
207
            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
208
            )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
209
            self.assertClose(rgb, image_ref_phong_dark, atol=0.05)
facebook-github-bot's avatar
facebook-github-bot committed
210
211
212

    def test_simple_sphere_elevated_camera(self):
        """
Patrick Labatut's avatar
Patrick Labatut committed
213
        Test output of phong and gouraud shading matches a reference image using
facebook-github-bot's avatar
facebook-github-bot committed
214
215
216
217
218
219
        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)

220
221
222
223
224
225
226
227
228
    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
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
    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
255
        half_half = (512.0 / 2.0, 512.0 / 2.0)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
256
257
258
259
260
        for cam_type in (PerspectiveCameras, OrthographicCameras):
            cameras = cam_type(
                device=device,
                R=R,
                T=T,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
261
262
                principal_point=(half_half,),
                focal_length=(half_half,),
Georgia Gkioxari's avatar
Georgia Gkioxari committed
263
                image_size=((512, 512),),
264
                in_ndc=False,
Georgia Gkioxari's avatar
Georgia Gkioxari committed
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
            )
            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
281
282
283
284
            if DEBUG:
                Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / f"{filename}_.png"
                )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
285
286
287
288

            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
289
290
    def test_simple_sphere_batched(self):
        """
Nikhila Ravi's avatar
Nikhila Ravi committed
291
        Test a mesh with vertex textures can be extended to form a batch, and
Nikhila Ravi's avatar
Nikhila Ravi committed
292
293
        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
294
        """
Nikhila Ravi's avatar
Nikhila Ravi committed
295
        batch_size = 5
facebook-github-bot's avatar
facebook-github-bot committed
296
297
        device = torch.device("cuda:0")

Nikhila Ravi's avatar
Nikhila Ravi committed
298
        # Init mesh with vertex textures.
facebook-github-bot's avatar
facebook-github-bot committed
299
300
301
        sphere_meshes = ico_sphere(5, device).extend(batch_size)
        verts_padded = sphere_meshes.verts_padded()
        faces_padded = sphere_meshes.faces_padded()
Nikhila Ravi's avatar
Nikhila Ravi committed
302
303
        feats = torch.ones_like(verts_padded, device=device)
        textures = TexturesVertex(verts_features=feats)
facebook-github-bot's avatar
facebook-github-bot committed
304
305
306
307
308
309
310
311
312
        sphere_meshes = Meshes(
            verts=verts_padded, faces=faces_padded, textures=textures
        )

        # Init rasterizer settings
        dist = torch.tensor([2.7]).repeat(batch_size).to(device)
        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
313
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
facebook-github-bot's avatar
facebook-github-bot committed
314
        raster_settings = RasterizationSettings(
Nikhila Ravi's avatar
Nikhila Ravi committed
315
            image_size=512, blur_radius=0.0, faces_per_pixel=4
facebook-github-bot's avatar
facebook-github-bot committed
316
317
318
319
        )

        # Init shader settings
        materials = Materials(device=device)
Nikhila Ravi's avatar
Nikhila Ravi committed
320
321
322
        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)
323
        blend_params = BlendParams(1e-4, 1e-4, (0, 0, 0))
facebook-github-bot's avatar
facebook-github-bot committed
324
325

        # Init renderer
326
        rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
Nikhila Ravi's avatar
Nikhila Ravi committed
327
328
329
        shader_tests = [
            ShaderTest(HardPhongShader, "phong", "hard_phong"),
            ShaderTest(SoftPhongShader, "phong", "soft_phong"),
Krzysztof Chalupka's avatar
Krzysztof Chalupka committed
330
            ShaderTest(SplatterPhongShader, "phong", "splatter_phong"),
Nikhila Ravi's avatar
Nikhila Ravi committed
331
332
333
334
335
336
337
            ShaderTest(HardGouraudShader, "gouraud", "hard_gouraud"),
            ShaderTest(HardFlatShader, "flat", "hard_flat"),
        ]
        for test in shader_tests:
            reference_name = test.reference_name
            debug_name = test.debug_name
            shader = test.shader(
338
339
340
341
342
                lights=lights,
                cameras=cameras,
                materials=materials,
                blend_params=blend_params,
            )
Nikhila Ravi's avatar
Nikhila Ravi committed
343
344
            renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
            images = renderer(sphere_meshes)
Nikhila Ravi's avatar
Nikhila Ravi committed
345
            image_ref = load_rgb_image(
Nikhila Ravi's avatar
Nikhila Ravi committed
346
347
                "test_simple_sphere_light_%s_%s.png"
                % (reference_name, type(cameras).__name__),
Georgia Gkioxari's avatar
Georgia Gkioxari committed
348
                DATA_DIR,
Nikhila Ravi's avatar
Nikhila Ravi committed
349
            )
Nikhila Ravi's avatar
Nikhila Ravi committed
350
351
            for i in range(batch_size):
                rgb = images[i, ..., :3].squeeze().cpu()
Nikhila Ravi's avatar
Nikhila Ravi committed
352
                if i == 0 and DEBUG:
Georgia Gkioxari's avatar
Georgia Gkioxari committed
353
                    filename = "DEBUG_simple_sphere_batched_%s_%s.png" % (
Nikhila Ravi's avatar
Nikhila Ravi committed
354
                        debug_name,
Georgia Gkioxari's avatar
Georgia Gkioxari committed
355
356
                        type(cameras).__name__,
                    )
Nikhila Ravi's avatar
Nikhila Ravi committed
357
358
359
                    Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                        DATA_DIR / filename
                    )
Nikhila Ravi's avatar
Nikhila Ravi committed
360
                self.assertClose(rgb, image_ref, atol=0.05)
facebook-github-bot's avatar
facebook-github-bot committed
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375

    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,
376
            clip_barycentric_coords=True,
facebook-github-bot's avatar
facebook-github-bot committed
377
378
379
        )

        # Init rasterizer settings
380
        R, T = look_at_view_transform(2.7, 0, 0)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
381
382
383
384
385
386
387
388
389
390
391
392
393
394
        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
395
            )
Georgia Gkioxari's avatar
Georgia Gkioxari committed
396
397
398
399
400
401
402
403
404
            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
405

Georgia Gkioxari's avatar
Georgia Gkioxari committed
406
407
408
409
            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
410

Georgia Gkioxari's avatar
Georgia Gkioxari committed
411
412
            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
413

Georgia Gkioxari's avatar
Georgia Gkioxari committed
414
415
416
417
418
419
            # 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
420
421
422

    def test_texture_map(self):
        """
423
424
        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
425
426
        """
        device = torch.device("cuda:0")
427
428

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

        # Load mesh + texture
Nikhila Ravi's avatar
Nikhila Ravi committed
431
432
433
434
435
436
437
438
439
        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
440
441

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

facebook-github-bot's avatar
facebook-github-bot committed
445
        raster_settings = RasterizationSettings(
Nikhila Ravi's avatar
Nikhila Ravi committed
446
            image_size=512, blur_radius=0.0, faces_per_pixel=1
facebook-github-bot's avatar
facebook-github-bot committed
447
448
449
450
451
        )

        # Init shader settings
        materials = Materials(device=device)
        lights = PointLights(device=device)
452
453
454
455

        # 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
456

457
458
459
460
461
        blend_params = BlendParams(
            sigma=1e-1,
            gamma=1e-4,
            background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
        )
facebook-github-bot's avatar
facebook-github-bot committed
462
463
        # Init renderer
        renderer = MeshRenderer(
464
            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
465
            shader=TexturedSoftPhongShader(
466
467
468
469
                lights=lights,
                cameras=cameras,
                materials=materials,
                blend_params=blend_params,
facebook-github-bot's avatar
facebook-github-bot committed
470
471
472
473
            ),
        )

        # Load reference image
Nikhila Ravi's avatar
Nikhila Ravi committed
474
        image_ref = load_rgb_image("test_texture_map_back.png", DATA_DIR)
facebook-github-bot's avatar
facebook-github-bot committed
475

Nikhila Ravi's avatar
Nikhila Ravi committed
476
477
478
479
480
        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
            images = renderer(mesh)
            rgb = images[0, ..., :3].squeeze().cpu()
facebook-github-bot's avatar
facebook-github-bot committed
481

Nikhila Ravi's avatar
Nikhila Ravi committed
482
483
484
485
486
487
488
489
490
491
            if DEBUG:
                Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / "DEBUG_texture_map_back.png"
                )

            # 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)
facebook-github-bot's avatar
facebook-github-bot committed
492
493

        # Check grad exists
494
        [verts] = mesh.verts_list()
facebook-github-bot's avatar
facebook-github-bot committed
495
        verts.requires_grad = True
496
        mesh2 = Meshes(verts=[verts], faces=mesh.faces_list(), textures=mesh.textures)
497
        images = renderer(mesh2)
facebook-github-bot's avatar
facebook-github-bot committed
498
499
        images[0, ...].sum().backward()
        self.assertIsNotNone(verts.grad)
500

501
502
503
504
505
        ##########################################
        # 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
506
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
507
508
509
510
511

        # 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
Nikhila Ravi's avatar
Nikhila Ravi committed
512
        image_ref = load_rgb_image("test_texture_map_front.png", DATA_DIR)
513

Nikhila Ravi's avatar
Nikhila Ravi committed
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
        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

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

            if DEBUG:
                Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / "DEBUG_texture_map_front.png"
                )

            # 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)
531

532
533
534
        #################################
        # Add blurring to rasterization
        #################################
535
        R, T = look_at_view_transform(2.7, 0, 180)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
536
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
537
538
539
540
541
        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,
542
            clip_barycentric_coords=True,
543
            perspective_correct=False,
544
545
546
        )

        # Load reference image
Nikhila Ravi's avatar
Nikhila Ravi committed
547
        image_ref = load_rgb_image("test_blurry_textured_rendering.png", DATA_DIR)
548

Nikhila Ravi's avatar
Nikhila Ravi committed
549
550
551
552
553
554
555
556
557
        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

            images = renderer(
                mesh.clone(),
                cameras=cameras,
                raster_settings=raster_settings,
                blend_params=blend_params,
558
            )
Nikhila Ravi's avatar
Nikhila Ravi committed
559
560
561
562
563
564
            rgb = images[0, ..., :3].squeeze().cpu()

            if DEBUG:
                Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / "DEBUG_blurry_textured_rendering.png"
                )
565

Nikhila Ravi's avatar
Nikhila Ravi committed
566
            self.assertClose(rgb, image_ref, atol=0.05)
567

568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
    def test_batch_uvs(self):
        """Test that two random tori with TexturesUV render the same as each individually."""
        torch.manual_seed(1)
        device = torch.device("cuda:0")
        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(
            sigma=1e-1,
            gamma=1e-4,
            background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
        )
        # Init renderer
        renderer = MeshRenderer(
            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
            shader=HardPhongShader(
                device=device, lights=lights, cameras=cameras, blend_params=blend_params
            ),
        )

        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)

648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
    def test_join_uvs(self):
        """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
664
        lights = AmbientLights(device=device)
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
        blend_params = BlendParams(
            sigma=1e-1,
            gamma=1e-4,
            background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
        )
        renderer = MeshRenderer(
            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
            shader=HardPhongShader(
                device=device, blend_params=blend_params, cameras=cameras, lights=lights
            ),
        )

        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
685
686
687
        verts_shifted3 = verts.clone()
        verts_shifted3 *= 0.5
        verts_shifted3[:, 1] -= 700
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
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

        [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)
733
734
735
736
737
            # 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])
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753

            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)

            image_ref = load_rgb_image(f"test_joinuvs{i}_final.png", DATA_DIR)
            map_ref = load_rgb_image(f"test_joinuvs{i}_map.png", DATA_DIR)

            if DEBUG:
                Image.fromarray((output.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / f"test_joinuvs{i}_final_.png"
                )
754
                Image.fromarray((merged.numpy() * 255).astype(np.uint8)).save(
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
                    DATA_DIR / f"test_joinuvs{i}_merged.png"
                )

                Image.fromarray((output1.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / f"test_joinuvs{i}_1.png"
                )
                Image.fromarray((output2.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / f"test_joinuvs{i}_2.png"
                )
                Image.fromarray((output3.numpy() * 255).astype(np.uint8)).save(
                    DATA_DIR / f"test_joinuvs{i}_3.png"
                )
                Image.fromarray(
                    (mesh.textures.maps_padded()[0].cpu().numpy() * 255).astype(
                        np.uint8
                    )
                ).save(DATA_DIR / f"test_joinuvs{i}_map_.png")
                Image.fromarray(
                    (mesh2.textures.maps_padded()[0].cpu().numpy() * 255).astype(
                        np.uint8
                    )
                ).save(DATA_DIR / f"test_joinuvs{i}_map2.png")
                Image.fromarray(
                    (mesh3.textures.maps_padded()[0].cpu().numpy() * 255).astype(
                        np.uint8
                    )
                ).save(DATA_DIR / f"test_joinuvs{i}_map3.png")

783
784
            self.assertClose(output, merged)
            self.assertClose(output, image_ref, atol=0.005)
785
786
            self.assertClose(mesh.textures.maps_padded()[0].cpu(), map_ref, atol=0.05)

787
788
789
790
791
792
793
794
795
796
797
798
    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,
        )
799
800
        self.assertEqual(a._num_faces_per_mesh, [100])
        self.assertEqual(b._num_faces_per_mesh, [50])
801
        c = a.join_batch([b]).join_scene()
802
803
804
        self.assertEqual(a._num_faces_per_mesh, [100])
        self.assertEqual(b._num_faces_per_mesh, [50])
        self.assertEqual(c._num_faces_per_mesh, [150])
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822

        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")

823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
    def test_join_verts(self):
        """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")
        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
849
        lights = AmbientLights(device=device)
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
        blend_params = BlendParams(
            sigma=1e-1,
            gamma=1e-4,
            background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
        )
        renderer = MeshRenderer(
            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
            shader=HardPhongShader(
                device=device, blend_params=blend_params, cameras=cameras, lights=lights
            ),
        )

        output = renderer(mesh)

        image_ref = load_rgb_image("test_joinverts_final.png", DATA_DIR)

        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)
            ).save(DATA_DIR / "test_joinverts_final_.png")
            Image.fromarray(
                (debugging_outputs[0][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
            ).save(DATA_DIR / "test_joinverts_1.png")
            Image.fromarray(
                (debugging_outputs[1][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
            ).save(DATA_DIR / "test_joinverts_2.png")

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

    def test_join_atlas(self):
        """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")
        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)
915
916
        self.assertEqual(textures1._num_faces_per_mesh, [len(faces)])
        self.assertEqual(textures2._num_faces_per_mesh, [len(faces)])
917
        mesh_joined = join_meshes_as_scene([mesh1, mesh2])
918
919
920
        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])
921
922
923
924
925

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

        raster_settings = RasterizationSettings(
926
927
928
929
            image_size=512,
            blur_radius=0.0,
            faces_per_pixel=1,
            perspective_correct=False,
930
931
        )

Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
932
        lights = AmbientLights(device=device)
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
        blend_params = BlendParams(
            sigma=1e-1,
            gamma=1e-4,
            background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
        )
        renderer = MeshRenderer(
            rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
            shader=HardPhongShader(
                device=device, blend_params=blend_params, cameras=cameras, lights=lights
            ),
        )

        output = renderer(mesh_joined)

        image_ref = load_rgb_image("test_joinatlas_final.png", DATA_DIR)

        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)
            ).save(DATA_DIR / "test_joinatlas_final_.png")
            Image.fromarray(
                (debugging_outputs[0][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
            ).save(DATA_DIR / "test_joinatlas_1.png")
            Image.fromarray(
                (debugging_outputs[1][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
            ).save(DATA_DIR / "test_joinatlas_2.png")

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

966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
    def test_joined_spheres(self):
        """
        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())
            )
990
        joined_sphere_mesh = join_meshes_as_scene(sphere_mesh_list)
Nikhila Ravi's avatar
Nikhila Ravi committed
991
992
        joined_sphere_mesh.textures = TexturesVertex(
            verts_features=torch.ones_like(joined_sphere_mesh.verts_padded())
993
994
995
996
        )

        # Init rasterizer settings
        R, T = look_at_view_transform(2.7, 0.0, 0.0)
Georgia Gkioxari's avatar
Georgia Gkioxari committed
997
        cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
998
        raster_settings = RasterizationSettings(
999
1000
1001
1002
            image_size=512,
            blur_radius=0.0,
            faces_per_pixel=1,
            perspective_correct=False,
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
        )

        # Init shader settings
        materials = Materials(device=device)
        lights = PointLights(device=device)
        lights.location = torch.tensor([0.0, 0.0, +2.0], device=device)[None]
        blend_params = BlendParams(1e-4, 1e-4, (0, 0, 0))

        # Init renderer
        rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
        shaders = {
            "phong": HardPhongShader,
            "gouraud": HardGouraudShader,
            "flat": HardFlatShader,
        }
        for (name, shader_init) in shaders.items():
            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
1035
1036
1037
1038

    def test_texture_map_atlas(self):
        """
        Test a mesh with a texture map as a per face atlas is loaded and rendered correctly.
Nikhila Ravi's avatar
Nikhila Ravi committed
1039
        Also check that the backward pass for texture atlas rendering is differentiable.
Nikhila Ravi's avatar
Nikhila Ravi committed
1040
1041
        """
        device = torch.device("cuda:0")
1042
1043

        obj_filename = TUTORIAL_DATA_DIR / "cow_mesh/cow.obj"
Nikhila Ravi's avatar
Nikhila Ravi committed
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053

        # 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
1054
        atlas = aux.texture_atlas
Nikhila Ravi's avatar
Nikhila Ravi committed
1055
1056
1057
        mesh = Meshes(
            verts=[verts],
            faces=[faces.verts_idx],
Nikhila Ravi's avatar
Nikhila Ravi committed
1058
            textures=TexturesAtlas(atlas=[atlas]),
Nikhila Ravi's avatar
Nikhila Ravi committed
1059
1060
1061
1062
        )

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

        raster_settings = RasterizationSettings(
Nikhila Ravi's avatar
Nikhila Ravi committed
1066
1067
1068
1069
            image_size=512,
            blur_radius=0.0,
            faces_per_pixel=1,
            cull_backfaces=True,
1070
            perspective_correct=False,
Nikhila Ravi's avatar
Nikhila Ravi committed
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
        )

        # Init shader settings
        materials = Materials(device=device, specular_color=((0, 0, 0),), shininess=0.0)
        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.
Nikhila Ravi's avatar
Nikhila Ravi committed
1082
        rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
Nikhila Ravi's avatar
Nikhila Ravi committed
1083
        renderer = MeshRenderer(
Nikhila Ravi's avatar
Nikhila Ravi committed
1084
            rasterizer=rasterizer,
Nikhila Ravi's avatar
Nikhila Ravi committed
1085
1086
1087
1088
            shader=HardPhongShader(lights=lights, cameras=cameras, materials=materials),
        )

        images = renderer(mesh)
Nikhila Ravi's avatar
Nikhila Ravi committed
1089
        rgb = images[0, ..., :3].squeeze()
Nikhila Ravi's avatar
Nikhila Ravi committed
1090
1091
1092
1093
1094

        # Load reference image
        image_ref = load_rgb_image("test_texture_atlas_8x8_back.png", DATA_DIR)

        if DEBUG:
Nikhila Ravi's avatar
Nikhila Ravi committed
1095
            Image.fromarray((rgb.detach().cpu().numpy() * 255).astype(np.uint8)).save(
Nikhila Ravi's avatar
Nikhila Ravi committed
1096
1097
1098
                DATA_DIR / "DEBUG_texture_atlas_8x8_back.png"
            )

Nikhila Ravi's avatar
Nikhila Ravi committed
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
        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,
            faces_per_pixel=5,
            cull_backfaces=True,
            clip_barycentric_coords=True,
        )
        images = renderer(mesh, raster_settings=raster_settings)
        images[0, ...].sum().backward()

        fragments = rasterizer(mesh, raster_settings=raster_settings)
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
1123
        # Some of the bary coordinates are outside the
Nikhila Ravi's avatar
Nikhila Ravi committed
1124
1125
1126
1127
        # [0, 1] range as expected because the blur is > 0
        self.assertTrue(fragments.bary_coords.ge(1.0).any())
        self.assertIsNotNone(atlas.grad)
        self.assertTrue(atlas.grad.sum().abs() > 0.0)
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187

    def test_simple_sphere_outside_zfar(self):
        """
        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
            )
            rasterizer = MeshRasterizer(
                cameras=cameras, raster_settings=raster_settings
            )
            blend_params = BlendParams(1e-4, 1e-4, (0, 0, 1.0))

            shader = SoftPhongShader(
                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_outside_zfar_%d.png" % int(zfar)

            # 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)
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238

    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
        R, T = look_at_view_transform(2.7, 0.0, 0.0)
        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
            )
            rasterizer = MeshRasterizer(raster_settings=raster_settings)
            blend_params = BlendParams(1e-4, 1e-4, (0, 0, 0))

            shader = HardPhongShader(
                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(
                "test_simple_sphere_light_phong_%s.png" % cam_type.__name__, DATA_DIR
            )
            self.assertClose(rgb, image_ref, atol=0.05)