openni_range_image_registration.cpp

Go to the documentation of this file.

/* \author Bastian Steder */

#include <iostream>
using namespace std;
#include <pcl/io/openni_grabber.h>
#include "pcl/range_image/range_image_planar.h"
#include "pcl/common/common_headers.h"
#include "pcl/visualization/range_image_visualizer.h"
#include "pcl/visualization/pcl_visualizer.h"
#include "pcl/console/parse.h"
#include <boost/filesystem.hpp> 

using namespace pcl;
using namespace pcl::visualization;

std::string device_id = "#1";
std::string output_prefix = "openni_icp";
std::string dir_name;

float angular_resolution = 0.5f;

float size_factor = 1.0;

int icp_search_radius_consecutive = 3;
int icp_num_iterations_consecutive = 20;
int icp_search_radius_keyframe = 1;
int icp_num_iterations_keyframe = 10;
float new_keyframe_overlap = 0.85;

float maximum_range = -1.0f;

int max_no_of_threads = 1;
bool capture = false;

bool show_color_image = 0,
     show_range_image = 1,
     show_registered_clouds = 1;

boost::mutex image_mutex;
pcl::PointCloud<pcl::PointXYZ>::ConstPtr point_cloud_ptr;
boost::shared_ptr<openni_wrapper::DepthImage> depth_image_ptr;
boost::shared_ptr<openni_wrapper::Image> color_image_ptr;

bool received_new_image_pair = false,
     received_new_depth_data = false;

int range_image_counter=0, key_frame_counter=0;

struct EventHelper
{
  void 
  depthAndColorCB (const boost::shared_ptr<openni_wrapper::Image>& color_image,
                   const boost::shared_ptr<openni_wrapper::DepthImage>& depth_image, float constant)
  {
    (void)constant;  // No warning
    if (!image_mutex.try_lock ())
      return;
    depth_image_ptr = depth_image;
    color_image_ptr = color_image;
    image_mutex.unlock ();
    received_new_image_pair = true;
  }
  
  void
  depth_image_cb (const boost::shared_ptr<openni_wrapper::DepthImage>& depth_image)
  {
    if (!image_mutex.try_lock ())
      return;
    depth_image_ptr = depth_image;
    image_mutex.unlock ();
    received_new_depth_data = true;
  }
};

Eigen::Affine3f last_pose, last_keyframe_pose, current_pose=Eigen::Affine3f::Identity();

