def test_quaternion_normalize(self): q = np.array([1, 0, 0, 0]) testing.assert_equal(quaternion_normalize(q), [1, 0, 0, 0]) q = np.array([1, 2, 3, 4]) testing.assert_almost_equal( quaternion_normalize(q), [0.18257419, 0.36514837, 0.54772256, 0.73029674])
def transform_coords(c1, c2, out=None): """Return Coordinates by applying c1 to c2 from the left Parameters ---------- c1 : skrobot.coordinates.Coordinates c2 : skrobot.coordinates.Coordinates Coordinates c3 : skrobot.coordinates.Coordinates or None Output argument. If this value is specified, the results will be in-placed. Returns ------- Coordinates(pos=translation, rot=q) : skrobot.coordinates.Coordinates new coordinates Examples -------- >>> from skrobot.coordinates import Coordinates >>> from skrobot.coordinates import transform_coords >>> from numpy import pi >>> c1 = Coordinates() >>> c2 = Coordinates() >>> c3 = transform_coords(c1, c2) >>> c3.translation array([0., 0., 0.]) >>> c3.rotation array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]) >>> c1 = Coordinates().translate([0.1, 0.2, 0.3]).rotate(pi / 3.0, 'x') >>> c2 = Coordinates().translate([0.3, -0.3, 0.1]).rotate(pi / 2.0, 'y') >>> c3 = transform_coords(c1, c2) >>> c3.translation array([ 0.4 , -0.03660254, 0.09019238]) >>> c3.rotation >>> c3.rotation array([[ 1.94289029e-16, 0.00000000e+00, 1.00000000e+00], [ 8.66025404e-01, 5.00000000e-01, -1.66533454e-16], [-5.00000000e-01, 8.66025404e-01, 2.77555756e-17]]) """ if out is None: out = Coordinates() elif not isinstance(out, Coordinates): raise TypeError("Input type should be skrobot.coordinates.Coordinates") out.translation = c1.translation + np.dot(c1.rotation, c2.translation) out.rotation = quaternion_normalize( quaternion_multiply(c1.quaternion, c2.quaternion)) return out
def __init__(self, qr=[1, 0, 0, 0], qd=[0, 0, 0, 0], enforce_unit_norm=False): if (isinstance(qd, list) or isinstance(qd, np.ndarray)) and \ len(qd) == 3: x, y, z = qd qr = quaternion_normalize(qr) qd = 0.5 * quaternion_multiply([0, x, y, z], qr) self.qr = qr self.qd = qd if enforce_unit_norm: norm = self.norm if not np.allclose(norm[0], [1]): raise ValueError("Dual quaternoin's norm " 'should be 1, but gives {}'.format(norm[0]))
def predict(self, rgb_img, depth, label_img, label, bboxes, intrinsic_matrix): num_points = self.num_points iteration = self.iteration batch_size = 1 lst = np.array(label.flatten(), dtype=np.int32) img_height, img_width, _ = rgb_img.shape cameramodel = cameramodels.PinholeCameraModel.from_intrinsic_matrix( intrinsic_matrix, img_height, img_width) translations = [] rotations = [] for idx in range(len(lst)): itemid = lst[idx] rmin, cmin, rmax, cmax = bboxes[idx] mask_depth = ma.getmaskarray(ma.masked_not_equal(depth, 0)) mask_label = ma.getmaskarray(ma.masked_equal(label_img, itemid)) mask = mask_label * mask_depth choose = mask[rmin:rmax, cmin:cmax].flatten().nonzero()[0] if len(choose) == 0: translations.append([]) rotations.append([]) continue if len(choose) > num_points: c_mask = np.zeros(len(choose), dtype=int) c_mask[:num_points] = 1 np.random.shuffle(c_mask) choose = choose[c_mask.nonzero()] else: choose = np.pad(choose, (0, num_points - len(choose)), 'wrap') xmap = np.array([[j for i in range(img_width)] for j in range(img_height)]) ymap = np.array([[i for i in range(img_width)] for j in range(img_height)]) depth_masked = depth[ rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype(np.float32) xmap_masked = xmap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype( np.float32) ymap_masked = ymap[rmin:rmax, cmin:cmax].flatten()[choose][:, np.newaxis].astype( np.float32) choose = np.array([choose]) cloud = cameramodel.batch_project_pixel_to_3d_ray( np.concatenate([ymap_masked, xmap_masked], axis=1), depth_masked) img_masked = np.array(rgb_img)[:, :, :3] img_masked = np.transpose(img_masked, (2, 0, 1)) img_masked = img_masked[:, rmin:rmax, cmin:cmax] with torch.no_grad(): cloud = torch.from_numpy(cloud.astype(np.float32)) choose = torch.LongTensor(choose.astype(np.int32)) img_masked = normalize( torch.from_numpy(img_masked.astype(np.float32))) index = torch.LongTensor([itemid - 1]) cloud = cloud.cuda() choose = choose.cuda() img_masked = img_masked.cuda() index = index.cuda() cloud = cloud.view(1, num_points, 3) img_masked = img_masked.view(1, 3, img_masked.size()[1], img_masked.size()[2]) pred_rot, pred_trans, pred_score, emb = self.estimator( img_masked, cloud, choose, index) pred_rot = pred_rot / torch.norm(pred_rot, dim=2).view( 1, num_points, 1) pred_score = pred_score.view(batch_size, num_points) _, which_max = torch.max(pred_score, 1) pred_trans = pred_trans.view(batch_size * num_points, 1, 3) points = cloud.view(batch_size * num_points, 1, 3) rotation = pred_rot[0][which_max[0]].view(-1).cpu().\ data.numpy() translation = (points + pred_trans)[which_max[0]].view(-1).\ cpu().data.numpy() for _ in range(iteration): T = torch.from_numpy(translation.astype(np.float32)).\ cuda().view(1, 3).\ repeat(num_points, 1).contiguous().\ view(1, num_points, 3) trans_matrix = np.eye(4) trans_matrix[:3, :3] = quaternion2matrix( quaternion_normalize(rotation)) R = torch.from_numpy(trans_matrix[:3, :3].astype( np.float32)).cuda().view(1, 3, 3) trans_matrix[0:3, 3] = translation new_cloud = torch.bmm((cloud - T), R).contiguous() refined_rot, refined_trans = self.refiner( new_cloud, emb, index) refined_rot = refined_rot.view(1, 1, -1) refined_rot = refined_rot / (torch.norm( refined_rot, dim=2).view(1, 1, 1)) rotation_2 = refined_rot.view(-1).cpu().data.numpy() translation_2 = refined_trans.view(-1).cpu().data.numpy() trans_matrix_2 = np.eye(4) trans_matrix_2[:3, :3] = quaternion2matrix( quaternion_normalize(rotation_2)) trans_matrix_2[0:3, 3] = translation_2 trans_matrix_final = np.dot(trans_matrix, trans_matrix_2) rotation_final = matrix2quaternion( trans_matrix_final[:3, :3]) translation_final = np.array([ trans_matrix_final[0][3], trans_matrix_final[1][3], trans_matrix_final[2][3] ]) rotation = rotation_final translation = translation_final translations.append(translation) rotations.append(quaternion_normalize(rotation)) return rotations, translations