Spaces:
Running
Running
| import torch | |
| import math | |
| def euler_to_rotation_matrix(euler): | |
| # http://euclideanspace.com/maths/geometry/rotations/conversions/eulerToMatrix/index.htm | |
| # Change coordinates system | |
| euler[:, 0] = -(euler[:, 0]-90) | |
| euler[:, 1] = -euler[:, 1] | |
| euler[:, 2] = -(euler[:, 2]+90) | |
| # Convert to radians | |
| rad = euler*(math.pi/180.0) | |
| cy = torch.cos(rad[:, 0]) | |
| sy = torch.sin(rad[:, 0]) | |
| cp = torch.cos(rad[:, 1]) | |
| sp = torch.sin(rad[:, 1]) | |
| cr = torch.cos(rad[:, 2]) | |
| sr = torch.sin(rad[:, 2]) | |
| # Init R matrix tensors | |
| working_device = None | |
| if euler.is_cuda: | |
| working_device = euler.device | |
| Ry = torch.zeros((euler.shape[0], 3, 3), device=working_device) | |
| Rp = torch.zeros((euler.shape[0], 3, 3), device=working_device) | |
| Rr = torch.zeros((euler.shape[0], 3, 3), device=working_device) | |
| # Yaw | |
| Ry[:, 0, 0] = cy | |
| Ry[:, 0, 2] = sy | |
| Ry[:, 1, 1] = 1. | |
| Ry[:, 2, 0] = -sy | |
| Ry[:, 2, 2] = cy | |
| # Pitch | |
| Rp[:, 0, 0] = cp | |
| Rp[:, 0, 1] = -sp | |
| Rp[:, 1, 0] = sp | |
| Rp[:, 1, 1] = cp | |
| Rp[:, 2, 2] = 1. | |
| # Roll | |
| Rr[:, 0, 0] = 1. | |
| Rr[:, 1, 1] = cr | |
| Rr[:, 1, 2] = -sr | |
| Rr[:, 2, 1] = sr | |
| Rr[:, 2, 2] = cr | |
| return torch.matmul(torch.matmul(Ry, Rp), Rr) | |
| def projectPoints(pts, rot, trl, cam_matrix): | |
| # Get working device | |
| working_device = None | |
| if pts.is_cuda: | |
| working_device = pts.device | |
| # Perspective projection model | |
| trl = trl.unsqueeze(2) | |
| extrinsics = torch.cat((rot, trl), 2) | |
| proj_matrix = torch.matmul(cam_matrix, extrinsics) | |
| # Homogeneous landmarks | |
| ones = torch.ones(pts.shape[:2], device=working_device, requires_grad=trl.requires_grad) | |
| ones = ones.unsqueeze(2) | |
| pts_hom = torch.cat((pts, ones), 2) | |
| # Project landmarks | |
| pts_proj = pts_hom.permute((0, 2, 1)) # Transpose | |
| pts_proj = torch.matmul(proj_matrix, pts_proj) | |
| pts_proj = pts_proj.permute((0, 2, 1)) | |
| pts_proj = pts_proj/pts_proj[:, :, 2].unsqueeze(2) # Lambda = 1 | |
| return pts_proj[:, :, :-1] | |