depth_image_proc register cover all angles with fill_upsampling_holes

depth_image_proc/CMakeLists.txt

@@ -54,3 +54,10 @@ install(TARGETS ${PROJECT_NAME}
 install(FILES nodelet_plugins.xml
         DESTINATION ${CATKIN_PACKAGE_SHARE_DESTINATION}
 )
+
+if(CATKIN_ENABLE_TESTING)
+  find_package(rostest REQUIRED)
+
+  add_rostest(test/test_register.test
+    DEPENDENCIES ${${PROJECT_NAME}_EXPORTED_TARGETS})
+endif()

depth_image_proc/package.xml

@@ -45,4 +45,7 @@
   <run_depend>tf2</run_depend>
   <run_depend>tf2_ros</run_depend>
   <run_depend>sensor_msgs</run_depend>
+
+  <test_depend>rostest</test_depend>
+  <test_depend>tf</test_depend>
 </package>

depth_image_proc/src/nodelets/register.cpp

@@ -195,6 +195,58 @@ void RegisterNodelet::imageCb(const sensor_msgs::ImageConstPtr& depth_image_msg,
   pub_registered_.publish(registered_msg, registered_info_msg);
 }
 
+template<typename T>
+bool transform_depth(
+    const Eigen::Affine3d& depth_to_rgb,
+    const double u, const double v,
+    const image_geometry::PinholeCameraModel& depth_model_,
+    const image_geometry::PinholeCameraModel& rgb_model_,
+    const double inv_depth_fx, const double inv_depth_fy,
+    const int width, const int height,
+    const double depth,
+    int& u_rgb, int& v_rgb,
+    T& new_depth)
+{
+  // TODO(lucasw) pulling these out shouldn't cost anything
+  const auto depth_cx = depth_model_.cx(), depth_cy = depth_model_.cy();
+  const auto depth_Tx = depth_model_.Tx(), depth_Ty = depth_model_.Ty();
+  const auto rgb_fx = rgb_model_.fx(), rgb_fy = rgb_model_.fy();
+  const auto rgb_cx = rgb_model_.cx(), rgb_cy = rgb_model_.cy();
+  const auto rgb_Tx = rgb_model_.Tx(), rgb_Ty = rgb_model_.Ty();
+
+  /// @todo Combine all operations into one matrix multiply on (u,v,d)
+  // Reproject (u,v,Z) to (X,Y,Z,1) in depth camera frame
+  Eigen::Vector4d xyz_depth;
+  xyz_depth << ((u - depth_cx)*depth - depth_Tx) * inv_depth_fx,
+               ((v - depth_cy)*depth - depth_Ty) * inv_depth_fy,
+               depth,
+               1;
+
+  // Transform to RGB camera frame
+  Eigen::Vector4d xyz_rgb = depth_to_rgb * xyz_depth;
+  new_depth = static_cast<T>(xyz_rgb.z());
+  // TODO(lucasw) is the intent to simulate what a real depth camera would see?  If so reject negative depth
+  // but if want to preserve as much data as possible it may make sense to pass through negative values
+  // (though don't overwrite positive values, use abs in the z buffer test)
+  if (new_depth < 0.0)
+  {
+    return false;
+  }
+
+  // Project to (u,v) in RGB image
+  const double inv_Z = 1.0 / new_depth;
+
+  u_rgb = (rgb_fx*xyz_rgb.x() + rgb_Tx)*inv_Z + rgb_cx + 0.5;
+  if (u_rgb < 0 || u_rgb >= width)
+    return false;
+
+  v_rgb = (rgb_fy*xyz_rgb.y() + rgb_Ty)*inv_Z + rgb_cy + 0.5;
+  if (v_rgb < 0 || v_rgb >= height)
+    return false;
+
+  return true;
+}
+
 template<typename T>
 void RegisterNodelet::convert(const sensor_msgs::ImageConstPtr& depth_msg,
                               const sensor_msgs::ImagePtr& registered_msg,
@@ -207,93 +259,100 @@ void RegisterNodelet::convert(const sensor_msgs::ImageConstPtr& depth_msg,
   DepthTraits<T>::initializeBuffer(registered_msg->data);
 
   // Extract all the parameters we need
-  double inv_depth_fx = 1.0 / depth_model_.fx();
-  double inv_depth_fy = 1.0 / depth_model_.fy();
-  double depth_cx = depth_model_.cx(), depth_cy = depth_model_.cy();
-  double depth_Tx = depth_model_.Tx(), depth_Ty = depth_model_.Ty();
-  double rgb_fx = rgb_model_.fx(), rgb_fy = rgb_model_.fy();
-  double rgb_cx = rgb_model_.cx(), rgb_cy = rgb_model_.cy();
-  double rgb_Tx = rgb_model_.Tx(), rgb_Ty = rgb_model_.Ty();
-
+  const double inv_depth_fx = 1.0 / depth_model_.fx();
+  const double inv_depth_fy = 1.0 / depth_model_.fy();
+
+  const int dst_width = static_cast<int>(registered_msg->width);
+  const int dst_height = static_cast<int>(registered_msg->height);
+
   // Transform the depth values into the RGB frame
   /// @todo When RGB is higher res, interpolate by rasterizing depth triangles onto the registered image  
   const T* depth_row = reinterpret_cast<const T*>(&depth_msg->data[0]);
-  int row_step = depth_msg->step / sizeof(T);
+  const int row_step = depth_msg->step / sizeof(T);
   T* registered_data = reinterpret_cast<T*>(&registered_msg->data[0]);
-  int raw_index = 0;
+  // TODO(lucasw) this isn't used
+  // int raw_index = 0;
   for (unsigned v = 0; v < depth_msg->height; ++v, depth_row += row_step)
   {
-    for (unsigned u = 0; u < depth_msg->width; ++u, ++raw_index)
+    for (unsigned u = 0; u < depth_msg->width; ++u  /* , ++raw_index */)
     {
-      T raw_depth = depth_row[u];
+      const T raw_depth = depth_row[u];
       if (!DepthTraits<T>::valid(raw_depth))
         continue;
-      
-      double depth = DepthTraits<T>::toMeters(raw_depth);
+
+      const double depth = DepthTraits<T>::toMeters(raw_depth);
 
       if (fill_upsampling_holes_ == false)
       {
-        /// @todo Combine all operations into one matrix multiply on (u,v,d)
-        // Reproject (u,v,Z) to (X,Y,Z,1) in depth camera frame
-        Eigen::Vector4d xyz_depth;
-        xyz_depth << ((u - depth_cx)*depth - depth_Tx) * inv_depth_fx,
-                     ((v - depth_cy)*depth - depth_Ty) * inv_depth_fy,
-                     depth,
-                     1;
-
-        // Transform to RGB camera frame
-        Eigen::Vector4d xyz_rgb = depth_to_rgb * xyz_depth;
-
-        // Project to (u,v) in RGB image
-        double inv_Z = 1.0 / xyz_rgb.z();
-        int u_rgb = (rgb_fx*xyz_rgb.x() + rgb_Tx)*inv_Z + rgb_cx + 0.5;
-        int v_rgb = (rgb_fy*xyz_rgb.y() + rgb_Ty)*inv_Z + rgb_cy + 0.5;
-
-        if (u_rgb < 0 || u_rgb >= (int)registered_msg->width ||
-            v_rgb < 0 || v_rgb >= (int)registered_msg->height)
+        int u_rgb;
+        int v_rgb;
+        T new_depth;
+        if (!transform_depth(depth_to_rgb, u, v,
+            depth_model_, rgb_model_,
+            inv_depth_fx, inv_depth_fy,
+            dst_width, dst_height,
+            depth,
+            u_rgb, v_rgb, new_depth)) {
           continue;
-
-        T& reg_depth = registered_data[v_rgb*registered_msg->width + u_rgb];
-        T  new_depth = DepthTraits<T>::fromMeters(xyz_rgb.z());
+        }
+
+        T& reg_depth = registered_data[v_rgb*dst_width + u_rgb];
         // Validity and Z-buffer checks
         if (!DepthTraits<T>::valid(reg_depth) || reg_depth > new_depth)
           reg_depth = new_depth;
+
       }
       else
       {
-        // Reproject (u,v,Z) to (X,Y,Z,1) in depth camera frame
-        Eigen::Vector4d xyz_depth_1, xyz_depth_2;
-        xyz_depth_1 << ((u-0.5f - depth_cx)*depth - depth_Tx) * inv_depth_fx,
-                       ((v-0.5f - depth_cy)*depth - depth_Ty) * inv_depth_fy,
-                       depth,
-                       1;
-        xyz_depth_2 << ((u+0.5f - depth_cx)*depth - depth_Tx) * inv_depth_fx,
-                       ((v+0.5f - depth_cy)*depth - depth_Ty) * inv_depth_fy,
-                       depth,
-                       1;
-
-        // Transform to RGB camera frame
-        Eigen::Vector4d xyz_rgb_1 = depth_to_rgb * xyz_depth_1;
-        Eigen::Vector4d xyz_rgb_2 = depth_to_rgb * xyz_depth_2;
-
-        // Project to (u,v) in RGB image
-        double inv_Z = 1.0 / xyz_rgb_1.z();
-        int u_rgb_1 = (rgb_fx*xyz_rgb_1.x() + rgb_Tx)*inv_Z + rgb_cx + 0.5;
-        int v_rgb_1 = (rgb_fy*xyz_rgb_1.y() + rgb_Ty)*inv_Z + rgb_cy + 0.5;
-        inv_Z = 1.0 / xyz_rgb_2.z();
-        int u_rgb_2 = (rgb_fx*xyz_rgb_2.x() + rgb_Tx)*inv_Z + rgb_cx + 0.5;
-        int v_rgb_2 = (rgb_fy*xyz_rgb_2.y() + rgb_Ty)*inv_Z + rgb_cy + 0.5;
-
-        if (u_rgb_1 < 0 || u_rgb_2 >= (int)registered_msg->width ||
-            v_rgb_1 < 0 || v_rgb_2 >= (int)registered_msg->height)
+        // TODO(lucasw) loop on two -0.5, 0.5 vectors, keep track of u_min/max and v_min/max
+        // as it goes
+        int u_rgb_min, u_rgb_max;
+        int v_rgb_min, v_rgb_max;
+        size_t count = 0;
+        T new_depth_sum = 0.0;
+        const std::array<double, 2> uv_offsets = {-0.5, 0.5};
+        for (const auto& v_offset : uv_offsets) {
+          for (const auto& u_offset : uv_offsets) {
+            int u_rgb;
+            int v_rgb;
+            T new_depth;
+            if (!transform_depth(depth_to_rgb,
+                u + u_offset, v + v_offset,
+                depth_model_, rgb_model_,
+                inv_depth_fx, inv_depth_fy,
+                dst_width, dst_height,
+                depth,
+                u_rgb, v_rgb, new_depth)) {
+              // TODO(lucasw) need to be able to handle partiall out of bounds squares here
+              continue;
+            }
+            if (count == 0) {
+              u_rgb_min = u_rgb;
+              u_rgb_max = u_rgb;
+              v_rgb_min = v_rgb;
+              v_rgb_max = v_rgb;
+            } else {
+              u_rgb_min = std::min(u_rgb, u_rgb_min);
+              u_rgb_max = std::max(u_rgb, u_rgb_max);
+              v_rgb_min = std::min(v_rgb, v_rgb_min);
+              v_rgb_max = std::max(v_rgb, v_rgb_max);
+            }
+            new_depth_sum += new_depth;
+            count++;
+          }
+        }
+
+        if (count == 0) {
           continue;
+        }
 
-        for (int nv=v_rgb_1; nv<=v_rgb_2; ++nv)
+        // fill in the square defined by uv range
+        const T new_depth = DepthTraits<T>::fromMeters(new_depth_sum / static_cast<double>(count));
+        for (int nv=v_rgb_min; nv<=v_rgb_max; ++nv)
         {
-          for (int nu=u_rgb_1; nu<=u_rgb_2; ++nu)
+          for (int nu=u_rgb_min; nu<=u_rgb_max; ++nu)
           {
-            T& reg_depth = registered_data[nv*registered_msg->width + nu];
-            T  new_depth = DepthTraits<T>::fromMeters(0.5*(xyz_rgb_1.z()+xyz_rgb_2.z()));
+            T& reg_depth = registered_data[nv*dst_width + nu];
             // Validity and Z-buffer checks
             if (!DepthTraits<T>::valid(reg_depth) || reg_depth > new_depth)
               reg_depth = new_depth;

depth_image_proc/test/test_register.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python
+
+import copy
+import unittest
+
+import numpy as np
+import rospy
+import rostest
+from cv_bridge import CvBridge
+from sensor_msgs.msg import CameraInfo
+from sensor_msgs.msg import Image
+
+
+class TestRegister(unittest.TestCase):
+    def test_register(self):
+        # publish a 2x2 float depth image with depths like
+        # [1.0, 2.0], [2.0, 1.0]
+        # publish a camera info with a aov of around 90 degrees
+        # publish two static tfs in the .test file, one shifted
+        # by 1.0 in x from the depth camera_info frame
+        # another by 2.0 in x, and 1.5 in z, with a 90 degree rotation
+        # looking back at those depth points
+        # publish another camera info, in series, have it shift from one frame to the other
+        # subscribe to the depth_image_proc node and verify output
+        # (make numbers come out even with 90 degree aov)
+
+        self.output_depth = rospy.get_param("~output_depth", 1.0)
+        self.expected_depth = rospy.get_param("~expected_depth", 7.0)
+        rospy.loginfo(f"published depth: {self.output_depth}, expected input depth: {self.expected_depth}")
+
+        self.cv_bridge = CvBridge()
+        self.depth_image_pub = rospy.Publisher("depth/image_rect", Image, queue_size=2)
+        self.depth_ci_pub = rospy.Publisher("depth/camera_info", CameraInfo, queue_size=2)
+        self.rgb_ci_pub = rospy.Publisher("rgb/camera_info", CameraInfo, queue_size=2)
+
+        self.received_msg = None
+
+        # TODO(lucasw) use time sync subscriber
+        self.depth_sub = rospy.Subscriber("depth_registered/image_rect", Image, self.depth_callback, queue_size=2)
+        self.ci_sub = rospy.Subscriber("depth_registered/camera_info", CameraInfo, self.ci_callback, queue_size=2)
+
+        ci_msg = CameraInfo()
+        wd = 3
+        ht = 3
+        ci_msg.header.frame_id = "station1"
+        ci_msg.height = ht
+        ci_msg.width = wd
+        ci_msg.distortion_model = "plumb_bob"
+
+        cx = 1.5
+        cy = 1.5
+        fx = 1.5
+        fy = 1.5
+        ci_msg.K = [fx, 0.0, cx,
+                    0.0, fy, cy,
+                    0.0, 0.0, 1.0]
+        ci_msg.R = [1, 0, 0,
+                    0, 1, 0,
+                    0, 0, 1]
+        ci_msg.P = [fx, 0.0, cx, 0.0,
+                    0.0, fy, cy, 0.0,
+                    0.0, 0.0, 1.0, 0.0]
+
+        depth = np.ones((ht, wd, 1), np.float32) * self.output_depth
+        rospy.loginfo(depth)
+
+        depth_msg = self.cv_bridge.cv2_to_imgmsg(depth, "32FC1")
+        rgb_ci_msg = copy.deepcopy(ci_msg)
+        rgb_ci_msg.header.frame_id = "station2"
+        depth_msg.header = ci_msg.header
+
+        for i in range(6):
+            rospy.loginfo(f"{i} waiting for register connections...")
+            num_im_sub = self.depth_image_pub.get_num_connections()
+            num_ci_sub = self.depth_ci_pub.get_num_connections()
+            num_rgb_ci_sub = self.rgb_ci_pub.get_num_connections()
+            rospy.sleep(1)
+            if num_im_sub > 0 and num_ci_sub > 0 and num_rgb_ci_sub > 0:
+                break
+
+        self.assertGreater(self.depth_image_pub.get_num_connections(), 0)
+        self.assertGreater(self.depth_ci_pub.get_num_connections(), 0)
+        self.assertGreater(self.rgb_ci_pub.get_num_connections(), 0)
+        rospy.sleep(1.0)
+
+        self.count = 0
+        rospy.loginfo("publishing depth and ci, wait for callbacks")
+        for i in range(4):
+            stamp = rospy.Time.now()
+            ci_msg.header.stamp = stamp
+            rgb_ci_msg.header.stamp = stamp
+            depth_msg.header.stamp = stamp
+            self.depth_image_pub.publish(depth_msg)
+            self.depth_ci_pub.publish(ci_msg)
+            self.rgb_ci_pub.publish(rgb_ci_msg)
+            rospy.sleep(0.1)
+
+        while i in range(6):
+            if self.count > 1:
+                break
+            rospy.sleep(1)
+        rospy.loginfo("done waiting")
+
+        self.assertIsNotNone(self.received_msg, "no depth message received")
+        registered_depth = self.cv_bridge.imgmsg_to_cv2(self.received_msg, "32FC1")
+        # the edges of the expected 3x3 image may be nans, but the center should be valid
+        valid_depth = registered_depth[1, 1]
+        self.assertAlmostEqual(valid_depth, self.expected_depth, places=1)
+
+    def depth_callback(self, msg):
+        rospy.loginfo("received depth")
+        self.received_msg = msg
+        self.count += 1
+
+    def ci_callback(self, msg):
+        rospy.logdebug("received camera info")
+
+
+if __name__ == "__main__":
+    node_name = 'test_register'
+    rospy.init_node(node_name)
+    # rostest.rosrun("depth_image_proc", node_name, sys.argv)
+    rostest.rosrun("depth_image_proc", node_name, TestRegister)