grab_box.cpp

Go to the documentation of this file.

#include <simple_object_capture/common.h>
#include <simple_object_capture/CaptureConfig.h>
#include <dynamic_reconfigure/server.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/visualization/pcl_visualizer.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/kdtree/kdtree.h>
#include <pcl/surface/mls.h>

#include <ros/ros.h>
#include <message_filters/subscriber.h>
#include <message_filters/time_synchronizer.h>
#include <message_filters/synchronizer.h>
#include <message_filters/sync_policies/approximate_time.h>
#include <sensor_msgs/PointCloud2.h>
#include <sensor_msgs/Image.h>
#include <sensor_msgs/CameraInfo.h>
#include <geometry_msgs/PoseStamped.h>

#include <opencv2/opencv.hpp>
#include <cv_bridge/cv_bridge.h> 

#include <boost/filesystem.hpp>
#include <rosbag/bag.h>

//NOTE: THIS FUNCTION REALLY SHOULDN"T BE REUSED
boost::shared_ptr<pcl::PointIndices> getIndexDifferenceAssumingSortedIndex(const pcl::PointIndices& input, int rows, int cols)
{
  boost::shared_ptr<pcl::PointIndices> output(new pcl::PointIndices());
  int index = 0;
  for(int i = 0; i < (int)input.indices.size() - 1; ++i)
  {
    if(input.indices[i] >= input.indices[i+1])
    {
      assert(false);
      printf("THE INDICES ARE NOT SORTED!!!!!!!!!!\n");
    }
  }

  int input_index = 0;
  for(int i = 0; i < rows; ++i)
  {
    for(int j = 0; j < cols; ++j)
    {
      if(input_index >= (int)input.indices.size())
        return output;

      if(input.indices[input_index] == index)
        input_index++;
      else
        output->indices.push_back(index);

      index++;
    }
  }

  return output;
}


const static float CUBE_SIZE = 0.4;

pcl::PointXYZ operator *(const Eigen::Affine3f& t, const pcl::PointXYZ& p)
{
  Eigen::Vector3f ep(p.x, p.y, p.z);
  ep = t * ep;
  return pcl::PointXYZ(ep[0], ep[1], ep[2]);
}

pcl::PointXYZRGB operator *(const Eigen::Affine3f& t, const pcl::PointXYZRGB& p)
{
  Eigen::Vector3f ep(p.x, p.y, p.z);
  ep = t * ep;
  pcl::PointXYZRGB ret;
  ret.x = ep[0];
  ret.y = ep[1];
  ret.z = ep[2];
  ret.rgb = p.rgb;
  return ret;
}

struct Box
{
  Eigen::Vector3f center;
  Eigen::Vector3f scale;
  Eigen::Affine3f transform;
};


class GrabBox
{
  public:
    typedef message_filters::sync_policies::ApproximateTime<sensor_msgs::Image, sensor_msgs::Image, sensor_msgs::CameraInfo, sensor_msgs::PointCloud2> ApproxTimeSync;

  GrabBox(std::string output_directory) :   viz_("Grab"), has_new_cloud_(false), 
                cloud_(new pcl::PointCloud<pcl::PointXYZRGB>()), 
                take_snapshot_(false), recompute_plane_(true), redraw_box_(false), output_directory_(output_directory)
  {
    //make sure that the directory we wish to save in exists
    if(!boost::filesystem::exists(output_directory_))
    {
      boost::filesystem::create_directory(output_directory_);
    }

    bag_.open(boost::filesystem::path(output_directory_ / "object_capture.bag").string(), rosbag::bagmode::Write);

    // hook up reconfigure callback
    dynamic_reconfigure::Server<simple_object_capture::CaptureConfig>::CallbackType reconfigure_cb(boost::bind(&GrabBox::reconfigure, this, _1, _2));
    dyn_conf_server_.setCallback(reconfigure_cb);

    ros::NodeHandle camera_in_nh("camera");
    rgb_sub.reset(new message_filters::Subscriber<sensor_msgs::Image>(camera_in_nh, "rgb/image_color", 2));
    depth_sub.reset(new message_filters::Subscriber<sensor_msgs::Image>(camera_in_nh, "depth_registered/image", 2));
    info_sub.reset(new message_filters::Subscriber<sensor_msgs::CameraInfo>(camera_in_nh, "rgb/camera_info", 2));
    cloud_sub.reset(new message_filters::Subscriber<sensor_msgs::PointCloud2>(camera_in_nh, "depth_registered/points", 2));

    sync_.reset(new message_filters::Synchronizer<ApproxTimeSync>(ApproxTimeSync(20), *rgb_sub, *depth_sub, *info_sub, *cloud_sub));
    sync_->registerCallback(boost::bind(&GrabBox::callback, this, _1, _2, _3, _4));
    
  }

