From 81fa739b92f0920f4d7f6e97ec103ae1497a1c74 Mon Sep 17 00:00:00 2001 From: RobinSchmid7 Date: Sun, 9 Jun 2024 22:04:12 +0200 Subject: [PATCH] Add distortion to input image function --- .../elevation_mapping_wrapper.hpp | 2 +- .../elevation_mapping.py | 2 ++ .../elevation_mapping_ros.py | 2 +- .../kernels/custom_image_kernels.py | 20 +++++++++++++++++-- .../src/elevation_mapping_ros.cpp | 7 +++++-- .../src/elevation_mapping_wrapper.cpp | 4 ++-- 6 files changed, 29 insertions(+), 8 deletions(-) diff --git a/elevation_mapping_cupy/include/elevation_mapping_cupy/elevation_mapping_wrapper.hpp b/elevation_mapping_cupy/include/elevation_mapping_cupy/elevation_mapping_wrapper.hpp index 116403e6..fb74574f 100644 --- a/elevation_mapping_cupy/include/elevation_mapping_cupy/elevation_mapping_wrapper.hpp +++ b/elevation_mapping_cupy/include/elevation_mapping_cupy/elevation_mapping_wrapper.hpp @@ -49,7 +49,7 @@ class ElevationMappingWrapper { void input(const RowMatrixXd& points, const std::vector& channels, const RowMatrixXd& R, const Eigen::VectorXd& t, const double positionNoise, const double orientationNoise); void input_image(const std::vector& multichannel_image, const std::vector& channels, const RowMatrixXd& R, - const Eigen::VectorXd& t, const RowMatrixXd& cameraMatrix, int height, int width); + const Eigen::VectorXd& t, const RowMatrixXd& cameraMatrix, const RowMatrixXd& D, int height, int width); void move_to(const Eigen::VectorXd& p, const RowMatrixXd& R); void clear(); void update_variance(); diff --git a/elevation_mapping_cupy/script/elevation_mapping_cupy/elevation_mapping.py b/elevation_mapping_cupy/script/elevation_mapping_cupy/elevation_mapping.py index 0ff32d9a..c9190725 100644 --- a/elevation_mapping_cupy/script/elevation_mapping_cupy/elevation_mapping.py +++ b/elevation_mapping_cupy/script/elevation_mapping_cupy/elevation_mapping.py @@ -466,6 +466,7 @@ def input_image( R: cp._core.core.ndarray, t: cp._core.core.ndarray, K: cp._core.core.ndarray, + D: cp._core.core.ndarray, image_height: int, image_width: int, ): @@ -514,6 +515,7 @@ def input_image( y1, z1, P.reshape(-1), + D.reshape(-1), image_height, image_width, self.center, diff --git a/elevation_mapping_cupy/script/elevation_mapping_cupy/elevation_mapping_ros.py b/elevation_mapping_cupy/script/elevation_mapping_cupy/elevation_mapping_ros.py index 3756ebd7..6e69861c 100644 --- a/elevation_mapping_cupy/script/elevation_mapping_cupy/elevation_mapping_ros.py +++ b/elevation_mapping_cupy/script/elevation_mapping_cupy/elevation_mapping_ros.py @@ -177,7 +177,7 @@ def image_callback(self, camera_msg, camera_info_msg, sub_key): assert np.all(np.array(camera_info_msg.D) == 0.0), "Undistortion not implemented" K = np.array(camera_info_msg.K, dtype=np.float32).reshape(3, 3) # process pointcloud - self._map.input_image(sub_key, semantic_img, R, t, K, camera_info_msg.height, camera_info_msg.width) + self._map.input_image(sub_key, semantic_img, R, t, K, D, camera_info_msg.height, camera_info_msg.width) self._image_process_counter += 1 def pointcloud_callback(self, msg, sub_key): diff --git a/elevation_mapping_cupy/script/elevation_mapping_cupy/kernels/custom_image_kernels.py b/elevation_mapping_cupy/script/elevation_mapping_cupy/kernels/custom_image_kernels.py index ce03ee8f..bf835c85 100644 --- a/elevation_mapping_cupy/script/elevation_mapping_cupy/kernels/custom_image_kernels.py +++ b/elevation_mapping_cupy/script/elevation_mapping_cupy/kernels/custom_image_kernels.py @@ -13,7 +13,7 @@ def image_to_map_correspondence_kernel(resolution, width, height, tolerance_z_co The function returns a kernel that can be used to perform the correspondence calculation. """ _image_to_map_correspondence_kernel = cp.ElementwiseKernel( - in_params="raw U map, raw U x1, raw U y1, raw U z1, raw U P, raw U image_height, raw U image_width, raw U center", + in_params="raw U map, raw U x1, raw U y1, raw U z1, raw U P, raw D, raw U image_height, raw U image_width, raw U center", out_params="raw U uv_correspondence, raw B valid_correspondence", preamble=string.Template( """ @@ -64,7 +64,23 @@ def image_to_map_correspondence_kernel(resolution, width, height, tolerance_z_co // filter point next to image plane if ((u < 0) || (v < 0) || (u >= image_width) || (v >= image_height)){ return; - } + } + + // Apply undistortion using distortion matrix D + float k1 = D[0]; + float k2 = D[1]; + float p1 = D[2]; + float p2 = D[3]; + float k3 = D[4]; + float x = (u - image_width / 2.0) / (image_width / 2.0); + float y = (v - image_height / 2.0) / (image_height / 2.0); + float r2 = x * x + y * y; + float radial_distortion = 1 + k1 * r2 + k2 * r2 * r2 + k3 * r2 * r2 * r2; + float x_corrected = x * radial_distortion + 2 * p1 * x * y + p2 * (r2 + 2 * x * x); + float y_corrected = y * radial_distortion + 2 * p2 * x * y + p1 * (r2 + 2 * y * y); + u = (x_corrected * (image_width / 2.0)) + (image_width / 2.0); + v = (y_corrected * (image_height / 2.0)) + (image_height / 2.0); + int y0_c = y0; int x0_c = x0; diff --git a/elevation_mapping_cupy/src/elevation_mapping_ros.cpp b/elevation_mapping_cupy/src/elevation_mapping_ros.cpp index 715afeb9..b6bbc2d1 100644 --- a/elevation_mapping_cupy/src/elevation_mapping_ros.cpp +++ b/elevation_mapping_cupy/src/elevation_mapping_ros.cpp @@ -388,6 +388,9 @@ void ElevationMappingNode::inputImage(const sensor_msgs::ImageConstPtr& image_ms // Extract camera matrix Eigen::Map> cameraMatrix(&camera_info_msg->K[0]); + // Extract distortion coefficients + Eigen::Map distortionCoeffs(camera_info_msg->D.data(), camera_info_msg->D.size()); + // Get pose of sensor in map frame tf::StampedTransform transformTf; std::string sensorFrameId = image_msg->header.frame_id; @@ -428,8 +431,8 @@ void ElevationMappingNode::inputImage(const sensor_msgs::ImageConstPtr& image_ms } // Pass image to pipeline - map_.input_image(multichannel_image, channels, transformationMapToSensor.rotation(), transformationMapToSensor.translation(), cameraMatrix, - image.rows, image.cols); + map_.input_image(multichannel_image, channels, transformationMapToSensor.rotation(), transformationMapToSensor.translation(), cameraMatrix, + distortionCoeffs, image.rows, image.cols); } void ElevationMappingNode::imageCallback(const sensor_msgs::ImageConstPtr& image_msg, diff --git a/elevation_mapping_cupy/src/elevation_mapping_wrapper.cpp b/elevation_mapping_cupy/src/elevation_mapping_wrapper.cpp index f66f9a79..b23207d6 100644 --- a/elevation_mapping_cupy/src/elevation_mapping_wrapper.cpp +++ b/elevation_mapping_cupy/src/elevation_mapping_wrapper.cpp @@ -178,10 +178,10 @@ void ElevationMappingWrapper::input(const RowMatrixXd& points, const std::vector } void ElevationMappingWrapper::input_image(const std::vector& multichannel_image, const std::vector& channels, const RowMatrixXd& R, - const Eigen::VectorXd& t, const RowMatrixXd& cameraMatrix, int height, int width) { + const Eigen::VectorXd& t, const RowMatrixXd& cameraMatrix, const RowMatrixXd& D, int height, int width) { py::gil_scoped_acquire acquire; map_.attr("input_image")(multichannel_image, channels, Eigen::Ref(R), Eigen::Ref(t), - Eigen::Ref(cameraMatrix), height, width); + Eigen::Ref(cameraMatrix), Eigen::Ref(D), height, width); } void ElevationMappingWrapper::move_to(const Eigen::VectorXd& p, const RowMatrixXd& R) {