Resolve "Добавить параметры в конфиг pose estimation узла"

rbs_perception/scripts/pose_estimation.py

@@ -3,14 +3,15 @@
   detection_service
   ROS 2 program for 6D Pose Estimation
 
-  @shalenikol release 0.1
+  @shalenikol release 0.2
 """
   
 # Import the necessary libraries
 import rclpy # Python library for ROS 2
 from rclpy.node import Node # Handles the creation of nodes
 from sensor_msgs.msg import Image # Image is the message type
-from geometry_msgs.msg import Quaternion
+from geometry_msgs.msg import Quaternion, TransformStamped
+from tf2_ros import TransformBroadcaster
 from cv_bridge import CvBridge # Package to convert between ROS and OpenCV Images
 import cv2 # OpenCV library
 from rbs_skill_interfaces.srv import DetectObject
@@ -23,9 +24,11 @@ import tempfile
 from pathlib import Path
 import numpy as np
 from ament_index_python.packages import get_package_share_directory
-import megapose
+# import megapose
 from megapose.scripts.run_inference_on_example import run_inference
 
+tf2_send_pose = True
+
 """
 # encoder for numpy array 
 def np_encoder(object):
@@ -44,49 +47,51 @@ class ImageSubscriber(Node):
       y = json.load(f)
 
     for name, val in y.items():
-      if name == "topicName":
-        self.topicName = val
+      if name == "nodeName":
+        self.nodeName = val
       elif name == "topicImage":
         self.topicImage = val
       elif name == "topicPubName":
         self.topicPubName = val
       elif name == "topicSrv":
         self.topicSrv = val
+      elif name == "tf2_send_pose":
+        self.tf2_send_pose = val
+      elif name == "camera_info":
+        self.K_, self.res_ = self._getCameraParam(val)
 
-  def _getCameraParam(self):
-    # {"K": [[924.855, 0.0, 320.0], [0.0, 923.076, 240.0], [0.0, 0.0, 1.0]], "resolution": [480, 640]}
+  def _getCameraParam(self, info):
     """
-    focal length: 25 mm (field of view: 38.1712°)
-    sensor type: Micro Four Thirds System (Стандарт MFT микро 4/3)
-    sensor_width: 17.3 mm
-    sensor_height: 13 mm
+    Returns the intrinsic matrix and resolution from the provided camera info.
     """
-    resolution = [480, 640] #np.array([480, 640])
-    intrinsic_matrix = [    #np.array([
-      [924.855, 0.0, resolution[1] / 2.0],
-      [0.0, 923.076, resolution[0] / 2.0],
-      [0.0, 0.0, 1.0]
-    ] #)
+    intrinsic_matrix = [
+        [info["fx"], 0.0, info["width"] / 2.0],
+        [0.0, info["fy"], info["height"] / 2.0],
+        [0.0, 0.0, 1.0]
+    ]
+    resolution = [info["height"], info["width"]]
     return intrinsic_matrix, resolution
 
   def __init__(self):
     """
     Class constructor to set up the node
     """
-    self.topicName = "image_sub2"
-    self.topicImage = "/ground_true/camera_node"
-    self.topicPubName = "pose6D_images"
-    self.topicSrv = "detect6Dpose"
+    self.nodeName = "image_sub2"
+    self.topicImage = "/outer_rgbd_camera/image"
+    self.topicPubName = self.nodeName + "/pose6D_images"
+    self.topicSrv = self.nodeName + "/detect6Dpose"
     self._InitService()
 
     self.tmpdir = tempfile.gettempdir()
     self.mytemppath = Path(self.tmpdir) / "rbs_per"
     self.mytemppath.mkdir(exist_ok=True)
-    self.K_, self.res_ = self._getCameraParam()
     #os.environ["MEGAPOSE_DATA_DIR"] = str(self.mytemppath)
     
     # Initiate the Node class's constructor and give it a name
-    super().__init__(self.topicName)
+    super().__init__(self.nodeName)
+
+    # Initialize the transform broadcaster
+    self.tf_broadcaster = TransformBroadcaster(self)
 
     self.subscription = None
     self.objName = ""
@@ -101,14 +106,15 @@ class ImageSubscriber(Node):
     self.service = self.create_service(DetectObject, self.topicSrv, self.service_callback)
    
   def service_callback(self, request, response):
-    self.get_logger().info(f"Incoming request for object detection ObjectInfo(name: {request.object.name}, mesh_path: {request.object.mesh_path})")
+    self.get_logger().info(f"Incoming request for pose estimation ObjectInfo(name: {request.object.name}, mesh_path: {request.object.mesh_path})")
 
     if not os.path.isfile(request.object.mesh_path):
       response.call_status = False
       response.error_msg = f"{request.object.mesh_path}: no such file"
       return response
     if request.object.id == -1:
-      self.subscription = None # ? сброс подписки
+      self.subscription = None # ? сброс подпискиpython -m megapose.scripts.download --example_data
+
       response.call_status = True
       return response
 
@@ -126,7 +132,7 @@ class ImageSubscriber(Node):
       output_fn = tPath / "object_data.json"
       output_json_dict = {
         "label": self.objName,
-        "bbox_modal": [2,2,self.res_[1]-1,self.res_[0]-1]
+        "bbox_modal": [2,2,self.res_[1]-4,self.res_[0]-4]
       }
       data = []
       data.append(output_json_dict)
@@ -146,7 +152,7 @@ class ImageSubscriber(Node):
       }
       data = []
       data.append(output_json_dict)
-      output_fn.write_text(json.dumps(data))
+      output_fn.write_text(json.dumps(output_json_dict))
 
       # Create the subscriber. This subscriber will receive an Image from the video_frames topic. The queue size is 3 messages.
       self.subscription = self.create_subscription(Image, self.topicImage, self.listener_callback, 3)
@@ -166,6 +172,29 @@ class ImageSubscriber(Node):
       data = "No result file: '" + str(f) + "'"
     return data
   
+  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 listener_callback(self, data):
     """
     Callback function.
@@ -182,16 +211,19 @@ class ImageSubscriber(Node):
     
     # 6D pose estimation
     self.get_logger().info(f"megapose: begin {self.cnt}")
+    print(self.objPath)
     run_inference(self.objPath,"megapose-1.0-RGB-multi-hypothesis")
 
     # опубликуем результат оценки позы
     data = self.load_result(self.objPath)
     if data[0] == "[":
       y = json.loads(data)[0]
-      pose = y['TWO']
+      pose = y["TWO"]
       quat = pose[0]
       #pose[1] - 3D перемещение
       self.publisher.publish(Quaternion(x=quat[1],y=quat[2],z=quat[3],w=quat[0]))
+      if tf2_send_pose:
+        self.tf_obj_pose(pose)
 
     self.get_logger().info(f"megapose: end {self.cnt}")