def matrix_from_q_position(q_position):
matrix = torch.zeros(q_position.shape[0], 4, 4, device=q_position.device)
matrix[:, :3, :3] = kornia.quaternion_to_rotation_matrix(
torch.nn.functional.normalize(q_position[:, 3:]))
matrix[:, :3, 3] = q_position[:, :3]
matrix[:, 3, 3] = 1
return matrix
def test_triplet_qma_xyzw(self, axis, device, dtype, atol, rtol):
array = [[0.0, 0.0, 0.0, 0.0]]
array[0][axis] = 1.0 # `0 + axis` should fail when WXYZ
quaternion = torch.tensor(array, device=device, dtype=dtype)
assert quaternion.shape[-1] == 4
with pytest.warns(UserWarning):
mm = kornia.quaternion_to_rotation_matrix(
quaternion, order=QuaternionCoeffOrder.XYZW)
assert mm.shape[-1] == 3
assert mm.shape[-2] == 3
angle_axis = kornia.rotation_matrix_to_angle_axis(mm)
assert angle_axis.shape[-1] == 3
angle_axis_expected = [[0.0, 0.0, 0.0]]
angle_axis_expected[0][axis] = kornia.pi
angle_axis_expected = torch.tensor(angle_axis_expected,
device=device,
dtype=dtype)
assert_allclose(angle_axis, angle_axis_expected, atol=atol, rtol=rtol)
with pytest.warns(UserWarning):
quaternion_hat = kornia.angle_axis_to_quaternion(
angle_axis, order=QuaternionCoeffOrder.XYZW)
assert_allclose(quaternion_hat, quaternion, atol=atol, rtol=rtol)
def test_jit(self, device):
@torch.jit.script
def op_script(input):
return kornia.quaternion_to_rotation_matrix(input)
quaternion = torch.tensor([0., 0., 1., 0.]).to(device)
actual = op_script(quaternion)
expected = kornia.quaternion_to_rotation_matrix(quaternion)
assert_allclose(actual, expected)
def test_z_rotation(self, device):
quaternion = torch.tensor([0., 0., 1., 0.]).to(device)
expected = torch.tensor([
[-1., 0., 0.],
[0., -1., 0.],
[0., 0., 1.],
]).to(device)
matrix = kornia.quaternion_to_rotation_matrix(quaternion)
assert_allclose(matrix, expected)
def test_back_and_forth(self, device):
matrix = torch.tensor([
[1., 0., 0.],
[0., 0., -1.],
[0., 1., 0.],
]).to(device)
quaternion = kornia.rotation_matrix_to_quaternion(matrix)
matrix_hat = kornia.quaternion_to_rotation_matrix(quaternion)
assert_allclose(matrix, matrix_hat)
def test_z_rotation(self, device, dtype):
quaternion = torch.tensor([0., 0., 1., 0.], device=device, dtype=dtype)
expected = torch.tensor([
[-1., 0., 0.],
[0., -1., 0.],
[0., 0., 1.],
], device=device, dtype=dtype)
matrix = kornia.quaternion_to_rotation_matrix(quaternion)
assert_allclose(matrix, expected, atol=1e-4, rtol=1e-4)
def test_x_rotation(self):
quaternion = torch.tensor([1., 0., 0., 0.])
expected = torch.tensor([
[1., 0., 0.],
[0., -1., 0.],
[0., 0., -1.],
])
matrix = kornia.quaternion_to_rotation_matrix(quaternion)
assert_allclose(matrix, expected)
def pose_reg(quat, tran, pose_px_ind, ind, gt_p03d, gt_p13d, gt_depth, max_obj, p03d_feat,img):
# solve the scale
alpha = torch.ones(quat.shape[0]).cuda()
for i in range(quat.shape[0]):
d1 = p03d_feat[i,-1]
d2 = gt_p03d[i,-1].view(-1)
alpha[i] = (d1*d2).sum()/(d1*d1).sum()
#pdb.set_trace()
#from utils.fusion import pcwrite
#pc1 = np.asarray(p03d_feat[0].T.cpu())
#pc2 = np.asarray(gt_p03d[0].view(3,-1).T.cpu())
#pc1 = pc1*np.asarray(alpha[i].cpu())
#pcwrite('/data/gengshay/0.ply',np.concatenate([pc1,pc1],-1))
#pcwrite('/data/gengshay/1.ply',np.concatenate([pc2,pc2],-1))
alpha = alpha.detach()
vis = torch.zeros_like(gt_depth)
quat = _transpose_and_gather_feat(quat, ind).view(-1,4)
tran = _transpose_and_gather_feat(tran, ind).view(-1,3)
gt_p03d = gt_p03d.permute(0,2,3,1)
gt_p13d = gt_p13d.permute(0,2,3,1)
gt_depth = gt_depth.permute(0,2,3,1)
loss = []
for it,obj_mask in enumerate(pose_px_ind):
imgid = it//max_obj
if len(obj_mask)>0:
p03d = gt_p03d[imgid][obj_mask]
p13d = gt_p13d[imgid][obj_mask]
depth =gt_depth[imgid][obj_mask]
# compute rigid transform
quatx = torch.nn.functional.normalize(quat[it],2,-1)
quatx = kornia.quaternion_to_rotation_matrix(quatx)
pred_p13d = quatx.matmul(p03d[:,:,None])[:,:,0]+tran[it] * alpha[imgid]
#pdb.set_trace()
#from utils.fusion import pcwrite
#pc1 = np.asarray(p03d.cpu())
#pc2 = np.asarray(pred_p13d.detach().cpu())
#pc3 = np.asarray(p13d.cpu())
#rgb = img[imgid][obj_mask].cpu()*255
#pcwrite('/data/gengshay/0.ply',np.concatenate([pc1,rgb],-1))
#pcwrite('/data/gengshay/1.ply',np.concatenate([pc2,rgb],-1))
#pcwrite('/data/gengshay/2.ply',np.concatenate([pc3,rgb],-1))
sub_loss = ((p13d - pred_p13d)/depth).abs()
loss.append( sub_loss.mean() )
# vis
sub_vis = torch.zeros_like(vis[0,0])
sub_vis[obj_mask] = sub_loss.mean(-1)
vis[imgid,0] += sub_vis
if len(loss)>0:
loss = torch.stack(loss).mean()
else:
loss = 0
return loss, vis
def rigid_transform(p03d,p13d,quat, tran,mask):
mask = torch.Tensor(mask).cuda()
for it in range(mask.max().int()):
obj_mask = mask==(it+1)
# compute rigid transform
quatx = torch.nn.functional.normalize(quat[it],2,-1)
quatx = kornia.quaternion_to_rotation_matrix(quatx)
p13d[obj_mask] = quatx.matmul(p03d[obj_mask][:,:,None])[:,:,0]+tran[it]
return p03d,p13d
def matrix_from_logq_position(logq_position):
quaternion = kornia.quaternion_log_to_exp(logq_position[:, 3:])
matrix = torch.zeros(logq_position.shape[0],
4,
4,
device=logq_position.device)
matrix[:, :3, :3] = kornia.quaternion_to_rotation_matrix(quaternion)
matrix[:, :3, 3] = logq_position[:, :3]
matrix[:, 3, 3] = 1
return matrix
def _build_extrinsic(extrinsic_matrix_params: dict) -> torch.Tensor:
quaternion = torch.tensor(extrinsic_matrix_params['orientation'])[[1, 2, 3, 0]]
quaternion[:3] *= -1
R = quaternion_to_rotation_matrix(quaternion)
t = torch.tensor(extrinsic_matrix_params['translation'])
tr = -torch.matmul(R, t)
return torch.cat([torch.cat([R, tr.unsqueeze(1)], dim=1)], dim=0,)
def ctdet_decode(heat, wh, reg=None, cat_spec_wh=False, K=100,quat=None,tran =None,p03d=None):
batch, cat, height, width = heat.size()
# heat = torch.sigmoid(heat)
# perform nms on heatmaps
heat = _nms(heat)
scores, inds, clses, ys, xs = _topk(heat, K=K)
if reg is not None:
reg = _transpose_and_gather_feat(reg, inds)
reg = reg.view(batch, K, 2)
xs = xs.view(batch, K, 1) + reg[:, :, 0:1]
ys = ys.view(batch, K, 1) + reg[:, :, 1:2]
else:
xs = xs.view(batch, K, 1)
ys = ys.view(batch, K, 1)
scores = scores.view(batch, K, 1)
pdist_ct = torch.cat([xs,ys],-1)
pdist_ind=(ys*width+xs).long()
pdist_pred = _transpose_and_gather_feat(wh, pdist_ind[:,:,0])
if quat is not None:
quat_pred = _transpose_and_gather_feat(quat, pdist_ind[:,:,0])
tran_pred = _transpose_and_gather_feat(tran, pdist_ind[:,:,0])
pdist_mask = (scores>0.1)[:,:,0]
contour_pred = np.zeros(wh.shape[2:])
mask_pred = np.zeros(wh.shape[2:])
angles = torch.range(0, 350, 10).cuda() / 180 * math.pi
bboxs = np.zeros((0,4))
p03d = p03d[0].permute(1,2,0)
p13d = p03d.clone()
if pdist_mask.sum()>0:
contour = distance2mask(pdist_ct[0][pdist_mask[0]], pdist_pred[0][pdist_mask[0]], angles, wh.shape[2:])
contour = np.asarray(contour.permute(0,2,1).cpu()[:,:,None],dtype=int)
contour_pred = cv2.drawContours(contour_pred, contour, -1,1,3)
mask_pred,bboxs = draw_masks(mask_pred, np.asarray(pdist_ct[0][pdist_mask[0]].cpu()), contour)
#pdb.set_trace()
if quat is not None:
quat_pred = quat_pred[0][pdist_mask[0]]
tran_pred = tran_pred[0][pdist_mask[0]]
#p03d,p13d = rigid_transform(p03d,p13d,quat_pred,tran_pred, mask_pred)
pred = np.concatenate([contour_pred, mask_pred],0)
rt = {}
rt['mask'] = pred
scores = np.asarray(scores[scores>0.1].cpu())
rt['bbox'] = np.concatenate([bboxs.reshape((-1,4)), scores[:,None]],-1)
if quat is not None:
rt['quat'] = np.asarray(kornia.quaternion_to_rotation_matrix(quat_pred).cpu())
rt['tran'] = np.asarray(tran_pred.cpu())
#rt['p03d'] = np.asarray(p03d.cpu())
#rt['p13d'] = np.asarray(p13d.cpu())
return rt
def test_triplet_aqm(self, axis, device, dtype, atol, rtol):
array = [[0.0, 0.0, 0.0]]
array[0][axis] = kornia.pi / 2.0
angle_axis = torch.tensor(array, device=device, dtype=dtype)
assert angle_axis.shape[-1] == 3
quaternion = kornia.angle_axis_to_quaternion(angle_axis, order=QuaternionCoeffOrder.WXYZ)
assert quaternion.shape[-1] == 4
rot_m = kornia.quaternion_to_rotation_matrix(quaternion, order=QuaternionCoeffOrder.WXYZ)
assert rot_m.shape[-1] == 3
assert rot_m.shape[-2] == 3
angle_axis_hat = kornia.rotation_matrix_to_angle_axis(rot_m)
assert_close(angle_axis_hat, angle_axis, atol=atol, rtol=rtol)
def loss_function(self, quat, trans, proj, lid_pts, target_pts):
R = kornia.quaternion_to_rotation_matrix(quat)
lid_pts = lid_pts.transpose(0, 1)
cam_pts = torch.matmul(R, lid_pts)
cam_pts = cam_pts.transpose(0, 1)
cam_pts += trans
pixel_index = kornia.geometry.project_points(cam_pts, proj[:3, :3])
pixel_index = pixel_index[:, torch.LongTensor(
[1, 0])] # Swap columns 1 and 0]
pixel_index = pixel_index.unsqueeze(1)
loss = torch.norm(pixel_index - target_pts, dim=2, p=2)
# loss = self.compute_dist(pixel_index,target_pts)
loss, _ = torch.min(loss, dim=1)
loss = torch.max(loss)
# criterion = SamplesLoss(loss="hausdorff",p=2,blur=1e-4)
# loss = criterion(pixel_index,target_pts)
return loss
def loss_function(self, quat, trans, proj, lid_pts, seg_mask):
R = kornia.quaternion_to_rotation_matrix(quat)
lid_pts = lid_pts.transpose(0, 1)
cam_pts = torch.matmul(R, lid_pts)
cam_pts = cam_pts.transpose(0, 1)
cam_pts += trans
pixel_index = kornia.geometry.project_points(cam_pts, proj[:3, :3])
pixel_index = pixel_index[:, torch.LongTensor(
[1, 0])] # Swap columns 1 and 0]
target_pts = torch.nonzero(seg_mask).float()
# pixel_index = pixel_index.unsqueeze(1)
# loss = torch.norm(pixel_index - target_pts, dim=2, p=2)
loss = self.compute_dist(pixel_index, target_pts)
loss, _ = torch.min(loss, dim=1)
loss = torch.max(loss)
return loss
def test_triplet_qma(self, axis, device, dtype, atol, rtol):
array = [[0.0, 0.0, 0.0, 0.0]]
array[0][1 + axis] = 1.0 # `1 + axis` this should fail when XYZW
quaternion = torch.tensor(array, device=device, dtype=dtype)
assert quaternion.shape[-1] == 4
mm = kornia.quaternion_to_rotation_matrix(quaternion, order=QuaternionCoeffOrder.WXYZ)
assert mm.shape[-1] == 3
assert mm.shape[-2] == 3
angle_axis = kornia.rotation_matrix_to_angle_axis(mm)
assert angle_axis.shape[-1] == 3
angle_axis_expected = [[0.0, 0.0, 0.0]]
angle_axis_expected[0][axis] = kornia.pi
angle_axis_expected = torch.tensor(angle_axis_expected, device=device, dtype=dtype)
assert_close(angle_axis, angle_axis_expected, atol=atol, rtol=rtol)
quaternion_hat = kornia.angle_axis_to_quaternion(angle_axis, order=QuaternionCoeffOrder.WXYZ)
assert_close(quaternion_hat, quaternion, atol=atol, rtol=rtol)
def test_triplet_aqm_xyzw(self, axis, device, dtype, atol, rtol):
array = [[0., 0., 0.]]
array[0][axis] = kornia.pi / 2.
angle_axis = torch.tensor(array, device=device, dtype=dtype)
assert angle_axis.shape[-1] == 3
with pytest.warns(UserWarning):
quaternion = kornia.angle_axis_to_quaternion(
angle_axis, order=QuaternionCoeffOrder.XYZW)
assert quaternion.shape[-1] == 4
with pytest.warns(UserWarning):
rot_m = kornia.quaternion_to_rotation_matrix(
quaternion, order=QuaternionCoeffOrder.XYZW)
assert rot_m.shape[-1] == 3
assert rot_m.shape[-2] == 3
angle_axis_hat = kornia.rotation_matrix_to_angle_axis(rot_m)
assert_allclose(angle_axis_hat, angle_axis, atol=atol, rtol=rtol)
def quaternion_to_rot_matrix(quaternion):
return kornia.quaternion_to_rotation_matrix(quaternion=quaternion)
#rotmat = cv2.Rodrigues(np.asarray([ 0.05*6.28, 0.10*6.28*(-0.5+i/nframes),0.]))[0] # y-axis temple
quat = kornia.rotation_matrix_to_quaternion(torch.Tensor(rotmat).cuda())
proj_cam = torch.zeros(1, 7).cuda()
depth = torch.zeros(1, 1).cuda()
proj_cam[:, 0] = 5 # focal=10
proj_cam[:, 1] = 0. # x translation = 0
proj_cam[:, 2] = 0. # y translation = 0
proj_cam[:, 3] = quat[3]
proj_cam[:, 4:] = quat[:3]
#depth[:,0] = 0.05 # for temple
depth[:, 0] = 1 # z translation (depth) =10 for spot
#depth[:,0] = 0.5 # z translation (depth) =10 for spot
#depth[:,0] = 200 # for kitti
# obj-cam transform
Rmat = kornia.quaternion_to_rotation_matrix(
torch.cat((-proj_cam[:, 4:], proj_cam[:, 3:4]), 1))
Tmat = torch.cat([proj_cam[:, 1:3], depth], 1)
verts = verts.matmul(Rmat) + Tmat[:, np.newaxis, :] # obj to cam transform
verts = torch.cat([verts, torch.ones_like(verts[:, :, 0:1])], dim=-1)
# pespective projection: x=fX/Z assuming px=py=0, normalization of Z
verts[:, :,
1] = verts[:, :, 1].clone() * proj_cam[:, :1] / verts[:, :,
2].clone()
verts[:, :,
0] = verts[:, :, 0].clone() * proj_cam[:, :1] / verts[:, :,
2].clone()
verts[:, :,
2] = ((verts[:, :, 2] - verts[:, :, 2].min()) /
(verts[:, :, 2].max() - verts[:, :, 2].min()) - 0.5).detach()
verts[:, :, 2] += 10
def test_smoke_batch(self, batch_size, device, dtype):
quaternion = torch.zeros(batch_size, 4, device=device, dtype=dtype)
matrix = kornia.quaternion_to_rotation_matrix(quaternion)
assert matrix.shape == (batch_size, 3, 3)
def op_script(input):
return kornia.quaternion_to_rotation_matrix(input)
