class PoseObservationPublisher(SimulatorPlugin): def __init__(self, multibody, camera_link, fov, near, far, noise_exp, topic_prefix='', frequency=30, debug=False): super(PoseObservationPublisher, self).__init__('PoseObservationPublisher') self.topic_prefix = topic_prefix self.publisher = rospy.Publisher('{}/pose_obs'.format(topic_prefix), PSAMsg, queue_size=1, tcp_nodelay=True) self.message_templates = {} self.multibody = multibody self.camera_link = camera_link self.fov = fov self.near = near self.far = far self.noise_exp = noise_exp self._enabled = True self.visualizer = ROSVisualizer('pose_obs_viz', 'world') if debug else None self._last_update = 1000 self._update_wait = 1.0 / frequency def post_physics_update(self, simulator, deltaT): """Implements post physics step behavior. :type simulator: BasicSimulator :type deltaT: float """ if not self._enabled: return self._last_update += deltaT if self._last_update >= self._update_wait: self._last_update = 0 cf_tuple = self.multibody.get_link_state(self.camera_link).worldFrame camera_frame = gm.frame3_quaternion(cf_tuple.position.x, cf_tuple.position.y, cf_tuple.position.z, *cf_tuple.quaternion) cov_proj = gm.rot_of(camera_frame)[:3, :3] inv_cov_proj = cov_proj.T out = PSAMsg() if self.visualizer is not None: self.visualizer.begin_draw_cycle() poses = [] for name, body in simulator.bodies.items(): if body == self.multibody: continue if isinstance(body, MultiBody): poses_to_process = [('{}/{}'.format(name, l), body.get_link_state(l).worldFrame) for l in body.links] else: poses_to_process = [(name, body.pose())] for pname, pose in poses_to_process: if not pname in self.message_templates: msg = PoseStampedMsg() msg.header.frame_id = pname self.message_templates[pname] = msg else: msg = self.message_templates[pname] obj_pos = gm.point3(*pose.position) c2o = obj_pos - gm.pos_of(camera_frame) dist = gm.norm(c2o) if dist < self.far and dist > self.near and gm.dot_product(c2o, gm.x_of(camera_frame)) > gm.cos(self.fov * 0.5) * dist: noise = 2 ** (self.noise_exp * dist) - 1 (n_quat, ) = np_random_quat_normal(1, 0, noise) (n_trans, ) = np_random_normal_offset(1, 0, noise) n_pose = pb.Transform(pb.Quaternion(*pose.quaternion), pb.Vector3(*pose.position)) *\ pb.Transform(pb.Quaternion(*n_quat), pb.Vector3(*n_trans[:3])) if self.visualizer is not None: poses.append(transform_to_matrix(n_pose)) msg.pose.position.x = n_pose.origin.x msg.pose.position.y = n_pose.origin.y msg.pose.position.z = n_pose.origin.z msg.pose.orientation.x = n_pose.rotation.x msg.pose.orientation.y = n_pose.rotation.y msg.pose.orientation.z = n_pose.rotation.z msg.pose.orientation.w = n_pose.rotation.w out.poses.append(msg) self.publisher.publish(out) if self.visualizer is not None: self.visualizer.draw_poses('debug', gm.se.eye(4), 0.1, 0.02, poses) self.visualizer.render() def disable(self, simulator): """Stops the execution of this plugin. :type simulator: BasicSimulator """ self._enabled = False self.publisher.unregister() def to_dict(self, simulator): """Serializes this plugin to a dictionary. :type simulator: BasicSimulator :rtype: dict """ return {'body': simulator.get_body_id(self.body.bId()), 'camera_link': self.camera_link, 'fov': self.fov, 'near': self.near, 'far': self.far, 'noise_exp': self.noise_exp, 'topic_prefix': self.topic_prefix} @classmethod def factory(cls, simulator, init_dict): body = simulator.get_body(init_dict['body']) if body is None: raise Exception('Body "{}" does not exist in the context of the given simulation.'.format(init_dict['body'])) return cls(body, init_dict['camera_link'], init_dict['fov'], init_dict['near'], init_dict['far'], init_dict['noise_exp'], init_dict['topic_prefix']) def reset(self, simulator): pass
vis = ROSVisualizer('axis_vis', 'world') az, ay = [Position(x) for x in 'ax az'.split(' ')] frame_rpy = frame3_rpy(0, ay, az, point3(0,0,0)) state = {ay: 0, az: 0} points = [point3(0,0,0) + get_rot_vector(frame_rpy.subs({ay: sin(v), az: cos(v)})) for v in [(3.14512 / 25) * x for x in range(51)]] vis.begin_draw_cycle('points') vis.draw_strip('points', se.eye(4), 0.03, points) vis.render('points') rospy.sleep(1) timer = Time() while not rospy.is_shutdown(): now = Time.now() if (now - timer).to_sec() >= 0.02: state[ay] = sin(now.to_sec()) state[az] = cos(now.to_sec()) frame = frame_rpy.subs(state) axis = get_rot_vector(frame) vis.begin_draw_cycle('visuals') vis.draw_poses('visuals', se.eye(4), 0.4, 0.02, [frame]) vis.draw_vector('visuals', pos_of(frame), axis) vis.render('visuals')
class ROSQPEManager(object): def __init__(self, tracker: Kineverse6DQPTracker, model=None, model_path=None, reference_frame='world', urdf_param='/qp_description_check', update_freq=30, observation_alias=None): self.tracker = tracker self.last_observation = 0 self.last_update = 0 self.reference_frame = reference_frame self.observation_aliases = {o: o for o in tracker.observation_names} if observation_alias is not None: for path, alias in observation_alias.items(): if path in self.observation_aliases: self.observation_aliases[alias] = path self.str_controls = {str(s) for s in self.tracker.get_controls()} self.vis = ROSVisualizer('~vis', reference_frame) if model is not None: self.broadcaster = ModelTFBroadcaster_URDF(urdf_param, model_path, model, Path('ekf')) else: self.broadcaster = None self.tf_buffer = tf2_ros.Buffer() self.listener = tf2_ros.TransformListener(self.tf_buffer) self.pub_state = rospy.Publisher('~state_estimate', ValueMapMsg, queue_size=1, tcp_nodelay=True) self.sub_observations = rospy.Subscriber('~observations', TransformStampedArrayMsg, callback=self.cb_obs, queue_size=1) self.sub_controls = rospy.Subscriber('~controls', ValueMapMsg, callback=self.cb_controls, queue_size=1) self.timer = rospy.Timer(rospy.Duration(1 / update_freq), self.cb_update) def cb_obs(self, transform_stamped_array_msg): ref_frames = {} num_valid_obs = 0 last_observation = {} for trans in transform_stamped_array_msg.transforms: if trans.child_frame_id not in self.observation_aliases: continue matrix = np_frame3_quaternion( trans.transform.translation.x, trans.transform.translation.y, trans.transform.translation.z, trans.transform.rotation.x, trans.transform.rotation.y, trans.transform.rotation.z, trans.transform.rotation.w) if trans.header.frame_id != self.reference_frame: if trans.header.frame_id not in ref_frames: try: ref_trans = self.tf_buffer.lookup_transform( self.reference_frame, trans.header.frame_id, rospy.Time(0)) np_ref_trans = np_frame3_quaternion( ref_trans.transform.translation.x, ref_trans.transform.translation.y, ref_trans.transform.translation.z, ref_trans.transform.rotation.x, ref_trans.transform.rotation.y, ref_trans.transform.rotation.z, ref_trans.transform.rotation.w) ref_frames[trans.header.frame_id] = np_ref_trans except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e: print( f'Exception raised while looking up {trans.header.frame_id} -> {self.reference_frame}:\n{e}' ) break else: np_ref_trans = ref_frames[trans.header.frame_id] matrix = np_ref_trans.dot(matrix) last_observation[self.observation_aliases[ trans.child_frame_id]] = matrix num_valid_obs += 1 # self.last_observation[trans.child_frame_id] = np_6d_pose_feature(trans.transform.translation.x, # trans.transform.translation.y, # trans.transform.translation.z, # trans.transform.rotation.x, # trans.transform.rotation.y, # trans.transform.rotation.z, # trans.transform.rotation.w) else: try: if num_valid_obs == 0: return self.tracker.process_observation(last_observation) self.last_observation += 1 except QPSolverException as e: print( f'Solver crashed during observation update. Skipping observation... Error:\n{e}' ) return self.vis.begin_draw_cycle('observations') # temp_poses = [gm.frame3_axis_angle(feature[3:] / np.sqrt(np.sum(feature[3:]**2)), np.sqrt(np.sum(feature[3:]**2)), feature[:3]) for feature in self.last_observation.values()] self.vis.draw_poses('observations', np.eye(4), 0.2, 0.01, last_observation.values()) # self.vis.draw_poses('observations', np.eye(4), 0.2, 0.01, temp_poses) self.vis.render('observations') def cb_controls(self, map_message): self.last_control = { gm.Symbol(str_symbol): v for str_symbol, v in zip(map_message.symbol, map_message.value) if str_symbol in self.str_controls } self.tracker.process_control(self.last_control) def cb_update(self, *args): if self.last_observation == self.last_update: return self.last_update = self.last_observation est_state = self.tracker.get_estimated_state() state_msg = ValueMapMsg() state_msg.header.stamp = rospy.Time.now() state_msg.symbol, state_msg.value = zip(*est_state.items()) state_msg.symbol = [str(s) for s in state_msg.symbol] self.pub_state.publish(state_msg) if self.broadcaster is not None: self.broadcaster.update_state(est_state) self.broadcaster.publish_state()
class ROSEKFManager(object): def __init__(self, tracker: Kineverse6DEKFTracker, model=None, model_path=None, urdf_param='/ekf_description_check'): self.tracker = tracker self.last_observation = {} self.last_update = None self.last_control = {s: 0.0 for s in self.tracker.controls} self.str_controls = {str(s) for s in self.tracker.controls} self.vis = ROSVisualizer('~vis', 'world') if model is not None: self.broadcaster = ModelTFBroadcaster_URDF(urdf_param, model_path, model, Path('ekf')) else: self.broadcaster = None self.pub_state = rospy.Publisher('~state_estimate', ValueMapMsg, queue_size=1, tcp_nodelay=True) self.sub_observations = rospy.Subscriber('~observations', TransformStampedArrayMsg, callback=self.cb_obs, queue_size=1) self.sub_controls = rospy.Subscriber('~controls', ValueMapMsg, callback=self.cb_controls, queue_size=1) def cb_obs(self, transform_stamped_array_msg): for trans in transform_stamped_array_msg.transforms: matrix = np_frame3_quaternion( trans.transform.translation.x, trans.transform.translation.y, trans.transform.translation.z, trans.transform.rotation.x, trans.transform.rotation.y, trans.transform.rotation.z, trans.transform.rotation.w) self.last_observation[trans.child_frame_id] = matrix # self.last_observation[trans.child_frame_id] = np_6d_pose_feature(trans.transform.translation.x, # trans.transform.translation.y, # trans.transform.translation.z, # trans.transform.rotation.x, # trans.transform.rotation.y, # trans.transform.rotation.z, # trans.transform.rotation.w) self.vis.begin_draw_cycle('observations') # temp_poses = [gm.frame3_axis_angle(feature[3:] / np.sqrt(np.sum(feature[3:]**2)), np.sqrt(np.sum(feature[3:]**2)), feature[:3]) for feature in self.last_observation.values()] self.vis.draw_poses('observations', np.eye(4), 0.2, 0.01, self.last_observation.values()) self.vis.render('observations') self.try_update() def cb_controls(self, map_message): self.last_control = { gm.Symbol(str_symbol): v for str_symbol, v in zip(map_message.symbol, map_message.value) if str_symbol in self.str_controls } self.try_update() def try_update(self): if min(p in self.last_observation for p in self.tracker.observation_names) is False: return if len(self.tracker.controls) > 0 and min( s in self.last_control for s in self.tracker.controls) is False: return now = rospy.Time.now() dt = 0.05 if self.last_update is None else (now - self.last_update).to_sec() self.tracker.process_update(self.last_observation, self.last_control, dt) est_state = self.tracker.get_estimated_state() if len(est_state) == 0: return state_msg = ValueMapMsg() state_msg.header.stamp = now state_msg.symbol, state_msg.value = zip(*est_state.items()) state_msg.symbol = [str(s) for s in state_msg.symbol] self.pub_state.publish(state_msg) # print('Performed update of estimate and published it.') # self.last_control = {} self.last_observation = {} if self.broadcaster is not None: self.broadcaster.update_state(self.tracker.get_estimated_state()) self.broadcaster.publish_state()