  ~GrabBox()
  {
    printf("Writing bag to disk\n");
    bag_.close();
  }


  void reconfigure(simple_object_capture::CaptureConfig &config, uint32_t level) 
  {
    ROS_DEBUG("Calling Dynamic reconfigure");
    
    boost::mutex::scoped_lock lock(cloud_mutex_);
    config_ = config;
    redraw_box_ = true;
  }


  void callback(const sensor_msgs::Image::ConstPtr& rgb, const sensor_msgs::Image::ConstPtr& depth, 
                const sensor_msgs::CameraInfo::ConstPtr& camera_info, const sensor_msgs::PointCloud2::ConstPtr& cloud)
  {
    ROS_DEBUG("Camera callback");

    boost::mutex::scoped_lock lock(cloud_mutex_);

    //copy over the data for the working thread
    pcl::fromROSMsg(*cloud, *cloud_);
    rgb_ = rgb;
    depth_ = depth;
    camera_info_ = camera_info;
    has_new_cloud_ = true;
  }

  template <class PointT>
  bool pointInBox(const Box& box, const PointT& pcl_pt)
  {
    if(isnan(pcl_pt.x) || isnan(pcl_pt.y) || isnan(pcl_pt.z))
      return false;

    Eigen::Vector3f pt(pcl_pt.x, pcl_pt.y, pcl_pt.z);
    Eigen::Vector3f pt_box = box.transform.inverse() * pt;

    return fabs(pt_box[0] - box.center[0]) < box.scale[0] / 2
      && fabs(pt_box[1] - box.center[1]) < box.scale[1] / 2
      && fabs(pt_box[2] - box.center[2]) < box.scale[2] / 2;
  }

  void drawCube(const pcl::PointXYZ& c1, const pcl::PointXYZ& c2, const pcl::PointXYZ& c3, const pcl::PointXYZ& c4,
                const pcl::PointXYZ& c5, const pcl::PointXYZ& c6, const pcl::PointXYZ& c7, const pcl::PointXYZ& c8,
                pcl::visualization::PCLVisualizer& viz)
  {
    pcl::PointCloud<pcl::PointXYZ> polygon;
    polygon.is_dense = true;
    polygon.width = 8;
    polygon.height = 1;

    polygon.points.push_back(c1);
    polygon.points.push_back(c2);
    polygon.points.push_back(c3);
    polygon.points.push_back(c4);

    //remove the old box from the viewer
    viz.removeShape("poly1");
    viz.removeShape("poly2");
    viz.removeShape("poly3");
    viz.removeShape("poly4");

    viz.addPolygon<pcl::PointXYZ>(polygon.makeShared(), 1.0, 0.0, 0.0, "poly1");

    polygon.points[0] = c4;
    polygon.points[1] = c3;
    polygon.points[2] = c7;
    polygon.points[3] = c8;
    viz.addPolygon<pcl::PointXYZ>(polygon.makeShared(), 1.0, 0.0, 0.0, "poly3");

    polygon.points[0] = c1;
    polygon.points[1] = c2;
    polygon.points[2] = c6;
    polygon.points[3] = c5;
    viz.addPolygon<pcl::PointXYZ>(polygon.makeShared(), 1.0, 0.0, 0.0, "poly4");

    polygon.points[0] = c5;
    polygon.points[1] = c6;
    polygon.points[2] = c7;
    polygon.points[3] = c8;
    viz.addPolygon<pcl::PointXYZ>(polygon.makeShared(), 0.0, 0.0, 1.0, "poly2");


  }

  template <class PointT>
  pcl::ModelCoefficients findPlaneCoefficients(const typename pcl::PointCloud<PointT>::ConstPtr cloud, const Eigen::Vector3f& axis, double eps_angle,
                                               pcl::PointIndices::Ptr inliers = boost::shared_ptr<pcl::PointIndices>())
  {
    if (!inliers)
      inliers.reset(new pcl::PointIndices);
    // Segments the plane from the cloud
    pcl::ModelCoefficients::Ptr coefficients(new pcl::ModelCoefficients);
    pcl::SACSegmentation<PointT> seg;
    seg.setMaxIterations(10000);
    seg.setOptimizeCoefficients(true);
    seg.setModelType(pcl::SACMODEL_PERPENDICULAR_PLANE);
    seg.setMethodType(pcl::SAC_RANSAC);
    seg.setDistanceThreshold(0.01);
    seg.setInputCloud(cloud);
    seg.setAxis(axis);
    seg.setEpsAngle(eps_angle);
    seg.segment(*inliers, *coefficients);

    if (inliers->indices.size() == 0)
    {
      PCL_ERROR("Could not estimate a plane (no inliers found)");
    }
    else
    {
      printf("Plane coefficients: (%f, %f, %f, %f)\n",
             coefficients->values[0], coefficients->values[1],
             coefficients->values[2], coefficients->values[3]);
    }
    return *coefficients;
  }

