def collate_batched_meshes(batch: List[Dict]): # pragma: no cover
"""
Take a list of objects in the form of dictionaries and merge them
into a single dictionary. This function can be used with a Dataset
object to create a torch.utils.data.Dataloader which directly
returns Meshes objects.
TODO: Add support for textures.
Args:
batch: List of dictionaries containing information about objects
in the dataset.
Returns:
collated_dict: Dictionary of collated lists. If batch contains both
verts and faces, a collated mesh batch is also returned.
"""
if batch is None or len(batch) == 0:
return None
collated_dict = {}
for k in batch[0].keys():
collated_dict[k] = [d[k] for d in batch]
collated_dict["mesh"] = None
if {"verts", "faces"}.issubset(collated_dict.keys()):
textures = None
if "textures" in collated_dict:
textures = TexturesAtlas(atlas=collated_dict["textures"])
collated_dict["mesh"] = Meshes(
verts=collated_dict["verts"],
faces=collated_dict["faces"],
textures=textures,
)
return collated_dict
def test_texture_map_atlas(self):
"""
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.
"""
device = torch.device("cuda:0")
obj_dir = Path(__file__).resolve().parent.parent / "docs/tutorials/data"
obj_filename = obj_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]),
)
# 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,
cull_backfaces=True,
perspective_correct=False,
)
# Init shader settings
materials = Materials(device=device, specular_color=((0, 0, 0),), shininess=0.0)
lights = PointLights(device=device)
# Place light behind the cow in world space. The front of
# the cow is facing the -z direction.
lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None]
# The HardPhongShader can be used directly with atlas textures.
rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
renderer = MeshRenderer(
rasterizer=rasterizer,
shader=HardPhongShader(lights=lights, cameras=cameras, materials=materials),
)
images = renderer(mesh)
rgb = images[0, ..., :3].squeeze()
# Load reference image
image_ref = load_rgb_image("test_texture_atlas_8x8_back.png", DATA_DIR)
if DEBUG:
Image.fromarray((rgb.detach().cpu().numpy() * 255).astype(np.uint8)).save(
DATA_DIR / "DEBUG_texture_atlas_8x8_back.png"
)
self.assertClose(rgb.cpu(), image_ref, atol=0.05)
# Check gradients are propagated
# correctly back to the texture atlas.
# Because of how texture sampling is implemented
# for the texture atlas it is not possible to get
# gradients back to the vertices.
atlas.requires_grad = True
mesh = Meshes(
verts=[verts],
faces=[faces.verts_idx],
textures=TexturesAtlas(atlas=[atlas]),
)
raster_settings = RasterizationSettings(
image_size=512,
blur_radius=0.0001,
faces_per_pixel=5,
cull_backfaces=True,
clip_barycentric_coords=True,
)
images = renderer(mesh, raster_settings=raster_settings)
images[0, ...].sum().backward()
fragments = rasterizer(mesh, raster_settings=raster_settings)
# Some of the bary coordinates are outisde the
# [0, 1] range as expected because the blur is > 0
self.assertTrue(fragments.bary_coords.ge(1.0).any())
self.assertIsNotNone(atlas.grad)
self.assertTrue(atlas.grad.sum().abs() > 0.0)
def test_join_atlas(self):
"""Meshes with TexturesAtlas joined into a scene"""
# Test the result of rendering two tori with separate textures.
# The expected result is consistent with rendering them each alone.
torch.manual_seed(1)
device = torch.device("cuda:0")
plain_torus = torus(r=1, R=4, sides=5, rings=6, device=device)
[verts] = plain_torus.verts_list()
verts_shifted1 = verts.clone()
verts_shifted1 *= 1.2
verts_shifted1[:, 0] += 4
verts_shifted1[:, 1] += 5
verts[:, 0] -= 4
verts[:, 1] -= 4
[faces] = plain_torus.faces_list()
map_size = 3
# Two random atlases.
# The averaging of the random numbers here is not consistent with the
# meaning of the atlases, but makes each face a bit smoother than
# if everything had a random color.
atlas1 = torch.rand(size=(faces.shape[0], map_size, map_size, 3), device=device)
atlas1[:, 1] = 0.5 * atlas1[:, 0] + 0.5 * atlas1[:, 2]
atlas1[:, :, 1] = 0.5 * atlas1[:, :, 0] + 0.5 * atlas1[:, :, 2]
atlas2 = torch.rand(size=(faces.shape[0], map_size, map_size, 3), device=device)
atlas2[:, 1] = 0.5 * atlas2[:, 0] + 0.5 * atlas2[:, 2]
atlas2[:, :, 1] = 0.5 * atlas2[:, :, 0] + 0.5 * atlas2[:, :, 2]
textures1 = TexturesAtlas(atlas=[atlas1])
textures2 = TexturesAtlas(atlas=[atlas2])
mesh1 = Meshes(verts=[verts], faces=[faces], textures=textures1)
mesh2 = Meshes(verts=[verts_shifted1], faces=[faces], textures=textures2)
mesh_joined = join_meshes_as_scene([mesh1, mesh2])
R, T = look_at_view_transform(18, 0, 0)
cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
raster_settings = RasterizationSettings(
image_size=512,
blur_radius=0.0,
faces_per_pixel=1,
perspective_correct=False,
)
lights = PointLights(
device=device,
ambient_color=((1.0, 1.0, 1.0),),
diffuse_color=((0.0, 0.0, 0.0),),
specular_color=((0.0, 0.0, 0.0),),
)
blend_params = BlendParams(
sigma=1e-1,
gamma=1e-4,
background_color=torch.tensor([1.0, 1.0, 1.0], device=device),
)
renderer = MeshRenderer(
rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings),
shader=HardPhongShader(
device=device, blend_params=blend_params, cameras=cameras, lights=lights
),
)
output = renderer(mesh_joined)
image_ref = load_rgb_image("test_joinatlas_final.png", DATA_DIR)
if DEBUG:
debugging_outputs = []
for mesh_ in [mesh1, mesh2]:
debugging_outputs.append(renderer(mesh_))
Image.fromarray(
(output[0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
).save(DATA_DIR / "test_joinatlas_final_.png")
Image.fromarray(
(debugging_outputs[0][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
).save(DATA_DIR / "test_joinatlas_1.png")
Image.fromarray(
(debugging_outputs[1][0, ..., :3].cpu().numpy() * 255).astype(np.uint8)
).save(DATA_DIR / "test_joinatlas_2.png")
result = output[0, ..., :3].cpu()
self.assertClose(result, image_ref, atol=0.05)
def test_texture_map_atlas(self):
"""
Test a mesh with a texture map as a per face atlas is loaded and rendered correctly.
"""
device = torch.device("cuda:0")
obj_dir = Path(
__file__).resolve().parent.parent / "docs/tutorials/data"
obj_filename = obj_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,
)
mesh = Meshes(
verts=[verts],
faces=[faces.verts_idx],
textures=TexturesAtlas(atlas=[aux.texture_atlas]),
)
# 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,
cull_backfaces=True)
# Init shader settings
materials = Materials(device=device,
specular_color=((0, 0, 0), ),
shininess=0.0)
lights = PointLights(device=device)
# Place light behind the cow in world space. The front of
# the cow is facing the -z direction.
lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None]
# The HardPhongShader can be used directly with atlas textures.
renderer = MeshRenderer(
rasterizer=MeshRasterizer(cameras=cameras,
raster_settings=raster_settings),
shader=HardPhongShader(lights=lights,
cameras=cameras,
materials=materials),
)
images = renderer(mesh)
rgb = images[0, ..., :3].squeeze().cpu()
# Load reference image
image_ref = load_rgb_image("test_texture_atlas_8x8_back.png", DATA_DIR)
if DEBUG:
Image.fromarray((rgb.numpy() * 255).astype(np.uint8)).save(
DATA_DIR / "DEBUG_texture_atlas_8x8_back.png")
self.assertClose(rgb, image_ref, atol=0.05)
