Merge pull request #98 from leggedrobotics/dev/nature_hiking/feat/dis…

…tortion

Image kernel distortion

elevation_mapping_cupy/config/setups/anymal/anymal_parameters.yaml

@@ -1,6 +1,6 @@
 #### Basic parameters ########
-resolution: 0.04                              # resolution in m.
-map_length: 8                           # map's size in m.
+resolution: 0.04                                # resolution in m.
+map_length: 8                                   # map's size in m.
 sensor_noise_factor: 0.05                       # point's noise is sensor_noise_factor*z^2 (z is distance from sensor).
 mahalanobis_thresh: 2.0                         # points outside this distance is outlier.
 outlier_variance: 0.01                          # if point is outlier, add this value to the cell.
@@ -9,7 +9,7 @@ max_drift: 0.1                                  # drift compensation happens onl
 drift_compensation_alpha: 0.1                   # drift compensation alpha for smoother update of drift compensation
 time_variance: 0.0001                           # add this value when update_variance is called.
 max_variance: 100.0                             # maximum variance for each cell.
-initial_variance: 1000.0                         # initial variance for each cell.
+initial_variance: 1000.0                        # initial variance for each cell.
 traversability_inlier: 0.9                      # cells with higher traversability are used for drift compensation.
 dilation_size: 3                                # dilation filter size before traversability filter.
 wall_num_thresh: 20                             # if there are more points than this value, only higher points than the current height are used to make the wall more sharp.

elevation_mapping_cupy/config/setups/anymal/anymal_sensor_parameter.yaml

@@ -18,12 +18,17 @@ image_channel_fusions:
 
 publishers:
   elevation_map_raw:
-    layers: ['elevation', 'traversability', 'variance', 'rgb']
+    layers: ['elevation', 'traversability', 'variance', 'rgb', 'upper_bound']
     basic_layers: ['elevation', 'traversability']
     fps: 5.0
 
+  elevation_map_recordable:
+    layers: ['elevation', 'traversability', 'variance', 'rgb']
+    basic_layers: ['elevation', 'traversability']
+    fps: 2.0
+
   filtered_elevation_map:
-    layers: ['inpaint', 'smooth', 'min_filter', 'erosion']
+    layers: ['inpaint', 'smooth', 'min_filter', 'upper_bound']
     basic_layers: ['inpaint']
     fps: 5.0
 

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 Eigen::VectorXd& D, int height, int width);
   void move_to(const Eigen::VectorXd& p, const RowMatrixXd& R);
   void clear();
   void update_variance();

elevation_mapping_cupy/script/elevation_mapping_cupy/elevation_mapping.py

@@ -228,7 +228,10 @@ def shift_map_z(self, delta_z):
     def compile_kernels(self):
         """Compile all kernels belonging to the elevation map."""
 
-        self.new_map = cp.zeros((self.elevation_map.shape[0], self.cell_n, self.cell_n), dtype=self.data_type,)
+        self.new_map = cp.zeros(
+            (self.elevation_map.shape[0], self.cell_n, self.cell_n),
+            dtype=self.data_type,
+        )
         self.traversability_input = cp.zeros((self.cell_n, self.cell_n), dtype=self.data_type)
         self.traversability_mask_dummy = cp.zeros((self.cell_n, self.cell_n), dtype=self.data_type)
         self.min_filtered = cp.zeros((self.cell_n, self.cell_n), dtype=self.data_type)
@@ -287,14 +290,18 @@ def compile_image_kernels(self):
                     np.zeros((self.cell_n, self.cell_n), dtype=np.bool_), dtype=np.bool_
                 )
                 self.uv_correspondence = cp.asarray(
-                    np.zeros((2, self.cell_n, self.cell_n), dtype=np.float32), dtype=np.float32,
+                    np.zeros((2, self.cell_n, self.cell_n), dtype=np.float32),
+                    dtype=np.float32,
                 )
                 # self.distance_correspondence = cp.asarray(
                 #     np.zeros((self.cell_n, self.cell_n), dtype=np.float32), dtype=np.float32
                 # )
                 # TODO tolerance_z_collision add parameter
                 self.image_to_map_correspondence_kernel = image_to_map_correspondence_kernel(
-                    resolution=self.resolution, width=self.cell_n, height=self.cell_n, tolerance_z_collision=0.10,
+                    resolution=self.resolution,
+                    width=self.cell_n,
+                    height=self.cell_n,
+                    tolerance_z_collision=0.10,
                 )
                 break
 
@@ -466,6 +473,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,
     ):
