def test_quaternion2matrix(self): testing.assert_array_equal(quaternion2matrix([1, 0, 0, 0]), np.eye(3)) testing.assert_almost_equal( quaternion2matrix([1.0 / np.sqrt(2), 1.0 / np.sqrt(2), 0, 0]), np.array([[1., 0., 0.], [0., 0., -1.], [0., 1., 0.]])) testing.assert_almost_equal( quaternion2matrix( normalize_vector( [1.0, 1 / np.sqrt(2), 1 / np.sqrt(2), 1 / np.sqrt(2)])), np.array([[0.2000000, -0.1656854, 0.9656854], [0.9656854, 0.2000000, -0.1656854], [-0.1656854, 0.9656854, 0.2000000]])) testing.assert_almost_equal( quaternion2matrix( normalize_vector( [1.0, -1 / np.sqrt(2), 1 / np.sqrt(2), -1 / np.sqrt(2)])), np.array([[0.2000000, 0.1656854, 0.9656854], [-0.9656854, 0.2000000, 0.1656854], [-0.1656854, -0.9656854, 0.2000000]])) testing.assert_almost_equal( quaternion2matrix([0.925754, 0.151891, 0.159933, 0.307131]), rotate_matrix( rotate_matrix(rotate_matrix(np.eye(3), 0.2, 'x'), 0.4, 'y'), 0.6, 'z'), decimal=5)
def rotation(self, rotation): """Set rotation of this coordinate This setter checkes the given rotation and set it this coordinate. Parameters ---------- rotation : list or numpy.ndarray we can take 3x3 rotation matrix or rpy angle [yaw, pitch, roll] or quaternion [w, x, y, z] order """ rotation = np.array(rotation) # Convert quaternions if rotation.shape == (4, ): self._q = np.array([q for q in rotation]) if np.abs(np.linalg.norm(self._q) - 1.0) > 1e-3: raise ValueError('Invalid quaternion. Must be ' 'norm 1.0, get {}'.format( np.linalg.norm(self._q))) rotation = quaternion2matrix(self._q) elif rotation.shape == (3, ): # Convert [yaw-pitch-roll] to rotation matrix self._q = rpy2quaternion(rotation) rotation = quaternion2matrix(self._q) else: self._q = matrix2quaternion(rotation) # Convert lists and tuples if type(rotation) in (list, tuple): rotation = np.array(rotation).astype(np.float32) _check_valid_rotation(rotation) self._rotation = rotation * 1.
def test_rotation(self): qr = normalize_vector(np.array([1, 2, 3, 4])) x, y, z = np.array([1, 2, 3]) qd = 0.5 * quaternion_multiply(np.array([0, x, y, z]), qr) dq = DualQuaternion(qr, qd) testing.assert_almost_equal(dq.rotation, quaternion2matrix(qr))
def test_matrix2quaternion(self): testing.assert_almost_equal(matrix2quaternion(np.eye(3)), np.array([1, 0, 0, 0])) m = rotate_matrix( rotate_matrix(rotate_matrix(np.eye(3), 0.2, 'x'), 0.4, 'y'), 0.6, 'z') testing.assert_almost_equal(quaternion2matrix(matrix2quaternion(m)), m) testing.assert_almost_equal(matrix2quaternion( np.array([[0.428571, 0.514286, -0.742857], [-0.857143, -0.028571, -0.514286], [-0.285714, 0.857143, 0.428571]])), normalize_vector(np.array([4, 3, -1, -3])), decimal=5)
def quaternion2vec(q, axis='x'): """Calculate axis vector from quaternion Parameters ---------- q : np.ndarray qauternion axis : str, optional axis x, y, z, by default 'x' Returns ------- vec : numpy.ndarray """ m = quaternion2matrix(q) return matrix2vec(m, axis=axis)
def rotation(self): """Return rotation matrix. Note that this property internally normalizes quaternion. Returns ------- quaternion2matrix(self.q) : numpy.ndarray 3x3 rotation matrix Examples -------- >>> from skrobot.coordinates.quaternion import Quaternion >>> q = Quaternion() >>> q #<Quaternion 0x12f1aa6a0 w: 1.0 x: 0.0 y: 0.0 z: 0.0> >>> q.rotation array([[1., 0., 0.], [0., 1., 0.], [0., 0., 1.]]) >>> q.q = [0, 1, 0, 0] >>> q #<Quaternion 0x12f1aa6a0 w: 0.0 x: 1.0 y: 0.0 z: 0.0> >>> q.rotation array([[ 1., 0., 0.], [ 0., -1., 0.], [ 0., 0., -1.]]) >>> q.q = [1, 2, 3, 4] >>> q #<Quaternion 0x12f1aa6a0 w: 1.0 x: 2.0 y: 3.0 z: 4.0> >>> q.rotation array([[-0.66666667, 0.13333333, 0.73333333], [ 0.66666667, -0.33333333, 0.66666667], [ 0.33333333, 0.93333333, 0.13333333]]) """ return quaternion2matrix(self.normalized.q)
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
def callback(self, camera_info_msg, img_raw_msg, img_msg, depth_msg): ymin = camera_info_msg.roi.y_offset xmin = camera_info_msg.roi.x_offset ymax = camera_info_msg.roi.y_offset + camera_info_msg.roi.height xmax = camera_info_msg.roi.x_offset + camera_info_msg.roi.width self.camera_model.roi = [ymin, xmin, ymax, xmax] self.camera_model.target_size = self.target_size bgr_raw = self.bridge.imgmsg_to_cv2(img_raw_msg, "bgr8") bgr = self.bridge.imgmsg_to_cv2(img_msg, "bgr8") cv_depth = self.bridge.imgmsg_to_cv2(depth_msg, "32FC1") if cv_depth is None or bgr is None: return remove_nan(cv_depth) cv_depth[cv_depth < self.depth_range[0]] = 0 cv_depth[cv_depth > self.depth_range[1]] = 0 bgr = cv2.resize(bgr, self.target_size) cv_depth = cv2.resize(cv_depth, self.target_size, interpolation=cv2.INTER_NEAREST) depth_bgr = colorize_depth(cv_depth, ignore_value=0) in_feature = cv_depth.copy().astype(np.float32) * 0.001 if self.config['use_bgr2gray']: gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY) gray = cv2.resize(gray, self.target_size)[..., None] / 255. normalized_depth = normalize_depth(cv_depth, self.depth_range[0], self.depth_range[1])[..., None] in_feature = np.concatenate((normalized_depth, gray), axis=2).astype(np.float32) if self.transform: in_feature = self.transform(in_feature) in_feature = in_feature.to(self.device) in_feature = in_feature.unsqueeze(0) confidence, depth, rotation = self.model(in_feature) confidence = confidence[0, 0:1, ...] confidence_np = confidence.cpu().detach().numpy().copy() * 255 confidence_np = confidence_np.transpose(1, 2, 0) confidence_np[confidence_np <= 0] = 0 confidence_np[confidence_np >= 255] = 255 confidence_img = confidence_np.astype(np.uint8) confidence_img = cv2.resize(confidence_img, self.target_size) heatmap = overlay_heatmap(bgr, confidence_img) axis_pred = bgr.copy() axis_pred_raw = bgr_raw.copy() hanging_points_pose_array = PoseArray() for i, (roi, roi_center) in enumerate( zip(self.model.rois_list[0], self.model.rois_center_list[0])): if roi.tolist() == [0, 0, 0, 0]: continue roi = roi.cpu().detach().numpy().copy() hanging_point_x = roi_center[0] hanging_point_y = roi_center[1] v = rotation[i].cpu().detach().numpy() v /= np.linalg.norm(v) rot = rotation_matrix_from_axis(v, [0, 1, 0], 'xy') q = matrix2quaternion(rot) hanging_point = np.array( self.camera_model.project_pixel_to_3d_ray( [int(hanging_point_x), int(hanging_point_y)])) if self.predict_depth: dep = depth[i].cpu().detach().numpy().copy() dep = unnormalize_depth(dep, self.depth_range[0], self.depth_range[1]) * 0.001 length = float(dep) / hanging_point[2] else: depth_roi = make_box(roi_center, width=self.depth_roi_size[1], height=self.depth_roi_size[0], img_shape=self.target_size, xywh=False) depth_roi_clip = cv_depth[depth_roi[0]:depth_roi[2], depth_roi[1]:depth_roi[3]] dep_roi_clip = depth_roi_clip[np.where( np.logical_and(self.depth_range[0] < depth_roi_clip, depth_roi_clip < self.depth_range[1]))] dep_roi_clip = np.median(dep_roi_clip) * 0.001 if dep_roi_clip == np.nan: continue length = float(dep_roi_clip) / hanging_point[2] hanging_point *= length hanging_point_pose = Pose() hanging_point_pose.position.x = hanging_point[0] hanging_point_pose.position.y = hanging_point[1] hanging_point_pose.position.z = hanging_point[2] hanging_point_pose.orientation.w = q[0] hanging_point_pose.orientation.x = q[1] hanging_point_pose.orientation.y = q[2] hanging_point_pose.orientation.z = q[3] hanging_points_pose_array.poses.append(hanging_point_pose) axis_pred_raw = cv2.rectangle( axis_pred_raw, (int(roi[0] * (xmax - xmin) / self.target_size[1] + xmin), int(roi[1] * (ymax - ymin) / self.target_size[0] + ymin)), (int(roi[2] * (xmax - xmin) / self.target_size[1] + xmin), int(roi[3] * (ymax - ymin) / self.target_size[0] + ymin)), (0, 255, 0), 1) try: axis_pred_raw = draw_axis(axis_pred_raw, quaternion2matrix(q), hanging_point, self.camera_model.full_K) except Exception: print('Fail to draw axis') axis_pred = self.camera_model.crop_image(axis_pred_raw, copy=True).astype(np.uint8) msg_out = self.bridge.cv2_to_imgmsg(heatmap, "bgr8") msg_out.header.stamp = depth_msg.header.stamp confidence_msg = self.bridge.cv2_to_imgmsg(confidence_img, "mono8") confidence_msg.header.stamp = depth_msg.header.stamp colorized_depth_msg = self.bridge.cv2_to_imgmsg(depth_bgr, "bgr8") colorized_depth_msg.header.stamp = depth_msg.header.stamp axis_pred_msg = self.bridge.cv2_to_imgmsg(axis_pred, "bgr8") axis_pred_msg.header.stamp = depth_msg.header.stamp axis_pred_raw_msg = self.bridge.cv2_to_imgmsg(axis_pred_raw, "bgr8") axis_pred_raw_msg.header.stamp = depth_msg.header.stamp hanging_points_pose_array.header = camera_info_msg.header self.pub.publish(msg_out) self.pub_confidence.publish(confidence_msg) self.pub_depth.publish(colorized_depth_msg) self.pub_axis.publish(axis_pred_msg) self.pub_axis_raw.publish(axis_pred_raw_msg) self.pub_hanging_points.publish(hanging_points_pose_array)
def step(self, dataloader, mode): print('Start {}'.format(mode)) # self.model = self.prev_model if mode == 'train': self.model.train() elif mode == 'val' or mode == 'test': self.model.eval() loss_sum = 0 confidence_loss_sum = 0 depth_loss_sum = 0 rotation_loss_sum = 0 rotation_loss_count = 0 for index, (hp_data, depth_image, camera_info_path, hp_data_gt, annotation_data) in tqdm.tqdm(enumerate(dataloader), total=len(dataloader), desc='{} epoch={}'.format( mode, self.epo), leave=False): # if index == 0: # self.model = self.prev_model self.cameramodel\ = cameramodels.PinholeCameraModel.from_yaml_file( camera_info_path[0]) self.cameramodel.target_size = self.target_size depth_image = hp_data.numpy().copy()[0, 0, ...] depth_image = np.nan_to_num(depth_image) depth_image = unnormalize_depth(depth_image, self.depth_range[0], self.depth_range[1]) hp_data = hp_data.to(self.device) depth_image_bgr = colorize_depth(depth_image, ignore_value=self.depth_range[0]) if mode == 'train': confidence, depth, rotation = self.model(hp_data) elif mode == 'val' or mode == 'test': with torch.no_grad(): confidence, depth, rotation = self.model(hp_data) confidence_np = confidence[0, ...].cpu().detach().numpy().copy() confidence_np[confidence_np >= 1] = 1. confidence_np[confidence_np <= 0] = 0. confidence_vis = cv2.cvtColor(confidence_np[0, ...] * 255, cv2.COLOR_GRAY2BGR) if mode != 'test': pos_weight = hp_data_gt.detach().numpy().copy() pos_weight = pos_weight[:, 0, ...] zeroidx = np.where(pos_weight < 0.5) nonzeroidx = np.where(pos_weight >= 0.5) pos_weight[zeroidx] = 0.5 pos_weight[nonzeroidx] = 1.0 pos_weight = torch.from_numpy(pos_weight) pos_weight = pos_weight.to(self.device) hp_data_gt = hp_data_gt.to(self.device) confidence_gt = hp_data_gt[:, 0:1, ...] rois_list_gt, rois_center_list_gt = find_rois(confidence_gt) criterion = HPNETLoss(self.use_coords).to(self.device) if self.model.rois_list is None or rois_list_gt is None: return None, None annotated_rois = annotate_rois(self.model.rois_list, rois_list_gt, annotation_data) confidence_loss, depth_loss, rotation_loss = criterion( confidence, hp_data_gt, pos_weight, depth, rotation, annotated_rois) if self.train_depth: loss = confidence_loss + rotation_loss + depth_loss else: loss = confidence_loss + rotation_loss if torch.isnan(loss): print('loss is nan!!') self.model = self.prev_model self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr, betas=(0.9, 0.999), eps=1e-10, weight_decay=0, amsgrad=False) self.optimizer.load_state_dict( self.prev_optimizer.state_dict()) continue else: self.prev_model = copy.deepcopy(self.model) self.prev_optimizer = copy.deepcopy(self.optimizer) if mode == 'train': self.optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(self.model.parameters(), 5) self.optimizer.step() axis_gt = depth_image_bgr.copy() confidence_gt_vis = cv2.cvtColor( confidence_gt[0, 0, ...].cpu().detach().numpy().copy() * 255, cv2.COLOR_GRAY2BGR) # Visualize gt axis and roi for roi, roi_c in zip(rois_list_gt[0], rois_center_list_gt[0]): if roi.tolist() == [0, 0, 0, 0]: continue roi = roi.cpu().detach().numpy().copy() cx = roi_c[0] cy = roi_c[1] depth_and_rotation_gt = get_value_gt([cx, cy], annotation_data[0]) rotation_gt = depth_and_rotation_gt[1:] depth_gt_val = depth_and_rotation_gt[0] unnormalized_depth_gt_val = unnormalize_depth( depth_gt_val, self.depth_range[0], self.depth_range[1]) hanging_point_pose = np.array( self.cameramodel.project_pixel_to_3d_ray( [int(cx), int(cy)])) \ * unnormalized_depth_gt_val * 0.001 if self.use_coords: rot = quaternion2matrix(rotation_gt), else: v = np.matmul(quaternion2matrix(rotation_gt), [1, 0, 0]) rot = rotation_matrix_from_axis(v, [0, 1, 0], 'xy') try: draw_axis(axis_gt, rot, hanging_point_pose, self.cameramodel.K) except Exception: print('Fail to draw axis') confidence_gt_vis = draw_roi(confidence_gt_vis, roi, val=depth_gt_val, gt=True) axis_gt = draw_roi(axis_gt, roi, val=depth_gt_val, gt=True) # Visualize pred axis and roi axis_pred = depth_image_bgr.copy() for i, (roi, roi_c) in enumerate( zip(self.model.rois_list[0], self.model.rois_center_list[0])): if roi.tolist() == [0, 0, 0, 0]: continue roi = roi.cpu().detach().numpy().copy() cx = roi_c[0] cy = roi_c[1] dep = depth[i].cpu().detach().numpy().copy() normalized_dep_pred = float(dep) dep = unnormalize_depth(dep, self.depth_range[0], self.depth_range[1]) confidence_vis = draw_roi(confidence_vis, roi, val=normalized_dep_pred) axis_pred = draw_roi(axis_pred, roi, val=normalized_dep_pred) if mode != 'test': if annotated_rois[i][2]: confidence_vis = draw_roi(confidence_vis, annotated_rois[i][0], val=annotated_rois[i][1][0], gt=True) axis_pred = draw_roi(axis_pred, annotated_rois[i][0], val=annotated_rois[i][1][0], gt=True) hanging_point_pose = np.array( self.cameramodel.project_pixel_to_3d_ray( [int(cx), int(cy)])) * float(dep * 0.001) if self.use_coords: # have not check this yet q = rotation[i].cpu().detach().numpy().copy() q /= np.linalg.norm(q) rot = quaternion2matrix(q) else: v = rotation[i].cpu().detach().numpy() v /= np.linalg.norm(v) rot = rotation_matrix_from_axis(v, [0, 1, 0], 'xy') try: draw_axis(axis_pred, rot, hanging_point_pose, self.cameramodel.K) except Exception: print('Fail to draw axis') axis_pred = cv2.cvtColor(axis_pred, cv2.COLOR_BGR2RGB) confidence_vis = cv2.cvtColor(confidence_vis, cv2.COLOR_BGR2RGB) if self.config['use_bgr']: if self.config['use_bgr2gray']: in_gray = hp_data.cpu().detach().numpy().copy()[0, 1:2, ...] * 255 in_gray = in_gray.transpose(1, 2, 0).astype(np.uint8) in_gray = cv2.cvtColor(in_gray, cv2.COLOR_GRAY2RGB) in_gray = in_gray.transpose(2, 0, 1) in_img = in_gray else: in_bgr = hp_data.cpu().detach().numpy().copy()[ 0, 3:, ...].transpose(1, 2, 0) in_rgb = cv2.cvtColor(in_bgr, cv2.COLOR_BGR2RGB).transpose( 2, 0, 1) in_img = in_rgb if mode != 'test': confidence_loss_sum += confidence_loss.item() axis_gt = cv2.cvtColor(axis_gt, cv2.COLOR_BGR2RGB) confidence_gt_vis = cv2.cvtColor(confidence_gt_vis, cv2.COLOR_BGR2RGB) if rotation_loss.item() > 0: depth_loss_sum += depth_loss.item() rotation_loss_sum += rotation_loss.item() loss_sum = loss_sum \ + confidence_loss.item() \ + rotation_loss.item() rotation_loss_count += 1 if np.mod(index, 1) == 0: print( 'epoch {}, {}/{},{} loss is confidence:{} rotation:{} depth:{}' .format( # noqa self.epo, index, len(dataloader), mode, confidence_loss.item(), rotation_loss.item(), depth_loss.item())) self.vis.images( [axis_gt.transpose(2, 0, 1), axis_pred.transpose(2, 0, 1)], win='{} axis'.format(mode), opts=dict(title='{} axis'.format(mode))) self.vis.images( [ confidence_gt_vis.transpose(2, 0, 1), confidence_vis.transpose(2, 0, 1) ], win='{}_confidence_roi'.format(mode), opts=dict(title='{} confidence(GT, Pred)'.format(mode))) if self.config['use_bgr']: self.vis.images([in_img], win='{} in_gray'.format(mode), opts=dict(title='{} in_gray'.format(mode))) else: if self.config['use_bgr']: self.vis.images( [ in_img, confidence_vis.transpose(2, 0, 1), axis_pred.transpose(2, 0, 1) ], win='{}-{}'.format(mode, index), opts=dict( title='{}-{} hanging_point_depth (pred)'.format( mode, index))) else: self.vis.images( [ confidence_vis.transpose(2, 0, 1), axis_pred.transpose(2, 0, 1) ], win='{}-{}'.format(mode, index), opts=dict( title='{}-{} hanging_point_depth (pred)'.format( mode, index))) if np.mod(index, 1000) == 0: save_file = osp.join( self.save_dir, 'hpnet_latestmodel_' + self.time_now + '.pt') print('save {}'.format(save_file)) torch.save(self.model.state_dict(), save_file, _use_new_zipfile_serialization=False) if mode != 'test': if len(dataloader) > 0: avg_confidence_loss\ = confidence_loss_sum / len(dataloader) if rotation_loss_count > 0: avg_rotation_loss\ = rotation_loss_sum / rotation_loss_count avg_depth_loss\ = depth_loss_sum / rotation_loss_count avg_loss\ = loss_sum / rotation_loss_count else: avg_rotation_loss = 1e10 avg_depth_loss = 1e10 avg_loss = 1e10 else: avg_loss = loss_sum avg_confidence_loss = confidence_loss_sum avg_rotation_loss = rotation_loss_sum avg_depth_loss = rotation_loss_sum self.vis.line(X=np.array([self.epo]), Y=np.array([avg_confidence_loss]), opts={'title': 'confidence'}, win='confidence loss', name='{}_confidence_loss'.format(mode), update='append') if rotation_loss_count > 0: self.vis.line(X=np.array([self.epo]), Y=np.array([avg_rotation_loss]), opts={'title': 'rotation loss'}, win='rotation loss', name='{}_rotation_loss'.format(mode), update='append') self.vis.line(X=np.array([self.epo]), Y=np.array([avg_depth_loss]), opts={'title': 'depth loss'}, win='depth loss', name='{}_depth_loss'.format(mode), update='append') self.vis.line(X=np.array([self.epo]), Y=np.array([avg_loss]), opts={'title': 'loss'}, win='loss', name='{}_loss'.format(mode), update='append') if mode == 'val': if np.mod(self.epo, self.save_model_interval) == 0: save_file = osp.join( self.save_dir, 'hpnet_latestmodel_' + self.time_now + '.pt') print('save {}'.format(save_file)) torch.save(self.model.state_dict(), save_file, _use_new_zipfile_serialization=False) if self.best_loss > avg_loss: print('update best model {} -> {}'.format( self.best_loss, avg_loss)) self.best_loss = avg_loss save_file = osp.join( self.save_dir, 'hpnet_bestmodel_' + self.time_now + '.pt') print('save {}'.format(save_file)) # For ros(python 2, torch 1.4) torch.save(self.model.state_dict(), save_file, _use_new_zipfile_serialization=False)
'--input-dir', '-i', type=str, help='input_dir', # default='/home/kosuke55/catkin_ws/src/pose_annotation_tool/annotation_obj') # noqa default='/media/kosuke55/SANDISK/meshdata/ycb_pouring_object_16/textured_urdf/annotation_obj') # noqa args = parser.parse_args() paths = list(sorted(Path(args.input_dir).glob('*.txt'))) for path in paths: input_file = str(path) print(input_file) contact_points = {'contact_points': [], 'labels': [], 'urdf_file': str(path.with_suffix('.urdf'))} with open(input_file)as f: for line in f.readlines(): values = [float(v) for v in line.split()] label = int(values[0]) pos = values[1:4] quaternion = values[4:] matrix = quaternion2matrix(quaternion) contact_points['contact_points'].append( np.vstack([pos, matrix]).tolist()) contact_points['labels'].append(label) print(line) save_json(str(path.with_suffix('.json')), contact_points) # print(lines)