def meshes_collate(batch): r"""Puts each data field into a tensor with outer dimension batch size""" elem = batch[0] elem_type = type(elem) if isinstance(elem, torch.Tensor): out = None if torch.utils.data.get_worker_info() is not None: # If we're in a background process, concatenate directly into a # shared memory tensor to avoid an extra copy numel = sum([x.numel() for x in batch]) storage = elem.storage()._new_shared(numel) out = elem.new(storage) return torch.stack(batch, 0, out=out) elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \ and elem_type.__name__ != 'string_': elem = batch[0] if elem_type.__name__ == 'ndarray': # array of string classes and object if collate.np_str_obj_array_pattern.search( elem.dtype.str) is not None: raise TypeError( collate.default_collate_err_msg_format.format(elem.dtype)) return meshes_collate([torch.as_tensor(b) for b in batch]) elif elem.shape == (): # scalars return torch.as_tensor(batch) elif isinstance(elem, float): return torch.tensor(batch, dtype=torch.float64) elif isinstance(elem, collate.int_classes): return torch.tensor(batch) elif isinstance(elem, collate.string_classes): return batch elif isinstance(elem, collate.container_abcs.Mapping): return {key: meshes_collate([d[key] for d in batch]) for key in elem} elif isinstance(elem, tuple) and hasattr(elem, '_fields'): # namedtuple return elem_type(*(meshes_collate(samples) for samples in zip(*batch))) elif isinstance(elem, collate.container_abcs.Sequence): transposed = zip(*batch) return [meshes_collate(samples) for samples in transposed] elif isinstance(elem, Meshes): return join_meshes_as_batch(batch) raise TypeError(collate.default_collate_err_msg_format.format(elem_type))
def test_batch_uvs(self): """Test that two random tori with TexturesUV render the same as each individually.""" torch.manual_seed(1) device = torch.device("cuda:0") plain_torus = torus(r=1, R=4, sides=10, rings=10, device=device) [verts] = plain_torus.verts_list() [faces] = plain_torus.faces_list() nocolor = torch.zeros((100, 100), device=device) color_gradient = torch.linspace(0, 1, steps=100, device=device) color_gradient1 = color_gradient[None].expand_as(nocolor) color_gradient2 = color_gradient[:, None].expand_as(nocolor) colors1 = torch.stack([nocolor, color_gradient1, color_gradient2], dim=2) colors2 = torch.stack([color_gradient1, color_gradient2, nocolor], dim=2) verts_uvs1 = torch.rand(size=(verts.shape[0], 2), device=device) verts_uvs2 = torch.rand(size=(verts.shape[0], 2), device=device) textures1 = TexturesUV( maps=[colors1], faces_uvs=[faces], verts_uvs=[verts_uvs1] ) textures2 = TexturesUV( maps=[colors2], faces_uvs=[faces], verts_uvs=[verts_uvs2] ) mesh1 = Meshes(verts=[verts], faces=[faces], textures=textures1) mesh2 = Meshes(verts=[verts], faces=[faces], textures=textures2) mesh_both = join_meshes_as_batch([mesh1, mesh2]) R, T = look_at_view_transform(10, 10, 0) cameras = FoVPerspectiveCameras(device=device, R=R, T=T) raster_settings = RasterizationSettings( image_size=128, blur_radius=0.0, faces_per_pixel=1 ) # Init shader settings lights = PointLights(device=device) lights.location = torch.tensor([0.0, 0.0, 2.0], device=device)[None] blend_params = BlendParams( sigma=1e-1, gamma=1e-4, background_color=torch.tensor([1.0, 1.0, 1.0], device=device), ) # Init renderer renderer = MeshRenderer( rasterizer=MeshRasterizer(cameras=cameras, raster_settings=raster_settings), shader=HardPhongShader( device=device, lights=lights, cameras=cameras, blend_params=blend_params ), ) outputs = [] for meshes in [mesh_both, mesh1, mesh2]: outputs.append(renderer(meshes)) if DEBUG: Image.fromarray( (outputs[0][0, ..., :3].cpu().numpy() * 255).astype(np.uint8) ).save(DATA_DIR / "test_batch_uvs0.png") Image.fromarray( (outputs[1][0, ..., :3].cpu().numpy() * 255).astype(np.uint8) ).save(DATA_DIR / "test_batch_uvs1.png") Image.fromarray( (outputs[0][1, ..., :3].cpu().numpy() * 255).astype(np.uint8) ).save(DATA_DIR / "test_batch_uvs2.png") Image.fromarray( (outputs[2][0, ..., :3].cpu().numpy() * 255).astype(np.uint8) ).save(DATA_DIR / "test_batch_uvs3.png") diff = torch.abs(outputs[0][0, ..., :3] - outputs[1][0, ..., :3]) Image.fromarray(((diff > 1e-5).cpu().numpy().astype(np.uint8) * 255)).save( DATA_DIR / "test_batch_uvs01.png" ) diff = torch.abs(outputs[0][1, ..., :3] - outputs[2][0, ..., :3]) Image.fromarray(((diff > 1e-5).cpu().numpy().astype(np.uint8) * 255)).save( DATA_DIR / "test_batch_uvs23.png" ) self.assertClose(outputs[0][0, ..., :3], outputs[1][0, ..., :3], atol=1e-5) self.assertClose(outputs[0][1, ..., :3], outputs[2][0, ..., :3], atol=1e-5)