void
saveFrame()
{
  if (!capture)
    return;
  
  stringstream filename_stream;
  filename_stream << setfill('0')<<setw(6)<<key_frame_counter;
  
  std::string info_filename = dir_name+"/"+filename_stream.str()+"_info.txt";
  std::ofstream info_file(info_filename.c_str());
  
  std::string depth_image_filename = dir_name+"/"+filename_stream.str()+"_depth_image.dat";
  const unsigned short* depth_map = depth_image_ptr->getDepthMetaData().Data();
  int width_depth_image = depth_image_ptr->getWidth (), height_depth_image = depth_image_ptr->getHeight ();
  int depth_image_size = width_depth_image*height_depth_image*sizeof(*depth_map);
  std::ofstream depth_image_file(depth_image_filename.c_str());
  depth_image_file.write((const char*)depth_map, depth_image_size);
  depth_image_file.close();
  
  info_file << "depth_image_filename "<<getFilenameWithoutPath (depth_image_filename)<<"\n"
            << "time_stamp_depth_image "  << depth_image_ptr->getTimeStamp () << "\n"
            << "frame_id_depth_image "  << depth_image_ptr->getFrameID () << "\n"
            << "width_depth_image "  << width_depth_image  << "\n"
            << "height_depth_image " << height_depth_image << "\n"
            << "focal_length_depth_image "  << depth_image_ptr->getFocalLength() << "\n";
  
  std::string color_image_filename = dir_name+"/"+filename_stream.str()+"_color_image.dat";
  const unsigned char* color_map = color_image_ptr->getMetaData().Data();
  int width_color_image = color_image_ptr->getWidth (),
      height_color_image = color_image_ptr->getHeight ();
  openni_wrapper::Image::Encoding encoding = color_image_ptr->getEncoding();
  int color_image_size = 0;
  if (encoding==openni_wrapper::Image::RGB)
    color_image_size = 3 * width_color_image*height_color_image;
  else if(encoding==openni_wrapper::Image::BAYER_GRBG)
    color_image_size = width_color_image*height_color_image;
  else if(encoding==openni_wrapper::Image::YUV422)
    color_image_size = 2 * width_color_image*height_color_image;
  std::ofstream color_image_file(color_image_filename.c_str());
  color_image_file.write((const char*)color_map, color_image_size);
  color_image_file.close();
  
  info_file << "color_image_filename "<<getFilenameWithoutPath (color_image_filename)<<"\n"
            << "time_stamp_color_image "  << color_image_ptr->getTimeStamp () << "\n"
            << "frame_id_color_image "  << color_image_ptr->getFrameID () << "\n"
            << "width_color_image "  << width_color_image  << "\n"
            << "height_color_image " << height_color_image << "\n";
  if (encoding==openni_wrapper::Image::RGB)
    info_file << "color_image_encoding RGB\n";
  else if(encoding==openni_wrapper::Image::BAYER_GRBG)
    info_file << "color_image_encoding BAYER_GRBG\n";
  else if(encoding==openni_wrapper::Image::YUV422)
    info_file << "color_image_encoding YUV422\n";
  
  Eigen::Affine3f pose_diff = Eigen::Affine3f::Identity();
  if (key_frame_counter > 1) {
    pose_diff = current_pose.inverse() * last_keyframe_pose;
  }
  Eigen::Vector3f translation = pose_diff.translation();
  Eigen::Quaternionf rotation = Eigen::Quaternionf(pose_diff.rotation());
  info_file << "pose_diff_translation_xyz "<<translation.x()<<" "<<translation.y()<<" "<<translation.z()<<"\n";
  info_file << "pose_diff_rotation_wxyz "<<rotation.w()<<" "<<rotation.x()<<" "<<rotation.y()<<" "<<rotation.z()<<"\n";

  
  info_file.close();
}

void
printUsage (const char* progName)
{
  cout << "\n\nUsage: "<<progName<<" [options] [scene.pcd] <model.pcl> [model_2.pcl] ... [model_n.pcl]\n\n"
       << "Options:\n"
       << "-------------------------------------------\n"
       << "-r <float>      angular resolution in degrees (default "<<angular_resolution<<")\n"
       << "-d <device_id>  set the device id (default \""<<device_id<<"\")\n"
       << "-t <int>        set the maximum number of threads for parallelization (default: "<<max_no_of_threads<<")\n"
       << "-c              save keyframes to disc\n"
       << "-m <float>      maximum range\n"
       << "-s <float>      size factor (default: "<<size_factor<<")\n"
       << "-h              this help\n"
       << "\n\n";
}