@@ -483,6 +491,7 @@ def input_image(
         Returns:
             None:
         """
+
         image = np.stack(image, axis=0)
         if len(image.shape) == 2:
             image = image[None]
@@ -492,9 +501,17 @@ def input_image(
         K = cp.asarray(K, dtype=self.data_type)
         R = cp.asarray(R, dtype=self.data_type)
         t = cp.asarray(t, dtype=self.data_type)
+        D = cp.asarray(D, dtype=self.data_type)
         image_height = cp.float32(image_height)
         image_width = cp.float32(image_width)
 
+        if len(D) < 4:
+            D = cp.zeros(5, dtype=self.data_type)
+        elif len(D) == 4:
+            D = cp.concatenate([D, cp.zeros(1, dtype=self.data_type)])
+        else:
+            D = D[:5]
+
         # Calculate transformation matrix
         P = cp.asarray(K @ cp.concatenate([R, t[:, None]], 1), dtype=np.float32)
         t_cam_map = -R.T @ t - self.center
@@ -505,7 +522,6 @@ def input_image(
 
         self.uv_correspondence *= 0
         self.valid_correspondence[:, :] = False
-        # self.distance_correspondence *= 0.0
 
         with self.map_lock:
             self.image_to_map_correspondence_kernel(
@@ -514,6 +530,8 @@ def input_image(
                 y1,
                 z1,
                 P.reshape(-1),
+                K.reshape(-1),
+                D.reshape(-1),
                 image_height,
                 image_width,
                 self.center,
@@ -522,7 +540,12 @@ def input_image(
                 size=int(self.cell_n * self.cell_n),
             )
             self.semantic_map.update_layers_image(
-                image, channels, self.uv_correspondence, self.valid_correspondence, image_height, image_width,
+                image,
+                channels,
+                self.uv_correspondence,
+                self.valid_correspondence,
+                image_height,
+                image_width,
             )
 
     def update_normal(self, dilated_map):
@@ -534,7 +557,10 @@ def update_normal(self, dilated_map):
         with self.map_lock:
             self.normal_map *= 0.0
             self.normal_filter_kernel(
-                dilated_map, self.elevation_map[2], self.normal_map, size=(self.cell_n * self.cell_n),
+                dilated_map,
+                self.elevation_map[2],
+                self.normal_map,
+                size=(self.cell_n * self.cell_n),
             )
 
     def process_map_for_publish(self, input_map, fill_nan=False, add_z=False, xp=cp):
@@ -580,7 +606,9 @@ def get_traversability(self):
             traversability layer
         """
         traversability = cp.where(
-            (self.elevation_map[2] + self.elevation_map[6]) > 0.5, self.elevation_map[3].copy(), cp.nan,
+            (self.elevation_map[2] + self.elevation_map[6]) > 0.5,
+            self.elevation_map[3].copy(),
+            cp.nan,
         )
         self.traversability_buffer[3:-3, 3:-3] = traversability[3:-3, 3:-3]
         traversability = self.traversability_buffer[1:-1, 1:-1]
@@ -602,7 +630,8 @@ def get_upper_bound(self):
         """
         if self.param.use_only_above_for_upper_bound:
             valid = cp.logical_or(
-                cp.logical_and(self.elevation_map[5] > 0.0, self.elevation_map[6] > 0.5), self.elevation_map[2] > 0.5,
+                cp.logical_and(self.elevation_map[5] > 0.0, self.elevation_map[6] > 0.5),
+                self.elevation_map[2] > 0.5,
             )
         else:
             valid = cp.logical_or(self.elevation_map[2] > 0.5, self.elevation_map[6] > 0.5)
@@ -618,7 +647,8 @@ def get_is_upper_bound(self):
         """
         if self.param.use_only_above_for_upper_bound:
             valid = cp.logical_or(
-                cp.logical_and(self.elevation_map[5] > 0.0, self.elevation_map[6] > 0.5), self.elevation_map[2] > 0.5,
+                cp.logical_and(self.elevation_map[5] > 0.0, self.elevation_map[6] > 0.5),
+                self.elevation_map[2] > 0.5,
             )
         else:
             valid = cp.logical_or(self.elevation_map[2] > 0.5, self.elevation_map[6] > 0.5)
@@ -779,7 +809,11 @@ def get_layer(self, name):
             return_map = self.semantic_map.semantic_map[idx]
         elif name in self.plugin_manager.layer_names:
             self.plugin_manager.update_with_name(
-                name, self.elevation_map, self.layer_names, self.semantic_map, self.base_rotation,
+                name,
+                self.elevation_map,
+                self.layer_names,
+                self.semantic_map,
+                self.base_rotation,
             )
             return_map = self.plugin_manager.get_map_with_name(name)
         else:
@@ -825,7 +859,10 @@ def get_polygon_traversability(self, polygon, result):
         else:
             t = cp.asarray(0.0, dtype=self.data_type)
         is_safe, un_polygon = is_traversable(
-            masked, self.param.safe_thresh, self.param.safe_min_thresh, self.param.max_unsafe_n,
+            masked,
+            self.param.safe_thresh,
+            self.param.safe_min_thresh,
+            self.param.max_unsafe_n,
         )
         untraversable_polygon_num = 0
         if un_polygon is not None:
@@ -881,7 +918,9 @@ def initialize_map(self, points, method="cubic"):
     t = xp.random.rand(3)
     print(R, t)
     param = Parameter(
-        use_chainer=False, weight_file="../config/weights.dat", plugin_config_file="../config/plugin_config.yaml",
+        use_chainer=False,
+        weight_file="../config/weights.dat",
+        plugin_config_file="../config/plugin_config.yaml",
     )
     param.additional_layers = ["rgb", "grass", "tree", "people"]
     param.fusion_algorithms = ["color", "class_bayesian", "class_bayesian", "class_bayesian"]

elevation_mapping_cupy/script/elevation_mapping_cupy/elevation_mapping_ros.py

@@ -174,10 +174,13 @@ def image_callback(self, camera_msg, camera_info_msg, sub_key):
         else:
             semantic_img = [semantic_img]
 
-        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)
+
+        assert np.all(np.array(camera_info_msg.D) == 0.0), "Undistortion not implemented"
+        D = np.array(camera_info_msg.D, dtype=np.float32).reshape(5, 1)
+
         # 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):

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 U K, raw U 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(
             """
@@ -22,7 +22,7 @@ def image_to_map_correspondence_kernel(resolution, width, height, tolerance_z_co
                 return layer * layer_n + idx;
             }
             __device__ bool is_inside_map(int x, int y) {
-                return (x >= 0 && y >= 0 && x<${width} && x<${height});
+                return (x >= 0 && y >= 0 && x<${width} && y<${height});
             }
             __device__ float get_l2_distance(int x0, int y0, int x1, int y1) {
                 float dx = x0-x1;
@@ -60,11 +60,35 @@ def image_to_map_correspondence_kernel(resolution, width, height, tolerance_z_co
             }
             u = u/d;
             v = v/d;
+
+            // Check if D is all zeros
+            bool is_D_zero = (D[0] == 0 && D[1] == 0 && D[2] == 0 && D[3] == 0 && D[4] == 0);
+
+            // Apply undistortion using distortion matrix D if not all zeros
+            if (!is_D_zero) {
+                float k1 = D[0];
+                float k2 = D[1];
+                float p1 = D[2];
+                float p2 = D[3];
+                float k3 = D[4];
+                float fx = K[0];
+                float fy = K[4];
+                float cx = K[2];
+                float cy = K[5];
+                float x = (u - cx) / fx;
+                float y = (v - cy) / fy;
+                float r2 = x * x + y * y;
+                float radial_distortion = 1 + k1 * r2 + k2 * r2 * r2 + k3 * r2 * r2 * r2;
+                float u_corrected = x * radial_distortion + 2 * p1 * x * y + p2 * (r2 + 2 * x * x);
+                float v_corrected = y * radial_distortion + 2 * p2 * x * y + p1 * (r2 + 2 * y * y);
+                u = fx * u_corrected + cx;
+                v = fy * v_corrected + cy;
+            }
             
             // filter point next to image plane
             if ((u < 0) || (v < 0) || (u >= image_width) || (v >= image_height)){
                 return;
-            } 
+            }
             
             int y0_c = y0;
             int x0_c = x0;

elevation_mapping_cupy/src/elevation_mapping_ros.cpp

@@ -388,6 +388,16 @@ 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::VectorXd distortionCoeffs;
+  if (!camera_info_msg->D.empty()) {
+    distortionCoeffs = Eigen::Map<const Eigen::VectorXd>(camera_info_msg->D.data(), camera_info_msg->D.size());
+  } else {
+    ROS_WARN("Distortion coefficients are empty.");
+    distortionCoeffs = Eigen::VectorXd::Zero(5);
+    // return;
+  }
+
   // Get pose of sensor in map frame
   tf::StampedTransform transformTf;
   std::string sensorFrameId = image_msg->header.frame_id;
@@ -428,8 +438,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,

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 Eigen::VectorXd& 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 Eigen::VectorXd>(D), height, width);
 }
 
 void ElevationMappingWrapper::move_to(const Eigen::VectorXd& p, const RowMatrixXd& R) {