#!/usr/bin/env python3 """ pose_estimation_lifecycle_node ROS 2 program for 6D Pose Estimation @shalenikol release 0.3 """ from typing import Optional import os import shutil import json import tempfile from pathlib import Path import rclpy from rclpy.lifecycle import Node from rclpy.lifecycle import Publisher from rclpy.lifecycle import State from rclpy.lifecycle import TransitionCallbackReturn from rclpy.timer import Timer from ament_index_python.packages import get_package_share_directory from sensor_msgs.msg import Image, CameraInfo from geometry_msgs.msg import Pose, TransformStamped from tf2_ros import TransformBroadcaster from cv_bridge import CvBridge # Package to convert between ROS and OpenCV Images import cv2 # OpenCV library from megapose.scripts.run_inference_on_example import ModelPreload, run_inference_rbs #, run_inference class PoseEstimator(Node): """Our lifecycle node.""" def _InitService(self): # Initialization service data p = os.path.join(get_package_share_directory("rbs_perception"), "config", "pose_estimation_config.json") # load config with open(p, "r") as f: y = json.load(f) for name, val in y.items(): if name == "nodeName": self.nodeName = val elif name == "topicImage": self.topicImage = val elif name == "topicCameraInfo": self.topicCameraInfo = val elif name == "topicDepth": self.topicDepth = val elif name == "publishDelay": self.publishDelay = val elif name == "topicSrv": self.topicSrv = val elif name == "tf2_send_pose": self.tf2_send_pose = val def __init__(self, node_name, **kwargs): """Construct the node.""" self._count: int = 0 self._pub: Optional[Publisher] = None self._timer: Optional[Timer] = None self._image_cnt: int = 0 self._sub = None self._sub_info = None self._sub_depth = None self._is_camerainfo = False self._K = [[0., 0., 0.], [0., 0., 0.], [0., 0., 0.]] self._res = [0, 0] self._pose = [[1., 0., 0., 0.], [0., 0., 0.]] self.tf2_send_pose = 0 self.megapose_model = None self.nodeName = node_name #"image_sub2" self.topicImage = "/outer_rgb_camera/image" self.topicCameraInfo = "/outer_rgb_camera/camera_info" self.topicDepth = "/outer_rgbd_camera/depth_image" self.publishDelay = 2.0 self.topicSrv = self.nodeName + "/detect6Dpose" self._InitService() self.tmpdir = tempfile.gettempdir() self.mytemppath = Path(self.tmpdir) / "rbs_per" self.mytemppath.mkdir(exist_ok=True) super().__init__(self.nodeName, **kwargs) self.declare_parameter("mesh_path", "") # Initialize the transform broadcaster self.tf_broadcaster = TransformBroadcaster(self) # Used to convert between ROS and OpenCV images self.br = CvBridge() self.objName = "" self.objMeshFile = "" self.objPath = "" def publish(self): """Publish a new message when enabled.""" self._count += 1 if self._pub is not None and self._pub.is_activated: # опубликуем результат оценки позы q = self._pose[0] t = self._pose[1] p = Pose() p.position.x = t[0] p.position.y = t[1] p.position.z = t[2] p.orientation.w = q[0] p.orientation.x = q[1] p.orientation.y = q[2] p.orientation.z = q[3] self._pub.publish(p) if self.tf2_send_pose: self.tf_obj_pose(self._pose) def tf_obj_pose(self,pose): """ Передача позиции объекта в tf2 """ t = TransformStamped() # assign pose to corresponding tf variables t.header.stamp = self.get_clock().now().to_msg() t.header.frame_id = 'world' t.child_frame_id = self.objName # coordinates tr = pose[1] t.transform.translation.x = tr[0] t.transform.translation.y = tr[1] t.transform.translation.z = tr[2] # rotation q = pose[0] t.transform.rotation.x = q[1] # 0 t.transform.rotation.y = q[2] # 1 t.transform.rotation.z = q[3] # 2 t.transform.rotation.w = q[0] # 3 # Send the transformation self.tf_broadcaster.sendTransform(t) def on_configure(self, state: State) -> TransitionCallbackReturn: """ Configure the node, after a configuring transition is requested. :return: The state machine either invokes a transition to the "inactive" state or stays in "unconfigured" depending on the return value. TransitionCallbackReturn.SUCCESS transitions to "inactive". TransitionCallbackReturn.FAILURE transitions to "unconfigured". TransitionCallbackReturn.ERROR or any uncaught exceptions to "errorprocessing" """ mesh_path = self.get_parameter("mesh_path").get_parameter_value().string_value if not os.path.isfile(mesh_path): self.get_logger().info("Parameter 'mesh_path' not set") return TransitionCallbackReturn.FAILURE data = os.path.basename(mesh_path) self.objName = os.path.splitext(data)[0] self.objMeshFile = mesh_path self.objPath = self.mytemppath / "examples" self.objPath.mkdir(exist_ok=True) self.objPath /= self.objName self.objPath.mkdir(exist_ok=True) tPath = self.objPath / "inputs" tPath.mkdir(exist_ok=True) tPath = self.objPath / "meshes" tPath.mkdir(exist_ok=True) tPath /= self.objName tPath.mkdir(exist_ok=True) shutil.copyfile(self.objMeshFile, str(tPath / (self.objName+".ply"))) # Create the subscribers. self._sub_info = self.create_subscription(CameraInfo, self.topicCameraInfo, self.listener_camera_info, 2) self._sub_depth = self.create_subscription(Image, self.topicDepth, self.listener_depth, 3) # Create the publisher. self._pub = self.create_lifecycle_publisher(Pose, self.topicSrv, 10) self._timer = self.create_timer(self.publishDelay, self.publish) # Preload Megapose model self.megapose_model = ModelPreload(self.objPath,"megapose-1.0-RGB-multi-hypothesis") self.get_logger().info('on_configure() is called.') return TransitionCallbackReturn.SUCCESS def on_activate(self, state: State) -> TransitionCallbackReturn: # Log self.get_logger().info('on_activate() is called.') # Create the main subscriber. self._sub = self.create_subscription(Image, self.topicImage, self.listener_callback, 3) return super().on_activate(state) def on_deactivate(self, state: State) -> TransitionCallbackReturn: # Log self.get_logger().info('on_deactivate() is called.') # Destroy the main subscriber. self.destroy_subscription(self._sub) return super().on_deactivate(state) def on_cleanup(self, state: State) -> TransitionCallbackReturn: """ Cleanup the node. :return: The state machine either invokes a transition to the "unconfigured" state or stays in "inactive" depending on the return value. TransitionCallbackReturn.SUCCESS transitions to "unconfigured". TransitionCallbackReturn.FAILURE transitions to "inactive". TransitionCallbackReturn.ERROR or any uncaught exceptions to "errorprocessing" """ # очистим параметры node_param = rclpy.parameter.Parameter("mesh_path", rclpy.Parameter.Type.STRING, "") all_node_param = [node_param] self.set_parameters(all_node_param) self._is_camerainfo = False self.destroy_timer(self._timer) self.destroy_publisher(self._pub) self.destroy_subscription(self._sub) self.destroy_subscription(self._sub_info) self.destroy_subscription(self._sub_depth) self.get_logger().info('on_cleanup() is called.') return TransitionCallbackReturn.SUCCESS def on_shutdown(self, state: State) -> TransitionCallbackReturn: """ Shutdown the node. :return: The state machine either invokes a transition to the "finalized" state or stays in the current state depending on the return value. TransitionCallbackReturn.SUCCESS transitions to "unconfigured". TransitionCallbackReturn.FAILURE transitions to "inactive". TransitionCallbackReturn.ERROR or any uncaught exceptions to "errorprocessing" """ self.destroy_timer(self._timer) self.destroy_publisher(self._pub) self.destroy_subscription(self._sub) self.destroy_subscription(self._sub_info) self.destroy_subscription(self._sub_depth) self.get_logger().info('on_shutdown() is called.') return TransitionCallbackReturn.SUCCESS def listener_camera_info(self, data): """ CameraInfo callback function. """ if self._is_camerainfo: # повторно инфо камеры не читаем return self._res = [data.height, data.width] k_ = data.k self._K = [ [k_[0], k_[1], k_[2]], [k_[3], k_[4], k_[5]], [k_[6], k_[7], k_[8]] ] tPath = self.objPath / "inputs" #{"label": "fork", "bbox_modal": [329, 189, 430, 270]} output_fn = tPath / "object_data.json" output_json_dict = { "label": self.objName, "bbox_modal": [2,2,self._res[1]-4,self._res[0]-4] } data = [] data.append(output_json_dict) output_fn.write_text(json.dumps(data)) #{"K": [[25.0, 0.0, 8.65], [0.0, 25.0, 6.5], [0.0, 0.0, 1.0]], "resolution": [480, 640]} output_fn = self.objPath / "camera_data.json" output_json_dict = { "K": self._K, "resolution": self._res } data = [] data.append(output_json_dict) output_fn.write_text(json.dumps(output_json_dict)) # установим признак получения инфо камеры self._is_camerainfo = True def listener_depth(self, data): """ Depth image callback function. """ #self.get_logger().info("Receiving depth image") # Convert ROS Image message to OpenCV image current_frame = self.br.imgmsg_to_cv2(data) # Save depth image for Megapose cv2.imwrite(str(self.objPath / "image_depth.png"), current_frame) def load_result(self, example_dir: Path, json_name = "object_data.json"): f = example_dir / "outputs" / json_name if os.path.isfile(f): data = f.read_text() else: data = "No result file: '" + str(f) + "'" return data def listener_callback(self, data): """ Image Callback function. """ if not self._is_camerainfo: self.get_logger().warning("No data from CameraInfo") return # Convert ROS Image message to OpenCV image current_frame = self.br.imgmsg_to_cv2(data) # Save image for Megapose cv2.imwrite(str(self.objPath / "image_rgb.png"), current_frame) self._image_cnt += 1 # 6D pose estimation self.get_logger().info(f"megapose: begin {self._image_cnt} {self.objPath}") #run_inference(self.objPath,"megapose-1.0-RGB-multi-hypothesis") run_inference_rbs(self.megapose_model) data = self.load_result(self.objPath) if data[0] == "[": y = json.loads(data)[0] self._pose = y["TWO"] self.get_logger().info(f"megapose: end {self._image_cnt}") def main(): rclpy.init() executor = rclpy.executors.SingleThreadedExecutor() lc_node = PoseEstimator("lc_pose_estimator") executor.add_node(lc_node) try: executor.spin() except (KeyboardInterrupt, rclpy.executors.ExternalShutdownException): lc_node.destroy_node() if __name__ == '__main__': main()