  void drawBox(const Box& box, pcl::visualization::PCLVisualizer& viz)
  {
    pcl::PointXYZ c1(box.scale[0]/2 + box.center[0], box.scale[1]/2 + box.center[1], -box.scale[2]/2 + box.center[2]);
    pcl::PointXYZ c2(-box.scale[0]/2 + box.center[0], box.scale[1]/2 + box.center[1], -box.scale[2]/2 + box.center[2]);
    pcl::PointXYZ c3(-box.scale[0]/2 + box.center[0], box.scale[1]/2 + box.center[1], box.scale[2]/2 + box.center[2]);
    pcl::PointXYZ c4(box.scale[0]/2 + box.center[0], box.scale[1]/2 + box.center[1], box.scale[2]/2 + box.center[2]);

    pcl::PointXYZ c5(box.scale[0]/2 + box.center[0], -box.scale[1]/2 + box.center[1], -box.scale[2]/2 + box.center[2]);
    pcl::PointXYZ c6(-box.scale[0]/2 + box.center[0], -box.scale[1]/2 + box.center[1], -box.scale[2]/2 + box.center[2]);
    pcl::PointXYZ c7(-box.scale[0]/2 + box.center[0], -box.scale[1]/2 + box.center[1], box.scale[2]/2 + box.center[2]);
    pcl::PointXYZ c8(box.scale[0]/2 + box.center[0], -box.scale[1]/2 + box.center[1], box.scale[2]/2 + box.center[2]);

    //make sure to transform the cube into the frame of the camera
    c1 = box.transform * c1;
    c2 = box.transform * c2;
    c3 = box.transform * c3;
    c4 = box.transform * c4;
    c5 = box.transform * c5;
    c6 = box.transform * c6;
    c7 = box.transform * c7;
    c8 = box.transform * c8;

    drawCube(c1, c2, c3, c4, c5, c6, c7, c8, viz);
  }

  void spin()
  {
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr viz_cloud;
    pcl::PointCloud<pcl::PointXYZRGB>::ConstPtr original_cloud;
    ros::Duration update_period(1.0);
    float box_angle = 0.0f;
    int snapshot_index = 0;

    Box box;
    //register a callback we'll use to do different operations
    viz_.registerKeyboardCallback(boost::bind(&GrabBox::keyPress, this, _1));
    
    while (!viz_.wasStopped())
    {
      ros::spinOnce();
      viz_.spinOnce();
      {
        if (has_new_cloud_)
        {
          // If the GUI locks up intermitantly, this is probably why
          {
            boost::mutex::scoped_lock lock(cloud_mutex_);
            box.scale = Eigen::Vector3f(config_.box_width, config_.box_depth, config_.box_height);
            box.center = Eigen::Vector3f(0, 0, box.scale[2]/2);

            viz_cloud.reset(new pcl::PointCloud<pcl::PointXYZRGB>(*cloud_));
            original_cloud = cloud_;
          }

          if(recompute_plane_)
          {
            object_frame_ = getObjectFrame(viz_cloud);
            box.transform = object_frame_ * Eigen::AngleAxisf(box_angle, Eigen::Vector3f(0, 0, 1));
            printf("Drawing the cube with angle %.2f\n", box_angle);
            drawBox(box, viz_);
            recompute_plane_ = false;
          }

          if(redraw_box_)
          {
            drawBox(box, viz_);
            redraw_box_ = false;
          }

          for(size_t i = 0; i < viz_cloud->points.size(); ++i)
          {
            if(!pointInBox(box, viz_cloud->points[i]))
            {
              viz_cloud->points[i].r = viz_cloud->points[i].r / 2;
              viz_cloud->points[i].g = viz_cloud->points[i].g / 2;
              viz_cloud->points[i].b = viz_cloud->points[i].b / 2;
            }
            else
            {
              viz_cloud->points[i].g = std::min(viz_cloud->points[i].g + 50, 255);
            }
          }

          if (!viz_.updatePointCloud(viz_cloud, "cloud"))
          {
            viz_.addPointCloud(viz_cloud, "cloud");
            viz_.resetCameraViewpoint("cloud");

          }
          has_new_cloud_ = false;
        }

        if (take_snapshot_)
        {
          // Removes the plane from the point cloud
          pcl::PointIndices::Ptr inliers(new pcl::PointIndices);
          pcl::ModelCoefficients coefficients = findPlaneCoefficients<pcl::PointXYZRGB>(original_cloud, Eigen::Vector3f(0, -cos(M_PI/8), -sin(M_PI/8)), M_PI/4, inliers);

          assert(inliers->indices.size() > 0);
          pcl::PointCloud<pcl::PointXYZRGB>::Ptr segmented (new pcl::PointCloud<pcl::PointXYZRGB>);
          boost::shared_ptr<pcl::PointIndices> outliers = getIndexDifferenceAssumingSortedIndex(*inliers, original_cloud->height, original_cloud->width);
          extractIndices(original_cloud, outliers, *segmented, true);

          // Finds the points inside the box
          pcl::PointCloud<pcl::PointXYZRGB>::Ptr box_cloud(new pcl::PointCloud<pcl::PointXYZRGB>());
          box_cloud->is_dense = true;
          box_cloud->width = 0;
          box_cloud->height = 1;
          Eigen::Affine3f sensor_frame = Eigen::Translation3f(original_cloud->sensor_origin_.head<3>()) * original_cloud->sensor_orientation_;
          Eigen::Affine3f box_sensor_frame = box.transform.inverse() * sensor_frame;
          box_cloud->sensor_origin_.head<3>() = box_sensor_frame.translation();
          box_cloud->sensor_origin_[3] = 0.f;
          box_cloud->sensor_orientation_ = (Eigen::Quaternionf)box_sensor_frame.rotation();

          //we want to create a masked image as we find points in the box
          cv::Mat_<uint8_t> mask = cv::Mat_<uint8_t>::zeros(original_cloud->height, original_cloud->width);

          //get a pointer to the beginning of the uint array
          uint8_t* mask_array = mask.ptr<uint8_t>(0);
          for(size_t i = 0; i < segmented->points.size(); ++i)
          {
            if(pointInBox(box, segmented->points[i]))
            {
              // Transforms the points into the box frame
              pcl::PointXYZRGB box_pt = box.transform.inverse() * segmented->points[i];
              box_cloud->points.push_back(box_pt);
              box_cloud->width++;

              //we want to mark the object as valid
              mask_array[outliers->indices[i]] = 255;
            }
          }

          //compute normals for the cloud
          pcl::PointCloud<pcl::PointXYZRGB>::Ptr downsampled_cloud = downsample<pcl::PointXYZRGB>(box_cloud, 0.005);
          pcl::PointCloud<pcl::Normal>::Ptr normals = compute_surface_normals<pcl::PointXYZRGB, pcl::Normal>(downsampled_cloud, 0.03);
          //quickViz<pcl::PointXYZRGB, pcl::Normal>(downsampled_cloud, normals, "cloud + normals");

          //concatenate the cloud for saving
          pcl::PointCloud<pcl::PointXYZRGBNormal> combined_box_cloud;
          pcl::concatenateFields(*normals, *downsampled_cloud, combined_box_cloud);
          

          char buf[1024];
          snprintf(buf, 1024, "snapshot%02d.pcd", snapshot_index);
          ++snapshot_index;
          pcl::io::savePCDFile<pcl::PointXYZRGBNormal>(boost::filesystem::path(output_directory_ / buf).string(), combined_box_cloud);

          box_angle += config_.angle_increment * M_PI / 180;
          if (box_angle > 2 * M_PI)
            box_angle -= 2 * M_PI;
          box.transform = object_frame_ * Eigen::AngleAxisf(box_angle, Eigen::Vector3f(0, 0, 1));
          printf("Re-drawing the cube with angle %.2f\n", box_angle);
          drawBox(box, viz_);

          //we'll also make sure to write the necessary topics
          bag_.write("camera/depth/camera_info", ros::Time::now(), camera_info_);
          bag_.write("camera/depth/image", ros::Time::now(), depth_);
          bag_.write("camera/rgb/camera_info", ros::Time::now(), camera_info_);
          bag_.write("camera/rgb/image_color", ros::Time::now(), rgb_);

          //Publish Mask
          cv_bridge::CvImage img_bridge;
          img_bridge.encoding = "mono8";
          img_bridge.header = camera_info_->header;
          img_bridge.image = mask;

          boost::shared_ptr<sensor_msgs::Image> img_msg = img_bridge.toImageMsg();
          bag_.write("camera/mask", ros::Time::now(), img_msg);

          //Publish Pose
          geometry_msgs::PoseStamped pose_msg;
          pose_msg.header = camera_info_->header;
          pose_msg.pose.position.x = box.transform.translation()[0];
          pose_msg.pose.position.y = box.transform.translation()[1];
          pose_msg.pose.position.z = box.transform.translation()[2];

          Eigen::Quaternionf q(box.transform.rotation());
          pose_msg.pose.orientation.w = q.w();
          pose_msg.pose.orientation.x = q.x();
          pose_msg.pose.orientation.y = q.y();
          pose_msg.pose.orientation.z = q.z();

          bag_.write("camera/pose", ros::Time::now(), pose_msg);

          take_snapshot_ = false;
        }
      }
    }
  }

