jarod_reynolds

experimental_lab3_q1

May 1st, 2021
722
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #!/usr/bin/python3
  2.  
  3. import matplotlib
  4. import matplotlib.pyplot as plt
  5. import rospy
  6. from sensor_msgs.msg import PointCloud2 as pc2
  7. import sensor_msgs.point_cloud2 as pc2
  8. from sensor_msgs.msg import LaserScan
  9. from nav_msgs.msg import Odometry
  10. from geometry_msgs.msg import Pose
  11. from laser_geometry import LaserProjection
  12. import open3d as o3d
  13. import numpy as np
  14. import collections
  15. import copy
  16. import time
  17.  
  18. # class OdomShiz:
  19. #     def __init__(self):
  20. #         self.OdomSubscriber = rospy.Subscriber("/odom", Odometry, self.odom_callback)
  21. #         self.odom_trajectory = {'x': [], 'y': []}
  22. #
  23. #     def odom_callback(self,data):
  24. #         # x_pos = data.pose.pose.position.x
  25. #         # y_pos = data.pose.pose.position.y
  26. #         # z_pos = data.pose.pose.position.z
  27. #
  28. #         self.odom_trajectory['x'].append(data.pose.pose.position.x)
  29. #         self.odom_trajectory['y'].append(data.pose.pose.position.y)
  30. #         self.update_odom_plot()
  31. #
  32. #     def update_odom_plot(self):
  33. #         odom_plot.plot(self.odom_trajectory['x'], self.odom_trajectory['y'])
  34. #         plt.draw()
  35. #         plt.pause(0.001)
  36. #         plt.show()
  37.  
  38.  
  39. class SickDawgPCz:
  40.     def __init__(self):
  41.         self.laserProj = LaserProjection()
  42.         self.laserSubscriber = rospy.Subscriber("/scan", LaserScan, self.laser_callback)
  43.         self.queue = collections.deque()
  44.         self.current_pos = np.array([0, 0, 0])
  45.         self.vis = o3d.visualization.Visualizer()
  46.         self.source_geometry = None
  47.         self.target_geometry = None
  48.         self.trajectory = {'x': [], 'y': []}
  49.         self.scan_count = 0
  50.         self.trans_init = None
  51.         # Odometry stuff
  52.         self.OdomSubscriber = rospy.Subscriber("/odom", Odometry, self.odom_callback)
  53.         self.odom_trajectory = {'x': [], 'y': []}
  54.  
  55.     def laser_callback(self, data):
  56.         pc = self.laser_scan_to_point_cloud(data)
  57.         # o3d.visualization.draw_geometries([pc])
  58.  
  59.         # Append cloud to the queue
  60.         self.queue.append(pc)
  61.         self.scan_count += 1
  62.  
  63.     def odom_callback(self, data):
  64.         self.odom_trajectory['x'].append(data.pose.pose.position.x)
  65.         self.odom_trajectory['y'].append(data.pose.pose.position.y)
  66.         # self.update_odom_plot()
  67.  
  68.     def laser_scan_to_point_cloud(self, scan):
  69.         cloud_out = self.laserProj.projectLaser(scan)
  70.  
  71.         # Get cloud data from ros_cloud
  72.         # Source: https://github.com/felixchenfy/open3d_ros_pointcloud_conversion/blob/master/lib_cloud_conversion_between_Open3D_and_ROS.py
  73.         field_names = [field.name for field in cloud_out.fields]
  74.         cloud_data = list(pc2.read_points(cloud_out, skip_nans=True, field_names=field_names))
  75.  
  76.         open3d_cloud = o3d.PointCloud()
  77.         xyz = [(x, y, z) for x, y, z, intensity, index in cloud_data]  # get xyz
  78.         open3d_cloud.points = o3d.Vector3dVector(np.array(xyz))
  79.  
  80.         return open3d_cloud
  81.  
  82.     def icp_registration(self):
  83.         # Check for 2 clouds in queue
  84.         if len(self.queue) < 2:
  85.             return
  86.  
  87.         # print("QUEUE LENGTH:" + str(len(self.queue)))
  88.         source = self.queue.popleft()
  89.         target = self.queue.popleft()
  90.         threshold = 0.02
  91.  
  92.         if self.trans_init is None:
  93.             self.trans_init = np.identity(4)
  94.  
  95.         reg_p2l = o3d.registration.registration_icp(
  96.             source, target, threshold, self.trans_init,
  97.             o3d.registration.TransformationEstimationPointToPoint(),
  98.             o3d.registration.ICPConvergenceCriteria(max_iteration=200))
  99.  
  100.         # Save this transformation to use in the next ICP
  101.         self.trans_init = reg_p2l.transformation[:]
  102.  
  103.         # print(reg_p2l)
  104.         # print(reg_p2l.transformation)
  105.         self.draw_registration_result(source, target, reg_p2l.transformation)
  106.         self.update_position(reg_p2l.transformation)
  107.  
  108.         # Add target back to queue
  109.         self.queue.appendleft(target)
  110.  
  111.         # Update trajectory
  112.         # Plotting is a bottleneck - only update the plot every 50 scans
  113.         # so the program can keep up with incoming scans
  114.         if self.scan_count % 10 == 0:
  115.             # pc_stuff.update_trajectory()
  116.             self.update_plots()
  117.  
  118.     def update_position(self, transform):
  119.         # Extract rotation and translation from transformation matrix
  120.         rot = transform[0:3, 0:3]
  121.         transl = transform[0:3, 3]
  122.         print(rot)
  123.         print(transl)
  124.         # print(self.current_pos)
  125.  
  126.         # Apply transformation to bot's current pose
  127.         new_pos = rot.dot(self.current_pos.T) + transl.T
  128.         self.current_pos = new_pos.T[:]
  129.  
  130.         # Add the new position to the trajectory
  131.         self.trajectory['x'].append(new_pos[0])
  132.         self.trajectory['y'].append(new_pos[1])
  133.  
  134.     def update_plots(self):
  135.         pc_plot.plot(self.trajectory['x'], self.trajectory['y'])
  136.         odom_plot.plot(self.odom_trajectory['x'], self.odom_trajectory['y'])
  137.         plt.draw()
  138.         plt.pause(0.001)
  139.         plt.show()
  140.  
  141.     # Helper visualisation function
  142.     # Source: http://www.open3d.org/docs/release/tutorial/pipelines/icp_registration.html
  143.     def draw_registration_result(self, source, target, transformation):
  144.  
  145.         source_temp = copy.deepcopy(source)
  146.         target_temp = copy.deepcopy(target)
  147.         source_temp.paint_uniform_color([1, 0.706, 0])
  148.         target_temp.paint_uniform_color([0, 0.651, 0.929])
  149.         source_temp.transform(transformation)
  150.         # o3d.visualization.draw_geometries([source_temp, target_temp]) # Blocking
  151.  
  152.         # Non-blocking visualisation
  153.         self.vis.remove_geometry(self.source_geometry)
  154.         self.vis.remove_geometry(self.target_geometry)
  155.         self.source_geometry = source_temp
  156.         self.target_geometry = target_temp
  157.         self.vis.add_geometry(self.source_geometry)
  158.         self.vis.add_geometry(self.target_geometry)
  159.         self.vis.poll_events()
  160.         self.vis.update_renderer()
  161.  
  162.  
  163. # Press the green button in the gutter to run the script.
  164. if __name__ == '__main__':
  165.     rospy.init_node("laser2PointCloud",  disable_signals=True)
  166.     pc_stuff = SickDawgPCz()
  167.     # odomOb = OdomShiz()
  168.  
  169.     pc_stuff.vis.create_window()
  170.  
  171.     # Initialise a plot for the trajectory
  172.     matplotlib.use('TKAgg')
  173.     fig,(pc_plot, odom_plot) = plt.subplots(1, 2)
  174.     plt.ion()
  175.     plt.show()
  176.     pc_plot.axis([-3, 1.5, -3, 1.5])
  177.     # odom_plot.axis([2, 6.5, -1.5, 3])
  178.     odom_plot.axis([-5, 5, -5, 5])
  179.     # plt.set(xlim=(-3, 1.5), ylim=(-3, 1.5))
  180.  
  181.     while True:
  182.         try:
  183.             pc_stuff.icp_registration()
  184.  
  185.         except KeyboardInterrupt as e:
  186.             print(e)
  187.             break
  188.     pc_stuff.vis.destroy_window()
  189.  
  190.     rospy.spin()
  191.  
RAW Paste Data

Adblocker detected! Please consider disabling it...

We've detected AdBlock Plus or some other adblocking software preventing Pastebin.com from fully loading.

We don't have any obnoxious sound, or popup ads, we actively block these annoying types of ads!

Please add Pastebin.com to your ad blocker whitelist or disable your adblocking software.

×