Reimplemented wrapper in cpp.

ros/rpt2D_wrapper_py/rpt2D_wrapper_py/rpt2D_wrapper.py

@@ -0,0 +1,191 @@
+import json
+import os
+import sys
+import threading
+import time
+
+import numpy as np
+import rclpy
+from cv_bridge import CvBridge
+from rclpy.qos import QoSHistoryPolicy, QoSProfile, QoSReliabilityPolicy
+from sensor_msgs.msg import Image
+from std_msgs.msg import String
+
+filepath = os.path.dirname(os.path.realpath(__file__)) + "/"
+sys.path.append(filepath + "../../../scripts/")
+import test_triangulate
+import utils_2d_pose
+
+# ==================================================================================================
+
+bridge = CvBridge()
+node = None
+publisher_pose = None
+
+cam_id = "camera01"
+img_input_topic = "/" + cam_id + "/pylon_ros2_camera_node/image_raw"
+pose_out_topic = "/poses/" + cam_id
+
+last_input_image = None
+kpt_model = None
+joint_names_2d = test_triangulate.joint_names_2d
+
+lock = threading.Lock()
+stop_flag = False
+
+# Model config
+min_bbox_score = 0.4
+min_bbox_area = 0.1 * 0.1
+batch_poses = False
+
+
+# ==================================================================================================
+
+
+def callback_images(image_data):
+    global last_input_image, lock
+
+    # Convert into cv images from image string
+    if image_data.encoding == "bayer_rggb8":
+        bayer_image = bridge.imgmsg_to_cv2(image_data, "bayer_rggb8")
+    elif image_data.encoding == "mono8":
+        bayer_image = bridge.imgmsg_to_cv2(image_data, "mono8")
+    elif image_data.encoding == "rgb8":
+        color_image = bridge.imgmsg_to_cv2(image_data, "rgb8")
+        bayer_image = test_triangulate.rgb2bayer(color_image)
+    else:
+        raise ValueError("Unknown image encoding:", image_data.encoding)
+
+    time_stamp = image_data.header.stamp.sec + image_data.header.stamp.nanosec / 1.0e9
+
+    with lock:
+        last_input_image = (bayer_image, time_stamp)
+
+
+# ==================================================================================================
+
+
+def callback_model():
+    global last_input_image, kpt_model, lock
+
+    ptime = time.time()
+    if last_input_image is None:
+        time.sleep(0.001)
+        return
+
+    # Collect inputs
+    images_2d = []
+    timestamps = []
+    with lock:
+        img = np.copy(last_input_image[0])
+        ts = last_input_image[1]
+        images_2d.append(img)
+        timestamps.append(ts)
+        last_input_image = None
+
+    # Predict 2D poses
+    images_2d = [test_triangulate.bayer2rgb(img) for img in images_2d]
+    poses_2d = utils_2d_pose.get_2d_pose(kpt_model, images_2d)
+    poses_2d = test_triangulate.update_keypoints(poses_2d, joint_names_2d)
+    poses_2d = poses_2d[0]
+
+    # Drop persons with no joints
+    poses_2d = np.asarray(poses_2d)
+    mask = np.sum(poses_2d[..., 2], axis=1) > 0
+    poses_2d = poses_2d[mask]
+
+    # Round poses
+    poses2D = [np.array(p).round(3).tolist() for p in poses_2d]
+
+    # Publish poses
+    ts_pose = time.time()
+    poses = {
+        "bodies2D": poses2D,
+        "joints": joint_names_2d,
+        "num_persons": len(poses2D),
+        "timestamps": {
+            "image": timestamps[0],
+            "pose": ts_pose,
+            "z-images-pose": ts_pose - timestamps[0],
+        },
+    }
+    publish(poses)
+
+    msg = "Detected persons: {} - Prediction time: {:.4f}s"
+    print(msg.format(poses["num_persons"], time.time() - ptime))
+
+
+# ==================================================================================================
+
+
+def callback_wrapper():
+    global stop_flag
+    while not stop_flag:
+        callback_model()
+        time.sleep(0.001)
+
+
+# ==================================================================================================
+
+
+def publish(data):
+    # Publish json data
+    msg = String()
+    msg.data = json.dumps(data)
+    publisher_pose.publish(msg)
+
+
+# ==================================================================================================
+
+
+def main():
+    global node, publisher_pose, kpt_model, stop_flag
+
+    # Start node
+    rclpy.init(args=sys.argv)
+    node = rclpy.create_node("rpt2D_wrapper")
+
+    # Quality of service settings
+    qos_profile = QoSProfile(
+        reliability=QoSReliabilityPolicy.RELIABLE,
+        history=QoSHistoryPolicy.KEEP_LAST,
+        depth=1,
+    )
+
+    # Create subscribers
+    _ = node.create_subscription(
+        Image,
+        img_input_topic,
+        callback_images,
+        qos_profile,
+    )
+
+    # Create publishers
+    publisher_pose = node.create_publisher(String, pose_out_topic, qos_profile)
+
+    # Load 2D pose model
+    whole_body = test_triangulate.whole_body
+    if any((whole_body[k] for k in whole_body)):
+        kpt_model = utils_2d_pose.load_wb_model()
+    else:
+        kpt_model = utils_2d_pose.load_model(min_bbox_score, min_bbox_area, batch_poses)
+
+    node.get_logger().info("Finished initialization of pose estimator")
+
+    # Start prediction thread
+    p1 = threading.Thread(target=callback_wrapper)
+    p1.start()
+
+    # Run ros update thread
+    rclpy.spin(node)
+
+    stop_flag = True
+    p1.join()
+    node.destroy_node()
+    rclpy.shutdown()
+
+
+# ==================================================================================================
+
+if __name__ == "__main__":
+    main()