  Eigen::Affine3f getObjectFrame(const pcl::PointCloud<pcl::PointXYZRGB>::ConstPtr cloud)
  {
    //find a plane
    pcl::ModelCoefficients coefficients = findPlaneCoefficients<pcl::PointXYZRGB>(cloud, Eigen::Vector3f(0, -cos(M_PI/8), -sin(M_PI/8)), M_PI/4);

    //we want to find the intersection of the z-axis with the plane. We'll use that as the center point for our cube
    double a = coefficients.values[0];
    double b = coefficients.values[1];
    double c = coefficients.values[2];
    double d = coefficients.values[3];

    double u = d / (c * -1.0);

    Eigen::Vector3f r = u * Eigen::Vector3f(0.0, 0.0, 1.0);
    printf("Intersection point: (%.2f, %.2f, %.2f)", r[0], r[1], r[2]);

    Eigen::Vector3f proj(0.0, 0.0, 0.6);
    Eigen::Vector3f normal(a, b, c);

    if(normal[2] > 0)
      normal = -1.0 * normal;

    normal.normalize();

    Eigen::Vector3f plane_pt = proj - (proj - r).dot(normal) * normal;
    printf("Plane point: (%.2f, %.2f, %.2f)", plane_pt[0], plane_pt[1], plane_pt[2]);

    //viz_.addSphere(pcl::PointXYZ(r[0], r[1], r[2]), 0.05, "sphere1");
    //viz_.addSphere(pcl::PointXYZ(plane_pt[0], plane_pt[1], plane_pt[2]), 0.05, "sphere2");

    float rot_angle = atan2(normal[1], normal[2]);
    Eigen::Affine3f rot(Eigen::AngleAxisf(rot_angle, Eigen::Vector3f(-1, 0, 0)));

    printf("Angle of rotation: %.2f\n", rot_angle);

    Eigen::Translation3f trans(plane_pt);

    return trans * rot;
  }

  void keyPress(const pcl::visualization::KeyboardEvent& e)
  {
    if(e.keyUp())
    {
      switch (e.getKeyCode())
      {
        case 32:  // SPACE
          take_snapshot_ = true;
          break;
        case 97: //a
          recompute_plane_ = true;
          break;
      }
      printf("Keycode %u\n", e.getKeyCode());
    }
  }

  dynamic_reconfigure::Server<simple_object_capture::CaptureConfig> dyn_conf_server_;
  simple_object_capture::CaptureConfig config_;

  pcl::visualization::PCLVisualizer viz_;
  boost::shared_ptr<message_filters::Synchronizer<ApproxTimeSync> > sync_;
  boost::shared_ptr<message_filters::Subscriber<sensor_msgs::Image> > rgb_sub, depth_sub;
  boost::shared_ptr<message_filters::Subscriber<sensor_msgs::CameraInfo> > info_sub;
  boost::shared_ptr<message_filters::Subscriber<sensor_msgs::PointCloud2> > cloud_sub;

  boost::mutex cloud_mutex_;
  bool has_new_cloud_;

  boost::shared_ptr<pcl::PointCloud<pcl::PointXYZRGB> > cloud_;
  boost::shared_ptr<const sensor_msgs::Image> rgb_, depth_;
  boost::shared_ptr<const sensor_msgs::CameraInfo> camera_info_;

  Eigen::Affine3f object_frame_;
  bool take_snapshot_, recompute_plane_, redraw_box_;

  boost::filesystem::path output_directory_;

  rosbag::Bag bag_;
};

// NEXT: keypress to trigger saving a point cloud

int main(int argc, char** argv)
{
  ros::init(argc, argv, "simple_object_capture");

  std::string output_dir = ".";
  if(argc > 1)
    output_dir = argv[1];

  GrabBox gb(output_dir);
  gb.spin();
  return 0;
}