int main (int argc, char** argv)
{
  // --------------------------------------
  // -----Parse Command Line Arguments-----
  // --------------------------------------
  if (pcl::console::find_argument (argc, argv, "-h") >= 0)
  {
    printUsage (argv[0]);
    return 0;
  }
  if (pcl::console::find_argument (argc, argv, "-c") >= 0)
  {
    capture = true;
    cout << "Switching capturing on.\n";
  }
  if (pcl::console::parse (argc, argv, "-d", device_id) >= 0)
    cout << "Using device id \""<<device_id<<"\".\n";
  if (pcl::console::parse (argc, argv, "-r", angular_resolution) >= 0)
    cout << "Setting angular resolution to "<<angular_resolution<<"deg.\n";
  if (pcl::console::parse (argc, argv, "-m", maximum_range) >= 0)
    cout << "Setting maximum range to "<<maximum_range<<"m.\n";
  angular_resolution = deg2rad (angular_resolution);
  if (pcl::console::parse (argc, argv, "-t", max_no_of_threads) >= 0)
    cout << "Setting maximum number of threads to "<< max_no_of_threads<<".\n";
  if (pcl::console::parse (argc, argv, "-s", size_factor) >= 0)
    cout << "Setting size factor to "<<size_factor<<".\n";
  
  float icp_max_distance_start_consecutive=size_factor*0.3, icp_max_distance_end_consecutive=size_factor*0.05;
  float icp_max_distance_start_keyframe=size_factor*0.15, icp_max_distance_end_keyframe=size_factor*0.05;
  float overlap_calculation_max_error = size_factor*0.25;
  
  RangeImage::max_no_of_threads = max_no_of_threads;
  
  RangeImageVisualizer range_image_widget ("Range Image"),
                       ir_image_widget ("IR Image"),
                       color_image_widget ("RGB Image");
  
  PCLVisualizer viewer ("3D Viewer");
  viewer.addCoordinateSystem (1.0f);
  viewer.setBackgroundColor (1, 1, 1);
  
  viewer.initCameraParameters ();
  viewer.camera_.pos[0] = 0.0;
  viewer.camera_.pos[1] = -0.3;
  viewer.camera_.pos[2] = -2.0;
  viewer.camera_.focal[0] = 0.0;
  viewer.camera_.focal[1] = 0.0+viewer.camera_.pos[1];
  viewer.camera_.focal[2] = 1.0;
  viewer.camera_.view[0] = 0.0;
  viewer.camera_.view[1] = -1.0;
  viewer.camera_.view[2] = 0.0;
  viewer.updateCamera();

  openni_wrapper::OpenNIDriver& driver = openni_wrapper::OpenNIDriver::getInstance ();
  if (driver.getNumberDevices () > 0)
  {
    for (unsigned deviceIdx = 0; deviceIdx < driver.getNumberDevices (); ++deviceIdx)
    {
      cout << "Device: " << deviceIdx + 1 << ", vendor: " << driver.getVendorName (deviceIdx)
           << ", product: " << driver.getProductName (deviceIdx) << ", connected: "
           << (int) driver.getBus (deviceIdx) << " @ " << (int) driver.getAddress (deviceIdx)
           << ", serial number: \'" << driver.getSerialNumber (deviceIdx) << "\'\n";
    }
  }
  else
  {
    cout << "\nNo devices connected.\n\n";
    return 1;
  }
  
  pcl::Grabber* interface = new pcl::OpenNIGrabber (device_id);
  EventHelper event_helper;
  
  //boost::function<void (const boost::shared_ptr<openni_wrapper::DepthImage>&) > f_depth_image =
    //boost::bind (&EventHelper::depth_image_cb, &event_helper, _1);
  //boost::signals2::connection c_depth_image = interface->registerCallback (f_depth_image);
  
  boost::function<void (const boost::shared_ptr<openni_wrapper::Image>&,
                        const boost::shared_ptr<openni_wrapper::DepthImage>&,
                        float) > f_image_pair =
    boost::bind (&EventHelper::depthAndColorCB, &event_helper, _1, _2, _3);
  boost::signals2::connection c_image_pair = interface->registerCallback (f_image_pair);
  
  cout << "Starting grabber\n";
  interface->start ();
  cout << "Done\n";
  
  boost::shared_ptr<RangeImagePlanar> range_image_ptr (new RangeImagePlanar);
  RangeImagePlanar& range_image = *range_image_ptr;
  boost::shared_ptr<RangeImagePlanar> last_range_image_ptr (new RangeImagePlanar);
  RangeImagePlanar& last_range_image = *last_range_image_ptr;
  boost::shared_ptr<RangeImagePlanar> last_keyframe_range_image_ptr (new RangeImagePlanar);
  RangeImagePlanar& last_keyframe_range_image = *last_keyframe_range_image_ptr;
  
  if (capture)
  {
    boost::posix_time::ptime current_time(boost::posix_time::second_clock::local_time());
    stringstream ss;
    boost::posix_time::time_facet* time_format = new boost::posix_time::time_facet("%Y_%m_%d__%H_%M_%S");
    ss.imbue(std::locale(ss.getloc(), time_format));
    ss << output_prefix << "_" << current_time;
    dir_name = ss.str();
    cout << "Creating directory \""<<dir_name<<"\" to save captured data.\n";
    boost::filesystem::create_directory(dir_name);
  }
  
  double last_50_frames_start_time = 0;
  
  while (true)
  {
    if (show_registered_clouds)
      viewer.spinOnce (10);
    if (show_color_image || show_range_image)
      ImageWidgetWX::spinOnce ();  // process GUI events
    boost::this_thread::sleep (boost::posix_time::microseconds (1000));
    
    bool got_new_range_image = false;
    if (!received_new_image_pair || !image_mutex.try_lock ())
      continue;
    
    received_new_image_pair = false;
    
    // Depth Image
    unsigned long time_stamp_depth_image = depth_image_ptr->getTimeStamp ();
    int frame_id_depth_image = depth_image_ptr->getFrameID ();
    const unsigned short* depth_map = depth_image_ptr->getDepthMetaData().Data();
    int width_depth_image = depth_image_ptr->getWidth (), height_depth_image = depth_image_ptr->getHeight ();
    float focal_length_x = depth_image_ptr->getFocalLength(), focal_length_y = focal_length_x;
    float original_angular_resolution = asinf (0.5f*float(width_depth_image)/float(focal_length_x)) /
                                              (0.5f*float(width_depth_image));
    float desired_angular_resolution = angular_resolution;
    std::swap(last_range_image, range_image);
    range_image.setDepthImage (depth_map, width_depth_image, height_depth_image,
                               width_depth_image/2, height_depth_image/2,
                               focal_length_x, focal_length_y, desired_angular_resolution);
    
    if (maximum_range > 0.0f)
    {
      for (int point_idx=0; point_idx<range_image.points.size(); ++point_idx)
        if (range_image.points[point_idx].range > maximum_range)
          range_image.points[point_idx].range = std::numeric_limits<float>::infinity ();
    }

    if (show_range_image)
      range_image_widget.setRangeImage (range_image, 0.0f, 10.0f);
    
    // Color Image
    unsigned long time_stamp_color_image = color_image_ptr->getTimeStamp ();
    int frame_id_color_image = color_image_ptr->getFrameID ();
    openni_wrapper::Image::Encoding encoding = color_image_ptr->getEncoding();
    //if (encoding==openni_wrapper::Image::RGB)
      //cout << "Image encoding is RGB.\n";
    //else if(encoding==openni_wrapper::Image::BAYER_GRBG)
      //cout << "Image encoding is BAYER_GRBG.\n";
    //else if(encoding==openni_wrapper::Image::YUV422)
      //cout << "Image encoding is YUV422.\n";
    const unsigned char* color_map = color_image_ptr->getMetaData().Data();
    int width_color_image = color_image_ptr->getWidth (), height_color_image = color_image_ptr->getHeight ();
    
    if (show_color_image)
    {
      if (encoding==openni_wrapper::Image::RGB)
        color_image_widget.setRGBImage(color_map, width_color_image, height_color_image);
      else if(encoding==openni_wrapper::Image::BAYER_GRBG)
        color_image_widget.setBayerGRBGImage(color_map, width_color_image, height_color_image);
    }
    
    if (!range_image.points.empty())
    {
      got_new_range_image = true;
      ++range_image_counter;
    }
    
    if (range_image_counter == 1)
    {
      last_keyframe_range_image = range_image;
      key_frame_counter = 1;
      last_pose = last_keyframe_pose = current_pose = Eigen::Affine3f::Identity();
      last_50_frames_start_time = getTime();
      cout << "k" << std::flush;
      saveFrame();
      image_mutex.unlock ();
      continue;
    }
    
    double icp_time = -getTime();
    Eigen::Affine3f consecutive_pose_diff_initial_guess = Eigen::Affine3f::Identity(),
                    keyframe_pose_diff_initial_guess = Eigen::Affine3f::Identity();
    last_pose = current_pose;
    Eigen::Affine3f consecutive_icp_pose_diff = range_image.doIcp(last_range_image,
                      consecutive_pose_diff_initial_guess, icp_search_radius_consecutive,
                      icp_max_distance_start_consecutive, icp_max_distance_end_consecutive,
                      icp_num_iterations_consecutive, 3, 1),
                    consecutive_icp_pose = last_pose * consecutive_icp_pose_diff.inverse();
    keyframe_pose_diff_initial_guess = consecutive_icp_pose.inverse() * last_keyframe_pose;
    
    Eigen::Affine3f keyframe_icp_pose_diff = range_image.doIcp(last_keyframe_range_image,
                                               keyframe_pose_diff_initial_guess, icp_search_radius_keyframe,
                                               icp_max_distance_start_keyframe, icp_max_distance_end_keyframe,
                                               icp_num_iterations_keyframe, 3, 1),
                    keyframe_icp_pose = last_keyframe_pose * keyframe_icp_pose_diff.inverse();
    //Eigen::Affine3f keyframe_icp_pose_diff = keyframe_pose_diff_initial_guess,
                    //keyframe_icp_pose = last_keyframe_pose * keyframe_icp_pose_diff.inverse();
    icp_time += getTime();
    //cout << PVARN(icp_time);

    double overlap_time = -getTime();
    // TODO: error range dependent
    float overlap = std::min(
      range_image.getOverlap(last_keyframe_range_image, keyframe_icp_pose_diff,
                             2, overlap_calculation_max_error, 3),
      last_keyframe_range_image.getOverlap(range_image, keyframe_icp_pose_diff.inverse(),
                                           1, overlap_calculation_max_error, 1));
    overlap_time += getTime();
    //cout << PVARN(overlap_time);
    
    current_pose = keyframe_icp_pose;
    
    if (overlap < new_keyframe_overlap)
    {
      cout << "k" << std::flush;
      ++key_frame_counter;
      saveFrame();
      last_keyframe_range_image = range_image;
      last_keyframe_pose = current_pose;
    }
    else
      cout << "." << std::flush;
    //cout << overlap << " " << std::flush;
    
    
    image_mutex.unlock ();
    
    if (range_image_counter%50 == 0)
    {
      if (range_image_counter > 1)
        cout << "\nFramerate: "<<50.0f / (getTime()-last_50_frames_start_time)<<"fps.\n";
      last_50_frames_start_time = getTime();
    }
    
    
    if (show_registered_clouds)
    {
      PointCloud<PointXYZ>::Ptr transformed_cloud_ptr(new PointCloud<PointXYZ>);
      PointCloud<PointXYZ>& transformed_cloud = *transformed_cloud_ptr;
      for (unsigned int i=0; i<range_image.points.size(); ++i)
      {
        if (range_image.isValid(i))
        {
          transformed_cloud.points.push_back(PointXYZ());
          transformed_cloud.points.back().getVector3fMap() = current_pose * range_image.points[i].getVector3fMap();
        }
      }
      
      PointCloudColorHandlerCustom<pcl::PointXYZ> color_handler_transformed_cloud (transformed_cloud_ptr, 0, 0, 0);
      if (!viewer.updatePointCloud<pcl::PointXYZ> (transformed_cloud_ptr, color_handler_transformed_cloud,
                                                   "transformed cloud"))
        viewer.addPointCloud<pcl::PointXYZ> (transformed_cloud_ptr, color_handler_transformed_cloud,
                                             "transformed cloud");
    }
  }
  
  interface->stop ();
}