def test_render_texture(self, uvs, faces, texture_maps, face_vertices_camera, face_vertices_image, face_camera_normals_z, height, width, dtype, device): batch_size = faces.shape[0] # attributes with uvs #uvs_with_mask = torch.nn.functional.pad(uvs, pad=(1, 0), value=1) #attributes = uvs_with_mask face_uvs = index_vertices_by_faces(uvs, faces) face_attributes = [ torch.ones((*face_uvs.shape[:-1], 1), device=device, dtype=dtype), face_uvs ] (texmask, texcoord), improb, imfaceidx = dibr_rasterization( height, width, face_vertices_camera[:, :, :, 2], face_vertices_image, face_attributes, face_camera_normals_z) texcolor = texture_mapping(texcoord, texture_maps, mode='bilinear') image = texcolor * texmask for bs in range(batch_size): image_gt = torch.from_numpy( np.array( Image.open(os.path.join(SAMPLE_DIR, f'texture_{bs}.png')))) image_gt = image_gt.to(device, dtype) / 255. assert torch.allclose(image[bs], image_gt, atol=1. / 255.0)
def test_all_grads(self, height, width, face_vertices_camera, face_vertices_image_zoom, face_vertex_colors, face_camera_normals_z, dtype, target_all_grads): _face_vertex_colors = face_vertex_colors.detach() _face_vertex_colors.requires_grad = True _face_vertices_image_zoom = face_vertices_image_zoom.detach() _face_vertices_image_zoom.requires_grad = True outputs = dibr_rasterization( 20, 30, face_vertices_camera.to(dtype)[:, :, :, 2], _face_vertices_image_zoom.to(dtype), _face_vertex_colors.to(dtype), face_camera_normals_z.to(dtype)) # gradients will be provided through the two differentiable outputs output = sum([torch.sum(output) for output in outputs]) output.backward() if dtype == torch.float: assert torch.allclose(target_all_grads[0], _face_vertices_image_zoom.grad, rtol=1e-5, atol=1e-4) else: # TODO(cfujitsang): non-determinism? assert torch.allclose(target_all_grads[0], _face_vertices_image_zoom.grad) assert torch.allclose(target_all_grads[1], _face_vertex_colors.grad)
def test_render_depths(self, face_vertices_camera, face_vertices_image, face_camera_normals_z, height, width, dtype, device): batch_size = face_vertices_camera.shape[0] # face_vertices_camera is of shape (num_batch, num_face, 9) num_batch, num_faces = face_vertices_camera.shape[:2] face_attributes = face_vertices_camera.reshape(num_batch, num_faces, 3, 3)[:, :, :, 2:3] # imfeat is interpolated features # improb is the soft mask # imfaceidx is the face index map, which pixel is covered by which face # it starts from 1, 0 is void. imfeat, improb, imfaceidx = dibr_rasterization( height, width, face_vertices_camera[:, :, :, 2], face_vertices_image, face_attributes, face_camera_normals_z) image = imfeat image_valid_region = image < 0 image_valid_values = image[image_valid_region] image_depth_norm = (image - image_valid_values.min()) / ( image_valid_values.max() - image_valid_values.min()) image = torch.where(image_valid_region, image_depth_norm, image) for bs in range(batch_size): image_gt = torch.from_numpy( np.array( Image.open(os.path.join(SAMPLE_DIR, f'depth_{bs}.png')))) image_gt = image_gt.to(device, dtype).unsqueeze(2) / 255. assert torch.allclose(image[bs], image_gt, atol=1 / 255.0)
def test_render_normal(self, face_vertices_camera, face_vertices_image, face_camera_normals_z, height, width, dtype, device): batch_size = face_vertices_camera.shape[0] face_normals_unit = face_normals(face_vertices_camera, unit=True) face_attributes = face_normals_unit.unsqueeze(-2).repeat(1, 1, 3, 1) # imfeat is interpolated features # improb is the soft mask # imfaceidx is the face index map, which pixel is covered by which face # it starts from 1, 0 is void. imfeat, improb, imfaceidx = dibr_rasterization( height, width, face_vertices_camera[:, :, :, 2], face_vertices_image, face_attributes, face_camera_normals_z) images = (imfeat + 1) / 2 images_gt = [ torch.from_numpy( np.array( Image.open( os.path.join(SAMPLE_DIR, f'vertex_normal_{bs}.png')))) for bs in range(batch_size) ] images_gt = torch.stack(images_gt, dim=0).to(device, dtype) / 255. if dtype == torch.double: num_pix_diff_tol = 8 else: num_pix_diff_tol = 0 num_pix_diff = torch.sum( ~torch.isclose(images, images_gt, atol=1. / 255.)) assert num_pix_diff <= num_pix_diff_tol
def test_render_vertex_colors(self, vertex_colors, faces, face_vertices_camera, face_vertices_image, face_camera_normals_z, height, width, dtype, device): batch_size = faces.shape[0] # face_vertex_colors attributes = vertex_colors face_attributes_idx = faces face_attributes = index_vertices_by_faces(attributes, face_attributes_idx) # imfeat is interpolated features # improb is the soft mask # imfaceidx is the face index map, which pixel is covered by which face # it starts from 1, 0 is void. imfeat, improb, imfaceidx = dibr_rasterization( height, width, face_vertices_camera[:, :, :, 2], face_vertices_image, face_attributes, face_camera_normals_z) image = imfeat images_gt = [ torch.from_numpy( np.array( Image.open( os.path.join(SAMPLE_DIR, f'vertex_color_{bs}.png')))) for bs in range(batch_size) ] images_gt = torch.stack(images_gt, dim=0).to(device, dtype) / 255. # the rendered soft mask is only tested here images_prob_gt = [ torch.from_numpy( np.array( Image.open(os.path.join(SAMPLE_DIR, f"image_prob_{bs}.png")))) for bs in range(batch_size) ] images_prob_gt = torch.stack(images_prob_gt, dim=0).to(device, dtype) / 255. # the rendered face_idx is only tested here images_face_idx_gt = [ torch.from_numpy( np.array( Image.open( os.path.join(SAMPLE_DIR, f"image_face_idx_{bs}.png")))) for bs in range(batch_size) ] images_face_idx_gt = \ (torch.stack(images_face_idx_gt, dim=0).to(device, torch.long) // 100) - 1 assert torch.allclose(image, images_gt, atol=1. / 255.) assert torch.allclose(improb, images_prob_gt, atol=1. / 255.0) if dtype == torch.double: num_pix_diff_tol = 4 else: num_pix_diff_tol = 0 num_pix_diff = torch.sum( ~torch.isclose(imfaceidx, images_face_idx_gt, atol=1. / 255.)) assert num_pix_diff <= num_pix_diff_tol
def render(self, **attributes): azimuths = attributes['azimuths'] elevations = attributes['elevations'] distances = attributes['distances'] batch_size = azimuths.shape[0] device = azimuths.device cam_proj = self.cam_proj.to(device) vertices = attributes['vertices'] textures = attributes['textures'] lights = attributes['lights'] faces = self.faces.to(device) face_uvs = self.face_uvs.to(device) num_faces = faces.shape[0] object_pos = torch.tensor([[0., 0., 0.]], dtype=torch.float, device=device).repeat(batch_size, 1) camera_up = torch.tensor([[0., 1., 0.]], dtype=torch.float, device=device).repeat(batch_size, 1) # camera_pos = torch.tensor([[0., 0., 4.]], dtype=torch.float, device=device).repeat(batch_size, 1) camera_pos = camera_position_from_spherical_angles(distances, elevations, azimuths, degrees=True) cam_transform = generate_transformation_matrix(camera_pos, object_pos, camera_up) face_vertices_camera, face_vertices_image, face_normals = \ prepare_vertices(vertices=vertices, faces=faces, camera_proj=cam_proj, camera_transform=cam_transform ) face_normals_unit = kal.ops.mesh.face_normals(face_vertices_camera, unit=True) face_normals_unit = face_normals_unit.unsqueeze(-2).repeat(1, 1, 3, 1) face_attributes = [ torch.ones((batch_size, num_faces, 3, 1), device=device), face_uvs.repeat(batch_size, 1, 1, 1), face_normals_unit ] image_features, soft_mask, face_idx = dibr_rasterization( self.image_size, self.image_size, face_vertices_camera[:, :, :, -1], face_vertices_image, face_attributes, face_normals[:, :, -1]) # image_features is a tuple in composed of the interpolated attributes of face_attributes # texture_coords, mask = image_features texmask, texcoord, imnormal = image_features texcolor = texture_mapping(texcoord, textures, mode='bilinear') coef = spherical_harmonic_lighting(imnormal, lights) image = texcolor * texmask * coef.unsqueeze(-1) + torch.ones_like(texcolor) * (1 - texmask) image = torch.clamp(image, 0, 1) render_img = image render_silhouttes = soft_mask[..., None] rgbs = torch.cat([render_img, render_silhouttes], axis=-1).permute(0, 3, 1, 2) attributes['face_normals'] = face_normals attributes['faces_image'] = face_vertices_image.mean(dim=2) attributes['visiable_faces'] = face_normals[:, :, -1] > 0.1 return rgbs, attributes
def target_face_vertex_colors_grad(self, height, width, face_vertices_camera, face_vertices_image, face_vertex_colors, face_camera_normals_z): _face_vertex_colors = face_vertex_colors.detach() _face_vertex_colors.requires_grad = True # This assume that test_vertex_colors_gradcheck pass outputs = dibr_rasterization(20, 30, face_vertices_camera[:, :, :, 2], face_vertices_image, _face_vertex_colors, face_camera_normals_z) # gradients will be provided through the two differentiable outputs output = sum([torch.sum(output) for output in outputs]) output.backward() return _face_vertex_colors.grad.clone()
def test_face_vertex_colors_grads(self, height, width, face_vertices_camera, face_vertices_image, face_vertex_colors, face_camera_normals_z, dtype, target_face_vertex_colors_grad): _face_vertex_colors = face_vertex_colors.detach() _face_vertex_colors.requires_grad = True outputs = dibr_rasterization( 20, 30, face_vertices_camera.to(dtype)[:, :, :, 2], face_vertices_image.to(dtype), _face_vertex_colors.to(dtype), face_camera_normals_z.to(dtype)) # gradients will be provided through the two differentiable outputs output = sum([torch.sum(output) for output in outputs]) output.backward() assert torch.allclose(_face_vertex_colors.grad, target_face_vertex_colors_grad)
def test_render_texture_with_light(self, uvs, faces, texture_maps, lights, face_vertices_camera, face_vertices_image, face_camera_normals_z, height, width, dtype, device): batch_size = faces.shape[0] # Note: in this example uv face is the same as mesh face # but they could be different face_uvs = index_vertices_by_faces(uvs, faces) # normal face_normals_unit = face_normals(face_vertices_camera, unit=True) face_normals_unit = face_normals_unit.unsqueeze(-2).repeat(1, 1, 3, 1) # merge them together face_attributes = [ torch.ones((*face_uvs.shape[:-1], 1), device=device, dtype=dtype), face_uvs, face_normals_unit ] (texmask, texcoord, imnormal), improb, imidx = dibr_rasterization( height, width, face_vertices_camera[:, :, :, 2], face_vertices_image, face_attributes, face_camera_normals_z) texcolor = texture_mapping(texcoord, texture_maps, mode='nearest') coef = spherical_harmonic_lighting(imnormal, lights) images = torch.clamp(texmask * texcolor * coef.unsqueeze(-1), 0, 1) if dtype == torch.double: num_pix_diff_tol = 74 # (over 2 x 256 x 512 x 3 pixels) else: num_pix_diff_tol = 0 images_gt = [ torch.from_numpy( np.array( Image.open( os.path.join(SAMPLE_DIR, f'texture_light_{bs}.png')))) for bs in range(batch_size) ] images_gt = torch.stack(images_gt, dim=0).to(device, dtype) / 255. num_pix_diff = torch.sum( ~torch.isclose(images, images_gt, atol=1. / 255.)) assert num_pix_diff <= num_pix_diff_tol
def test_optimize_vertex_position(self, vertices, faces, vertex_colors, vertices_image, camera_rot, camera_trans, camera_proj, height, width, dtype, device): batch_size = faces.shape[0] # face_vertex_colors camera_rot = camera_rot.to(device, dtype) camera_rot.requires_grad = False camera_trans = camera_trans.to(device, dtype) camera_trans.requires_grad = False camera_proj = camera_proj.to(device, dtype) camera_proj.requires_grad = False face_attributes = index_vertices_by_faces( vertex_colors.to(device, dtype), faces) vertices = vertices.to(device, dtype).clone().detach() vertices.requires_grad = False moved_vertices = vertices.to(device, dtype).clone() moved_vertices[:, 0, :2] += 0.4 moved_vertices = moved_vertices.detach() moved_vertices.requires_grad = True images_gt = [ torch.from_numpy( np.array( Image.open( os.path.join(SAMPLE_DIR, f'vertex_color_{bs}.png')))) for bs in range(batch_size) ] images_gt = torch.stack(images_gt, dim=0).to(device, dtype) / 255. with torch.no_grad(): moved_vertices_camera = rotate_translate_points( moved_vertices, camera_rot, camera_trans) moved_vertices_image = perspective_camera(moved_vertices_camera, camera_proj) # test that the vertex are far enough to fail the test. assert not torch.allclose( moved_vertices_image, vertices_image, atol=1e-2, rtol=1e-2) with torch.no_grad(): face_moved_vertices_camera, face_moved_vertices_image, face_moved_normals = \ prepare_vertices(moved_vertices, faces, camera_proj, camera_rot, camera_trans) face_moved_normals_z = face_moved_normals[:, :, 2] imfeat, _, _ = dibr_rasterization( height, width, face_moved_vertices_camera[:, :, :, 2], face_moved_vertices_image, face_attributes, face_moved_normals_z) original_loss = torch.mean(torch.abs(imfeat - images_gt)) # test that the loss is high enough assert original_loss > 0.01 optimizer = torch.optim.Adam([moved_vertices], lr=5e-3) for i in range(100): optimizer.zero_grad() face_moved_vertices_camera, face_moved_vertices_image, face_moved_normals = \ prepare_vertices(moved_vertices, faces, camera_proj, camera_rot, camera_trans) face_moved_normals_z = face_moved_normals[:, :, 2] imfeat, _, _ = dibr_rasterization( height, width, face_moved_vertices_camera[:, :, :, 2], face_moved_vertices_image, face_attributes, face_moved_normals_z) loss = torch.mean(torch.abs(imfeat - images_gt)) loss.backward() optimizer.step() moved_vertices_camera = rotate_translate_points( moved_vertices, camera_rot, camera_trans) moved_vertices_image = perspective_camera(moved_vertices_camera, camera_proj) # test that the loss went down assert loss < 0.001 # We only test on image plan since we don't change camera angle during training we don't expect depth to be correct. # We could probably fine-tune the test to have a lower tolerance (TODO: cfujitsang) assert torch.allclose(moved_vertices_image, vertices_image, atol=1e-2, rtol=1e-2)