├── .gitignore ├── CMakeLists.txt ├── LICENSE ├── README.md ├── assets ├── README.md ├── spiral_stairs_mapping.gif └── straight_stairs_mapping.gif ├── config └── config_param.yaml ├── data └── dataset.txt ├── msg ├── MapManager.msg └── Polygon3D.msg ├── package.xml └── src ├── main.py ├── modules ├── __init__.py ├── image_processes_cupy.py ├── mapmanager.py └── ros_node.py └── read_dataset.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Ignore Python cache files and folders 2 | __pycache__/ 3 | *.py[cod] 4 | 5 | -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.0.2) 2 | project(polygon_mapping) 3 | 4 | find_package(catkin REQUIRED COMPONENTS 5 | roscpp 6 | std_msgs 7 | geometry_msgs 8 | message_generation 9 | ) 10 | 11 | add_message_files( 12 | FILES 13 | Polygon3D.msg 14 | MapManager.msg 15 | ) 16 | 17 | generate_messages( 18 | DEPENDENCIES 19 | std_msgs 20 | geometry_msgs 21 | ) 22 | 23 | catkin_package( 24 | CATKIN_DEPENDS message_runtime roscpp std_msgs geometry_msgs 25 | ) 26 | 27 | include_directories( 28 | ${catkin_INCLUDE_DIRS} 29 | ) 30 | 31 | 32 | 33 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 BTFrontier 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Real-time and efficient polygonal mapping designed for humanoid robots. 2 | 3 | ## Demo 4 | ### straight_stairs_mapping 5 | ![straight_stairs_mapping](assets/straight_stairs_mapping.gif) 6 | 7 | ### spiral_stairs_mapping 8 | ![spiral_stairs_mapping](assets/spiral_stairs_mapping.gif) 9 | 10 | ## Install 11 | 12 | ### Requirements 13 | - Tested on `Ubuntu 20.04 / ROS Noetic` 14 | 15 | ### Dependencies 16 | This package depends on 17 | - pyrealsense2 18 | - cupy 19 | - ... 20 | 21 | ### Installation procedure 22 | It is assumed that ROS is installed. 23 | 24 | 1. Clone to your catkin_ws 25 | ```bash 26 | mkdir -p catkin_ws/src 27 | cd catkin_ws/src 28 | git clone https://github.com/BTFrontier/polygon_mapping.git 29 | 30 | ``` 31 | 32 | 2. Install dependent packages. 33 | 34 | 35 | 3. Build a package. 36 | ```bash 37 | cd catkin_ws 38 | catkin build polygon_mapping 39 | ``` 40 | 41 | ## Datasets 42 | You can download the test dataset [here](https://1drv.ms/f/c/1e83680b5fbc1ae4/Et2MgY6eCHRMpczAZAwRXBUBvlHg70gRJopAoxf9fdi9vg?e=DQmDKZ). 43 | Once downloaded, extract the `dataset_39` and `dataset_71` folders into the following directory: 44 | ```bash 45 | catkin_ws/src/polygon_mapping/data 46 | ``` 47 | Now, you can run the test case with the following command: 48 | ```bash 49 | rosrun polygon_mapping read_dataset.py 50 | ``` 51 | You can modify the dataset directory in `read_dataset.py` to use the dataset of your choice. During the testing process, intermediate images will be saved in the `processed` subdirectory within the dataset directory, allowing you to review the results later. 52 | 53 | ## Running on a Real Robot 54 | 55 | In addition to the datasets, the system can be run on your own depth camera. By default, the code retrieves depth images from a RealSense camera. If you're using the L515, you can configure additional LiDAR parameters. If you're using another depth camera, you may need to modify the depth image input accordingly. Before running the mapping program, you can configure various parameters in the `config/config_param.yaml` file. 56 | 57 | ### External Odometry 58 | 59 | First, ensure that the robot has an odometry system or another method for estimating its own state. Then, using the relative pose of the depth camera to the odometry, you need to send the depth camera's pose in the odometry coordinate frame via ROS tf in real-time. The configuration for the odometry frame and depth camera frame listener can be modified in the `config/config_param.yaml` file. 60 | 61 | ### Camera Parameters 62 | 63 | This section includes parameters such as the serial number of the RealSense depth camera (important when using multiple cameras), depth image resolution, camera intrinsic parameters, etc. For the L515, adjusting several parameters can help improve the quality of the depth map. 64 | 65 | ### Algorithm Parameters 66 | 67 | This includes image processing parameters and RANSAC parameters. These parameters influence the quality and real-time performance of the output and can be adjusted based on your needs. 68 | 69 | ### Running the Command 70 | 71 | Once all parameters are configured, you can run the mapping program with the following command: 72 | 73 | ```bash 74 | rosrun polygon_mapping main.py 75 | ``` 76 | 77 | ## Citing 78 | 79 | If you find this code useful in your research, please consider citing our paper: 80 | [Real-Time Polygonal Semantic Mapping for Humanoid Robot Stair Climbing](https://arxiv.org/abs/2411.01919) 81 | 82 | ```bash 83 | @inproceedings{bin2024real, 84 | title={Real-Time Polygonal Semantic Mapping for Humanoid Robot Stair Climbing}, 85 | author={Bin, Teng and Yao, Jianming and Lam, Tin Lun and Zhang, Tianwei}, 86 | booktitle={2024 IEEE-RAS 23rd International Conference on Humanoid Robots (Humanoids)}, 87 | pages={866--873}, 88 | year={2024}, 89 | organization={IEEE} 90 | } 91 | ``` 92 | -------------------------------------------------------------------------------- /assets/README.md: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /assets/spiral_stairs_mapping.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BTFrontier/polygon_mapping/109c4fe22e80daa897d7511d4f627941e2a28bac/assets/spiral_stairs_mapping.gif -------------------------------------------------------------------------------- /assets/straight_stairs_mapping.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BTFrontier/polygon_mapping/109c4fe22e80daa897d7511d4f627941e2a28bac/assets/straight_stairs_mapping.gif -------------------------------------------------------------------------------- /config/config_param.yaml: -------------------------------------------------------------------------------- 1 | #### Basic parameters ######## 2 | 3 | save_data: True ## Whether to save the dataset 4 | 5 | #### Camera Parameters ########## 6 | camera_serial_id: 'f1231454' # Serial number of the depth camera used for terrain mapping 7 | depth_width: 320 8 | depth_height: 240 9 | L515_pre_processing_sharpening: 3.0 # Internal sharpening of the depth image by the camera 10 | L515_laser_power: 100.0 # Laser power 11 | L515_confidence_threshold: 2.0 # Confidence threshold (meaning unclear) 12 | 13 | ## Depth camera intrinsics 14 | depth_fx: 228.686 15 | depth_fy: 230.078 16 | depth_cx: 157.748 17 | depth_cy: 123.968 18 | 19 | rgb_width: 640 # RGB image is used for display only 20 | rgb_height: 480 21 | 22 | #### Image Processing ########### 23 | 24 | anisotropic_diffusion_num_iter: 60 # Number of iterations for anisotropic diffusion filtering 25 | anisotropic_diffusion_kappa: 134 26 | anisotropic_diffusion_gamma: 2 27 | 28 | use_sharpening: True # Sharpen normal map before edge detection 29 | 30 | Canny_thresh1: 50 31 | Canny_thresh2: 100 # thresh2 should be greater than thresh1 32 | 33 | closing_kernel: 13 # Size of the closing operation kernel for closing overly blurred edges 34 | 35 | ### RANSAC Parameters 36 | ran_max_trials: 20 ## Number of RANSAC fitting attempts 37 | ran_min_samples: 3 ## Minimum of 3 points required to fit a plane in RANSAC 38 | ran_residual_threshold: 0.01 ## Inlier distance threshold (m) 39 | ran_outlier_threshold: 0.05 ## Outlier tolerance ratio 40 | 41 | #### Map Manager ######### 42 | 43 | min_area_thresh: 5000 # Polygons with pixel area below this threshold will be ignored 44 | merge_dis_thresh: 0.05 # Distance threshold for merging polygons (m) 45 | 46 | map_frame: 'world' # The map frame used by the odometry source 47 | camera_frame: 'camera2_depth' 48 | -------------------------------------------------------------------------------- /data/dataset.txt: -------------------------------------------------------------------------------- 1 | 2 | you can download datasets from the following link 3 | 4 | https://1drv.ms/f/c/1e83680b5fbc1ae4/Et2MgY6eCHRMpczAZAwRXBUBvlHg70gRJopAoxf9fdi9vg?e=DQmDKZ 5 | -------------------------------------------------------------------------------- /msg/MapManager.msg: -------------------------------------------------------------------------------- 1 | std_msgs/Header header 2 | Polygon3D[] polygons 3 | -------------------------------------------------------------------------------- /msg/Polygon3D.msg: -------------------------------------------------------------------------------- 1 | std_msgs/Header header 2 | geometry_msgs/Point32[] vertices 3 | geometry_msgs/Vector3 normal 4 | float32[4] plane_eq 5 | float32 area 6 | -------------------------------------------------------------------------------- /package.xml: -------------------------------------------------------------------------------- 1 | 2 | 3 | polygon_mapping 4 | 0.0.0 5 | The polygon_mapping package 6 | BTFrontier 7 | TODO 8 | catkin 9 | roscpp 10 | std_msgs 11 | geometry_msgs 12 | message_generation 13 | roscpp 14 | std_msgs 15 | geometry_msgs 16 | message_runtime 17 | 18 | -------------------------------------------------------------------------------- /src/main.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import time 4 | import rospy 5 | import rospkg 6 | import numpy as np 7 | import copy 8 | import cv2 9 | import pyrealsense2 as rs 10 | import os 11 | import yaml 12 | 13 | from modules.mapmanager import Polygon3D, MapManager 14 | from modules.image_processes_cupy import process_polygon, project_points_to_plane, depth_to_point_cloud_region, fit_plane_ransac, anisotropic_diffusion, depth_to_normals, add_border_conditionally 15 | from modules.ros_node import MappingNode 16 | 17 | def nothing(x): 18 | """Placeholder function for trackbar.""" 19 | pass 20 | 21 | def load_config(config_file): 22 | """Load configuration from a YAML file.""" 23 | with open(config_file, 'r') as file: 24 | config = yaml.safe_load(file) 25 | return config 26 | 27 | def create_dataset_folder(base_path): 28 | """Create a unique dataset folder and subfolders for RGB and depth images.""" 29 | dataset_index = 1 30 | while True: 31 | dataset_name = f"dataset_{dataset_index}" 32 | dataset_path = os.path.join(base_path, dataset_name) 33 | if not os.path.exists(dataset_path): 34 | os.makedirs(os.path.join(dataset_path, "rgb"), exist_ok=True) 35 | os.makedirs(os.path.join(dataset_path, "depth"), exist_ok=True) 36 | return dataset_path 37 | dataset_index += 1 38 | 39 | def set_cv_config(thresh1=50, thresh2=100, num_iter=60, kappa=134, gamma=2): 40 | """Initialize OpenCV windows and trackbars for image processing parameters.""" 41 | cv2.namedWindow('Filtered Image with Largest Contour') 42 | cv2.createTrackbar('Canny thresh1', 'Filtered Image with Largest Contour', thresh1, max(thresh1, 1000), nothing) 43 | cv2.createTrackbar('Canny thresh2', 'Filtered Image with Largest Contour', thresh2, max(thresh2, 1000), nothing) 44 | 45 | cv2.namedWindow('smoothed_depth_img') 46 | cv2.createTrackbar('num_iter', 'smoothed_depth_img', num_iter, max(num_iter, 360), nothing) 47 | cv2.createTrackbar('kappa', 'smoothed_depth_img', kappa, max(kappa, 500), nothing) 48 | cv2.createTrackbar('gamma', 'smoothed_depth_img', gamma, max(gamma, 80), nothing) 49 | 50 | def main(): 51 | # Initialize ROS package manager 52 | rospack = rospkg.RosPack() 53 | package_path = rospack.get_path('polygon_mapping') # Get package path 54 | config_path = os.path.join(package_path, 'config', 'config_param.yaml') # Construct path to config file 55 | config_param = load_config(config_path) # Load configuration 56 | 57 | try: 58 | # Initialize map manager with z-axis merging threshold 59 | map_manager = MapManager(z_threshold=config_param['merge_dis_thresh']) 60 | node = MappingNode(map_manager, map_frame=config_param['map_frame'], camera_frame=config_param['camera_frame']) 61 | time.sleep(1) # Wait 3 seconds for listener setup 62 | 63 | # Initialize and configure RealSense L515 64 | pipeline = rs.pipeline() 65 | config = rs.config() 66 | config.enable_device(config_param['camera_serial_id']) 67 | device = config.resolve(pipeline).get_device() 68 | depth_sensor = device.first_depth_sensor() 69 | 70 | # Set depth sensor options 71 | preset_index = 0 72 | depth_sensor.set_option(rs.option.visual_preset, preset_index) 73 | 74 | # Enable color and depth streams 75 | config.enable_stream(rs.stream.color, config_param['rgb_width'], config_param['rgb_height'], rs.format.bgr8, 30) 76 | config.enable_stream(rs.stream.depth, config_param['depth_width'], config_param['depth_height'], rs.format.z16, 30) 77 | fx, fy, cx, cy = config_param['depth_fx'], config_param['depth_fy'], config_param['depth_cx'], config_param['depth_cy'] 78 | 79 | # Start pipeline streaming 80 | pipeline.start(config) 81 | 82 | # Set additional RealSense options if supported 83 | if depth_sensor.supports(rs.option.pre_processing_sharpening): 84 | depth_sensor.set_option(rs.option.pre_processing_sharpening, config_param['L515_pre_processing_sharpening']) 85 | print("Pre-processing sharpening set to ", config_param['L515_pre_processing_sharpening']) 86 | 87 | if depth_sensor.supports(rs.option.laser_power): 88 | depth_sensor.set_option(rs.option.laser_power, config_param['L515_laser_power']) 89 | print("Laser power set to ", config_param['L515_laser_power']) 90 | 91 | if depth_sensor.supports(rs.option.confidence_threshold): 92 | depth_sensor.set_option(rs.option.confidence_threshold, config_param['L515_confidence_threshold']) 93 | print("Confidence threshold set to ", config_param['L515_confidence_threshold']) 94 | 95 | # Set up OpenCV configurations 96 | set_cv_config(config_param['Canny_thresh1'], config_param['Canny_thresh2'], config_param['anisotropic_diffusion_num_iter'], config_param['anisotropic_diffusion_kappa'], config_param['anisotropic_diffusion_gamma']) 97 | closing_kernel = config_param['closing_kernel'] 98 | 99 | # Prepare dataset folder and file paths if data saving is enabled 100 | if config_param['save_data']: 101 | image_index = 1 102 | dataset_root = os.path.join(package_path, 'data') 103 | dataset_path = create_dataset_folder(dataset_root) 104 | os.makedirs(os.path.join(dataset_path, "rgb"), exist_ok=True) 105 | os.makedirs(os.path.join(dataset_path, "depth"), exist_ok=True) 106 | t_file_path = os.path.join(dataset_path, "transformation_matrix.txt") 107 | 108 | # Main loop 109 | while True: 110 | frames = pipeline.wait_for_frames() 111 | depth_frame = frames.get_depth_frame() 112 | color_frame = frames.get_color_frame() 113 | color_image = np.asanyarray(color_frame.get_data()) 114 | if not depth_frame: 115 | continue 116 | 117 | time_start_loop = time.time() 118 | depth_image = np.asanyarray(depth_frame.get_data()) 119 | timestamp = time.time() 120 | 121 | T_lock = copy.deepcopy(node.T) # Capture transformation matrix 122 | 123 | # Apply anisotropic diffusion to smooth depth image 124 | time_start = time.time() 125 | smoothed_depth_img = anisotropic_diffusion(depth_image) 126 | print('Anisotropic diffusion time cost', 1000 * (time.time() - time_start), 'ms') 127 | 128 | # Compute surface normals from the depth image 129 | normals = depth_to_normals(smoothed_depth_img, fx, fy, cx, cy) 130 | 131 | # Retrieve Canny threshold values from trackbars 132 | thresh1 = cv2.getTrackbarPos('Canny thresh1', 'Filtered Image with Largest Contour') 133 | thresh2 = cv2.getTrackbarPos('Canny thresh2', 'Filtered Image with Largest Contour') 134 | 135 | # Apply sharpening filter to enhance edges if enabled 136 | kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]]) 137 | if config_param['use_sharpening']: 138 | filtered_img = cv2.filter2D(normals, -1, kernel) 139 | else: 140 | filtered_img = normals 141 | 142 | # Add border to the filtered image 143 | filtered_img = add_border_conditionally(filtered_img) 144 | # Perform edge detection 145 | edges = cv2.Canny(filtered_img, thresh1, thresh2) 146 | 147 | # Use morphological closing to fill gaps in edges 148 | kernel = np.ones((closing_kernel, closing_kernel), np.uint8) 149 | edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel) 150 | 151 | # Find contours in the edge-detected image 152 | contours, hierarchy = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) 153 | min_area_thresh = config_param['min_area_thresh'] 154 | new_polygon_list = [] 155 | 156 | for i, contour in enumerate(contours): 157 | # Skip contours without parent contours 158 | if hierarchy[0][i][3] == -1: 159 | continue 160 | 161 | # Skip contours below area threshold 162 | if cv2.contourArea(contour) < min_area_thresh: 163 | continue 164 | 165 | # Simplify contours using polygon approximation 166 | epsilon = 0.02 * cv2.arcLength(contour, True) 167 | approx = cv2.approxPolyDP(contour, epsilon, True) 168 | if cv2.contourArea(approx) < min_area_thresh: 169 | continue 170 | 171 | # Draw contour and vertices 172 | cv2.drawContours(filtered_img, [approx], 0, (0, 255, 0), 2) 173 | for point in approx: 174 | cv2.circle(filtered_img, tuple(point[0]), 5, (0, 0, 255), -1) 175 | 176 | polygon_points = depth_to_point_cloud_region(smoothed_depth_img, approx.reshape(-1, 2), fx, fy, cx, cy) 177 | if polygon_points.size == 0: 178 | print("No valid points in the polygon region") 179 | continue 180 | 181 | try: 182 | # Fit a plane using RANSAC 183 | time_start = time.time() 184 | normal, d = fit_plane_ransac( 185 | polygon_points, 186 | max_trials=config_param['ran_max_trials'], 187 | min_samples=config_param['ran_min_samples'], 188 | residual_threshold=config_param['ran_residual_threshold'], 189 | outlier_threshold=config_param['ran_outlier_threshold'] 190 | ) 191 | print('fit_plane_ransac time cost', 1000 * (time.time() - time_start), 'ms') 192 | 193 | projected_points = project_points_to_plane((*normal, d), fx, fy, cx, cy, approx.reshape(-1, 2)) 194 | 195 | # Process polygon and transform to world coordinates 196 | angle, world_vertices = process_polygon(projected_points, normal, T_lock) 197 | if world_vertices is not None: 198 | new_polygon_list.append(world_vertices) 199 | 200 | except ValueError as e: 201 | print("Error in fitting plane:", e) 202 | 203 | # Add polygons to the map manager and estimate z-drift 204 | map_manager.add_polygon_list(new_polygon_list) 205 | 206 | # Display the polygon map 207 | if map_manager.polygons: 208 | map_img = map_manager.plot_polygons() 209 | cv2.imshow('Polygon Map', map_img) 210 | 211 | # Save images and transformation matrix if data saving is enabled 212 | if config_param['save_data']: 213 | rgb_filename = os.path.join(dataset_path, "rgb", f"{image_index:06d}.png") 214 | cv2.imwrite(rgb_filename, color_image) 215 | 216 | depth_filename = os.path.join(dataset_path, "depth", f"{image_index:06d}.png") 217 | cv2.imwrite(depth_filename, depth_image) 218 | 219 | with open(t_file_path, 'a') as t_file: 220 | t_file.write(f"{image_index},{T_lock.flatten()},{timestamp}\n") 221 | image_index += 1 222 | 223 | # Display results 224 | cv2.imshow('smoothed_depth_img', cv2.applyColorMap(cv2.convertScaleAbs(smoothed_depth_img, alpha=0.03), cv2.COLORMAP_JET)) 225 | cv2.imshow('normals Image', normals) 226 | cv2.imshow('Color Image', color_image) 227 | cv2.imshow('Edges', edges) 228 | cv2.imshow('Filtered Image with Largest Contour', filtered_img) 229 | print('Loop time cost', 1000 * (time.time() - time_start_loop), 'ms') 230 | 231 | if cv2.waitKey(1) & 0xFF == ord('q'): 232 | break 233 | 234 | rospy.spin() 235 | 236 | except rospy.ROSInterruptException: 237 | pass 238 | 239 | finally: 240 | # Ensure pipeline stops and all OpenCV windows close 241 | pipeline.stop() 242 | cv2.destroyAllWindows() 243 | 244 | if config_param.get('save_data', False): 245 | localtime = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) 246 | map_manager.save_polygons_to_file(os.path.join(dataset_path, f'{localtime}_polygons.txt')) 247 | cv2.imwrite(os.path.join(dataset_path, f'{localtime}_map_img.jpg'), map_img) 248 | 249 | if __name__ == '__main__': 250 | main() 251 | -------------------------------------------------------------------------------- /src/modules/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/BTFrontier/polygon_mapping/109c4fe22e80daa897d7511d4f627941e2a28bac/src/modules/__init__.py -------------------------------------------------------------------------------- /src/modules/image_processes_cupy.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import numpy as np 4 | import copy 5 | import cv2 6 | import pyrealsense2 as rs 7 | from skimage import morphology 8 | import cupy as cp 9 | import scipy.ndimage 10 | 11 | def transform_point(point, T): 12 | """ Transform a point from camera coordinates to world coordinates. """ 13 | point_h = np.append(point, 1) # Convert point to homogeneous coordinates 14 | transformed_point_h = np.dot(T, point_h) 15 | return transformed_point_h[:3] 16 | 17 | def transform_vector(vector, T): 18 | """ Transform a vector from camera coordinates to world coordinates. """ 19 | # Only consider the rotation part, so we take the upper-left 3x3 part of T 20 | R = T[:3, :3] 21 | transformed_vector = np.dot(R, vector) 22 | if transformed_vector[2] < 0: 23 | transformed_vector = -transformed_vector 24 | return transformed_vector 25 | 26 | def angle_between_vectors(v1, v2): 27 | """ Calculate the angle (in radians) between two vectors. """ 28 | dot_product = np.dot(v1, v2) 29 | norms = np.linalg.norm(v1) * np.linalg.norm(v2) 30 | cos_angle = dot_product / norms 31 | angle = np.arccos(np.clip(cos_angle, -1.0, 1.0)) 32 | return angle 33 | 34 | def process_polygon(vertices, normal, T): 35 | """ 36 | Args: 37 | - vertices: Polygon vertices in camera coordinates, Nx3 numpy array 38 | - normal: Normalized normal vector of the polygon plane, 3-element numpy array 39 | 40 | Returns: 41 | - angle: Angle between the normal vector and the z-axis in world coordinates (degrees) 42 | - world_vertices: Polygon vertices in world coordinates 43 | """ 44 | 45 | # Transform the normal vector from camera coordinates to world coordinates 46 | world_normal = transform_vector(normal, T) 47 | 48 | # Calculate the angle between the normal vector and the z-axis in world coordinates 49 | world_z_axis = np.array([0, 0, 1]) 50 | angle_radians = angle_between_vectors(world_normal, world_z_axis) 51 | angle_degrees = np.degrees(angle_radians) 52 | 53 | # Check if the angle is less than 25° 54 | thresh_angle = 25 55 | if angle_degrees < thresh_angle: 56 | # Transform vertices from camera coordinates to world coordinates 57 | world_vertices = np.array([transform_point(vertex, T) for vertex in vertices]) 58 | return angle_degrees, world_vertices 59 | else: 60 | return angle_degrees, None 61 | 62 | def project_points_to_plane(plane_eq, fx, fy, cx, cy, uv_coords): 63 | """ 64 | Projects points from image coordinates (u, v) to a specified plane in camera coordinates. 65 | 66 | Args: 67 | - plane_eq (tuple): Coefficients (A, B, C, D) of the plane equation Ax + By + Cz + D = 0. 68 | - fx, fy, cx, cy: Camera intrinsic parameters. 69 | - uv_coords (np.ndarray): Array of image coordinates (u, v) with shape (N, 2). 70 | 71 | Returns: 72 | - np.ndarray: Array of 3D coordinates (x, y, z) projected onto the plane with shape (N, 3). 73 | """ 74 | A, B, C, D = plane_eq 75 | points_3d = [] 76 | 77 | for (u, v) in uv_coords: 78 | # Convert image coordinates to normalized camera coordinates 79 | x_norm = (u - cx) / fx 80 | y_norm = (v - cy) / fy 81 | 82 | # Calculate the z-coordinate using the plane equation 83 | z = -D / (A * x_norm + B * y_norm + C) 84 | x = x_norm * z 85 | y = y_norm * z 86 | 87 | points_3d.append([x, y, z]) 88 | 89 | return np.array(points_3d) 90 | 91 | def depth_to_point_cloud_region(depth_map, polygon_vertices, fx, fy, cx, cy): 92 | """Generate a point cloud for a region specified by polygon vertices in the depth map.""" 93 | height, width = depth_map.shape 94 | mask = np.zeros(depth_map.shape, dtype=np.uint8) 95 | cv2.fillPoly(mask, [polygon_vertices], 1) 96 | 97 | mask = cp.asarray(mask) 98 | depth_map = cp.asarray(depth_map) 99 | 100 | indices = cp.where(mask == 1) 101 | z = 0.00025 * depth_map[indices] 102 | i, j = indices[1], indices[0] 103 | x = (i - cx) * z / fx 104 | y = (j - cy) * z / fy 105 | point_cloud = cp.stack((x, y, z), axis=-1) 106 | return point_cloud 107 | 108 | def fit_plane_ransac(points, max_trials=100, min_samples=3, residual_threshold=0.005, outlier_threshold=0.05): 109 | """Estimate a plane from points using RANSAC with Cupy for GPU acceleration.""" 110 | best_model = None 111 | best_inliers = 0 112 | 113 | # Pre-generate random samples 114 | sample_indices = cp.random.choice(points.shape[0], (max_trials, min_samples), replace=False) 115 | sample_points = points[sample_indices] 116 | 117 | normals = cp.zeros((max_trials, 3)) 118 | ds = cp.zeros(max_trials) 119 | num_inliers = cp.zeros(max_trials) 120 | 121 | for i in range(max_trials): 122 | if sample_points[i].shape[0] < 3: 123 | continue 124 | 125 | p1, p2, p3 = sample_points[i][:3] 126 | v1 = p2 - p1 127 | v2 = p3 - p1 128 | normal = cp.cross(v1, v2) 129 | normal = normal / cp.linalg.norm(normal) 130 | 131 | d = -cp.dot(normal, p1) 132 | 133 | normals[i] = normal 134 | ds[i] = d 135 | 136 | distances = cp.abs(cp.dot(points, normals.T) + ds) 137 | inliers = distances < residual_threshold 138 | num_inliers = cp.sum(inliers, axis=0) 139 | 140 | best_index = cp.argmax(num_inliers) 141 | best_inliers = num_inliers[best_index] 142 | 143 | if best_inliers == 0: 144 | raise ValueError("Could not find a valid plane") 145 | 146 | best_model = (normals[best_index], ds[best_index]) 147 | 148 | normal, d = best_model 149 | outlier_ratio = 1 - (best_inliers / points.shape[0]) 150 | 151 | if outlier_ratio > outlier_threshold: 152 | raise ValueError(f"Outlier ratio too high: {outlier_ratio:.2f}, threshold: {outlier_threshold}") 153 | 154 | if d > 0: 155 | normal = -normal 156 | d = -d 157 | 158 | return cp.asnumpy(normal), cp.asnumpy(d) 159 | 160 | def anisotropic_diffusion(depth_img, num_iter=20, kappa=50, gamma=0.1, option=1): 161 | """Perform anisotropic diffusion on a depth image using GPU with Cupy.""" 162 | num_iter = cv2.getTrackbarPos('num_iter', 'smoothed_depth_img') 163 | kappa = cv2.getTrackbarPos('kappa', 'smoothed_depth_img') 164 | gamma = 0.1 * cv2.getTrackbarPos('gamma', 'smoothed_depth_img') 165 | img = cp.array(depth_img, dtype=cp.float32) 166 | h, w = img.shape 167 | 168 | for i in range(num_iter): 169 | # Compute gradients 170 | deltaN = cp.zeros_like(img) 171 | deltaS = cp.zeros_like(img) 172 | deltaE = cp.zeros_like(img) 173 | deltaW = cp.zeros_like(img) 174 | 175 | deltaN[1:, :] = img[:-1, :] - img[1:, :] 176 | deltaS[:-1, :] = img[1:, :] - img[:-1, :] 177 | deltaE[:, :-1] = img[:, 1:] - img[:, :-1] 178 | deltaW[:, 1:] = img[:, :-1] - img[:, 1:] 179 | 180 | # Compute diffusion coefficients 181 | if option == 1: 182 | cN = cp.exp(-(deltaN / kappa) ** 2) 183 | cS = cp.exp(-(deltaS / kappa) ** 2) 184 | cE = cp.exp(-(deltaE / kappa) ** 2) 185 | cW = cp.exp(-(deltaW / kappa) ** 2) 186 | elif option == 2: 187 | cN = 1. / (1. + (deltaN / kappa) ** 2) 188 | cS = 1. / (1. + (deltaS / kappa) ** 2) 189 | cE = 1. / (1. + (deltaE / kappa) ** 2) 190 | cW = 1. / (1. + (deltaW / kappa) ** 2) 191 | 192 | # Update image 193 | img += gamma * (cN * deltaN + cS * deltaS + cE * deltaE + cW * deltaW) 194 | 195 | return cp.asnumpy(img) 196 | 197 | def depth_to_normals(depth_img, fx, fy, cx, cy): 198 | """Calculate surface normals from a depth image using the Sobel operator.""" 199 | du = cv2.Sobel(depth_img, cv2.CV_64F, 1, 0, ksize=5) 200 | dv = cv2.Sobel(depth_img, cv2.CV_64F, 0, 1, ksize=5) 201 | 202 | X_u = np.ones_like(depth_img, dtype=np.float64) / fx 203 | Y_v = np.ones_like(depth_img, dtype=np.float64) / fy 204 | Z_u = du 205 | Z_v = dv 206 | 207 | # Calculate normal vector components 208 | nx = Y_v * Z_u 209 | ny = Z_v * X_u 210 | nz = -(X_u * Y_v) 211 | 212 | # Normalize the normal vectors 213 | norm = np.sqrt(nx**2 + ny**2 + nz**2 + 1e-8) 214 | nx /= norm 215 | ny /= norm 216 | nz /= norm 217 | 218 | # Handle invalid depth (e.g., zero depth) 219 | mask = depth_img == 0 220 | nx[mask] = 0 221 | ny[mask] = 0 222 | nz[mask] = 0 223 | 224 | # Map normal vector components to RGB color space in the range [0, 255] 225 | normals_rgb = cv2.merge(((nx + 1) / 2 * 255, (ny + 1) / 2 * 255, (nz + 1) / 2 * 255)).astype(np.uint8) 226 | return normals_rgb 227 | 228 | def add_border_conditionally(image): 229 | """Add a border conditionally around the image if certain criteria are met.""" 230 | rows, cols, _ = image.shape 231 | new_image = image.copy() 232 | 233 | # Create masks to check for border conditions 234 | mask_top = np.any(image[0, :, :] != [127, 127, 127], axis=-1) 235 | mask_bottom = np.any(image[-1, :, :] != [127, 127, 127], axis=-1) 236 | mask_left = np.any(image[:, 0, :] != [127, 127, 127], axis=-1) 237 | mask_right = np.any(image[:, -1, :] != [127, 127, 127], axis=-1) 238 | 239 | # Apply conditional border 240 | new_image[0, mask_top] = [255, 255, 255] 241 | new_image[-1, mask_bottom] = [255, 255, 255] 242 | new_image[mask_left, 0] = [255, 255, 255] 243 | new_image[mask_right, -1] = [255, 255, 255] 244 | 245 | return new_image 246 | -------------------------------------------------------------------------------- /src/modules/mapmanager.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import time 4 | import numpy as np 5 | import copy 6 | from shapely.geometry import Polygon as ShapelyPolygon, MultiPolygon 7 | from shapely.ops import unary_union 8 | import cv2 9 | from scipy.spatial import ConvexHull 10 | from scipy.optimize import minimize 11 | 12 | ########################################################################## Map Manager ########################################### 13 | class Polygon3D: 14 | def __init__(self, vertices): 15 | self.vertices = np.array(vertices) 16 | self.normal = self.compute_normal() 17 | self.plane_eq = self.compute_plane_equation() 18 | self.area = self.compute_3d_area() 19 | self.xy_projection = ShapelyPolygon(self.vertices[:, :2]) 20 | if not self.xy_projection.is_valid: 21 | raise ValueError(f"Invalid polygon with vertices: {vertices}") 22 | 23 | def compute_normal(self): 24 | if len(self.vertices) < 3: 25 | return np.array([np.nan, np.nan, np.nan]) 26 | 27 | p1, p2, p3 = self.vertices[:3] 28 | v1 = p2 - p1 29 | v2 = p3 - p1 30 | normal = np.cross(v1, v2) 31 | 32 | if np.linalg.norm(normal) == 0: 33 | return np.array([np.nan, np.nan, np.nan]) 34 | 35 | normal = normal / np.linalg.norm(normal) 36 | 37 | # Ensure the z component of the normal is non-negative 38 | if normal[2] < 0: 39 | normal = -normal 40 | 41 | return normal 42 | 43 | def compute_plane_equation(self): 44 | if np.any(np.isnan(self.normal)): 45 | return (np.nan, np.nan, np.nan, np.nan) 46 | 47 | point = self.vertices[0] 48 | a, b, c = self.normal 49 | d = -np.dot(self.normal, point) 50 | return (a, b, c, d) 51 | 52 | def compute_3d_area(self): 53 | if len(self.vertices) < 3: 54 | return np.nan 55 | 56 | total_area = 0.0 57 | origin = self.vertices[0] 58 | 59 | for i in range(1, len(self.vertices) - 1): 60 | v1 = self.vertices[i] - origin 61 | v2 = self.vertices[i + 1] - origin 62 | cross_product = np.cross(v1, v2) 63 | triangle_area = np.linalg.norm(cross_product) / 2 64 | total_area += triangle_area 65 | 66 | return total_area 67 | 68 | def merge_with(self, other): 69 | all_vertices = np.vstack((self.vertices, other.vertices)) 70 | hull = ConvexHull(all_vertices[:, :2]) 71 | hull_vertices = all_vertices[hull.vertices] 72 | xy_hull_vertices = hull_vertices[:, :2] 73 | 74 | new_normal = (self.normal * self.area + other.normal * other.area) / (self.area + other.area) 75 | 76 | if np.linalg.norm(new_normal) == 0: 77 | return self # Do not merge if normal is invalid 78 | 79 | new_normal /= np.linalg.norm(new_normal) 80 | 81 | # Compute d directly by averaging distances to the plane 82 | all_distances = np.dot(all_vertices, new_normal) 83 | d_opt = -np.mean(all_distances) 84 | 85 | new_vertices = np.column_stack((xy_hull_vertices, np.zeros(len(xy_hull_vertices)))) 86 | for i, vertex in enumerate(new_vertices): 87 | xy_vertex = vertex[:2] 88 | new_vertices[i, 2] = (-d_opt - new_normal[0] * xy_vertex[0] - new_normal[1] * xy_vertex[1]) / new_normal[2] 89 | 90 | return Polygon3D(new_vertices) 91 | 92 | def compute_z_for_xy(self, xy): 93 | a, b, c, d = self.plane_eq 94 | x, y = xy 95 | return (-d - a * x - b * y) / c 96 | 97 | def to_string(self): 98 | return f"Vertices: {self.vertices.tolist()}\nNormal: {self.normal.tolist()}\nPlane equation: {self.plane_eq}\nArea: {self.area}\n" 99 | 100 | class MapManager: 101 | def __init__(self, z_threshold=0.02): 102 | self.polygons = [] # Stores existing polygons 103 | self.z_threshold = z_threshold # Distance threshold for Z-axis merging 104 | self.z_drift = 0.0 # Estimated Z-axis drift 105 | self.single_frame = [] 106 | 107 | def add_polygon_list(self, new_polygon_list): 108 | # 1. Convert incoming polygon vertices to Polygon3D objects 109 | new_polygons = [] 110 | for vertices in new_polygon_list: 111 | try: 112 | new_polygon = Polygon3D(vertices) 113 | if np.any(np.isnan(new_polygon.normal)): 114 | print("Skipping invalid polygon with vertices:", vertices) 115 | continue 116 | new_polygons.append(new_polygon) 117 | except ValueError as e: 118 | print(e) 119 | 120 | # 2. Initially compensate new polygons with the last z_drift 121 | compensated_new_polygons = [self.apply_z_drift(polygon, self.z_drift) for polygon in new_polygons] 122 | 123 | # 3. Find overlapping parts between new and existing polygons to prepare for z_drift optimization 124 | merge_candidates = [] 125 | for new_polygon in compensated_new_polygons: 126 | for existing_polygon in self.polygons: 127 | if self.is_overlap_in_xy(new_polygon, existing_polygon): 128 | z_distance = abs(self.compute_average_z_distance(new_polygon, existing_polygon)) 129 | if z_distance <= self.z_threshold: 130 | merge_candidates.append((new_polygon, existing_polygon)) 131 | break 132 | 133 | # 4. If there are overlapping polygons, optimize z_drift 134 | if merge_candidates: 135 | print("1 self.z_drift (before optimization):", self.z_drift) 136 | z_drift_increment = self.optimize_z_drift(merge_candidates) # Obtain increment 137 | self.z_drift += z_drift_increment # Update z_drift incrementally 138 | print("2 self.z_drift (after optimization):", self.z_drift) 139 | 140 | # 5. Re-compensate new polygons with the updated z_drift 141 | compensated_new_polygons = [self.apply_z_drift(polygon, self.z_drift) for polygon in new_polygons] 142 | 143 | # 6. Merge compensated polygons into the map manager 144 | for new_polygon in compensated_new_polygons: 145 | self.merge_polygon(new_polygon) 146 | 147 | def apply_z_drift(self, polygon, z_drift): 148 | """Apply Z-axis drift compensation to a polygon""" 149 | compensated_vertices = polygon.vertices.copy() 150 | compensated_vertices[:, 2] += z_drift # Compensate only on Z-axis 151 | return Polygon3D(compensated_vertices) 152 | 153 | def merge_polygon(self, new_polygon): 154 | """Merge the new polygon with existing polygons""" 155 | try: 156 | if np.any(np.isnan(new_polygon.normal)): 157 | print("Skipping invalid polygon with vertices:", new_polygon.vertices) 158 | return 159 | 160 | merged = False 161 | while True: 162 | to_merge = None 163 | for i, existing_polygon in enumerate(self.polygons): 164 | if self.is_overlap_in_xy(new_polygon, existing_polygon): 165 | z_distance = abs(self.compute_average_z_distance(new_polygon, existing_polygon)) 166 | if z_distance <= self.z_threshold: 167 | to_merge = existing_polygon 168 | break 169 | 170 | if to_merge is not None: 171 | new_polygon = new_polygon.merge_with(to_merge) 172 | self.polygons.remove(to_merge) 173 | merged = True 174 | else: 175 | break 176 | 177 | self.polygons.append(new_polygon) 178 | if not merged: 179 | self.merge_polygons() 180 | except ValueError as e: 181 | print(e) 182 | 183 | def is_overlap_in_xy(self, poly1, poly2): 184 | """Check if two polygons overlap in the XY plane""" 185 | return poly1.xy_projection.intersects(poly2.xy_projection) 186 | 187 | def compute_average_z_distance(self, poly1, poly2): 188 | """Calculate the average Z-axis height difference over the overlap region in XY plane""" 189 | intersected_points = poly1.xy_projection.intersection(poly2.xy_projection) 190 | if intersected_points.is_empty: 191 | return float('inf') 192 | 193 | points = [] 194 | if isinstance(intersected_points, MultiPolygon): 195 | for sub_poly in intersected_points: 196 | points.extend(sub_poly.exterior.coords) 197 | else: 198 | points = list(intersected_points.exterior.coords) 199 | 200 | points = np.array(points) 201 | 202 | if points.size == 0: 203 | return float('inf') 204 | 205 | z_distances = [] 206 | for point in points: 207 | xy = point[:2] 208 | z1 = poly1.compute_z_for_xy(xy) 209 | z2 = poly2.compute_z_for_xy(xy) 210 | z_distances.append(abs(z1 - z2)) 211 | 212 | return np.mean(z_distances) 213 | 214 | def optimize_z_drift(self, merge_candidates): 215 | """Optimize Z-axis drift to minimize height difference for overlapping polygon pairs""" 216 | 217 | def cost_function(z_drift_increment): 218 | total_cost = 0.0 219 | for new_polygon, existing_polygon in merge_candidates: 220 | # Apply z_drift_increment to the new polygon 221 | compensated_polygon = self.apply_z_drift(new_polygon, self.z_drift + z_drift_increment) 222 | # Calculate Z height difference in overlap region 223 | z_distance = self.compute_average_z_distance(compensated_polygon, existing_polygon) 224 | # Weight by area of new polygon 225 | area = new_polygon.area 226 | # Cost = area-weighted square of height difference 227 | total_cost += area * z_distance ** 2 228 | return total_cost 229 | 230 | # Optimize z_drift increment using scipy.optimize.minimize 231 | result = minimize(cost_function, 0.0, method='BFGS') 232 | return result.x[0] # Return optimized z_drift increment 233 | 234 | def merge_polygons(self): 235 | """Merge all overlapping polygons in the existing set""" 236 | i = 0 237 | while i < len(self.polygons): 238 | polygon = self.polygons[i] 239 | j = i + 1 240 | while j < len(self.polygons): 241 | other_polygon = self.polygons[j] 242 | if self.is_overlap_in_xy(polygon, other_polygon): 243 | z_distance = abs(self.compute_average_z_distance(polygon, other_polygon)) 244 | if z_distance <= self.z_threshold: 245 | merged_polygon = polygon.merge_with(other_polygon) 246 | self.polygons[i] = merged_polygon 247 | self.polygons.pop(j) 248 | # Re-check the merged polygon 249 | j = i + 1 250 | else: 251 | j += 1 252 | else: 253 | j += 1 254 | i += 1 255 | 256 | def plot_polygons(self): 257 | img = np.ones((600, 600, 3), dtype=np.uint8) * 255 258 | all_points = np.vstack([poly.vertices[:, :2] for poly in self.polygons]) 259 | min_x, min_y = np.min(all_points, axis=0) 260 | max_x, max_y = np.max(all_points, axis=0) 261 | scale_x = 500 / (max_x - min_x) 262 | scale_y = 500 / (max_y - min_y) 263 | scale = min(scale_x, scale_y) 264 | offset_x = (600 - (max_x - min_x) * scale) / 2 265 | offset_y = (600 - (max_y - min_y) * scale) / 2 266 | 267 | for poly in self.polygons: 268 | pts = np.array(poly.xy_projection.exterior.coords, np.float32) 269 | pts[:, 0] = (pts[:, 0] - min_x) * scale + offset_x 270 | pts[:, 1] = (max_y - pts[:, 1]) * scale + offset_y # Flip y coordinate 271 | pts = pts.astype(np.int32) 272 | pts = pts.reshape((-1, 1, 2)) 273 | cv2.polylines(img, [pts], isClosed=True, color=(0, 255, 0), thickness=2) 274 | return img 275 | 276 | def save_polygons_to_file(self, filename): 277 | with open(filename, 'w') as f: 278 | i = 0 279 | for polygon in self.polygons: 280 | f.write(str(i) + "\n") 281 | i += 1 282 | f.write(polygon.to_string() + "\n") 283 | -------------------------------------------------------------------------------- /src/modules/ros_node.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import rospy 4 | import tf2_ros 5 | import numpy as np 6 | from geometry_msgs.msg import TransformStamped 7 | from nav_msgs.msg import Path 8 | import tf.transformations 9 | import copy 10 | 11 | from polygon_mapping.msg import Polygon3D as Polygon3D_msg 12 | from polygon_mapping.msg import MapManager as MapManager_msg 13 | from geometry_msgs.msg import Point32, Vector3 14 | 15 | from visualization_msgs.msg import Marker 16 | from geometry_msgs.msg import Point 17 | from geometry_msgs.msg import PoseStamped 18 | import tf.transformations as tf_trans 19 | 20 | ############################################################### Listen to TF and publish map ######################################################## 21 | class MappingNode: 22 | def __init__(self, map_manager, map_frame='world', camera_frame='camera2_depth', frequency=30.0): 23 | # Initialize the ROS node 24 | rospy.init_node('polygon_node', anonymous=True) 25 | 26 | self.T = np.eye(4) 27 | 28 | # Create TF buffer and listener 29 | self.tfBuffer = tf2_ros.Buffer() 30 | self.listener = tf2_ros.TransformListener(self.tfBuffer) 31 | 32 | self.map_frame = map_frame 33 | self.camera_frame = camera_frame 34 | 35 | # Set callback function for transformations 36 | self.set_transform_callback() 37 | 38 | # Connect MapManager to ROS 39 | self.map_manager = map_manager 40 | 41 | self.frequency = frequency 42 | self.pub = rospy.Publisher('/map_manager', MapManager_msg, queue_size=5) 43 | 44 | # Create a publisher for visualization markers 45 | self.previous_marker_count = 0 # Track the previous number of markers published 46 | self.marker_pub = rospy.Publisher('/polygon_markers', Marker, queue_size=20) 47 | 48 | # Initialize pose publishers 49 | self.pose_pub = rospy.Publisher('/camera_pose', PoseStamped, queue_size=1) 50 | self.path_pub = rospy.Publisher('/camera_path', Path, queue_size=10) 51 | self.path = Path() 52 | self.path.header.frame_id = "world" 53 | 54 | # Raw (uncompensated) pose publisher 55 | self.pose_raw_pub = rospy.Publisher('/camera_pose_raw', PoseStamped, queue_size=1) 56 | self.path_raw_pub = rospy.Publisher('/camera_raw_path', Path, queue_size=10) 57 | self.path_raw = Path() 58 | self.path_raw.header.frame_id = "world" 59 | 60 | self.timer = rospy.Timer(rospy.Duration(1.0 / self.frequency), self.publish_polygons) 61 | 62 | def set_transform_callback(self): 63 | self.timer = rospy.Timer(rospy.Duration(1.0 / 30.0), self.transform_callback) 64 | 65 | def transform_callback(self, event): 66 | try: 67 | # Retrieve the transform from /camera2_link to /world 68 | trans = self.tfBuffer.lookup_transform(self.map_frame, self.camera_frame, rospy.Time(0)) 69 | self.T = self.transform_to_matrix(trans) 70 | # rospy.loginfo("Transform from /camera2_link to /world: \n%s" % T) 71 | except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException): 72 | rospy.logwarn("TF lookup failed") 73 | 74 | def transform_to_matrix(self, trans): 75 | # Convert TransformStamped to a 4x4 transformation matrix 76 | translation = trans.transform.translation 77 | rotation = trans.transform.rotation 78 | 79 | q = [rotation.x, rotation.y, rotation.z, rotation.w] 80 | T = tf.transformations.quaternion_matrix(q) 81 | T[0, 3] = translation.x 82 | T[1, 3] = translation.y 83 | T[2, 3] = translation.z 84 | 85 | return T 86 | 87 | def publish_polygons(self, event): 88 | # Publish the global map 89 | msg = MapManager_msg() 90 | for polygon in self.map_manager.polygons: 91 | poly_msg = Polygon3D_msg() 92 | for vertex in polygon.vertices: 93 | point = Point32() 94 | point.x, point.y, point.z = vertex 95 | poly_msg.vertices.append(point) 96 | 97 | normal = Vector3() 98 | normal.x, normal.y, normal.z = polygon.normal 99 | poly_msg.normal = normal 100 | 101 | poly_msg.plane_eq = polygon.plane_eq 102 | poly_msg.area = polygon.area 103 | 104 | msg.polygons.append(poly_msg) 105 | self.pub.publish(msg) 106 | 107 | T_lock = copy.deepcopy(self.T) 108 | self.publish_camera_pose(T_lock) 109 | 110 | self.publish_markers() 111 | 112 | def publish_markers(self): 113 | # Assume self.map_manager.polygons is updated 114 | for i, polygon in enumerate(self.map_manager.polygons): 115 | marker = Marker() 116 | marker.header.frame_id = "world" 117 | marker.header.stamp = rospy.Time.now() 118 | marker.ns = "polygons" 119 | marker.id = i 120 | marker.type = Marker.LINE_STRIP 121 | marker.action = Marker.ADD 122 | marker.lifetime = rospy.Duration(0.5) # Marker automatically disappears after 0.5 seconds if not updated 123 | marker.pose.orientation.w = 1.0 124 | marker.scale.x = 0.05 125 | marker.color.r = 0.0 126 | marker.color.g = 1.0 127 | marker.color.b = 0.0 128 | marker.color.a = 1.0 129 | 130 | # Add polygon vertices 131 | for vertex in polygon.vertices: 132 | point = Point() 133 | point.x = vertex[0] 134 | point.y = vertex[1] 135 | point.z = vertex[2] 136 | marker.points.append(point) 137 | if len(polygon.vertices) > 2: 138 | point = Point() 139 | point.x = polygon.vertices[0][0] 140 | point.y = polygon.vertices[0][1] 141 | point.z = polygon.vertices[0][2] 142 | marker.points.append(point) # Close the polygon 143 | 144 | self.marker_pub.publish(marker) 145 | 146 | def publish_camera_pose(self, T): 147 | # Extract translation vector 148 | translation = T[0:3, 3] 149 | 150 | # Extract rotation matrix and convert to quaternion, keeping yaw 151 | rotation_matrix = T[0:3, 0:3] 152 | euler_angles = tf_trans.euler_from_matrix(rotation_matrix, 'sxyz') 153 | yaw = euler_angles[2] + 3.14159 / 2 154 | 155 | # Set roll and pitch to 0, reconstruct quaternion 156 | quaternion = tf_trans.quaternion_from_euler(0, 0, yaw) 157 | 158 | # Apply drift compensation to Z axis 159 | translation[2] += self.map_manager.z_drift 160 | 161 | # Create PoseStamped message 162 | pose_msg = PoseStamped() 163 | pose_msg.header.stamp = rospy.Time.now() 164 | pose_msg.header.frame_id = "world" 165 | 166 | # Set position 167 | pose_msg.pose.position.x = translation[0] 168 | pose_msg.pose.position.y = translation[1] 169 | pose_msg.pose.position.z = translation[2] 170 | 171 | # Set orientation (quaternion) 172 | pose_msg.pose.orientation.x = quaternion[0] 173 | pose_msg.pose.orientation.y = quaternion[1] 174 | pose_msg.pose.orientation.z = quaternion[2] 175 | pose_msg.pose.orientation.w = quaternion[3] 176 | 177 | # Publish the camera pose 178 | self.pose_pub.publish(pose_msg) 179 | 180 | # Update the path 181 | self.path.poses.append(pose_msg) 182 | self.path.header.stamp = rospy.Time.now() 183 | self.path_pub.publish(self.path) 184 | 185 | pose_raw_msg = copy.deepcopy(pose_msg) 186 | pose_raw_msg.pose.position.z -= self.map_manager.z_drift 187 | self.pose_raw_pub.publish(pose_raw_msg) 188 | 189 | # Update the raw path 190 | self.path_raw.poses.append(pose_raw_msg) 191 | self.path_raw.header.stamp = rospy.Time.now() 192 | self.path_raw_pub.publish(self.path_raw) 193 | -------------------------------------------------------------------------------- /src/read_dataset.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import time 4 | import rospy 5 | import rospkg 6 | import tf2_ros 7 | import numpy as np 8 | from geometry_msgs.msg import TransformStamped 9 | from nav_msgs.msg import Path 10 | import tf.transformations 11 | from shapely.geometry import Polygon as ShapelyPolygon, MultiPolygon 12 | from shapely.ops import unary_union 13 | import cv2 14 | from scipy.spatial import ConvexHull 15 | import copy 16 | from polygon_mapping.msg import Polygon3D as Polygon3D_msg 17 | from polygon_mapping.msg import MapManager as MapManager_msg 18 | from geometry_msgs.msg import Point32, Vector3 19 | from visualization_msgs.msg import Marker 20 | from geometry_msgs.msg import Point, PoseStamped 21 | import tf.transformations as tf_trans 22 | 23 | from modules.mapmanager import Polygon3D, MapManager 24 | from modules.image_processes_cupy import ( 25 | process_polygon, project_points_to_plane, depth_to_point_cloud_region, 26 | fit_plane_ransac, anisotropic_diffusion, depth_to_normals, add_border_conditionally 27 | ) 28 | 29 | import pyrealsense2 as rs 30 | from skimage import morphology 31 | import os 32 | import re 33 | import glob 34 | 35 | # Global transformation matrix 36 | T = np.eye(4) 37 | T_lock = np.eye(4) 38 | 39 | class PosePublishNode: 40 | def __init__(self, map_manager, frequency=30.0): 41 | # Initialize the ROS node 42 | rospy.init_node('polygon_node', anonymous=True) 43 | global T 44 | 45 | # Create TF buffer and listener 46 | self.tfBuffer = tf2_ros.Buffer() 47 | self.listener = tf2_ros.TransformListener(self.tfBuffer) 48 | 49 | # Link MapManager to ROS 50 | self.map_manager = map_manager 51 | self.frequency = frequency 52 | self.pub = rospy.Publisher('/map_manager', MapManager_msg, queue_size=5) 53 | self.marker_pub = rospy.Publisher('/polygon_markers', Marker, queue_size=100) 54 | 55 | # Initialize pose and path publishers 56 | self.pose_pub = rospy.Publisher('/camera_pose', PoseStamped, queue_size=1) 57 | self.path_pub = rospy.Publisher('/camera_path', Path, queue_size=10) 58 | self.path = Path() 59 | self.path.header.frame_id = "world" 60 | 61 | self.pose_raw_pub = rospy.Publisher('/camera_pose_raw', PoseStamped, queue_size=1) 62 | self.path_raw_pub = rospy.Publisher('/camera_raw_path', Path, queue_size=10) 63 | self.path_raw = Path() 64 | self.path_raw.header.frame_id = "world" 65 | 66 | # Drift counters for tracking positive and negative drift during stair climbing 67 | self.positive_drift_cont = 0 68 | self.nagetive_drift_cont = 0 69 | 70 | # Timer for polygon publishing 71 | self.timer = rospy.Timer(rospy.Duration(1.0 / self.frequency), self.publish_polygons) 72 | 73 | def publish_polygons(self, event): 74 | """Publish global map with polygons.""" 75 | msg = MapManager_msg() 76 | for polygon in self.map_manager.polygons: 77 | poly_msg = Polygon3D_msg() 78 | for vertex in polygon.vertices: 79 | point = Point32(x=vertex[0], y=vertex[1], z=vertex[2]) 80 | poly_msg.vertices.append(point) 81 | 82 | normal = Vector3(x=polygon.normal[0], y=polygon.normal[1], z=polygon.normal[2]) 83 | poly_msg.normal = normal 84 | poly_msg.plane_eq = polygon.plane_eq 85 | poly_msg.area = polygon.area 86 | msg.polygons.append(poly_msg) 87 | 88 | self.pub.publish(msg) 89 | global T 90 | T_lock = copy.deepcopy(T) 91 | self.publish_camera_pose(T_lock) 92 | self.publish_markers() 93 | 94 | def publish_markers(self): 95 | """Publish markers for polygons.""" 96 | for i, polygon in enumerate(self.map_manager.polygons): 97 | marker = Marker() 98 | marker.header.frame_id = "world" 99 | marker.header.stamp = rospy.Time.now() 100 | marker.ns = "polygons" 101 | marker.id = i 102 | marker.type = Marker.LINE_STRIP 103 | marker.action = Marker.ADD 104 | marker.lifetime = rospy.Duration(0.5) 105 | marker.pose.orientation.w = 1.0 106 | marker.scale.x = 0.05 107 | marker.color.r = 0.0 108 | marker.color.g = 1.0 109 | marker.color.b = 0.0 110 | marker.color.a = 1.0 111 | 112 | 113 | for vertex in polygon.vertices: 114 | point = Point() 115 | point.x = vertex[0] 116 | point.y = vertex[1] 117 | point.z = vertex[2] 118 | marker.points.append(point) 119 | if len(polygon.vertices) > 2: 120 | point = Point() 121 | point.x = polygon.vertices[0][0] 122 | point.y = polygon.vertices[0][1] 123 | point.z = polygon.vertices[0][2] 124 | marker.points.append(point) 125 | 126 | self.marker_pub.publish(marker) 127 | 128 | def publish_camera_pose(self, T): 129 | # Extract translation vector 130 | translation = T[0:3, 3] 131 | 132 | # Extract rotation matrix and convert it to a quaternion 133 | rotation_matrix = T[0:3, 0:3] 134 | euler_angles = tf_trans.euler_from_matrix(rotation_matrix, 'sxyz') 135 | yaw = euler_angles[2] + 3.14159 / 2 # Retain the yaw angle 136 | 137 | # Set roll and pitch to 0, reconstruct quaternion 138 | quaternion = tf_trans.quaternion_from_euler(0, 0, yaw) 139 | 140 | # Apply drift compensation on the Z-axis 141 | translation[2] += self.map_manager.z_drift 142 | 143 | if self.map_manager.z_drift > 0: 144 | self.positive_drift_cont += 1 145 | else: 146 | self.nagetive_drift_cont += 1 147 | 148 | # Create PoseStamped message 149 | pose_msg = PoseStamped() 150 | pose_msg.header.stamp = rospy.Time.now() 151 | pose_msg.header.frame_id = "world" # World coordinate frame 152 | 153 | # Set position 154 | pose_msg.pose.position.x = translation[0] 155 | pose_msg.pose.position.y = translation[1] 156 | pose_msg.pose.position.z = translation[2] 157 | 158 | # Set orientation (quaternion) 159 | pose_msg.pose.orientation.x = quaternion[0] 160 | pose_msg.pose.orientation.y = quaternion[1] 161 | pose_msg.pose.orientation.z = quaternion[2] 162 | pose_msg.pose.orientation.w = quaternion[3] 163 | 164 | # Publish camera pose 165 | self.pose_pub.publish(pose_msg) 166 | 167 | # Update path 168 | self.path.poses.append(pose_msg) 169 | self.path.header.stamp = rospy.Time.now() 170 | self.path_pub.publish(self.path) 171 | 172 | pose_raw_msg = copy.deepcopy(pose_msg) 173 | 174 | # Remove drift compensation for raw pose 175 | pose_raw_msg.pose.position.z -= self.map_manager.z_drift 176 | self.pose_raw_pub.publish(pose_raw_msg) 177 | 178 | # Update raw path 179 | self.path_raw.poses.append(pose_raw_msg) 180 | self.path_raw.header.stamp = rospy.Time.now() 181 | self.path_raw_pub.publish(self.path_raw) 182 | 183 | 184 | def nothing(x): 185 | """Placeholder function for trackbar.""" 186 | pass 187 | 188 | def read_transformation_matrix(t_file_path): 189 | """Read transformation matrices and timestamps from a file.""" 190 | T_matrices, timestamps = [], [] 191 | current_data = "" 192 | with open(t_file_path, 'r') as t_file: 193 | for line in t_file: 194 | current_data += line.strip() + " " 195 | numbers = re.findall(r'-?\d+\.\d+e[+-]?\d+|-?\d+\.\d+|-?\d+', current_data) 196 | if len(numbers) == 18: 197 | idx = int(numbers[0]) 198 | T = np.array(numbers[1:17]).astype(float).reshape(4, 4) 199 | timestamp = float(numbers[-1]) 200 | T_matrices.append(T) 201 | timestamps.append(timestamp) 202 | current_data = "" 203 | return T_matrices, timestamps 204 | 205 | def set_cv_config(thresh1=50, thresh2=100, num_iter=60, kappa=134, gamma=2): 206 | """Set up OpenCV windows and trackbars for image processing parameters.""" 207 | cv2.namedWindow('Filtered Image with Largest Contour') 208 | cv2.createTrackbar('Canny thresh1', 'Filtered Image with Largest Contour', thresh1, max(thresh1, 1000), nothing) 209 | cv2.createTrackbar('Canny thresh2', 'Filtered Image with Largest Contour', thresh2, max(thresh2, 1000), nothing) 210 | 211 | cv2.namedWindow('smoothed_depth_img') 212 | cv2.createTrackbar('num_iter', 'smoothed_depth_img', num_iter, max(num_iter, 360), nothing) 213 | cv2.createTrackbar('kappa', 'smoothed_depth_img', kappa, max(kappa, 500), nothing) 214 | cv2.createTrackbar('gamma', 'smoothed_depth_img', gamma, max(gamma, 80), nothing) 215 | 216 | def main(): 217 | try: 218 | # Create map manager 219 | map_manager = MapManager(z_threshold=0.05) 220 | node = PosePublishNode(map_manager) 221 | set_cv_config() 222 | 223 | global T 224 | fx, fy, cx, cy = 228.686, 230.078, 157.748, 123.968 # Example camera intrinsics 225 | 226 | # Initialize ROS package and get dataset paths 227 | rospack = rospkg.RosPack() 228 | package_path = rospack.get_path('polygon_mapping') 229 | dataset_path = os.path.join(package_path, 'data/', 'dataset_39/') 230 | rgb_images = sorted(glob.glob(dataset_path + 'rgb/*.png')) 231 | depth_images = sorted(glob.glob(dataset_path + 'depth/*.png')) 232 | 233 | # Load transformation matrices and timestamps 234 | T_file_path = os.path.join(dataset_path, 'transformation_matrix.txt') 235 | T_matrices, timestamps = read_transformation_matrix(T_file_path) 236 | 237 | os.makedirs(os.path.join(dataset_path, "processed","canny"), exist_ok=True) 238 | os.makedirs(os.path.join(dataset_path, "processed","depth"), exist_ok=True) 239 | os.makedirs(os.path.join(dataset_path, "processed","normal_raw"), exist_ok=True) 240 | os.makedirs(os.path.join(dataset_path, "processed","normal_smoothed"), exist_ok=True) 241 | os.makedirs(os.path.join(dataset_path, "processed","polygon_extracted"), exist_ok=True) 242 | os.makedirs(os.path.join(dataset_path, "processed","polygon_map"), exist_ok=True) 243 | 244 | # Process each image 245 | for idx, (T_matrix, timestamp) in enumerate(zip(T_matrices, timestamps)): 246 | rgb_image_path = os.path.join(dataset_path, 'rgb', f'{str(idx + 1).zfill(6)}.png') 247 | depth_image_path = os.path.join(dataset_path, 'depth', f'{str(idx + 1).zfill(6)}.png') 248 | 249 | color_image = cv2.imread(rgb_image_path) 250 | depth_image = cv2.imread(depth_image_path, cv2.IMREAD_UNCHANGED) 251 | T = T_matrix 252 | T_lock = copy.deepcopy(T) 253 | 254 | time_start_loop = time.time() 255 | 256 | # Smooth depth image using anisotropic diffusion 257 | time_start = time.time() 258 | smoothed_depth_img = anisotropic_diffusion(depth_image) 259 | print('Anisotropic diffusion time cost', 1000 * (time.time() - time_start), 'ms') 260 | 261 | # Save processed depth and normal images 262 | cv2.imwrite(dataset_path + "processed/depth/" + f'{str(idx + 1).zfill(6)}.png', 263 | cv2.applyColorMap(cv2.convertScaleAbs(smoothed_depth_img, alpha=0.03), cv2.COLORMAP_JET)) 264 | normals = depth_to_normals(smoothed_depth_img, fx, fy, cx, cy) 265 | cv2.imwrite(dataset_path + "processed/normal_smoothed/" + f'{str(idx + 1).zfill(6)}.png', normals) 266 | normals_raw = depth_to_normals(depth_image, fx, fy, cx, cy) 267 | cv2.imwrite(dataset_path + "processed/normal_raw/" + f'{str(idx + 1).zfill(6)}.png', normals_raw) 268 | 269 | # Retrieve Canny threshold values from trackbars 270 | thresh1 = cv2.getTrackbarPos('Canny thresh1', 'Filtered Image with Largest Contour') 271 | thresh2 = cv2.getTrackbarPos('Canny thresh2', 'Filtered Image with Largest Contour') 272 | 273 | # Apply sharpening filter for edge enhancement 274 | kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]]) 275 | filtered_img = cv2.filter2D(normals, -1, kernel) 276 | filtered_img = add_border_conditionally(filtered_img) 277 | edges = cv2.Canny(filtered_img, thresh1, thresh2) 278 | 279 | # Apply morphological closing to fill edge gaps 280 | kernel = np.ones((13, 13), np.uint8) 281 | edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel) 282 | 283 | # Find and process contours 284 | contours, hierarchy = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) 285 | min_area = 5000 286 | new_polygon_list = [] 287 | 288 | for i, contour in enumerate(contours): 289 | if hierarchy[0][i][3] == -1: 290 | continue 291 | if cv2.contourArea(contour) < min_area: 292 | continue 293 | 294 | hull = contour 295 | epsilon = 0.02 * cv2.arcLength(hull, True) 296 | approx = cv2.approxPolyDP(hull, epsilon, True) 297 | if cv2.contourArea(approx) < min_area: 298 | continue 299 | 300 | cv2.drawContours(filtered_img, [approx], 0, (0, 255, 0), 2) 301 | for point in approx: 302 | cv2.circle(filtered_img, tuple(point[0]), 5, (0, 0, 255), -1) 303 | 304 | polygon_points = depth_to_point_cloud_region(smoothed_depth_img, approx.reshape(-1, 2), fx, fy, cx, cy) 305 | if polygon_points.size == 0: 306 | print("No valid points in the polygon region") 307 | 308 | try: 309 | normal, d = fit_plane_ransac( 310 | polygon_points, max_trials=20, min_samples=3, residual_threshold=0.01 311 | ) 312 | projected_points = project_points_to_plane((*normal, d), fx, fy, cx, cy, approx.reshape(-1, 2)) 313 | angle, world_vertices = process_polygon(projected_points, normal, T_lock) 314 | if world_vertices is not None: 315 | new_polygon_list.append(world_vertices) 316 | except ValueError as e: 317 | print("Error in fitting plane:", e) 318 | 319 | map_manager.add_polygon_list(new_polygon_list) 320 | 321 | if map_manager.polygons: 322 | map_img = map_manager.plot_polygons() 323 | cv2.imshow('Polygon Map', map_img) 324 | cv2.imwrite(dataset_path + "processed/polygon_map/" + f'{str(idx + 1).zfill(6)}.png', map_img) 325 | 326 | cv2.imshow('smoothed_depth_img', cv2.applyColorMap(cv2.convertScaleAbs(smoothed_depth_img, alpha=0.03), cv2.COLORMAP_JET)) 327 | cv2.imshow('normals Image', normals) 328 | cv2.imshow('Color Image', color_image) 329 | cv2.imshow('Edges', edges) 330 | cv2.imwrite(dataset_path + "processed/canny/" + f'{str(idx + 1).zfill(6)}.png', edges) 331 | cv2.imshow('Filtered Image with Largest Contour', filtered_img) 332 | cv2.imwrite(dataset_path + "processed/polygon_extracted/" + f'{str(idx + 1).zfill(6)}.png', filtered_img) 333 | print('Loop time cost', 1000 * (time.time() - time_start_loop), 'ms') 334 | 335 | if cv2.waitKey(1) & 0xFF == ord('q'): 336 | break 337 | 338 | rospy.spin() 339 | cv2.destroyAllWindows() 340 | except rospy.ROSInterruptException: 341 | pass 342 | 343 | if __name__ == '__main__': 344 | main() 345 | --------------------------------------------------------------------------------