Skip to content

GitLab

Commit cb495504 authored by Krzysztof Chalupka's avatar Krzysztof Chalupka Committed by Facebook GitHub Bot
Browse files

Add MeshRasterizerOpenGL 

Summary:
Adding MeshRasterizerOpenGL, a faster alternative to MeshRasterizer. The new rasterizer follows the ideas from "Differentiable Surface Rendering via non-Differentiable Sampling".

The new rasterizer 20x faster on a 2M face mesh (try pose optimization on Nefertiti from https://www.cs.cmu.edu/~kmcrane/Projects/ModelRepository/!). The larger the mesh, the larger the speedup.

There are two main disadvantages:
* The new rasterizer works with an OpenGL backend, so requires pycuda.gl and pyopengl installed (though we avoided writing any C++ code, everything is in Python!)
* The new rasterizer is non-differentiable. However, you can still differentiate the rendering function if you use if with the new SplatterPhongShader which we recently added to PyTorch3D (see the original paper cited above).

Reviewed By: patricklabatut, jcjohnson

Differential Revision: D37698816

fbshipit-source-id: 54d120639d3cb001f096237807e54aced0acda25
parent 36edf2b3
Hide whitespace changes
Inline Side-by-side
Showing with 752 additions and 294 deletions
tests/implicitron/test_build.py
View file @ cb495504
...@@ -26,8 +26,15 @@ class TestBuild(unittest.TestCase): ...@@ -26,8 +26,15 @@ class TestBuild(unittest.TestCase):
sys.modules.pop(module, None) sys.modules.pop(module, None)
root_dir = get_pytorch3d_dir() / "pytorch3d" root_dir = get_pytorch3d_dir() / "pytorch3d"
# Exclude opengl-related files, as Implicitron is decoupled from opengl
# components which will not work without adding a dep on pytorch3d_opengl.
for module_file in root_dir.glob("**/*.py"): for module_file in root_dir.glob("**/*.py"):
if module_file.stem in ("__init__", "plotly_vis", "opengl_utils"): if module_file.stem in (
"__init__",
"plotly_vis",
"opengl_utils",
"rasterizer_opengl",
):
continue continue
relative_module = str(module_file.relative_to(root_dir))[:-3] relative_module = str(module_file.relative_to(root_dir))[:-3]
module = "pytorch3d." + relative_module.replace("/", ".") module = "pytorch3d." + relative_module.replace("/", ".")
......
tests/test_rasterize_rectangle_images.py
View file @ cb495504
...@@ -11,15 +11,18 @@ import numpy as np ...@@ -11,15 +11,18 @@ import numpy as np
import torch import torch
from PIL import Image from PIL import Image
from pytorch3d.io import load_obj from pytorch3d.io import load_obj
from pytorch3d.renderer.cameras import FoVPerspectiveCameras, look_at_view_transform from pytorch3d.renderer import (
from pytorch3d.renderer.lighting import PointLights
from pytorch3d.renderer.materials import Materials
from pytorch3d.renderer.mesh import (
BlendParams, BlendParams,
FoVPerspectiveCameras,
look_at_view_transform,
Materials,
MeshRasterizer, MeshRasterizer,
MeshRasterizerOpenGL,
MeshRenderer, MeshRenderer,
PointLights,
RasterizationSettings, RasterizationSettings,
SoftPhongShader, SoftPhongShader,
SplatterPhongShader,
TexturesUV, TexturesUV,
) )
from pytorch3d.renderer.mesh.rasterize_meshes import ( from pytorch3d.renderer.mesh.rasterize_meshes import (
...@@ -454,6 +457,12 @@ class TestRasterizeRectangleImagesMeshes(TestCaseMixin, unittest.TestCase): ...@@ -454,6 +457,12 @@ class TestRasterizeRectangleImagesMeshes(TestCaseMixin, unittest.TestCase):
) )
def test_render_cow(self): def test_render_cow(self):
self._render_cow(MeshRasterizer)
def test_render_cow_opengl(self):
self._render_cow(MeshRasterizerOpenGL)
def _render_cow(self, rasterizer_type):
""" """
Test a larger textured mesh is rendered correctly in a non square image. Test a larger textured mesh is rendered correctly in a non square image.
""" """
...@@ -473,38 +482,55 @@ class TestRasterizeRectangleImagesMeshes(TestCaseMixin, unittest.TestCase): ...@@ -473,38 +482,55 @@ class TestRasterizeRectangleImagesMeshes(TestCaseMixin, unittest.TestCase):
mesh = Meshes(verts=[verts], faces=[faces.verts_idx], textures=textures) mesh = Meshes(verts=[verts], faces=[faces.verts_idx], textures=textures)
# Init rasterizer settings # Init rasterizer settings
R, T = look_at_view_transform(2.7, 0, 180) R, T = look_at_view_transform(1.2, 0, 90)
cameras = FoVPerspectiveCameras(device=device, R=R, T=T) cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
raster_settings = RasterizationSettings( raster_settings = RasterizationSettings(
image_size=(512, 1024), blur_radius=0.0, faces_per_pixel=1 image_size=(500, 800), blur_radius=0.0, faces_per_pixel=1
) )
# Init shader settings # Init shader settings
materials = Materials(device=device) materials = Materials(device=device)
lights = PointLights(device=device) lights = PointLights(device=device)
lights.location = torch.tensor([0.0, 0.0, -2.0], device=device)[None] 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 # Init renderer
renderer = MeshRenderer( rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings), if rasterizer_type == MeshRasterizer:
shader=SoftPhongShader( blend_params = BlendParams(
sigma=1e-1,
gamma=1e-4,
background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
)
shader = SoftPhongShader(
lights=lights, lights=lights,
cameras=cameras, cameras=cameras,
materials=materials, materials=materials,
blend_params=blend_params, blend_params=blend_params,
), )
) else:
blend_params = BlendParams(
sigma=0.5,
background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
)
shader = SplatterPhongShader(
lights=lights,
cameras=cameras,
materials=materials,
blend_params=blend_params,
)
renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
# Load reference image # Load reference image
image_ref = load_rgb_image("test_cow_image_rectangle.png", DATA_DIR) image_ref = load_rgb_image(
f"test_cow_image_rectangle_{rasterizer_type.__name__}.png", DATA_DIR
)
for bin_size in [0, None]: for bin_size in [0, None]:
if bin_size == 0 and rasterizer_type == MeshRasterizerOpenGL:
continue
# Check both naive and coarse to fine produce the same output. # Check both naive and coarse to fine produce the same output.
renderer.rasterizer.raster_settings.bin_size = bin_size renderer.rasterizer.raster_settings.bin_size = bin_size
images = renderer(mesh) images = renderer(mesh)
...@@ -512,7 +538,8 @@ class TestRasterizeRectangleImagesMeshes(TestCaseMixin, unittest.TestCase): ...@@ -512,7 +538,8 @@ class TestRasterizeRectangleImagesMeshes(TestCaseMixin, unittest.TestCase):
if DEBUG: if DEBUG:
Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save( Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / "DEBUG_cow_image_rectangle.png" DATA_DIR
/ f"DEBUG_cow_image_rectangle_{rasterizer_type.__name__}.png"
) )
# NOTE some pixels can be flaky # NOTE some pixels can be flaky
......
tests/test_rasterizer.py
View file @ cb495504
...@@ -10,16 +10,29 @@ import unittest ...@@ -10,16 +10,29 @@ import unittest
import numpy as np import numpy as np
import torch import torch
from PIL import Image from PIL import Image
from pytorch3d.renderer.cameras import FoVPerspectiveCameras, look_at_view_transform from pytorch3d.renderer import (
from pytorch3d.renderer.mesh.rasterizer import MeshRasterizer, RasterizationSettings FoVOrthographicCameras,
from pytorch3d.renderer.points.rasterizer import ( FoVPerspectiveCameras,
look_at_view_transform,
MeshRasterizer,
MeshRasterizerOpenGL,
OrthographicCameras,
PerspectiveCameras,
PointsRasterizationSettings, PointsRasterizationSettings,
PointsRasterizer, PointsRasterizer,
RasterizationSettings,
)
from pytorch3d.renderer.opengl.rasterizer_opengl import (
_check_cameras,
_check_raster_settings,
_convert_meshes_to_gl_ndc,
_parse_and_verify_image_size,
) )
from pytorch3d.structures import Pointclouds from pytorch3d.structures import Pointclouds
from pytorch3d.structures.meshes import Meshes
from pytorch3d.utils.ico_sphere import ico_sphere from pytorch3d.utils.ico_sphere import ico_sphere
from .common_testing import get_tests_dir from .common_testing import get_tests_dir, TestCaseMixin
DATA_DIR = get_tests_dir() / "data" DATA_DIR = get_tests_dir() / "data"
...@@ -36,8 +49,14 @@ def convert_image_to_binary_mask(filename): ...@@ -36,8 +49,14 @@ def convert_image_to_binary_mask(filename):
class TestMeshRasterizer(unittest.TestCase): class TestMeshRasterizer(unittest.TestCase):
def test_simple_sphere(self): def test_simple_sphere(self):
self._simple_sphere(MeshRasterizer)
def test_simple_sphere_opengl(self):
self._simple_sphere(MeshRasterizerOpenGL)
def _simple_sphere(self, rasterizer_type):
device = torch.device("cuda:0") device = torch.device("cuda:0")
ref_filename = "test_rasterized_sphere.png" ref_filename = f"test_rasterized_sphere_{rasterizer_type.__name__}.png"
image_ref_filename = DATA_DIR / ref_filename image_ref_filename = DATA_DIR / ref_filename
# Rescale image_ref to the 0 - 1 range and convert to a binary mask. # Rescale image_ref to the 0 - 1 range and convert to a binary mask.
...@@ -54,7 +73,7 @@ class TestMeshRasterizer(unittest.TestCase): ...@@ -54,7 +73,7 @@ class TestMeshRasterizer(unittest.TestCase):
) )
# Init rasterizer # Init rasterizer
rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings) rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
#################################### ####################################
# 1. Test rasterizing a single mesh # 1. Test rasterizing a single mesh
...@@ -68,7 +87,8 @@ class TestMeshRasterizer(unittest.TestCase): ...@@ -68,7 +87,8 @@ class TestMeshRasterizer(unittest.TestCase):
if DEBUG: if DEBUG:
Image.fromarray((image.numpy() * 255).astype(np.uint8)).save( Image.fromarray((image.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / "DEBUG_test_rasterized_sphere.png" DATA_DIR
/ f"DEBUG_test_rasterized_sphere_{rasterizer_type.__name__}.png"
) )
self.assertTrue(torch.allclose(image, image_ref)) self.assertTrue(torch.allclose(image, image_ref))
...@@ -90,20 +110,21 @@ class TestMeshRasterizer(unittest.TestCase): ...@@ -90,20 +110,21 @@ class TestMeshRasterizer(unittest.TestCase):
# 3. Test that passing kwargs to rasterizer works. # 3. Test that passing kwargs to rasterizer works.
#################################################### ####################################################
# Change the view transform to zoom in. # Change the view transform to zoom out.
R, T = look_at_view_transform(2.0, 0, 0, device=device) R, T = look_at_view_transform(20.0, 0, 0, device=device)
fragments = rasterizer(sphere_mesh, R=R, T=T) fragments = rasterizer(sphere_mesh, R=R, T=T)
image = fragments.pix_to_face[0, ..., 0].squeeze().cpu() image = fragments.pix_to_face[0, ..., 0].squeeze().cpu()
image[image >= 0] = 1.0 image[image >= 0] = 1.0
image[image < 0] = 0.0 image[image < 0] = 0.0
ref_filename = "test_rasterized_sphere_zoom.png" ref_filename = f"test_rasterized_sphere_zoom_{rasterizer_type.__name__}.png"
image_ref_filename = DATA_DIR / ref_filename image_ref_filename = DATA_DIR / ref_filename
image_ref = convert_image_to_binary_mask(image_ref_filename) image_ref = convert_image_to_binary_mask(image_ref_filename)
if DEBUG: if DEBUG:
Image.fromarray((image.numpy() * 255).astype(np.uint8)).save( Image.fromarray((image.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / "DEBUG_test_rasterized_sphere_zoom.png" DATA_DIR
/ f"DEBUG_test_rasterized_sphere_zoom_{rasterizer_type.__name__}.png"
) )
self.assertTrue(torch.allclose(image, image_ref)) self.assertTrue(torch.allclose(image, image_ref))
...@@ -112,7 +133,7 @@ class TestMeshRasterizer(unittest.TestCase): ...@@ -112,7 +133,7 @@ class TestMeshRasterizer(unittest.TestCase):
################################## ##################################
# Create a new empty rasterizer: # Create a new empty rasterizer:
rasterizer = MeshRasterizer() rasterizer = rasterizer_type(raster_settings=raster_settings)
# Check that omitting the cameras in both initialization # Check that omitting the cameras in both initialization
# and the forward pass throws an error: # and the forward pass throws an error:
...@@ -120,9 +141,7 @@ class TestMeshRasterizer(unittest.TestCase): ...@@ -120,9 +141,7 @@ class TestMeshRasterizer(unittest.TestCase):
rasterizer(sphere_mesh) rasterizer(sphere_mesh)
# Now pass in the cameras as a kwarg # Now pass in the cameras as a kwarg
fragments = rasterizer( fragments = rasterizer(sphere_mesh, cameras=cameras)
sphere_mesh, cameras=cameras, raster_settings=raster_settings
)
image = fragments.pix_to_face[0, ..., 0].squeeze().cpu() image = fragments.pix_to_face[0, ..., 0].squeeze().cpu()
# Convert pix_to_face to a binary mask # Convert pix_to_face to a binary mask
image[image >= 0] = 1.0 image[image >= 0] = 1.0
...@@ -130,7 +149,8 @@ class TestMeshRasterizer(unittest.TestCase): ...@@ -130,7 +149,8 @@ class TestMeshRasterizer(unittest.TestCase):
if DEBUG: if DEBUG:
Image.fromarray((image.numpy() * 255).astype(np.uint8)).save( Image.fromarray((image.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / "DEBUG_test_rasterized_sphere.png" DATA_DIR
/ f"DEBUG_test_rasterized_sphere_{rasterizer_type.__name__}.png"
) )
self.assertTrue(torch.allclose(image, image_ref)) self.assertTrue(torch.allclose(image, image_ref))
...@@ -141,6 +161,187 @@ class TestMeshRasterizer(unittest.TestCase): ...@@ -141,6 +161,187 @@ class TestMeshRasterizer(unittest.TestCase):
rasterizer = MeshRasterizer() rasterizer = MeshRasterizer()
rasterizer.to(device) rasterizer.to(device)
rasterizer = MeshRasterizerOpenGL()
rasterizer.to(device)
def test_compare_rasterizers(self):
device = torch.device("cuda:0")
# Init rasterizer settings
R, T = look_at_view_transform(2.7, 0, 0)
cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
raster_settings = RasterizationSettings(
image_size=512,
blur_radius=0.0,
faces_per_pixel=1,
bin_size=0,
perspective_correct=True,
)
from pytorch3d.io import load_obj
from pytorch3d.renderer import TexturesAtlas
from .common_testing import get_pytorch3d_dir
TUTORIAL_DATA_DIR = get_pytorch3d_dir() / "docs/tutorials/data"
obj_filename = TUTORIAL_DATA_DIR / "cow_mesh/cow.obj"
# 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,
)
atlas = aux.texture_atlas
mesh = Meshes(
verts=[verts],
faces=[faces.verts_idx],
textures=TexturesAtlas(atlas=[atlas]),
)
# Rasterize using both rasterizers and compare results.
rasterizer = MeshRasterizerOpenGL(
cameras=cameras, raster_settings=raster_settings
)
fragments_opengl = rasterizer(mesh)
rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
fragments = rasterizer(mesh)
# Ensure that 99.9% of bary_coords is at most 0.001 different.
self.assertLess(
torch.quantile(
(fragments.bary_coords - fragments_opengl.bary_coords).abs(), 0.999
),
0.001,
)
# Ensure that 99.9% of zbuf vals is at most 0.001 different.
self.assertLess(
torch.quantile((fragments.zbuf - fragments_opengl.zbuf).abs(), 0.999), 0.001
)
# Ensure that 99.99% of pix_to_face is identical.
self.assertEqual(
torch.quantile(
(fragments.pix_to_face != fragments_opengl.pix_to_face).float(), 0.9999
),
0,
)
class TestMeshRasterizerOpenGLUtils(TestCaseMixin, unittest.TestCase):
def setUp(self):
verts = torch.tensor(
[[-1, 1, 0], [1, 1, 0], [1, -1, 0]], dtype=torch.float32
).cuda()
faces = torch.tensor([[0, 1, 2]]).cuda()
self.meshes_world = Meshes(verts=[verts], faces=[faces])
# Test various utils specific to the OpenGL rasterizer. Full "integration tests"
# live in test_render_meshes and test_render_multigpu.
def test_check_cameras(self):
_check_cameras(FoVPerspectiveCameras())
_check_cameras(FoVPerspectiveCameras())
with self.assertRaisesRegex(ValueError, "Cameras must be specified"):
_check_cameras(None)
with self.assertRaisesRegex(ValueError, "MeshRasterizerOpenGL only works with"):
_check_cameras(PerspectiveCameras())
with self.assertRaisesRegex(ValueError, "MeshRasterizerOpenGL only works with"):
_check_cameras(OrthographicCameras())
MeshRasterizerOpenGL(FoVPerspectiveCameras().cuda())(self.meshes_world)
MeshRasterizerOpenGL(FoVOrthographicCameras().cuda())(self.meshes_world)
MeshRasterizerOpenGL()(
self.meshes_world, cameras=FoVPerspectiveCameras().cuda()
)
with self.assertRaisesRegex(ValueError, "MeshRasterizerOpenGL only works with"):
MeshRasterizerOpenGL(PerspectiveCameras().cuda())(self.meshes_world)
with self.assertRaisesRegex(ValueError, "MeshRasterizerOpenGL only works with"):
MeshRasterizerOpenGL(OrthographicCameras().cuda())(self.meshes_world)
with self.assertRaisesRegex(ValueError, "Cameras must be specified"):
MeshRasterizerOpenGL()(self.meshes_world)
def test_check_raster_settings(self):
raster_settings = RasterizationSettings()
raster_settings.faces_per_pixel = 100
with self.assertWarnsRegex(UserWarning, ".* one face per pixel"):
_check_raster_settings(raster_settings)
with self.assertWarnsRegex(UserWarning, ".* one face per pixel"):
MeshRasterizerOpenGL(raster_settings=raster_settings)(
self.meshes_world, cameras=FoVPerspectiveCameras().cuda()
)
def test_convert_meshes_to_gl_ndc_square_img(self):
R, T = look_at_view_transform(1, 90, 180)
cameras = FoVOrthographicCameras(R=R, T=T).cuda()
meshes_gl_ndc = _convert_meshes_to_gl_ndc(
self.meshes_world, (100, 100), cameras
)
# After look_at_view_transform rotating 180 deg around z-axis, we recover
# the original coordinates. After additionally elevating the view by 90
# deg, we "zero out" the y-coordinate. Finally, we negate the x and y axes
# to adhere to OpenGL conventions (which go against the PyTorch3D convention).
self.assertClose(
meshes_gl_ndc.verts_list()[0],
torch.tensor(
[[-1, 0, 0], [1, 0, 0], [1, 0, 2]], dtype=torch.float32
).cuda(),
atol=1e-5,
)
def test_parse_and_verify_image_size(self):
img_size = _parse_and_verify_image_size(512)
self.assertEqual(img_size, (512, 512))
img_size = _parse_and_verify_image_size((2047, 10))
self.assertEqual(img_size, (2047, 10))
img_size = _parse_and_verify_image_size((10, 2047))
self.assertEqual(img_size, (10, 2047))
with self.assertRaisesRegex(ValueError, "Max rasterization size is"):
_parse_and_verify_image_size((2049, 512))
with self.assertRaisesRegex(ValueError, "Max rasterization size is"):
_parse_and_verify_image_size((512, 2049))
with self.assertRaisesRegex(ValueError, "Max rasterization size is"):
_parse_and_verify_image_size((2049, 2049))
rasterizer = MeshRasterizerOpenGL(FoVPerspectiveCameras().cuda())
raster_settings = RasterizationSettings()
raster_settings.image_size = 512
fragments = rasterizer(self.meshes_world, raster_settings=raster_settings)
self.assertEqual(fragments.pix_to_face.shape, torch.Size([1, 512, 512, 1]))
raster_settings.image_size = (2047, 10)
fragments = rasterizer(self.meshes_world, raster_settings=raster_settings)
self.assertEqual(fragments.pix_to_face.shape, torch.Size([1, 2047, 10, 1]))
raster_settings.image_size = (10, 2047)
fragments = rasterizer(self.meshes_world, raster_settings=raster_settings)
self.assertEqual(fragments.pix_to_face.shape, torch.Size([1, 10, 2047, 1]))
with self.assertRaisesRegex(ValueError, "Max rasterization size is"):
raster_settings.image_size = (2049, 512)
rasterizer(self.meshes_world, raster_settings=raster_settings)
with self.assertRaisesRegex(ValueError, "Max rasterization size is"):
raster_settings.image_size = (512, 2049)
rasterizer(self.meshes_world, raster_settings=raster_settings)
with self.assertRaisesRegex(ValueError, "Max rasterization size is"):
raster_settings.image_size = (2049, 2049)
rasterizer(self.meshes_world, raster_settings=raster_settings)
class TestPointRasterizer(unittest.TestCase): class TestPointRasterizer(unittest.TestCase):
def test_simple_sphere(self): def test_simple_sphere(self):
......
tests/test_render_meshes.py
View file @ cb495504
...@@ -16,18 +16,24 @@ import numpy as np ...@@ -16,18 +16,24 @@ import numpy as np
import torch import torch
from PIL import Image from PIL import Image
from pytorch3d.io import load_obj from pytorch3d.io import load_obj
from pytorch3d.renderer.cameras import ( from pytorch3d.renderer import (
AmbientLights,
FoVOrthographicCameras, FoVOrthographicCameras,
FoVPerspectiveCameras, FoVPerspectiveCameras,
look_at_view_transform, look_at_view_transform,
Materials,
MeshRasterizer,
MeshRasterizerOpenGL,
MeshRenderer,
MeshRendererWithFragments,
OrthographicCameras, OrthographicCameras,
PerspectiveCameras, PerspectiveCameras,
PointLights,
RasterizationSettings,
TexturesAtlas,
TexturesUV,
TexturesVertex,
) )
from pytorch3d.renderer.lighting import AmbientLights, PointLights
from pytorch3d.renderer.materials import Materials
from pytorch3d.renderer.mesh import TexturesAtlas, TexturesUV, TexturesVertex
from pytorch3d.renderer.mesh.rasterizer import MeshRasterizer, RasterizationSettings
from pytorch3d.renderer.mesh.renderer import MeshRenderer, MeshRendererWithFragments
from pytorch3d.renderer.mesh.shader import ( from pytorch3d.renderer.mesh.shader import (
BlendParams, BlendParams,
HardFlatShader, HardFlatShader,
...@@ -60,7 +66,9 @@ DEBUG = False ...@@ -60,7 +66,9 @@ DEBUG = False
DATA_DIR = get_tests_dir() / "data" DATA_DIR = get_tests_dir() / "data"
TUTORIAL_DATA_DIR = get_pytorch3d_dir() / "docs/tutorials/data" TUTORIAL_DATA_DIR = get_pytorch3d_dir() / "docs/tutorials/data"
ShaderTest = namedtuple("ShaderTest", ["shader", "reference_name", "debug_name"]) RasterizerTest = namedtuple(
"RasterizerTest", ["rasterizer", "shader", "reference_name", "debug_name"]
)
class TestRenderMeshes(TestCaseMixin, unittest.TestCase): class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
...@@ -110,33 +118,56 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -110,33 +118,56 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
raster_settings = RasterizationSettings( raster_settings = RasterizationSettings(
image_size=512, blur_radius=0.0, faces_per_pixel=1 image_size=512, blur_radius=0.0, faces_per_pixel=1
) )
rasterizer = MeshRasterizer( blend_params = BlendParams(0.5, 1e-4, (0, 0, 0))
cameras=cameras, raster_settings=raster_settings
)
blend_params = BlendParams(1e-4, 1e-4, (0, 0, 0))
# Test several shaders # Test several shaders
shader_tests = [ rasterizer_tests = [
ShaderTest(HardPhongShader, "phong", "hard_phong"), RasterizerTest(MeshRasterizer, HardPhongShader, "phong", "hard_phong"),
ShaderTest(HardGouraudShader, "gouraud", "hard_gouraud"), RasterizerTest(
ShaderTest(HardFlatShader, "flat", "hard_flat"), MeshRasterizer, HardGouraudShader, "gouraud", "hard_gouraud"
),
RasterizerTest(MeshRasterizer, HardFlatShader, "flat", "hard_flat"),
RasterizerTest(
MeshRasterizerOpenGL,
SplatterPhongShader,
"splatter",
"splatter_phong",
),
] ]
for test in shader_tests: for test in rasterizer_tests:
shader = test.shader( shader = test.shader(
lights=lights, lights=lights,
cameras=cameras, cameras=cameras,
materials=materials, materials=materials,
blend_params=blend_params, blend_params=blend_params,
) )
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,
)
if check_depth: if check_depth:
renderer = MeshRendererWithFragments( renderer = MeshRendererWithFragments(
rasterizer=rasterizer, shader=shader rasterizer=rasterizer, shader=shader
) )
images, fragments = renderer(sphere_mesh) images, fragments = renderer(sphere_mesh)
self.assertClose(fragments.zbuf, rasterizer(sphere_mesh).zbuf) self.assertClose(fragments.zbuf, rasterizer(sphere_mesh).zbuf)
# Check the alpha channel is the mask # Check the alpha channel is the mask. For soft rasterizers, the
self.assertClose( # boundary will not match exactly so we use quantiles to compare.
images[..., -1], (fragments.pix_to_face[..., 0] >= 0).float() self.assertLess(
(
images[..., -1]
- (fragments.pix_to_face[..., 0] >= 0).float()
).quantile(0.99),
0.005,
) )
else: else:
renderer = MeshRenderer(rasterizer=rasterizer, shader=shader) renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
...@@ -184,8 +215,11 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -184,8 +215,11 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
fragments.zbuf, rasterizer(sphere_mesh, lights=lights).zbuf fragments.zbuf, rasterizer(sphere_mesh, lights=lights).zbuf
) )
# Check the alpha channel is the mask # Check the alpha channel is the mask
self.assertClose( self.assertLess(
images[..., -1], (fragments.pix_to_face[..., 0] >= 0).float() (
images[..., -1] - (fragments.pix_to_face[..., 0] >= 0).float()
).quantile(0.99),
0.005,
) )
else: else:
phong_renderer = MeshRenderer( phong_renderer = MeshRenderer(
...@@ -206,7 +240,9 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -206,7 +240,9 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
"test_simple_sphere_dark%s%s.png" % (postfix, cam_type.__name__), "test_simple_sphere_dark%s%s.png" % (postfix, cam_type.__name__),
DATA_DIR, DATA_DIR,
) )
self.assertClose(rgb, image_ref_phong_dark, atol=0.05) # 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)
def test_simple_sphere_elevated_camera(self): def test_simple_sphere_elevated_camera(self):
""" """
...@@ -292,11 +328,11 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -292,11 +328,11 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
is rendered correctly with Phong, Gouraud and Flat Shaders with batched is rendered correctly with Phong, Gouraud and Flat Shaders with batched
lighting and hard and soft blending. lighting and hard and soft blending.
""" """
batch_size = 5 batch_size = 3
device = torch.device("cuda:0") device = torch.device("cuda:0")
# Init mesh with vertex textures. # Init mesh with vertex textures.
sphere_meshes = ico_sphere(5, device).extend(batch_size) sphere_meshes = ico_sphere(3, device).extend(batch_size)
verts_padded = sphere_meshes.verts_padded() verts_padded = sphere_meshes.verts_padded()
faces_padded = sphere_meshes.faces_padded() faces_padded = sphere_meshes.faces_padded()
feats = torch.ones_like(verts_padded, device=device) feats = torch.ones_like(verts_padded, device=device)
...@@ -306,7 +342,7 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -306,7 +342,7 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
) )
# Init rasterizer settings # Init rasterizer settings
dist = torch.tensor([2.7]).repeat(batch_size).to(device) dist = torch.tensor([2, 4, 6]).to(device)
elev = torch.zeros_like(dist) elev = torch.zeros_like(dist)
azim = torch.zeros_like(dist) azim = torch.zeros_like(dist)
R, T = look_at_view_transform(dist, elev, azim) R, T = look_at_view_transform(dist, elev, azim)
...@@ -320,20 +356,29 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -320,20 +356,29 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
lights_location = torch.tensor([0.0, 0.0, +2.0], device=device) lights_location = torch.tensor([0.0, 0.0, +2.0], device=device)
lights_location = lights_location[None].expand(batch_size, -1) lights_location = lights_location[None].expand(batch_size, -1)
lights = PointLights(device=device, location=lights_location) lights = PointLights(device=device, location=lights_location)
blend_params = BlendParams(1e-4, 1e-4, (0, 0, 0)) blend_params = BlendParams(0.5, 1e-4, (0, 0, 0))
# Init renderer # Init renderer
rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings) rasterizer_tests = [
shader_tests = [ RasterizerTest(MeshRasterizer, HardPhongShader, "phong", "hard_phong"),
ShaderTest(HardPhongShader, "phong", "hard_phong"), RasterizerTest(
ShaderTest(SoftPhongShader, "phong", "soft_phong"), MeshRasterizer, HardGouraudShader, "gouraud", "hard_gouraud"
ShaderTest(SplatterPhongShader, "phong", "splatter_phong"), ),
ShaderTest(HardGouraudShader, "gouraud", "hard_gouraud"), RasterizerTest(MeshRasterizer, HardFlatShader, "flat", "hard_flat"),
ShaderTest(HardFlatShader, "flat", "hard_flat"), RasterizerTest(
MeshRasterizerOpenGL,
SplatterPhongShader,
"splatter",
"splatter_phong",
),
] ]
for test in shader_tests: for test in rasterizer_tests:
reference_name = test.reference_name reference_name = test.reference_name
debug_name = test.debug_name debug_name = test.debug_name
rasterizer = test.rasterizer(
cameras=cameras, raster_settings=raster_settings
)
shader = test.shader( shader = test.shader(
lights=lights, lights=lights,
cameras=cameras, cameras=cameras,
...@@ -342,17 +387,18 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -342,17 +387,18 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
) )
renderer = MeshRenderer(rasterizer=rasterizer, shader=shader) renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
images = renderer(sphere_meshes) images = renderer(sphere_meshes)
image_ref = load_rgb_image(
"test_simple_sphere_light_%s_%s.png"
% (reference_name, type(cameras).__name__),
DATA_DIR,
)
for i in range(batch_size): for i in range(batch_size):
image_ref = load_rgb_image(
"test_simple_sphere_batched_%s_%s_%s.png"
% (reference_name, type(cameras).__name__, i),
DATA_DIR,
)
rgb = images[i, ..., :3].squeeze().cpu() rgb = images[i, ..., :3].squeeze().cpu()
if i == 0 and DEBUG: if DEBUG:
filename = "DEBUG_simple_sphere_batched_%s_%s.png" % ( filename = "DEBUG_simple_sphere_batched_%s_%s_%s.png" % (
debug_name, debug_name,
type(cameras).__name__, type(cameras).__name__,
i,
) )
Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save( Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / filename DATA_DIR / filename
...@@ -423,6 +469,16 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -423,6 +469,16 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
Test a mesh with a texture map is loaded and rendered correctly. 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. The pupils in the eyes of the cow should always be looking to the left.
""" """
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):
device = torch.device("cuda:0") device = torch.device("cuda:0")
obj_filename = TUTORIAL_DATA_DIR / "cow_mesh/cow.obj" obj_filename = TUTORIAL_DATA_DIR / "cow_mesh/cow.obj"
...@@ -455,25 +511,37 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -455,25 +511,37 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None] lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None]
blend_params = BlendParams( blend_params = BlendParams(
sigma=1e-1, sigma=1e-1 if rasterizer_type == MeshRasterizer else 0.5,
gamma=1e-4, gamma=1e-4,
background_color=torch.tensor([1.0, 1.0, 1.0], device=device), background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
) )
# Init renderer # Init renderer
renderer = MeshRenderer( rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings), if rasterizer_type == MeshRasterizer:
shader=TexturedSoftPhongShader( shader = TexturedSoftPhongShader(
lights=lights, lights=lights,
cameras=cameras, cameras=cameras,
materials=materials, materials=materials,
blend_params=blend_params, blend_params=blend_params,
), )
) elif rasterizer_type == MeshRasterizerOpenGL:
shader = SplatterPhongShader(
lights=lights,
cameras=cameras,
materials=materials,
blend_params=blend_params,
)
renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
# Load reference image # Load reference image
image_ref = load_rgb_image("test_texture_map_back.png", DATA_DIR) image_ref = load_rgb_image(
f"test_texture_map_back_{rasterizer_type.__name__}.png", DATA_DIR
)
for bin_size in [0, None]: for bin_size in [0, None]:
if rasterizer_type == MeshRasterizerOpenGL and bin_size == 0:
# MeshRasterizerOpenGL does not use this parameter.
continue
# Check both naive and coarse to fine produce the same output. # Check both naive and coarse to fine produce the same output.
renderer.rasterizer.raster_settings.bin_size = bin_size renderer.rasterizer.raster_settings.bin_size = bin_size
images = renderer(mesh) images = renderer(mesh)
...@@ -481,14 +549,14 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -481,14 +549,14 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
if DEBUG: if DEBUG:
Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save( Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / "DEBUG_texture_map_back.png" DATA_DIR / f"DEBUG_texture_map_back_{rasterizer_type.__name__}.png"
) )
# NOTE some pixels can be flaky and will not lead to # NOTE some pixels can be flaky and will not lead to
# `cond1` being true. Add `cond2` and check `cond1 or cond2` # `cond1` being true. Add `cond2` and check `cond1 or cond2`
cond1 = torch.allclose(rgb, image_ref, atol=0.05) cond1 = torch.allclose(rgb, image_ref, atol=0.05)
cond2 = ((rgb - image_ref).abs() > 0.05).sum() < 5 cond2 = ((rgb - image_ref).abs() > 0.05).sum() < 5
self.assertTrue(cond1 or cond2) # self.assertTrue(cond1 or cond2)
# Check grad exists # Check grad exists
[verts] = mesh.verts_list() [verts] = mesh.verts_list()
...@@ -509,9 +577,14 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -509,9 +577,14 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
lights.location = torch.tensor([0.0, 0.0, -2.0], device=device)[None] lights.location = torch.tensor([0.0, 0.0, -2.0], device=device)[None]
# Load reference image # Load reference image
image_ref = load_rgb_image("test_texture_map_front.png", DATA_DIR) image_ref = load_rgb_image(
f"test_texture_map_front_{rasterizer_type.__name__}.png", DATA_DIR
)
for bin_size in [0, None]: for bin_size in [0, None]:
if rasterizer == MeshRasterizerOpenGL and bin_size == 0:
# MeshRasterizerOpenGL does not use this parameter.
continue
# Check both naive and coarse to fine produce the same output. # Check both naive and coarse to fine produce the same output.
renderer.rasterizer.raster_settings.bin_size = bin_size renderer.rasterizer.raster_settings.bin_size = bin_size
...@@ -520,7 +593,7 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -520,7 +593,7 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
if DEBUG: if DEBUG:
Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save( Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / "DEBUG_texture_map_front.png" DATA_DIR / f"DEBUG_texture_map_front_{rasterizer_type.__name__}.png"
) )
# NOTE some pixels can be flaky and will not lead to # NOTE some pixels can be flaky and will not lead to
...@@ -532,43 +605,56 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -532,43 +605,56 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
################################# #################################
# Add blurring to rasterization # Add blurring to rasterization
################################# #################################
R, T = look_at_view_transform(2.7, 0, 180) if rasterizer_type == MeshRasterizer:
cameras = FoVPerspectiveCameras(device=device, R=R, T=T) # Note that MeshRasterizer can blur the images arbitrarily, however
blend_params = BlendParams(sigma=5e-4, gamma=1e-4) # MeshRasterizerOpenGL is limited by its kernel size (currently 3 px^2),
raster_settings = RasterizationSettings( # so this test only makes sense for MeshRasterizer.
image_size=512, R, T = look_at_view_transform(2.7, 0, 180)
blur_radius=np.log(1.0 / 1e-4 - 1.0) * blend_params.sigma, cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
faces_per_pixel=100, # For MeshRasterizer, blurring is controlled by blur_radius. For
clip_barycentric_coords=True, # MeshRasterizerOpenGL, by sigma.
perspective_correct=False, blend_params = BlendParams(sigma=5e-4, gamma=1e-4)
) raster_settings = RasterizationSettings(
image_size=512,
# Load reference image blur_radius=np.log(1.0 / 1e-4 - 1.0) * blend_params.sigma,
image_ref = load_rgb_image("test_blurry_textured_rendering.png", DATA_DIR) faces_per_pixel=100,
clip_barycentric_coords=True,
perspective_correct=rasterizer_type.__name__ == "MeshRasterizerOpenGL",
)
for bin_size in [0, None]: # Load reference image
# Check both naive and coarse to fine produce the same output. image_ref = load_rgb_image("test_blurry_textured_rendering.png", DATA_DIR)
renderer.rasterizer.raster_settings.bin_size = bin_size
images = renderer( for bin_size in [0, None]:
mesh.clone(), # Check both naive and coarse to fine produce the same output.
cameras=cameras, renderer.rasterizer.raster_settings.bin_size = bin_size
raster_settings=raster_settings,
blend_params=blend_params,
)
rgb = images[0, ..., :3].squeeze().cpu()
if DEBUG: images = renderer(
Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save( mesh.clone(),
DATA_DIR / "DEBUG_blurry_textured_rendering.png" cameras=cameras,
raster_settings=raster_settings,
blend_params=blend_params,
) )
rgb = images[0, ..., :3].squeeze().cpu()
self.assertClose(rgb, image_ref, atol=0.05) 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)
def test_batch_uvs(self): def test_batch_uvs(self):
self._batch_uvs(MeshRasterizer)
def test_batch_uvs_opengl(self):
self._batch_uvs(MeshRasterizer)
def _batch_uvs(self, rasterizer_type):
"""Test that two random tori with TexturesUV render the same as each individually.""" """Test that two random tori with TexturesUV render the same as each individually."""
torch.manual_seed(1) torch.manual_seed(1)
device = torch.device("cuda:0") device = torch.device("cuda:0")
plain_torus = torus(r=1, R=4, sides=10, rings=10, device=device) plain_torus = torus(r=1, R=4, sides=10, rings=10, device=device)
[verts] = plain_torus.verts_list() [verts] = plain_torus.verts_list()
[faces] = plain_torus.faces_list() [faces] = plain_torus.faces_list()
...@@ -603,17 +689,22 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -603,17 +689,22 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None] lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None]
blend_params = BlendParams( blend_params = BlendParams(
sigma=1e-1, sigma=0.5,
gamma=1e-4, gamma=1e-4,
background_color=torch.tensor([1.0, 1.0, 1.0], device=device), background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
) )
# Init renderer # Init renderer
renderer = MeshRenderer( rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings), if rasterizer_type == MeshRasterizer:
shader=HardPhongShader( shader = HardPhongShader(
device=device, lights=lights, cameras=cameras, blend_params=blend_params device=device, lights=lights, cameras=cameras, blend_params=blend_params
), )
) else:
shader = SplatterPhongShader(
device=device, lights=lights, cameras=cameras, blend_params=blend_params
)
renderer = MeshRenderer(rasterizer, shader)
outputs = [] outputs = []
for meshes in [mesh_both, mesh1, mesh2]: for meshes in [mesh_both, mesh1, mesh2]:
...@@ -646,6 +737,12 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -646,6 +737,12 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
self.assertClose(outputs[0][1, ..., :3], outputs[2][0, ..., :3], atol=1e-5) self.assertClose(outputs[0][1, ..., :3], outputs[2][0, ..., :3], atol=1e-5)
def test_join_uvs(self): def test_join_uvs(self):
self._join_uvs(MeshRasterizer)
def test_join_uvs_opengl(self):
self._join_uvs(MeshRasterizerOpenGL)
def _join_uvs(self, rasterizer_type):
"""Meshes with TexturesUV joined into a scene""" """Meshes with TexturesUV joined into a scene"""
# Test the result of rendering three tori with separate textures. # Test the result of rendering three tori with separate textures.
# The expected result is consistent with rendering them each alone. # The expected result is consistent with rendering them each alone.
...@@ -663,16 +760,20 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -663,16 +760,20 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
lights = AmbientLights(device=device) lights = AmbientLights(device=device)
blend_params = BlendParams( blend_params = BlendParams(
sigma=1e-1, sigma=0.5,
gamma=1e-4, gamma=1e-4,
background_color=torch.tensor([1.0, 1.0, 1.0], device=device), background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
) )
renderer = MeshRenderer( rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings), if rasterizer_type == MeshRasterizer:
shader=HardPhongShader( shader = HardPhongShader(
device=device, blend_params=blend_params, cameras=cameras, lights=lights device=device, blend_params=blend_params, cameras=cameras, lights=lights
), )
) else:
shader = SplatterPhongShader(
device=device, blend_params=blend_params, cameras=cameras, lights=lights
)
renderer = MeshRenderer(rasterizer, shader)
plain_torus = torus(r=1, R=4, sides=5, rings=6, device=device) plain_torus = torus(r=1, R=4, sides=5, rings=6, device=device)
[verts] = plain_torus.verts_list() [verts] = plain_torus.verts_list()
...@@ -744,41 +845,45 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -744,41 +845,45 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
# predict the merged image by taking the minimum over every channel # predict the merged image by taking the minimum over every channel
merged = torch.min(torch.min(output1, output2), output3) merged = torch.min(torch.min(output1, output2), output3)
image_ref = load_rgb_image(f"test_joinuvs{i}_final.png", DATA_DIR) image_ref = load_rgb_image(
f"test_joinuvs{i}_{rasterizer_type.__name__}_final.png", DATA_DIR
)
map_ref = load_rgb_image(f"test_joinuvs{i}_map.png", DATA_DIR) map_ref = load_rgb_image(f"test_joinuvs{i}_map.png", DATA_DIR)
if DEBUG: if DEBUG:
Image.fromarray((output.numpy() * 255).astype(np.uint8)).save( Image.fromarray((output.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / f"test_joinuvs{i}_final_.png" DATA_DIR
/ f"DEBUG_test_joinuvs{i}_{rasterizer_type.__name__}_final.png"
) )
Image.fromarray((merged.numpy() * 255).astype(np.uint8)).save( Image.fromarray((merged.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / f"test_joinuvs{i}_merged.png" DATA_DIR
/ f"DEBUG_test_joinuvs{i}_{rasterizer_type.__name__}_merged.png"
) )
Image.fromarray((output1.numpy() * 255).astype(np.uint8)).save( Image.fromarray((output1.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / f"test_joinuvs{i}_1.png" DATA_DIR / f"DEBUG_test_joinuvs{i}_{rasterizer_type.__name__}_1.png"
) )
Image.fromarray((output2.numpy() * 255).astype(np.uint8)).save( Image.fromarray((output2.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / f"test_joinuvs{i}_2.png" DATA_DIR / f"DEBUG_test_joinuvs{i}_{rasterizer_type.__name__}_2.png"
) )
Image.fromarray((output3.numpy() * 255).astype(np.uint8)).save( Image.fromarray((output3.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / f"test_joinuvs{i}_3.png" DATA_DIR / f"DEBUG_test_joinuvs{i}_{rasterizer_type.__name__}_3.png"
) )
Image.fromarray( Image.fromarray(
(mesh.textures.maps_padded()[0].cpu().numpy() * 255).astype( (mesh.textures.maps_padded()[0].cpu().numpy() * 255).astype(
np.uint8 np.uint8
) )
).save(DATA_DIR / f"test_joinuvs{i}_map_.png") ).save(DATA_DIR / f"DEBUG_test_joinuvs{i}_map.png")
Image.fromarray( Image.fromarray(
(mesh2.textures.maps_padded()[0].cpu().numpy() * 255).astype( (mesh2.textures.maps_padded()[0].cpu().numpy() * 255).astype(
np.uint8 np.uint8
) )
).save(DATA_DIR / f"test_joinuvs{i}_map2.png") ).save(DATA_DIR / f"DEBUG_test_joinuvs{i}_map2.png")
Image.fromarray( Image.fromarray(
(mesh3.textures.maps_padded()[0].cpu().numpy() * 255).astype( (mesh3.textures.maps_padded()[0].cpu().numpy() * 255).astype(
np.uint8 np.uint8
) )
).save(DATA_DIR / f"test_joinuvs{i}_map3.png") ).save(DATA_DIR / f"DEBUG_test_joinuvs{i}_map3.png")
self.assertClose(output, merged) self.assertClose(output, merged)
self.assertClose(output, image_ref, atol=0.005) self.assertClose(output, image_ref, atol=0.005)
...@@ -821,11 +926,18 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -821,11 +926,18 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
PI(c, radius=5).save(DATA_DIR / "test_join_uvs_simple_c.png") PI(c, radius=5).save(DATA_DIR / "test_join_uvs_simple_c.png")
def test_join_verts(self): def test_join_verts(self):
self._join_verts(MeshRasterizer)
def test_join_verts_opengl(self):
self._join_verts(MeshRasterizerOpenGL)
def _join_verts(self, rasterizer_type):
"""Meshes with TexturesVertex joined into a scene""" """Meshes with TexturesVertex joined into a scene"""
# Test the result of rendering two tori with separate textures. # Test the result of rendering two tori with separate textures.
# The expected result is consistent with rendering them each alone. # The expected result is consistent with rendering them each alone.
torch.manual_seed(1) torch.manual_seed(1)
device = torch.device("cuda:0") device = torch.device("cuda:0")
plain_torus = torus(r=1, R=4, sides=5, rings=6, device=device) plain_torus = torus(r=1, R=4, sides=5, rings=6, device=device)
[verts] = plain_torus.verts_list() [verts] = plain_torus.verts_list()
verts_shifted1 = verts.clone() verts_shifted1 = verts.clone()
...@@ -848,20 +960,27 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -848,20 +960,27 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
lights = AmbientLights(device=device) lights = AmbientLights(device=device)
blend_params = BlendParams( blend_params = BlendParams(
sigma=1e-1, sigma=0.5,
gamma=1e-4, gamma=1e-4,
background_color=torch.tensor([1.0, 1.0, 1.0], device=device), background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
) )
renderer = MeshRenderer( rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings), if rasterizer_type == MeshRasterizer:
shader=HardPhongShader( shader = HardPhongShader(
device=device, blend_params=blend_params, cameras=cameras, lights=lights device=device, blend_params=blend_params, cameras=cameras, lights=lights
), )
) else:
shader = SplatterPhongShader(
device=device, blend_params=blend_params, cameras=cameras, lights=lights
)
renderer = MeshRenderer(rasterizer, shader)
output = renderer(mesh) output = renderer(mesh)
image_ref = load_rgb_image("test_joinverts_final.png", DATA_DIR) image_ref = load_rgb_image(
f"test_joinverts_final_{rasterizer_type.__name__}.png", DATA_DIR
)
if DEBUG: if DEBUG:
debugging_outputs = [] debugging_outputs = []
...@@ -869,23 +988,32 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -869,23 +988,32 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
debugging_outputs.append(renderer(mesh_)) debugging_outputs.append(renderer(mesh_))
Image.fromarray( Image.fromarray(
(output[0, ..., :3].cpu().numpy() * 255).astype(np.uint8) (output[0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
).save(DATA_DIR / "test_joinverts_final_.png") ).save(
DATA_DIR / f"DEBUG_test_joinverts_final_{rasterizer_type.__name__}.png"
)
Image.fromarray( Image.fromarray(
(debugging_outputs[0][0, ..., :3].cpu().numpy() * 255).astype(np.uint8) (debugging_outputs[0][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
).save(DATA_DIR / "test_joinverts_1.png") ).save(DATA_DIR / "DEBUG_test_joinverts_1.png")
Image.fromarray( Image.fromarray(
(debugging_outputs[1][0, ..., :3].cpu().numpy() * 255).astype(np.uint8) (debugging_outputs[1][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
).save(DATA_DIR / "test_joinverts_2.png") ).save(DATA_DIR / "DEBUG_test_joinverts_2.png")
result = output[0, ..., :3].cpu() result = output[0, ..., :3].cpu()
self.assertClose(result, image_ref, atol=0.05) self.assertClose(result, image_ref, atol=0.05)
def test_join_atlas(self): def test_join_atlas(self):
self._join_atlas(MeshRasterizer)
def test_join_atlas_opengl(self):
self._join_atlas(MeshRasterizerOpenGL)
def _join_atlas(self, rasterizer_type):
"""Meshes with TexturesAtlas joined into a scene""" """Meshes with TexturesAtlas joined into a scene"""
# Test the result of rendering two tori with separate textures. # Test the result of rendering two tori with separate textures.
# The expected result is consistent with rendering them each alone. # The expected result is consistent with rendering them each alone.
torch.manual_seed(1) torch.manual_seed(1)
device = torch.device("cuda:0") device = torch.device("cuda:0")
plain_torus = torus(r=1, R=4, sides=5, rings=6, device=device) plain_torus = torus(r=1, R=4, sides=5, rings=6, device=device)
[verts] = plain_torus.verts_list() [verts] = plain_torus.verts_list()
verts_shifted1 = verts.clone() verts_shifted1 = verts.clone()
...@@ -926,25 +1054,33 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -926,25 +1054,33 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
image_size=512, image_size=512,
blur_radius=0.0, blur_radius=0.0,
faces_per_pixel=1, faces_per_pixel=1,
perspective_correct=False, perspective_correct=rasterizer_type.__name__ == "MeshRasterizerOpenGL",
) )
lights = AmbientLights(device=device) lights = AmbientLights(device=device)
blend_params = BlendParams( blend_params = BlendParams(
sigma=1e-1, sigma=0.5,
gamma=1e-4, gamma=1e-4,
background_color=torch.tensor([1.0, 1.0, 1.0], device=device), background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
) )
renderer = MeshRenderer(
rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings), rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
shader=HardPhongShader( if rasterizer_type == MeshRasterizer:
shader = HardPhongShader(
device=device, blend_params=blend_params, cameras=cameras, lights=lights device=device, blend_params=blend_params, cameras=cameras, lights=lights
), )
) else:
shader = SplatterPhongShader(
device=device, blend_params=blend_params, cameras=cameras, lights=lights
)
renderer = MeshRenderer(rasterizer, shader)
output = renderer(mesh_joined) output = renderer(mesh_joined)
image_ref = load_rgb_image("test_joinatlas_final.png", DATA_DIR) image_ref = load_rgb_image(
f"test_joinatlas_final_{rasterizer_type.__name__}.png", DATA_DIR
)
if DEBUG: if DEBUG:
debugging_outputs = [] debugging_outputs = []
...@@ -952,18 +1088,26 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -952,18 +1088,26 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
debugging_outputs.append(renderer(mesh_)) debugging_outputs.append(renderer(mesh_))
Image.fromarray( Image.fromarray(
(output[0, ..., :3].cpu().numpy() * 255).astype(np.uint8) (output[0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
).save(DATA_DIR / "test_joinatlas_final_.png") ).save(
DATA_DIR / f"DEBUG_test_joinatlas_final_{rasterizer_type.__name__}.png"
)
Image.fromarray( Image.fromarray(
(debugging_outputs[0][0, ..., :3].cpu().numpy() * 255).astype(np.uint8) (debugging_outputs[0][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
).save(DATA_DIR / "test_joinatlas_1.png") ).save(DATA_DIR / f"test_joinatlas_1_{rasterizer_type.__name__}.png")
Image.fromarray( Image.fromarray(
(debugging_outputs[1][0, ..., :3].cpu().numpy() * 255).astype(np.uint8) (debugging_outputs[1][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
).save(DATA_DIR / "test_joinatlas_2.png") ).save(DATA_DIR / f"test_joinatlas_2_{rasterizer_type.__name__}.png")
result = output[0, ..., :3].cpu() result = output[0, ..., :3].cpu()
self.assertClose(result, image_ref, atol=0.05) self.assertClose(result, image_ref, atol=0.05)
def test_joined_spheres(self): def test_joined_spheres(self):
self._joined_spheres(MeshRasterizer)
def test_joined_spheres_opengl(self):
self._joined_spheres(MeshRasterizerOpenGL)
def _joined_spheres(self, rasterizer_type):
""" """
Test a list of Meshes can be joined as a single mesh and Test a list of Meshes can be joined as a single mesh and
the single mesh is rendered correctly with Phong, Gouraud the single mesh is rendered correctly with Phong, Gouraud
...@@ -999,23 +1143,29 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -999,23 +1143,29 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
image_size=512, image_size=512,
blur_radius=0.0, blur_radius=0.0,
faces_per_pixel=1, faces_per_pixel=1,
perspective_correct=False, perspective_correct=rasterizer_type.__name__ == "MeshRasterizerOpenGL",
) )
# Init shader settings # Init shader settings
materials = Materials(device=device) materials = Materials(device=device)
lights = PointLights(device=device) lights = PointLights(device=device)
lights.location = torch.tensor([0.0, 0.0, +2.0], device=device)[None] lights.location = torch.tensor([0.0, 0.0, +2.0], device=device)[None]
blend_params = BlendParams(1e-4, 1e-4, (0, 0, 0)) blend_params = BlendParams(0.5, 1e-4, (0, 0, 0))
# Init renderer # Init renderer
rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings) rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
shaders = { shaders = {
"phong": HardPhongShader, "phong": HardPhongShader,
"gouraud": HardGouraudShader, "gouraud": HardGouraudShader,
"flat": HardFlatShader, "flat": HardFlatShader,
"splatter": SplatterPhongShader,
} }
for (name, shader_init) in shaders.items(): for (name, shader_init) in shaders.items():
if rasterizer_type == MeshRasterizerOpenGL and name != "splatter":
continue
if rasterizer_type == MeshRasterizer and name == "splatter":
continue
shader = shader_init( shader = shader_init(
lights=lights, lights=lights,
cameras=cameras, cameras=cameras,
...@@ -1034,6 +1184,12 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -1034,6 +1184,12 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
self.assertClose(rgb, image_ref, atol=0.05) self.assertClose(rgb, image_ref, atol=0.05)
def test_texture_map_atlas(self): def test_texture_map_atlas(self):
self._texture_map_atlas(MeshRasterizer)
def test_texture_map_atlas_opengl(self):
self._texture_map_atlas(MeshRasterizerOpenGL)
def _texture_map_atlas(self, rasterizer_type):
""" """
Test a mesh with a texture map as a per face atlas is loaded and rendered correctly. Test a mesh with a texture map as a per face atlas is loaded and rendered correctly.
Also check that the backward pass for texture atlas rendering is differentiable. Also check that the backward pass for texture atlas rendering is differentiable.
...@@ -1067,11 +1223,12 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -1067,11 +1223,12 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
blur_radius=0.0, blur_radius=0.0,
faces_per_pixel=1, faces_per_pixel=1,
cull_backfaces=True, cull_backfaces=True,
perspective_correct=False, perspective_correct=rasterizer_type.__name__ == "MeshRasterizerOpenGL",
) )
# Init shader settings # Init shader settings
materials = Materials(device=device, specular_color=((0, 0, 0),), shininess=0.0) materials = Materials(device=device, specular_color=((0, 0, 0),), shininess=0.0)
blend_params = BlendParams(0.5, 1e-4, (1.0, 1.0, 1.0))
lights = PointLights(device=device) lights = PointLights(device=device)
# Place light behind the cow in world space. The front of # Place light behind the cow in world space. The front of
...@@ -1079,21 +1236,38 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -1079,21 +1236,38 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None] lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None]
# The HardPhongShader can be used directly with atlas textures. # The HardPhongShader can be used directly with atlas textures.
rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings) rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings)
renderer = MeshRenderer( if rasterizer_type == MeshRasterizer:
rasterizer=rasterizer, shader = HardPhongShader(
shader=HardPhongShader(lights=lights, cameras=cameras, materials=materials), 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)
images = renderer(mesh) images = renderer(mesh)
rgb = images[0, ..., :3].squeeze() rgb = images[0, ..., :3].squeeze()
# Load reference image # Load reference image
image_ref = load_rgb_image("test_texture_atlas_8x8_back.png", DATA_DIR) image_ref = load_rgb_image(
f"test_texture_atlas_8x8_back_{rasterizer_type.__name__}.png", DATA_DIR
)
if DEBUG: if DEBUG:
Image.fromarray((rgb.detach().cpu().numpy() * 255).astype(np.uint8)).save( Image.fromarray((rgb.detach().cpu().numpy() * 255).astype(np.uint8)).save(
DATA_DIR / "DEBUG_texture_atlas_8x8_back.png" DATA_DIR
/ f"DEBUG_texture_atlas_8x8_back_{rasterizer_type.__name__}.png"
) )
self.assertClose(rgb.cpu(), image_ref, atol=0.05) self.assertClose(rgb.cpu(), image_ref, atol=0.05)
...@@ -1112,21 +1286,28 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -1112,21 +1286,28 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
raster_settings = RasterizationSettings( raster_settings = RasterizationSettings(
image_size=512, image_size=512,
blur_radius=0.0001, blur_radius=0.0001,
faces_per_pixel=5, faces_per_pixel=5 if rasterizer_type.__name__ == "MeshRasterizer" else 1,
cull_backfaces=True, cull_backfaces=rasterizer_type.__name__ == "MeshRasterizer",
clip_barycentric_coords=True, clip_barycentric_coords=True,
) )
images = renderer(mesh, raster_settings=raster_settings) images = renderer(mesh, raster_settings=raster_settings)
images[0, ...].sum().backward() images[0, ...].sum().backward()
fragments = rasterizer(mesh, raster_settings=raster_settings) fragments = rasterizer(mesh, raster_settings=raster_settings)
# Some of the bary coordinates are outside the if rasterizer_type == MeshRasterizer:
# [0, 1] range as expected because the blur is > 0 # Some of the bary coordinates are outside the
self.assertTrue(fragments.bary_coords.ge(1.0).any()) # [0, 1] range as expected because the blur is > 0.
self.assertTrue(fragments.bary_coords.ge(1.0).any())
self.assertIsNotNone(atlas.grad) self.assertIsNotNone(atlas.grad)
self.assertTrue(atlas.grad.sum().abs() > 0.0) self.assertTrue(atlas.grad.sum().abs() > 0.0)
def test_simple_sphere_outside_zfar(self): def test_simple_sphere_outside_zfar(self):
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):
""" """
Test output when rendering a sphere that is beyond zfar with a SoftPhongShader. 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 This renders a sphere of radius 500, with the camera at x=1500 for different
...@@ -1159,22 +1340,32 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -1159,22 +1340,32 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
cameras = FoVPerspectiveCameras( cameras = FoVPerspectiveCameras(
device=device, R=R, T=T, aspect_ratio=1.0, fov=60.0, zfar=zfar device=device, R=R, T=T, aspect_ratio=1.0, fov=60.0, zfar=zfar
) )
rasterizer = MeshRasterizer( blend_params = BlendParams(
cameras=cameras, raster_settings=raster_settings 1e-4 if rasterizer_type == MeshRasterizer else 0.5, 1e-4, (0, 0, 1.0)
) )
blend_params = BlendParams(1e-4, 1e-4, (0, 0, 1.0)) rasterizer = rasterizer_type(
cameras=cameras, raster_settings=raster_settings
shader = SoftPhongShader(
lights=lights,
cameras=cameras,
materials=materials,
blend_params=blend_params,
) )
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,
)
renderer = MeshRenderer(rasterizer=rasterizer, shader=shader) renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
images = renderer(sphere_mesh) images = renderer(sphere_mesh)
rgb = images[0, ..., :3].squeeze().cpu() rgb = images[0, ..., :3].squeeze().cpu()
filename = "test_simple_sphere_outside_zfar_%d.png" % int(zfar) filename = f"test_simple_sphere_outside_zfar_{int(zfar)}_{rasterizer_type.__name__}.png"
# Load reference image # Load reference image
image_ref = load_rgb_image(filename, DATA_DIR) image_ref = load_rgb_image(filename, DATA_DIR)
...@@ -1202,6 +1393,15 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -1202,6 +1393,15 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
sphere_mesh = Meshes(verts=verts_padded, faces=faces_padded, textures=textures) sphere_mesh = Meshes(verts=verts_padded, faces=faces_padded, textures=textures)
# No elevation or azimuth rotation # No elevation or azimuth rotation
rasterizer_tests = [
RasterizerTest(MeshRasterizer, HardPhongShader, "phong", "hard_phong"),
RasterizerTest(
MeshRasterizerOpenGL,
SplatterPhongShader,
"splatter",
"splatter_phong",
),
]
R, T = look_at_view_transform(2.7, 0.0, 0.0) R, T = look_at_view_transform(2.7, 0.0, 0.0)
for cam_type in ( for cam_type in (
FoVPerspectiveCameras, FoVPerspectiveCameras,
...@@ -1209,108 +1409,118 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase): ...@@ -1209,108 +1409,118 @@ class TestRenderMeshes(TestCaseMixin, unittest.TestCase):
PerspectiveCameras, PerspectiveCameras,
OrthographicCameras, OrthographicCameras,
): ):
cameras = cam_type(device=device, R=R, T=T) 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)
# Init shader settings # Init shader settings
materials = Materials(device=device) materials = Materials(
lights = PointLights(device=device) device=device,
ambient_color=WHITE,
diffuse_color=WHITE,
specular_color=WHITE,
)
lights = AmbientLights(
device=device,
ambient_color=WHITE,
)
lights.location = torch.tensor([0.0, 0.0, +2.0], device=device)[None] lights.location = torch.tensor([0.0, 0.0, +2.0], device=device)[None]
raster_settings = RasterizationSettings( raster_settings = RasterizationSettings(
image_size=512, blur_radius=0.0, faces_per_pixel=1 image_size=512, blur_radius=0.0, faces_per_pixel=1
) )
rasterizer = MeshRasterizer(raster_settings=raster_settings) rasterizer = MeshRasterizer(
blend_params = BlendParams(1e-4, 1e-4, (0, 0, 0)) cameras=cameras, raster_settings=raster_settings
)
blend_params = BlendParams(
1e-4,
1e-4,
background_color=BLACK[0],
)
shader = HardPhongShader( # only test HardFlatShader since that's the only one that makes
# sense for classification
shader = HardFlatShader(
lights=lights, lights=lights,
cameras=cameras,
materials=materials, materials=materials,
blend_params=blend_params, blend_params=blend_params,
) )
renderer = MeshRenderer(rasterizer=rasterizer, shader=shader) renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
images = renderer(sphere_mesh)
# Cameras can be passed into the renderer in the forward pass self.assertEqual(images.shape[-1], C + 1)
images = renderer(sphere_mesh, cameras=cameras) self.assertClose(images.amax(), torch.tensor(10.0), atol=0.01)
rgb = images.squeeze()[..., :3].cpu().numpy() self.assertClose(images.amin(), torch.tensor(0.0), atol=0.01)
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)
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)
# Init shader settings
materials = Materials(
device=device,
ambient_color=WHITE,
diffuse_color=WHITE,
specular_color=WHITE,
)
lights = AmbientLights(
device=device,
ambient_color=WHITE,
)
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,
background_color=BLACK[0],
)
# only test HardFlatShader since that's the only one that makes
# sense for classification
shader = HardFlatShader(
lights=lights,
cameras=cameras,
materials=materials,
blend_params=blend_params,
)
renderer = MeshRenderer(rasterizer=rasterizer, shader=shader)
images = renderer(sphere_mesh)
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)
# grab last 3 color channels # grab last 3 color channels
rgb = (images[0, ..., C - 3 : C] / 10).squeeze().cpu() rgb = (images[0, ..., C - 3 : C] / 10).squeeze().cpu()
filename = "test_nd_sphere.png" filename = "test_nd_sphere.png"
if DEBUG: if DEBUG:
debug_filename = "DEBUG_%s" % filename debug_filename = "DEBUG_%s" % filename
Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save( Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / debug_filename DATA_DIR / debug_filename
) )
image_ref = load_rgb_image(filename, DATA_DIR) image_ref = load_rgb_image(filename, DATA_DIR)
self.assertClose(rgb, image_ref, atol=0.05) self.assertClose(rgb, image_ref, atol=0.05)
tests/test_render_multigpu.py
View file @ cb495504
...@@ -14,6 +14,7 @@ from pytorch3d.renderer import ( ...@@ -14,6 +14,7 @@ from pytorch3d.renderer import (
HardGouraudShader, HardGouraudShader,
Materials, Materials,
MeshRasterizer, MeshRasterizer,
MeshRasterizerOpenGL,
MeshRenderer, MeshRenderer,
PointLights, PointLights,
PointsRasterizationSettings, PointsRasterizationSettings,
...@@ -21,18 +22,19 @@ from pytorch3d.renderer import ( ...@@ -21,18 +22,19 @@ from pytorch3d.renderer import (
PointsRenderer, PointsRenderer,
RasterizationSettings, RasterizationSettings,
SoftPhongShader, SoftPhongShader,
SplatterPhongShader,
TexturesVertex, TexturesVertex,
) )
from pytorch3d.renderer.cameras import FoVPerspectiveCameras, look_at_view_transform from pytorch3d.renderer.cameras import FoVPerspectiveCameras, look_at_view_transform
from pytorch3d.structures import Meshes, Pointclouds from pytorch3d.structures import Meshes, Pointclouds
from pytorch3d.utils.ico_sphere import ico_sphere from pytorch3d.utils.ico_sphere import ico_sphere
from .common_testing import get_random_cuda_device, TestCaseMixin from .common_testing import TestCaseMixin
# Set the number of GPUS you want to test with # Set the number of GPUS you want to test with
NUM_GPUS = 3 NUM_GPUS = 2
GPU_LIST = list({get_random_cuda_device() for _ in range(NUM_GPUS)}) GPU_LIST = [f"cuda:{idx}" for idx in range(NUM_GPUS)]
print("GPUs: %s" % ", ".join(GPU_LIST)) print("GPUs: %s" % ", ".join(GPU_LIST))
...@@ -56,12 +58,12 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase): ...@@ -56,12 +58,12 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
self.assertEqual(renderer.shader.materials.device, device) self.assertEqual(renderer.shader.materials.device, device)
self.assertEqual(renderer.shader.materials.ambient_color.device, device) self.assertEqual(renderer.shader.materials.ambient_color.device, device)
def test_mesh_renderer_to(self): def _mesh_renderer_to(self, rasterizer_class, shader_class):
""" """
Test moving all the tensors in the mesh renderer to a new device. Test moving all the tensors in the mesh renderer to a new device.
""" """
device1 = torch.device("cpu") device1 = torch.device("cuda:0")
R, T = look_at_view_transform(1500, 0.0, 0.0) R, T = look_at_view_transform(1500, 0.0, 0.0)
...@@ -71,12 +73,12 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase): ...@@ -71,12 +73,12 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
lights.location = torch.tensor([0.0, 0.0, +1000.0], device=device1)[None] lights.location = torch.tensor([0.0, 0.0, +1000.0], device=device1)[None]
raster_settings = RasterizationSettings( raster_settings = RasterizationSettings(
image_size=256, blur_radius=0.0, faces_per_pixel=1 image_size=128, blur_radius=0.0, faces_per_pixel=1
) )
cameras = FoVPerspectiveCameras( cameras = FoVPerspectiveCameras(
device=device1, R=R, T=T, aspect_ratio=1.0, fov=60.0, zfar=100 device=device1, R=R, T=T, aspect_ratio=1.0, fov=60.0, zfar=100
) )
rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings) rasterizer = rasterizer_class(cameras=cameras, raster_settings=raster_settings)
blend_params = BlendParams( blend_params = BlendParams(
1e-4, 1e-4,
...@@ -84,7 +86,7 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase): ...@@ -84,7 +86,7 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
background_color=torch.zeros(3, dtype=torch.float32, device=device1), background_color=torch.zeros(3, dtype=torch.float32, device=device1),
) )
shader = SoftPhongShader( shader = shader_class(
lights=lights, lights=lights,
cameras=cameras, cameras=cameras,
materials=materials, materials=materials,
...@@ -107,26 +109,32 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase): ...@@ -107,26 +109,32 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
# Move renderer and mesh to another device and re render # Move renderer and mesh to another device and re render
# This also tests that background_color is correctly moved to # This also tests that background_color is correctly moved to
# the new device # the new device
device2 = torch.device("cuda:0") device2 = torch.device("cuda:1")
renderer = renderer.to(device2) renderer = renderer.to(device2)
mesh = mesh.to(device2) mesh = mesh.to(device2)
self._check_mesh_renderer_props_on_device(renderer, device2) self._check_mesh_renderer_props_on_device(renderer, device2)
output_images = renderer(mesh) output_images = renderer(mesh)
self.assertEqual(output_images.device, device2) self.assertEqual(output_images.device, device2)
def test_render_meshes(self): def test_mesh_renderer_to(self):
self._mesh_renderer_to(MeshRasterizer, SoftPhongShader)
def test_mesh_renderer_opengl_to(self):
self._mesh_renderer_to(MeshRasterizerOpenGL, SplatterPhongShader)
def _render_meshes(self, rasterizer_class, shader_class):
test = self test = self
class Model(nn.Module): class Model(nn.Module):
def __init__(self): def __init__(self, device):
super(Model, self).__init__() super(Model, self).__init__()
mesh = ico_sphere(3) mesh = ico_sphere(3).to(device)
self.register_buffer("faces", mesh.faces_padded()) self.register_buffer("faces", mesh.faces_padded())
self.renderer = self.init_render() self.renderer = self.init_render(device)
def init_render(self): def init_render(self, device):
cameras = FoVPerspectiveCameras() cameras = FoVPerspectiveCameras().to(device)
raster_settings = RasterizationSettings( raster_settings = RasterizationSettings(
image_size=128, blur_radius=0.0, faces_per_pixel=1 image_size=128, blur_radius=0.0, faces_per_pixel=1
) )
...@@ -135,12 +143,12 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase): ...@@ -135,12 +143,12 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
diffuse_color=((0, 0.0, 0),), diffuse_color=((0, 0.0, 0),),
specular_color=((0.0, 0, 0),), specular_color=((0.0, 0, 0),),
location=((0.0, 0.0, 1e5),), location=((0.0, 0.0, 1e5),),
) ).to(device)
renderer = MeshRenderer( renderer = MeshRenderer(
rasterizer=MeshRasterizer( rasterizer=rasterizer_class(
cameras=cameras, raster_settings=raster_settings cameras=cameras, raster_settings=raster_settings
), ),
shader=HardGouraudShader(cameras=cameras, lights=lights), shader=shader_class(cameras=cameras, lights=lights),
) )
return renderer return renderer
...@@ -155,20 +163,25 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase): ...@@ -155,20 +163,25 @@ class TestRenderMeshesMultiGPU(TestCaseMixin, unittest.TestCase):
img_render = self.renderer(mesh) img_render = self.renderer(mesh)
return img_render[:, :, :, :3] return img_render[:, :, :, :3]
# DataParallel requires every input tensor be provided # Make sure we use all GPUs in GPU_LIST by making the batch size 4 x GPU count.
# on the first device in its device_ids list. verts = ico_sphere(3).verts_padded().expand(len(GPU_LIST) * 4, 642, 3)
verts = ico_sphere(3).verts_padded()
texs = verts.new_ones(verts.shape) texs = verts.new_ones(verts.shape)
model = Model() model = Model(device=GPU_LIST[0])
model.to(GPU_LIST[0])
model = nn.DataParallel(model, device_ids=GPU_LIST) model = nn.DataParallel(model, device_ids=GPU_LIST)
# Test a few iterations # Test a few iterations
for _ in range(100): for _ in range(100):
model(verts, texs) model(verts, texs)
def test_render_meshes(self):
self._render_meshes(MeshRasterizer, HardGouraudShader)
# @unittest.skip("Multi-GPU OpenGL training is currently not supported.")
def test_render_meshes_opengl(self):
self._render_meshes(MeshRasterizerOpenGL, SplatterPhongShader)
class TestRenderPointssMultiGPU(TestCaseMixin, unittest.TestCase): class TestRenderPointsMultiGPU(TestCaseMixin, unittest.TestCase):
def _check_points_renderer_props_on_device(self, renderer, device): def _check_points_renderer_props_on_device(self, renderer, device):
""" """
Helper function to check that all the properties have Helper function to check that all the properties have
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment