register_box.cpp

Go to the documentation of this file.

#include <simple_object_capture/common.h>

#include <pcl/registration/icp.h>
#include <pcl/registration/registration.h>
#include <pcl/registration/transformation_estimation.h>
#include <pcl/registration/transformation_estimation_point_to_plane.h>
#include <pcl/registration/transformation_estimation_point_to_plane_lls.h>
#include <pcl/point_types.h>
#include <pcl/io/pcd_io.h>
#include <pcl/features/normal_3d.h>
#include <pcl/features/integral_image_normal.h>
#include <pcl/sample_consensus/method_types.h>
#include <pcl/sample_consensus/model_types.h>
#include <pcl/segmentation/sac_segmentation.h>
#include <pcl/surface/mls.h>
#include <pcl/filters/radius_outlier_removal.h>


#include <rosbag/bag.h>
#include <rosbag/query.h>
#include <rosbag/view.h>
#include <string>
#include <sstream>
#include <fstream>
#include <sensor_msgs/CameraInfo.h>
#include <sensor_msgs/Image.h>
#include <geometry_msgs/PoseStamped.h>
#include <image_geometry/pinhole_camera_model.h>
#include <sensor_msgs/image_encodings.h>
#include <boost/filesystem.hpp>

typedef pcl::registration::TransformationEstimationPointToPlaneLLS<pcl::PointXYZRGBNormal, pcl::PointXYZRGBNormal> PointToPlane;

namespace enc = sensor_msgs::image_encodings;

template <class T>
std::vector<typename T::ConstPtr> getAllMsgs(const rosbag::Bag& bag, const std::string& topic)
{
  std::vector<typename T::ConstPtr> res;

  rosbag::View view(bag, rosbag::TopicQuery(topic));
  BOOST_FOREACH(rosbag::MessageInstance const m, view)
  {
    typename T::ConstPtr i = m.instantiate<T>();
    res.push_back(i);
  }
  return res;
}


template <class P>
void pcdWrite(typename pcl::PointCloud<P>::Ptr cloud, const std::string& name, int count, const boost::filesystem::path& folder = "")
{
  std::stringstream s;

  // create folder if needed
  if (folder != "")
  {
    s << folder.string() << "/";
    if (!boost::filesystem::is_directory(folder))
      boost::filesystem::create_directory(folder);
  }

  pcl::PCDWriter write;
  s << name;
  if (count >= 0)
  {
    s << "_";
    if (count < 10)
      s << "00";
    else if (count < 100)
      s << "0";
    s << count;
  }
  s << ".pcd";
  write.write(s.str(), *cloud);
}



int main(int argc, char** argv)
{
  // usage
  if (argc != 3)
  {
    printf("Usage: ./register_box bagfile output_folder\n");
    return -1;
  }

  // open bag file
  rosbag::Bag bag;
  std::string bagfile(argv[1]);
  std::string folder(argv[2]);
  printf("Opening bag file %s\n", bagfile.c_str());
  bag.open(bagfile, rosbag::bagmode::Read);

  // get messages from bag file
  std::vector<sensor_msgs::CameraInfo::ConstPtr> depth_info = getAllMsgs<sensor_msgs::CameraInfo>(bag, "camera/depth/camera_info");
  std::vector<sensor_msgs::Image::ConstPtr> depth_image = getAllMsgs<sensor_msgs::Image>(bag, "camera/depth/image");
  std::vector<sensor_msgs::Image::ConstPtr> rgb_image = getAllMsgs<sensor_msgs::Image>(bag, "camera/rgb/image_color");
  std::vector<sensor_msgs::Image::ConstPtr> mask = getAllMsgs<sensor_msgs::Image>(bag, "camera/mask");
  std::vector<geometry_msgs::PoseStamped::ConstPtr> pose = getAllMsgs<geometry_msgs::PoseStamped>(bag, "camera/pose");
  assert(depth_info.size() == depth_image.size());
  assert(depth_info.size() == rgb_image.size());
  assert(depth_info.size() == mask.size());
  assert(depth_info.size() == pose.size());
  printf("Read %zu pointcloud snapshots from bag file\n", depth_info.size());
  bag.close();

  // create pointclouds from messages
  std::vector<pcl::PointCloud<pcl::PointXYZRGB>::Ptr> full_clouds;
  std::vector<pcl::PointCloud<pcl::PointXYZRGBNormal>::Ptr> normal_clouds;
  
  printf("Importing all pointclouds\n");
  for (unsigned i=0; i<depth_info.size(); i++)
  {
    printf("  - Importing pointcloud %d / %d\n", (int)i+1, (int)depth_info.size());

    if (depth_image[i]->encoding != enc::TYPE_32FC1)
    {
      printf("Expected data of type [%s] for depth image, got [%s]", enc::TYPE_32FC1.c_str(), depth_image[i]->encoding.c_str());
      return -1;
    }

    // Update the camera model 
    image_geometry::PinholeCameraModel model;
    model.fromCameraInfo(depth_info[i]);
    float center_x = model.cx();
    float center_y = model.cy();
    float constant_x = 1.0 / model.fx();
    float constant_y = 1.0 / model.fy();
    float bad_point = std::numeric_limits<float>::quiet_NaN ();

    // create initial pointlcoud  
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr cloud (new pcl::PointCloud<pcl::PointXYZRGB>);
    cloud->header = depth_image[i]->header;
    cloud->height = depth_image[i]->height;
    cloud->width  = depth_image[i]->width;
    cloud->is_dense = false;
    cloud->points.resize (cloud->height * cloud->width);

    // create the pointcloud while applying the mask
    int idx = 0;
    pcl::PointCloud<pcl::PointXYZRGB>::iterator pt_iter = cloud->begin ();
    const float* depth_data = reinterpret_cast<const float*>(&depth_image[i]->data[0]);
    const uint8_t* mask_data  = reinterpret_cast<const uint8_t*>(&mask[i]->data[0]);
    const uint8_t* image_data  = reinterpret_cast<const uint8_t*>(&rgb_image[i]->data[0]);

    for (int v = 0; v < (int)cloud->height; ++v)
    {
      for (int u = 0; u < (int)cloud->width; ++u, ++idx, ++pt_iter)
      {
        pcl::PointXYZRGB& pt = *pt_iter;

        // Check for invalid measurements
        if (mask_data[idx] == 0 || depth_data[idx] == 0)
        {
          // not valid
          pt.x = pt.y = pt.z = bad_point;
          continue;
        }

        // Fill in XYZ
        pt.x = (u - center_x) * depth_data[idx] * constant_x;
        pt.y = (v - center_y) * depth_data[idx] * constant_y;
        pt.z = depth_data[idx];

        pt.r = image_data[3*idx+0];
        pt.g = image_data[3*idx+1];
        pt.b = image_data[3*idx+2];
      }
    }

    // transform pointcloud with prior
    transform_inverse<pcl::PointXYZRGB>(cloud, cloud, pose[i]->pose);

    // downsample cloud
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr down_cloud = downsample<pcl::PointXYZRGB>(cloud, 0.005);

    //compute normals for the cloud
    pcl::PointCloud<pcl::Normal>::Ptr normals = compute_surface_normals<pcl::PointXYZRGB, pcl::Normal>(down_cloud, 0.03);
    pcl::PointCloud<pcl::PointXYZRGBNormal>::Ptr normal_cloud(new pcl::PointCloud<pcl::PointXYZRGBNormal>);
    pcl::concatenateFields(*normals, *down_cloud, *normal_cloud);

    // write pointcloud to file for debugging
    pcdWrite<pcl::PointXYZRGBNormal>(normal_cloud, "normal", i, folder + "_debug");
    pcdWrite<pcl::PointXYZRGB>(cloud, "full", i, folder + "_debug");

    // store result
    normal_clouds.push_back(normal_cloud);
    full_clouds.push_back(cloud);
  }


  // create point-to-plane icp object
  printf("Registring all pointclouds\n");
  pcl::PointCloud<pcl::PointXYZRGBNormal>::Ptr target(new pcl::PointCloud<pcl::PointXYZRGBNormal>(*(normal_clouds[0])));
  pcl::PointCloud<pcl::PointXYZRGB>::Ptr full(new pcl::PointCloud<pcl::PointXYZRGB>(*(full_clouds[0])));
  pcl::IterativeClosestPoint<pcl::PointXYZRGBNormal, pcl::PointXYZRGBNormal> icp;
  icp.setTransformationEstimation(pcl::registration::TransformationEstimation<pcl::PointXYZRGBNormal, pcl::PointXYZRGBNormal>::Ptr(new PointToPlane));

  for (unsigned i=1; i<normal_clouds.size(); i++)
  {
    printf("  - Registring pointcloud %d to clouds 1-%d\n", (int)i+1, (int)i);
    // smooth cloud and recompute normals
    pcl::PointCloud<pcl::PointXYZRGBNormal>::Ptr target_smooth(new pcl::PointCloud<pcl::PointXYZRGBNormal>);
    pcl::MovingLeastSquares<pcl::PointXYZRGBNormal, pcl::PointXYZRGBNormal> mls;
    mls.setInputCloud(target);
    mls.setSearchRadius(0.03);
    mls.setSqrGaussParam(0.03*0.03);
    mls.setPolynomialOrder(2);
    mls.setPolynomialFit(true);
    mls.setOutputNormals(target_smooth);
    mls.reconstruct(*target_smooth);

    icp.setInputTarget(target_smooth);
    icp.setInputCloud(normal_clouds[i]);
    pcl::PointCloud<pcl::PointXYZRGBNormal> Final;
    icp.align(Final);
    
    // build up cloud to register against
    pcl::transformPointCloudWithNormals(*(normal_clouds[i]), *(normal_clouds[i]), icp.getFinalTransformation());
    *target += *(normal_clouds[i]);

    // build up full cloud
    pcl::transformPointCloud(*(full_clouds[i]), *(full_clouds[i]), icp.getFinalTransformation());
    *full += *(full_clouds[i]);

    // write result to disk
    pcdWrite<pcl::PointXYZRGBNormal>(target, "target", i, folder + "_debug");
    pcdWrite<pcl::PointXYZRGBNormal>(target_smooth, "target_smooth", i, folder + "_debug");
    pcdWrite<pcl::PointXYZRGB>(full, "target_full", i, folder + "_debug");
  } 

  printf("Downsampling and smoothing result\n");
  // downsample cloud
  pcl::PointCloud<pcl::PointXYZRGB>::Ptr full_down = downsample<pcl::PointXYZRGB>(full, 0.002);
  pcdWrite<pcl::PointXYZRGB>(full_down, "final_down", -1, folder + "_debug");
  
  //compute normals for the cloud
  pcl::PointCloud<pcl::Normal>::Ptr tmp_normals = compute_surface_normals<pcl::PointXYZRGB, pcl::Normal>(full_down, 0.03);
  pcl::PointCloud<pcl::PointXYZRGBNormal>::Ptr full_normals(new pcl::PointCloud<pcl::PointXYZRGBNormal>);
  pcl::concatenateFields(*tmp_normals, *full_down, *full_normals);

  // smooth cloud and recompute normals
  pcl::PointCloud<pcl::PointXYZRGBNormal>::Ptr full_smooth(new pcl::PointCloud<pcl::PointXYZRGBNormal>);
  pcl::MovingLeastSquares<pcl::PointXYZRGBNormal, pcl::PointXYZRGBNormal> mls;
  mls.setInputCloud(full_normals);
  mls.setSearchRadius(0.03);
  mls.setSqrGaussParam(0.03*0.03);
  mls.setPolynomialOrder(2);
  mls.setPolynomialFit(true);
  mls.setOutputNormals(full_smooth);
  mls.reconstruct(*full_smooth);
  pcdWrite<pcl::PointXYZRGBNormal>(full_smooth, "final_smooth", -1, folder + "_debug");

  // remove outliers
  pcl::PointCloud<pcl::PointXYZRGBNormal>::Ptr full_outliers(new pcl::PointCloud<pcl::PointXYZRGBNormal>);
  pcl::RadiusOutlierRemoval<pcl::PointXYZRGBNormal> ror;
  ror.setRadiusSearch(0.01);
  ror.setInputCloud(full_smooth);
  ror.setMinNeighborsInRadius(60);
  ror.filter(*full_outliers);  
  pcdWrite<pcl::PointXYZRGBNormal>(full_outliers, "final_smooth_outliers", -1, folder + "_debug");



  // write output for RoboEarth
  // --------------------------

  // write dense pointclouds of measurements
  for (unsigned i=0; i<full_clouds.size(); i++)
    pcdWrite<pcl::PointXYZRGB>(full_clouds[i], "dense_face", i, folder);
  pcdWrite<pcl::PointXYZRGBNormal>(full_outliers, "merged", -1, folder);

  std::ofstream f((folder + "/meta.xml").c_str());
  f << "<?xml version=\"1.0\" encoding=\"UTF-8\"?>" << std::endl 
    << "<model>" << std::endl
    << "  <!-- name of the model, as it is stored in the database. i.e. 'object class'.'object name' -->" << std::endl
    << "  <name>" << folder << "</name>" << std::endl
    << std::endl
    << "  <!-- object model type, for pointcloud models 'kinect_pcl' -->" << std::endl
    << "  <type>kinect_pcl</type>" << std::endl
    << std::endl
    << "  <!-- number of dense_face_xxx.pcd files -->" << std::endl
    << "  <faces>" << full_clouds.size() << "</faces>" << std::endl
    << std::endl
    << "  <!-- this is only for compatibility with re_vision's models -->" << std::endl
    << "  <dimensions> <scale>1</scale> </dimensions>" << std::endl
    << std::endl
    << "</model>" << std::endl;
  f.close();

  std::stringstream cmd;
  cmd << "zip -r " << folder << ".zip " << folder;
  int res = system(cmd.str().c_str());
  if (res != 0)
    printf("Failed to create zip file called %s.zip\n", folder.c_str());

  printf("Finished!\nThe result is stored in %s.zip\n", folder.c_str());
  return 0;
}