# Copyright (c) Meta Platforms, Inc. and affiliates. # 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. import re import unittest from itertools import product import numpy as np import torch from PIL import Image from pytorch3d.io import load_obj from pytorch3d.renderer import ( BlendParams, FoVPerspectiveCameras, look_at_view_transform, Materials, MeshRasterizer, MeshRenderer, PointLights, RasterizationSettings, SoftPhongShader, SplatterPhongShader, TexturesUV, ) from pytorch3d.renderer.mesh.rasterize_meshes import ( rasterize_meshes, rasterize_meshes_python, ) from pytorch3d.renderer.mesh.rasterizer import Fragments from pytorch3d.renderer.opengl import MeshRasterizerOpenGL from pytorch3d.renderer.points import ( AlphaCompositor, PointsRasterizationSettings, PointsRasterizer, PointsRenderer, ) from pytorch3d.renderer.points.rasterize_points import ( rasterize_points, rasterize_points_python, ) from pytorch3d.renderer.points.rasterizer import PointFragments from pytorch3d.structures import Meshes, Pointclouds from pytorch3d.transforms.transform3d import Transform3d from pytorch3d.utils import torus from .common_testing import ( get_pytorch3d_dir, get_tests_dir, load_rgb_image, TestCaseMixin, usesOpengl, ) DEBUG = False DATA_DIR = get_tests_dir() / "data" # Verts/Faces for a simple mesh with two faces. verts0 = torch.tensor( [ [-0.7, -0.70, 1.0], [0.0, -0.1, 1.0], [0.7, -0.7, 1.0], [-0.7, 0.1, 1.0], [0.0, 0.7, 1.0], [0.7, 0.1, 1.0], ], dtype=torch.float32, ) faces0 = torch.tensor([[1, 0, 2], [4, 3, 5]], dtype=torch.int64) # Points for a simple point cloud. Get the vertices from a # torus and apply rotations such that the points are no longer # symmerical in X/Y. torus_mesh = torus(r=0.25, R=1.0, sides=5, rings=2 * 5) t = ( Transform3d() .rotate_axis_angle(angle=90, axis="Y") .rotate_axis_angle(angle=45, axis="Z") .scale(0.3) ) torus_points = t.transform_points(torus_mesh.verts_padded()).squeeze() def _save_debug_image(idx, image_size, bin_size, blur): """ Save a mask image from the rasterization output for debugging. """ H, W = image_size # Save out the last image for debugging rgb = (idx[-1, ..., :3].cpu() > -1).squeeze() suffix = "square" if H == W else "non_square" filename = "%s_bin_size_%s_blur_%.3f_%dx%d.png" filename = filename % (suffix, str(bin_size), blur, H, W) if DEBUG: filename = "DEBUG_%s" % filename Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(DATA_DIR / filename) class TestRasterizeRectangleImagesErrors(TestCaseMixin, unittest.TestCase): def test_mesh_image_size_arg(self): meshes = Meshes(verts=[verts0], faces=[faces0]) with self.assertRaisesRegex(ValueError, re.escape("tuple/list of (H, W)")): rasterize_meshes( meshes, (100, 200, 3), 0.0001, faces_per_pixel=1, ) with self.assertRaisesRegex(ValueError, "sizes must be greater than 0"): rasterize_meshes( meshes, (0, 10), 0.0001, faces_per_pixel=1, ) with self.assertRaisesRegex(ValueError, "sizes must be integers"): rasterize_meshes( meshes, (100.5, 120.5), 0.0001, faces_per_pixel=1, ) def test_points_image_size_arg(self): points = Pointclouds([verts0]) with self.assertRaisesRegex(ValueError, re.escape("tuple/list of (H, W)")): rasterize_points( points, (100, 200, 3), 0.0001, points_per_pixel=1, ) with self.assertRaisesRegex(ValueError, "sizes must be greater than 0"): rasterize_points( points, (0, 10), 0.0001, points_per_pixel=1, ) with self.assertRaisesRegex(ValueError, "sizes must be integers"): rasterize_points( points, (100.5, 120.5), 0.0001, points_per_pixel=1, ) class TestRasterizeRectangleImagesMeshes(TestCaseMixin, unittest.TestCase): @staticmethod def _clone_mesh(verts0, faces0, device, batch_size): """ Helper function to detach and clone the verts/faces. This is needed in order to set up the tensors for gradient computation in different tests. """ verts = verts0.detach().clone() verts.requires_grad = True meshes = Meshes(verts=[verts], faces=[faces0]) meshes = meshes.to(device).extend(batch_size) return verts, meshes def _rasterize(self, meshes, image_size, bin_size, blur): """ Simple wrapper around the rasterize function to return the fragment data. """ face_idxs, zbuf, bary_coords, pix_dists = rasterize_meshes( meshes, image_size, blur, faces_per_pixel=1, bin_size=bin_size, ) return Fragments( pix_to_face=face_idxs, zbuf=zbuf, bary_coords=bary_coords, dists=pix_dists, ) @staticmethod def _save_debug_image(fragments, image_size, bin_size, blur): """ Save a mask image from the rasterization output for debugging. """ H, W = image_size # Save out the last image for debugging rgb = (fragments.pix_to_face[-1, ..., :3].cpu() > -1).squeeze() suffix = "square" if H == W else "non_square" filename = "triangle_%s_bin_size_%s_blur_%.3f_%dx%d.png" filename = filename % (suffix, str(bin_size), blur, H, W) if DEBUG: filename = "DEBUG_%s" % filename Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save( DATA_DIR / filename ) def _check_fragments(self, frag_1, frag_2): """ Helper function to check that the tensors in the Fragments frag_1 and frag_2 are the same. """ self.assertClose(frag_1.pix_to_face, frag_2.pix_to_face) self.assertClose(frag_1.dists, frag_2.dists) self.assertClose(frag_1.bary_coords, frag_2.bary_coords) self.assertClose(frag_1.zbuf, frag_2.zbuf) def _compare_square_with_nonsq( self, image_size, blur, device, verts0, faces0, nonsq_fragment_gradtensor_list, batch_size=1, ): """ Calculate the output from rasterizing a square image with the minimum of (H, W). Then compare this with the same square region in the non square image. The input mesh faces given by faces0 and verts0 are contained within the [-1, 1] range of the image so all the relevant pixels will be within the square region. `nonsq_fragment_gradtensor_list` is a list of fragments and verts grad tensors from rasterizing non square images. """ # Rasterize the square version of the image H, W = image_size S = min(H, W) verts_square, meshes_sq = self._clone_mesh(verts0, faces0, device, batch_size) square_fragments = self._rasterize( meshes_sq, image_size=(S, S), bin_size=0, blur=blur ) # Save debug image _save_debug_image(square_fragments.pix_to_face, (S, S), 0, blur) # Extract the values in the square image which are non zero. square_mask = square_fragments.pix_to_face > -1 square_dists = square_fragments.dists[square_mask] square_zbuf = square_fragments.zbuf[square_mask] square_bary = square_fragments.bary_coords[square_mask] # Retain gradients on the output of fragments to check # intermediate values with the non square outputs. square_fragments.dists.retain_grad() square_fragments.bary_coords.retain_grad() square_fragments.zbuf.retain_grad() # Calculate gradient for the square image torch.manual_seed(231) grad_zbuf = torch.randn_like(square_zbuf) grad_dist = torch.randn_like(square_dists) grad_bary = torch.randn_like(square_bary) loss0 = ( (grad_dist * square_dists).sum() + (grad_zbuf * square_zbuf).sum() + (grad_bary * square_bary).sum() ) loss0.backward() # Now compare against the non square outputs provided # in the nonsq_fragment_gradtensor_list list for fragments, grad_tensor, _name in nonsq_fragment_gradtensor_list: # Check that there are the same number of non zero pixels # in both the square and non square images. non_square_mask = fragments.pix_to_face > -1 self.assertEqual(non_square_mask.sum().item(), square_mask.sum().item()) # Check dists, zbuf and bary match the square image non_square_dists = fragments.dists[non_square_mask] non_square_zbuf = fragments.zbuf[non_square_mask] non_square_bary = fragments.bary_coords[non_square_mask] self.assertClose(square_dists, non_square_dists) self.assertClose(square_zbuf, non_square_zbuf) self.assertClose( square_bary, non_square_bary, atol=2e-7, ) # Retain gradients to compare values with outputs from # square image fragments.dists.retain_grad() fragments.bary_coords.retain_grad() fragments.zbuf.retain_grad() loss1 = ( (grad_dist * non_square_dists).sum() + (grad_zbuf * non_square_zbuf).sum() + (grad_bary * non_square_bary).sum() ) loss1.sum().backward() # Get the non zero values in the intermediate gradients # and compare with the values from the square image non_square_grad_dists = fragments.dists.grad[non_square_mask] non_square_grad_bary = fragments.bary_coords.grad[non_square_mask] non_square_grad_zbuf = fragments.zbuf.grad[non_square_mask] self.assertClose( non_square_grad_dists, square_fragments.dists.grad[square_mask], ) self.assertClose( non_square_grad_bary, square_fragments.bary_coords.grad[square_mask], ) self.assertClose( non_square_grad_zbuf, square_fragments.zbuf.grad[square_mask], ) # Finally check the gradients of the input vertices for # the square and non square case self.assertClose(verts_square.grad, grad_tensor.grad, rtol=3e-4, atol=5e-3) def test_gpu(self): """ Test that the output of rendering non square images gives the same result as square images. i.e. the dists, zbuf, bary are all the same for the square region which is present in both images. """ # Test both cases: (W > H), (H > W) as well as the case where # H and W are not integer multiples of each other (i.e. float aspect ratio) image_sizes = [(64, 128), (128, 64), (128, 256), (256, 128), (600, 1110)] devices = ["cuda:0"] blurs = [0.0, 0.001] batch_sizes = [1, 4] test_cases = product(image_sizes, blurs, devices, batch_sizes) for image_size, blur, device, batch_size in test_cases: # Initialize the verts grad tensor and the meshes objects verts_nonsq_naive, meshes_nonsq_naive = self._clone_mesh( verts0, faces0, device, batch_size ) verts_nonsq_binned, meshes_nonsq_binned = self._clone_mesh( verts0, faces0, device, batch_size ) # Get the outputs for both naive and coarse to fine rasterization fragments_naive = self._rasterize( meshes_nonsq_naive, image_size, blur=blur, bin_size=0, ) fragments_binned = self._rasterize( meshes_nonsq_binned, image_size, blur=blur, bin_size=None, ) # Save out debug images if needed _save_debug_image(fragments_naive.pix_to_face, image_size, 0, blur) _save_debug_image(fragments_binned.pix_to_face, image_size, None, blur) # Check naive and binned fragments give the same outputs self._check_fragments(fragments_naive, fragments_binned) # Here we want to compare the square image with the naive and the # coarse to fine methods outputs nonsq_fragment_gradtensor_list = [ (fragments_naive, verts_nonsq_naive, "naive"), (fragments_binned, verts_nonsq_binned, "coarse-to-fine"), ] self._compare_square_with_nonsq( image_size, blur, device, verts0, faces0, nonsq_fragment_gradtensor_list, batch_size, ) def test_cpu(self): """ Test that the output of rendering non square images gives the same result as square images. i.e. the dists, zbuf, bary are all the same for the square region which is present in both images. In this test we compare between the naive C++ implementation and the naive python implementation as the Coarse/Fine method is not fully implemented in C++ """ # Test both when (W > H) and (H > W). # Using smaller image sizes here as the Python rasterizer is really slow. image_sizes = [(32, 64), (64, 32), (60, 110)] devices = ["cpu"] blurs = [0.0, 0.001] batch_sizes = [1] test_cases = product(image_sizes, blurs, devices, batch_sizes) for image_size, blur, device, batch_size in test_cases: # Initialize the verts grad tensor and the meshes objects verts_nonsq_naive, meshes_nonsq_naive = self._clone_mesh( verts0, faces0, device, batch_size ) verts_nonsq_python, meshes_nonsq_python = self._clone_mesh( verts0, faces0, device, batch_size ) # Compare Naive CPU with Python as Coarse/Fine rasteriztation # is not implemented for CPU fragments_naive = self._rasterize( meshes_nonsq_naive, image_size, bin_size=0, blur=blur ) face_idxs, zbuf, bary_coords, pix_dists = rasterize_meshes_python( meshes_nonsq_python, image_size, blur, faces_per_pixel=1, ) fragments_python = Fragments( pix_to_face=face_idxs, zbuf=zbuf, bary_coords=bary_coords, dists=pix_dists, ) # Save debug images if DEBUG is set to true at the top of the file. _save_debug_image(fragments_naive.pix_to_face, image_size, 0, blur) _save_debug_image(fragments_python.pix_to_face, image_size, "python", blur) # List of non square outputs to compare with the square output nonsq_fragment_gradtensor_list = [ (fragments_naive, verts_nonsq_naive, "naive"), (fragments_python, verts_nonsq_python, "python"), ] self._compare_square_with_nonsq( image_size, blur, device, verts0, faces0, nonsq_fragment_gradtensor_list, batch_size, ) def test_render_cow(self): self._render_cow(MeshRasterizer) @usesOpengl 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. """ device = torch.device("cuda:0") obj_dir = get_pytorch3d_dir() / "docs/tutorials/data" obj_filename = obj_dir / "cow_mesh/cow.obj" # Load mesh + texture 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) # Init rasterizer settings R, T = look_at_view_transform(1.2, 0, 90) cameras = FoVPerspectiveCameras(device=device, R=R, T=T) raster_settings = RasterizationSettings( image_size=(500, 800), blur_radius=0.0, faces_per_pixel=1 ) # Init shader settings materials = Materials(device=device) lights = PointLights(device=device) lights.location = torch.tensor([0.0, 0.0, -2.0], device=device)[None] # Init renderer rasterizer = rasterizer_type(cameras=cameras, raster_settings=raster_settings) if rasterizer_type == MeshRasterizer: 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, cameras=cameras, materials=materials, 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 image_ref = load_rgb_image( f"test_cow_image_rectangle_{rasterizer_type.__name__}.png", DATA_DIR ) 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. renderer.rasterizer.raster_settings.bin_size = bin_size images = renderer(mesh) rgb = images[0, ..., :3].squeeze().cpu() if DEBUG: Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save( DATA_DIR / f"DEBUG_cow_image_rectangle_{rasterizer_type.__name__}.png" ) # NOTE some pixels can be flaky cond1 = torch.allclose(rgb, image_ref, atol=0.05) self.assertTrue(cond1) class TestRasterizeRectangleImagesPointclouds(TestCaseMixin, unittest.TestCase): @staticmethod def _clone_pointcloud(verts0, device, batch_size): """ Helper function to detach and clone the verts. This is needed in order to set up the tensors for gradient computation in different tests. """ verts = verts0.detach().clone() verts.requires_grad = True pointclouds = Pointclouds(points=[verts]) pointclouds = pointclouds.to(device).extend(batch_size) return verts, pointclouds def _rasterize(self, meshes, image_size, bin_size, blur): """ Simple wrapper around the rasterize function to return the fragment data. """ idxs, zbuf, dists = rasterize_points( meshes, image_size, blur, points_per_pixel=1, bin_size=bin_size, ) return PointFragments( idx=idxs, zbuf=zbuf, dists=dists, ) def _check_fragments(self, frag_1, frag_2): """ Helper function to check that the tensors in the Fragments frag_1 and frag_2 are the same. """ self.assertClose(frag_1.idx, frag_2.idx) self.assertClose(frag_1.dists, frag_2.dists) self.assertClose(frag_1.zbuf, frag_2.zbuf) def _compare_square_with_nonsq( self, image_size, blur, device, points, nonsq_fragment_gradtensor_list, batch_size=1, ): """ Calculate the output from rasterizing a square image with the minimum of (H, W). Then compare this with the same square region in the non square image. The input points are contained within the [-1, 1] range of the image so all the relevant pixels will be within the square region. `nonsq_fragment_gradtensor_list` is a list of fragments and verts grad tensors from rasterizing non square images. """ # Rasterize the square version of the image H, W = image_size S = min(H, W) points_square, pointclouds_sq = self._clone_pointcloud( points, device, batch_size ) square_fragments = self._rasterize( pointclouds_sq, image_size=(S, S), bin_size=0, blur=blur ) # Save debug image _save_debug_image(square_fragments.idx, (S, S), 0, blur) # Extract the values in the square image which are non zero. square_mask = square_fragments.idx > -1 square_dists = square_fragments.dists[square_mask] square_zbuf = square_fragments.zbuf[square_mask] # Retain gradients on the output of fragments to check # intermediate values with the non square outputs. square_fragments.dists.retain_grad() square_fragments.zbuf.retain_grad() # Calculate gradient for the square image torch.manual_seed(231) grad_zbuf = torch.randn_like(square_zbuf) grad_dist = torch.randn_like(square_dists) loss0 = (grad_dist * square_dists).sum() + (grad_zbuf * square_zbuf).sum() loss0.backward() # Now compare against the non square outputs provided # in the nonsq_fragment_gradtensor_list list for fragments, grad_tensor, _name in nonsq_fragment_gradtensor_list: # Check that there are the same number of non zero pixels # in both the square and non square images. non_square_mask = fragments.idx > -1 self.assertEqual(non_square_mask.sum().item(), square_mask.sum().item()) # Check dists, zbuf and bary match the square image non_square_dists = fragments.dists[non_square_mask] non_square_zbuf = fragments.zbuf[non_square_mask] self.assertClose(square_dists, non_square_dists) self.assertClose(square_zbuf, non_square_zbuf) # Retain gradients to compare values with outputs from # square image fragments.dists.retain_grad() fragments.zbuf.retain_grad() loss1 = (grad_dist * non_square_dists).sum() + ( grad_zbuf * non_square_zbuf ).sum() loss1.sum().backward() # Get the non zero values in the intermediate gradients # and compare with the values from the square image non_square_grad_dists = fragments.dists.grad[non_square_mask] non_square_grad_zbuf = fragments.zbuf.grad[non_square_mask] self.assertClose( non_square_grad_dists, square_fragments.dists.grad[square_mask], ) self.assertClose( non_square_grad_zbuf, square_fragments.zbuf.grad[square_mask], ) # Finally check the gradients of the input vertices for # the square and non square case self.assertClose(points_square.grad, grad_tensor.grad, rtol=2e-4) def test_gpu(self): """ Test that the output of rendering non square images gives the same result as square images. i.e. the dists, zbuf, idx are all the same for the square region which is present in both images. """ # Test both cases: (W > H), (H > W) as well as the case where # H and W are not integer multiples of each other (i.e. float aspect ratio) image_sizes = [(64, 128), (128, 64), (128, 256), (256, 128), (600, 1110)] devices = ["cuda:0"] blurs = [5e-2] batch_sizes = [1, 4] test_cases = product(image_sizes, blurs, devices, batch_sizes) for image_size, blur, device, batch_size in test_cases: # Initialize the verts grad tensor and the meshes objects verts_nonsq_naive, pointcloud_nonsq_naive = self._clone_pointcloud( torus_points, device, batch_size ) verts_nonsq_binned, pointcloud_nonsq_binned = self._clone_pointcloud( torus_points, device, batch_size ) # Get the outputs for both naive and coarse to fine rasterization fragments_naive = self._rasterize( pointcloud_nonsq_naive, image_size, blur=blur, bin_size=0, ) fragments_binned = self._rasterize( pointcloud_nonsq_binned, image_size, blur=blur, bin_size=None, ) # Save out debug images if needed _save_debug_image(fragments_naive.idx, image_size, 0, blur) _save_debug_image(fragments_binned.idx, image_size, None, blur) # Check naive and binned fragments give the same outputs self._check_fragments(fragments_naive, fragments_binned) # Here we want to compare the square image with the naive and the # coarse to fine methods outputs nonsq_fragment_gradtensor_list = [ (fragments_naive, verts_nonsq_naive, "naive"), (fragments_binned, verts_nonsq_binned, "coarse-to-fine"), ] self._compare_square_with_nonsq( image_size, blur, device, torus_points, nonsq_fragment_gradtensor_list, batch_size, ) def test_cpu(self): """ Test that the output of rendering non square images gives the same result as square images. i.e. the dists, zbuf, idx are all the same for the square region which is present in both images. In this test we compare between the naive C++ implementation and the naive python implementation as the Coarse/Fine method is not fully implemented in C++ """ # Test both when (W > H) and (H > W). # Using smaller image sizes here as the Python rasterizer is really slow. image_sizes = [(32, 64), (64, 32), (60, 110)] devices = ["cpu"] blurs = [5e-2] batch_sizes = [1] test_cases = product(image_sizes, blurs, devices, batch_sizes) for image_size, blur, device, batch_size in test_cases: # Initialize the verts grad tensor and the meshes objects verts_nonsq_naive, pointcloud_nonsq_naive = self._clone_pointcloud( torus_points, device, batch_size ) verts_nonsq_python, pointcloud_nonsq_python = self._clone_pointcloud( torus_points, device, batch_size ) # Compare Naive CPU with Python as Coarse/Fine rasteriztation # is not implemented for CPU fragments_naive = self._rasterize( pointcloud_nonsq_naive, image_size, bin_size=0, blur=blur ) idxs, zbuf, pix_dists = rasterize_points_python( pointcloud_nonsq_python, image_size, blur, points_per_pixel=1, ) fragments_python = PointFragments( idx=idxs, zbuf=zbuf, dists=pix_dists, ) # Save debug images if DEBUG is set to true at the top of the file. _save_debug_image(fragments_naive.idx, image_size, 0, blur) _save_debug_image(fragments_python.idx, image_size, "python", blur) # List of non square outputs to compare with the square output nonsq_fragment_gradtensor_list = [ (fragments_naive, verts_nonsq_naive, "naive"), (fragments_python, verts_nonsq_python, "python"), ] self._compare_square_with_nonsq( image_size, blur, device, torus_points, nonsq_fragment_gradtensor_list, batch_size, ) def test_render_pointcloud(self): """ Test a textured point cloud is rendered correctly in a non square image. """ device = torch.device("cuda:0") pointclouds = Pointclouds( points=[torus_points * 2.0], features=torch.ones_like(torus_points[None, ...]), ).to(device) R, T = look_at_view_transform(2.7, 0.0, 0.0) cameras = FoVPerspectiveCameras(device=device, R=R, T=T) raster_settings = PointsRasterizationSettings( image_size=(512, 1024), radius=5e-2, points_per_pixel=1 ) rasterizer = PointsRasterizer(cameras=cameras, raster_settings=raster_settings) compositor = AlphaCompositor() renderer = PointsRenderer(rasterizer=rasterizer, compositor=compositor) # Load reference image image_ref = load_rgb_image("test_pointcloud_rectangle_image.png", DATA_DIR) 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(pointclouds) rgb = images[0, ..., :3].squeeze().cpu() if DEBUG: Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save( DATA_DIR / "DEBUG_pointcloud_rectangle_image.png" ) # NOTE some pixels can be flaky cond1 = torch.allclose(rgb, image_ref, atol=0.05) self.assertTrue(cond1)