Apply will also undistort the depth image.

Improved speed, there was still a double conversion in one if statement.

Author: Thiemo Wiedemeyer

examples/protonect/Protonect.cpp

@@ -124,7 +124,7 @@ int main(int argc, char *argv[])
 
   libfreenect2::SyncMultiFrameListener listener(libfreenect2::Frame::Color | libfreenect2::Frame::Ir | libfreenect2::Frame::Depth);
   libfreenect2::FrameMap frames;
-  libfreenect2::Frame registered(512, 424, 3);
+  libfreenect2::Frame undistorted(512, 424, 4), registered(512, 424, 3);
 
   dev->setColorFrameListener(&listener);
   dev->setIrAndDepthFrameListener(&listener);
@@ -146,7 +146,9 @@ int main(int argc, char *argv[])
     cv::imshow("ir", cv::Mat(ir->height, ir->width, CV_32FC1, ir->data) / 20000.0f);
     cv::imshow("depth", cv::Mat(depth->height, depth->width, CV_32FC1, depth->data) / 4500.0f);
 
-    registration->apply(rgb,depth,&registered);
+    registration->apply(rgb,depth,&undistorted,&registered);
+
+    cv::imshow("undistorted", cv::Mat(undistorted.height, undistorted.width, CV_32FC1, undistorted.data) / 4500.0f);
     cv::imshow("registered", cv::Mat(registered.height, registered.width, CV_8UC3, registered.data));
 
     int key = cv::waitKey(1);

examples/protonect/include/libfreenect2/registration.h

@@ -44,19 +44,19 @@ class LIBFREENECT2_API Registration
   void apply(int dx, int dy, float dz, float& cx, float &cy) const;
 
   // undistort/register a whole image
-  void apply(const Frame* rgb, const Frame* depth, Frame* registered) const;
+  void apply(const Frame* rgb, const Frame* depth, Frame* undistorted, Frame* registered) const;
 
 private:
-  void undistort_depth(int dx, int dy, float& mx, float& my);
-  void depth_to_color(float mx, float my, float& rx, float& ry);
+  void distort(int mx, int my, float& dx, float& dy) const;
+  void depth_to_color(float mx, float my, float& rx, float& ry) const;
 
   Freenect2Device::IrCameraParams depth;
   Freenect2Device::ColorCameraParams color;
 
-  float undistort_map[512 * 424 * 2];
+  int distort_map[512 * 424];
   float depth_to_color_map_x[512 * 424];
   float depth_to_color_map_y[512 * 424];
-  int depth_to_color_map_i[512 * 424];
+  int depth_to_color_map_yi[512 * 424];
 };
 
 } /* namespace libfreenect2 */

examples/protonect/src/registration.cpp

@@ -39,26 +39,23 @@ namespace libfreenect2
 static const float depth_q = 0.01;
 static const float color_q = 0.002199;
 
-void Registration::undistort_depth(int x, int y, float& mx, float& my)
+void Registration::distort(int mx, int my, float& x, float& y) const
 {
-  float dx = ((float)x - depth.cx) / depth.fx;
-  float dy = ((float)y - depth.cy) / depth.fy;
-
-  float ps = (dx * dx) + (dy * dy);
-  float qs = ((ps * depth.k3 + depth.k2) * ps + depth.k1) * ps + 1.0;
-  for (int i = 0; i < 9; i++) {
-    float qd = ps / (qs * qs);
-    qs = ((qd * depth.k3 + depth.k2) * qd + depth.k1) * qd + 1.0;
-  }
-
-  mx = dx / qs;
-  my = dy / qs;
+  float dx = ((float)mx - depth.cx) / depth.fx;
+  float dy = ((float)my - depth.cy) / depth.fy;
+  float dx2 = dx * dx;
+  float dy2 = dy * dy;
+  float r2 = dx2 + dy2;
+  float dxdy2 = 2 * dx * dy;
+  float kr = 1 + ((depth.k3 * r2 + depth.k2) * r2 + depth.k1) * r2;
+  x = depth.fx * (dx * kr + depth.p2 * (r2 + 2 * dx2) + depth.p1 * dxdy2) + depth.cx;
+  y = depth.fy * (dy * kr + depth.p1 * (r2 + 2 * dy2) + depth.p2 * dxdy2) + depth.cy;
 }
 
-void Registration::depth_to_color(float mx, float my, float& rx, float& ry)
+void Registration::depth_to_color(float mx, float my, float& rx, float& ry) const
 {
-  mx *= depth.fx * depth_q;
-  my *= depth.fy * depth_q;
+  mx = (mx - depth.cx) * depth_q;
+  my = (my - depth.cy) * depth_q;
 
   float wx =
     (mx * mx * mx * color.mx_x3y0) + (my * my * my * color.mx_x0y3) +
@@ -86,71 +83,90 @@ void Registration::apply( int dx, int dy, float dz, float& cx, float &cy) const
   cx = rx * color.fx + color.cx;
 }
 
-void Registration::apply(const Frame* rgb, const Frame* depth, Frame* registered) const
+void Registration::apply(const Frame *rgb, const Frame *depth, Frame *undistorted, Frame *registered) const
 {
-  if (!depth || !rgb || !registered ||
-      depth->width != 512 || depth->height != 424 || depth->bytes_per_pixel != 4 ||
+  if (!undistorted || !rgb || !registered ||
       rgb->width != 1920 || rgb->height != 1080 || rgb->bytes_per_pixel != 3 ||
+      depth->width != 512 || depth->height != 424 || depth->bytes_per_pixel != 4 ||
+      undistorted->width != 512 || undistorted->height != 424 || undistorted->bytes_per_pixel != 4 ||
       registered->width != 512 || registered->height != 424 || registered->bytes_per_pixel != 3)
     return;
 
-  const float *depth_raw = (float*)depth->data;
+  const float *depth_data = (float*)depth->data;
+  float *undistorted_data = (float*)undistorted->data;
+  unsigned char *registered_data = registered->data;
+  const int *map_dist = distort_map;
   const float *map_x = depth_to_color_map_x;
-  const int *map_i = depth_to_color_map_i;
-  unsigned char *registered_raw = registered->data;
+  const int *map_yi = depth_to_color_map_yi;
   const int size_depth = 512 * 424;
   const int size_color = 1920 * 1080 * 3;
 
-  for (int i = 0; i < size_depth; ++i, ++registered_raw, ++map_x, ++map_i, ++depth_raw) {
-    const float z_raw = *depth_raw;
+  for (int i = 0; i < size_depth; ++i, ++registered_data, ++undistorted_data, ++map_dist, ++map_x, ++map_yi) {
+    const int index = *map_dist;
+
+    if(index < 0){
+      *undistorted_data = 0;
+      *registered_data = 0;
+      *++registered_data = 0;
+      *++registered_data = 0;
+      continue;
+    }
+
+    const float z_raw = depth_data[index];
+    *undistorted_data = z_raw;
 
-    if (z_raw == 0.0) {
-      *registered_raw = 0;
-      *++registered_raw = 0;
-      *++registered_raw = 0;
+    if (z_raw <= 0.0f) {
+      *registered_data = 0;
+      *++registered_data = 0;
+      *++registered_data = 0;
       continue;
     }
 
     const float rx = (*map_x + (color.shift_m / z_raw)) * color.fx + color.cx;
     const int cx = rx + 0.5f; // same as round for positive numbers
-    const int cy = *map_i;
+    const int cy = *map_yi;
     const int c_off = cx * 3 + cy;
 
     if (c_off < 0 || c_off > size_color || rx < -0.5f) {
-      *registered_raw = 0;
-      *++registered_raw = 0;
-      *++registered_raw = 0;
+      *registered_data = 0;
+      *++registered_data = 0;
+      *++registered_data = 0;
       continue;
     }
 
     const unsigned char *rgb_data = rgb->data + c_off;
-    *registered_raw = *rgb_data;
-    *++registered_raw = *++rgb_data;
-    *++registered_raw = *++rgb_data;
+    *registered_data = *rgb_data;
+    *++registered_data = *++rgb_data;
+    *++registered_data = *++rgb_data;
   }
 }
 
 Registration::Registration(Freenect2Device::IrCameraParams depth_p, Freenect2Device::ColorCameraParams rgb_p):
   depth(depth_p), color(rgb_p)
 {
   float mx, my;
+  int ix, iy, index;
   float rx, ry;
-  float *it_undist = undistort_map;
+  int *map_dist = distort_map;
   float *map_x = depth_to_color_map_x;
   float *map_y = depth_to_color_map_y;
-  int *map_i = depth_to_color_map_i;
+  int *map_yi = depth_to_color_map_yi;
 
   for (int y = 0; y < 424; y++) {
     for (int x = 0; x < 512; x++) {
-
-      undistort_depth(x,y,mx,my);
-      *it_undist++ = mx;
-      *it_undist++ = my;
-
-      depth_to_color(mx,my,rx,ry);
+      distort(x,y,mx,my);
+      ix = roundf(mx);
+      iy = roundf(my);
+      if(ix < 0 || ix >= 512 || iy < 0 || iy >= 424)
+        index = -1;
+      else
+        index = iy * 512 + ix;
+      *map_dist++ = index;
+
+      depth_to_color(x,y,rx,ry);
       *map_x++ = rx;
       *map_y++ = ry;
-      *map_i++ = round(ry) * 1920 * 3;
+      *map_yi++ = roundf(ry) * 1920 * 3;
     }
   }
 }