提交测试

examples/recipes/camera/cameras_differentiable.py

@@ -0,0 +1,65 @@
+# ====================================================================================================================
+# The following snippet demonstrates how cameras can be used for optimizing specific extrinsic / intrinsic parameters
+# ====================================================================================================================
+
+import torch
+import torch.optim as optim
+from kaolin.render.camera import Camera
+
+# Create simple perspective camera
+cam = Camera.from_args(
+    eye=torch.tensor([4.0, 4.0, 4.0]),
+    at=torch.tensor([0.0, 0.0, 0.0]),
+    up=torch.tensor([0.0, 1.0, 0.0]),
+    width=800, height=600, focal_x=300.0
+)
+
+# When requires_grad is on, the camera will automatically switch to differentiation friendly backend
+# (implicitly calling cam.switch_backend('matrix_6dof_rotation') )
+cam.requires_grad_(True)
+
+# Constraint camera to optimize only fov and camera position (cannot rotate)
+ext_mask, int_mask = cam.gradient_mask('t', 'focal_x', 'focal_y')
+ext_params, int_params = cam.parameters()
+ext_params.register_hook(lambda grad: grad * ext_mask.float())
+grad_scale = 1e5    # Used to move the projection matrix elements faster
+int_params.register_hook(lambda grad: grad * int_mask.float() * grad_scale)
+
+# Make the camera a bit noisy
+# Currently can't copy the camera here after requires_grad is true because we're still missing a camera.detach() op
+target = Camera.from_args(
+    eye=torch.tensor([4.0, 4.0, 4.0]),
+    at=torch.tensor([0.0, 0.0, 0.0]),
+    up=torch.tensor([0.0, 1.0, 0.0]),
+    width=800, height=600, focal_x=300.0
+)
+target.t = target.t + torch.randn_like(target.t)
+target.focal_x = target.focal_x + torch.randn_like(target.focal_x)
+target.focal_y = target.focal_y + torch.randn_like(target.focal_y)
+target_mat = target.view_projection_matrix()
+
+# Save for later so we have some comparison of what changed
+initial_view = cam.view_matrix().detach().clone()
+initial_proj = cam.projection_matrix().detach().clone()
+
+# Train a few steps
+optimizer = optim.SGD(cam.parameters(), lr=0.1, momentum=0.9)
+for idx in range(10):
+    view_proj = cam.view_projection_matrix()
+    optimizer.zero_grad()
+    loss = torch.nn.functional.mse_loss(target_mat, view_proj)
+    loss.backward()
+    optimizer.step()
+    print(f'Iteration {idx}:')
+    print(f'Loss: {loss.item()}')
+    print(f'Extrinsics: {cam.extrinsics.parameters()}')
+    print(f'Intrinsics: {cam.intrinsics.parameters()}')
+
+# Projection matrix grads are much smaller as they're scaled by the view-frustum dimensions..
+print(f'View matrix before: {initial_view}')
+print(f'View matrix after: {cam.view_matrix()}')
+print(f'Projection matrix before: {initial_proj}')
+print(f'Projection matrix after: {cam.projection_matrix()}')
+
+print('Did the camera change?')
+print(not torch.allclose(cam, target))