Add distortion to input image function

leggedrobotics · Jun 9, 2024 · 81fa739 · 81fa739
1 parent 0126df1
commit 81fa739
Show file tree

Hide file tree

Showing 6 changed files with 29 additions and 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
@@ -49,7 +49,7 @@ class ElevationMappingWrapper {
   void input(const RowMatrixXd& points, const std::vector<std::string>& channels, const RowMatrixXd& R, const Eigen::VectorXd& t,
              const double positionNoise, const double orientationNoise);
   void input_image(const std::vector<ColMatrixXf>& multichannel_image, const std::vector<std::string>& 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
@@ -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
@@ -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
@@ -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
@@ -388,6 +388,9 @@ void ElevationMappingNode::inputImage(const sensor_msgs::ImageConstPtr& image_ms
   // Extract camera matrix
   Eigen::Map<const Eigen::Matrix<double, 3, 3, Eigen::RowMajor>> cameraMatrix(&camera_info_msg->K[0]);
 
+  // Extract distortion coefficients
+  Eigen::Map<const Eigen::VectorXd> 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
@@ -178,10 +178,10 @@ void ElevationMappingWrapper::input(const RowMatrixXd& points, const std::vector
 }
 
 void ElevationMappingWrapper::input_image(const std::vector<ColMatrixXf>& multichannel_image, const std::vector<std::string>& 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<const RowMatrixXd>(R), Eigen::Ref<const Eigen::VectorXd>(t),
-                           Eigen::Ref<const RowMatrixXd>(cameraMatrix), height, width);
+                           Eigen::Ref<const RowMatrixXd>(cameraMatrix), Eigen::Ref<const RowMatrixXd>(D), height, width);
 }
 
 void ElevationMappingWrapper::move_to(const Eigen::VectorXd& p, const RowMatrixXd& R) {