3

Home

wxchen edited this page 2023-05-10 01:14:05 +08:00

Welcome to the Wiki.

#!/usr/bin/env python

import rospy
from sensor_msgs.msg import Image, PointCloud2, PointField, CameraInfo
from cv_bridge import CvBridge
import numpy as np
import open3d as o3d

from std_msgs.msg import Header
from sensor_msgs.msg import PointCloud2, PointField
import sensor_msgs.point_cloud2 as pc2

class ImageAndPointCloudProcessor:
    def __init__(self):
        self.bridge = CvBridge()
        self.image_sub = rospy.Subscriber('/camera/color/image_raw', Image, self.image_callback)
        self.pointcloud_sub = rospy.Subscriber('/camera/depth/color/points', PointCloud2, self.pointcloud_callback)
        self.camera_info_sub = rospy.Subscriber('/camera/color/camera_info', CameraInfo, self.camera_info_callback)
        self.pointcloud_pub = rospy.Publisher('/custom/pointcloud', PointCloud2, queue_size=10)
        self.box = [0, 0, 640, 480]  # default box
        self.camera_intrinsics = None

    def image_callback(self, msg):
        try:
            cv_image = self.bridge.imgmsg_to_cv2(msg, desired_encoding='bgr8')
        except Exception as e:
            rospy.logerr(e)
            return

        # Crop image
        xmin, ymin, xmax, ymax = self.box
        cropped_image = cv_image[ymin:ymax, xmin:xmax]

        # Process cropped image here
        # ...

    def camera_info_callback(self, msg):
        self.camera_intrinsics = np.array(msg.K).reshape((3, 3))

    def pointcloud_callback(self, msg):
        if self.camera_intrinsics is None:
            return

        try:
            pointcloud = o3d.geometry.PointCloud()
            pointcloud.points = o3d.utility.Vector3dVector(np.array(list(msg.data)).reshape(-1, 4)[:, :3])
        except Exception as e:
            rospy.logerr(e)
            return

        # Crop pointcloud
        xmin, ymin, xmax, ymax = self.box
        x_indices = np.logical_and(pointcloud.points[:, 0] >= xmin, pointcloud.points[:, 0] < xmax)
        y_indices = np.logical_and(pointcloud.points[:, 1] >= ymin, pointcloud.points[:, 1] < ymax)
        indices = np.logical_and(x_indices, y_indices)
        pointcloud_cropped = pointcloud.select_by_index(np.where(indices)[0])

        # Plane fitting
        plane_model, inliers = pointcloud_cropped.segment_plane(distance_threshold=0.01, ransac_n=3)
        inlier_cloud = pointcloud_cropped.select_by_index(inliers)
        outlier_cloud = pointcloud_cropped.select_by_index(inliers, invert=True)

        # Output rotation angle
        plane_normal = np.array(plane_model[:3])
        z_axis = np.array([0, 0, 1])
        rotation_axis = np.cross(z_axis, plane_normal)
        rotation_angle = np.arccos(np.dot(z_axis, plane_normal) / (np.linalg.norm(z_axis) * np.linalg.norm(plane_normal)))
        rotation_euler = np.array([rotation_angle, rotation_axis[0], rotation_axis[1], rotation_axis[2]])

        # Publish red pointcloud for the plane
        header = Header()
        header.stamp = rospy.Time.now()
        header.frame_id = msg.header.frame_id
        fields = [PointField('x', 0, PointField.FLOAT32, 1),
                  PointField('y', 4, PointField.FLOAT32, 1),
                  PointField('z', 8, PointField.FLOAT32, 1),
                  PointField('rgb', 12, PointField.FLOAT32, 1)]
        
        # Create a list of points representing the plane
        plane_points = []
        for p in inlier_cloud.points:
            u = (p[0] - self.camera_intrinsics[0, 2]) / self.camera_intrinsics[0, 0]
            v = (p[1] - self.camera_intrinsics[1, 2]) / self.camera_intrinsics[1, 1]
            plane_points.append([p[0], p[1], p[2], np.float32(0xff0000)])
        
        # Create and publish the pointcloud message
        msg = pc2.create_cloud(header, fields, plane_points)
        self.pointcloud_pub.publish(msg)

def main():
    rospy.init_node('image_and_pointcloud_processor')
    processor = ImageAndPointCloudProcessor()
    processor.box = [200, 200, 440, 320]  # example box
    rospy.spin()

if __name__ == '__main__':
    main()