def vertices_image(self, vertices_camera, camera_proj): return perspective_camera(vertices_camera, camera_proj)
def vertices_images_zoom(self, vertices_camera, camera_proj_zoom): # the two faces are fully covering the camera return perspective_camera(vertices_camera, camera_proj_zoom)
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)