object_tracker.cc

Go to the documentation of this file.

/*
 * object_tracker.cc
 * Mac Mason <mmason@willowgarage.com>
 *
 * This is the tied-for-second (with the plane tracker) stage in our pipeline.
 * We listen for cluster objects from the filter (which, at the moment,
 * publishes only the currently-largest cluster in the scene) and then
 * associate them together in useful ways. We make various things visible,
 * too.
 */

#include <cassert>
#include <limits>
#include <vector>
#include <sstream>

#include <boost/foreach.hpp>
#define foreach BOOST_FOREACH
#include <boost/thread.hpp>

#include "ros/ros.h"
#include "tf/transform_listener.h"
#include "tf/transform_broadcaster.h"
#include "semanticmodel/blob_store.h"
#include "pcl16/point_types.h"
#include "pcl16_ros/point_cloud.h"
#include "pcl16_ros/transforms.h"
#include "pcl16/surface/convex_hull.h"
#include "pcl16/segmentation/extract_clusters.h"
#include "pcl16/segmentation/sac_segmentation.h"
#include "pcl16/filters/extract_indices.h"
#include "pcl16/sample_consensus/method_types.h"
#include "pcl16/sample_consensus/model_types.h"
#include "image_geometry/pinhole_camera_model.h"
#include "opencv2/opencv.hpp"

#include "semanticmodel/centroid.hh"
#include "semanticmodel/miniatureoccupancygrid.hh"
#include "semanticmodel/Blobs.h"
#include "semanticmodel/Segment.h"
#include "semanticmodel/blob.hh"
#include "semanticmodel/GetCollectionNamespace.h"

#include "visualization_msgs/MarkerArray.h"

#include "dynamic_reconfigure/server.h"
#include "semanticmodel/ObjectTrackerConfig.h"

using namespace std;
using visualization_msgs::Marker;
using visualization_msgs::MarkerArray;
using semanticmodel::Blob;
using semanticmodel::Blobs;
using semanticmodel::Segment;
using semanticmodel::GetCollectionNamespace;
using sensor_msgs::PointCloud2;
using semanticmodel::MiniatureOccupancyGrid;

typedef pcl16::PointXYZRGB Point;
typedef pcl16::PointCloud<Point> PointCloud;

semanticmodel::ObjectTrackerConfig config;
void config_cb(const semanticmodel::ObjectTrackerConfig& _config, uint32_t i)
{
  // See the note in filter.cc's version of config_cb about race conditions.
  config = _config;
}

// The data we're storing.
vector<Blob*> objects;

// Interface to db
boost::shared_ptr<semanticmodel::BlobStore> blob_store;

// Our occupancy grid. Note that we keep just one for the duration; that means
// we're going to track this over the duration of the run.
MiniatureOccupancyGrid OG(0.05);

// The (randomly generated) colors we're using to show the object clouds.
// These are stored as [r g b r g b r g b], all in the range [0, 255]. There
// should be three times as many of these as there are elements of objects.
vector<int> colors;

cv::RNG prng(2);

/* Talking to the outside world. */
ros::Subscriber cluster_sub;
ros::Subscriber full_cloud_sub;  // Used for populating our OG.
ros::Publisher all_objects_pub;  // The objects (as one point cloud).
ros::Publisher flat_frames_publisher;  // Push out 2D frames.
ros::Publisher viz_pub; // The objects (as vis markers)
ros::Publisher og_pub;  // Our occupancy grid.

// These have to be pointers, or we'll build a node before ::init.
boost::shared_ptr<tf::TransformListener> listener; 
boost::shared_ptr<tf::TransformBroadcaster> broadcaster; 

void push_one_frame_segment(const Blob* blob, 
                            const sensor_msgs::Image& rgb_img,
                            const sensor_msgs::CameraInfo& rgb_info,
                            const sensor_msgs::Image& depth_img,
                            const sensor_msgs::CameraInfo& depth_info);

// Originally part of Segmenter; this is likely better.
void cluster_objects(const PointCloud::ConstPtr& aboves,
                     std::vector<PointCloud::Ptr>& clusters)
{
  if (!aboves)
    return;
  if (aboves->points.size() < 100)
    return;
  if (!ros::ok())
    return;

  clusters.clear();
  pcl16::search::KdTree<Point>::Ptr tree(new pcl16::search::KdTree<Point>());
  tree->setInputCloud(aboves);

  pcl16::EuclideanClusterExtraction<Point> euclid;
  std::vector<pcl16::PointIndices> indices;
  euclid.setClusterTolerance(config.cluster_tolerance);
  euclid.setMinClusterSize(config.min_cluster_size);
  euclid.setMaxClusterSize(config.max_cluster_size);
  euclid.setSearchMethod(tree);
  euclid.setInputCloud(aboves);
  ROS_DEBUG_NAMED("cluster", 
                  "About to cluster %zu points", aboves->points.size());
  euclid.extract(indices);

  // Fill in our output vector.
  pcl16::ExtractIndices<Point> extractor;
  foreach(pcl16::PointIndices& index, indices)
  {
    if (!ros::ok())
      return;
    PointCloud::Ptr new_cluster(new PointCloud());
    extractor.setInputCloud(aboves);
    pcl16::PointIndices::Ptr indexptr(new pcl16::PointIndices(index));
    extractor.setIndices(indexptr); // This is SO STUPID. Egad.
    extractor.filter(*new_cluster);
    clusters.push_back(new_cluster);
  }
}

void deplane_objects(std::vector<PointCloud::Ptr>& clusters)
{
  if (!ros::ok())
    return;
  // Fill this in with those that pass the test; then do a swap.
  std::vector<PointCloud::Ptr> newclusters;

  // Shared stuff.
  pcl16::SACSegmentation<Point> segmenter;
  segmenter.setOptimizeCoefficients(true);
  segmenter.setModelType(pcl16::SACMODEL_PLANE);
  segmenter.setMethodType(pcl16::SAC_RANSAC);
  segmenter.setDistanceThreshold(config.plane_tolerance);
  segmenter.setMaxIterations(config.object_ransac_iterations);

  foreach(PointCloud::Ptr& cluster, clusters)
  {
    if (cluster->points.size() == 0)
      continue;
    if (!ros::ok())
      return;

    pcl16::ModelCoefficients::Ptr model(new pcl16::ModelCoefficients());
    pcl16::PointIndices::Ptr indices(new pcl16::PointIndices());
    segmenter.setInputCloud(cluster);
    segmenter.segment(*indices, *model);

    double fraction = (float)indices->indices.size() / cluster->points.size();
//    ROS_FATAL("Fit %zu indices (%f%%), eq: [%f %f %f %f]",
//             indices->indices.size(),
//             100 * fraction,
//             model->values[0],
//             model->values[1],
//             model->values[2],
//             model->values[3]);

    if (fraction <= config.object_plane_percentage)
    {
      newclusters.push_back(cluster);
      ROS_DEBUG_NAMED("object_tracker", "Kept cluster with %zu points.", 
                      indices->indices.size());
    }
    else
    {
      ROS_DEBUG_NAMED("object_tracker", "Rejected cluster with %zu points.", 
                      indices->indices.size());
    }
  }

  // And we're done.
  std::swap(clusters, newclusters);
}

// Use this cloud to update our occupancy grid.
void fc_cb(const PointCloud::ConstPtr& input)
{
  if (!ros::ok())
    return;

  PointCloud::Ptr canonical(new PointCloud());
  if (!pcl16_ros::transformPointCloud(config.canonical_frame, 
                                    *input, 
                                    *canonical, 
                                    *listener))
    return;

  // Step 1: fill in the grid.
  boost::scoped_ptr<boost::mutex::scoped_lock> lock(OG.get_lock());
  foreach(const Point& p, *input)
  {
    // Because +z is "up" in /map, we can just ignore the other coordinates.
    if (p.z >= config.occ_trick_low && p.z <= config.occ_trick_high)
    {
      OG.set(p.x, p.y, MiniatureOccupancyGrid::OCCUPIED);
    }
  }
  nav_msgs::OccupancyGrid og;
  OG.to_msg(og);
  og_pub.publish(og);
}

// Perform the occupancy grid trick to throw away those objects we consider
// unlikely. input is both an input and output argument, as you would expect
// by now.
void occupancy_trick(PointCloud::Ptr& input)
{
  if (!ros::ok())
    return;

  /*
   * We're going to need to walk input twice. The first time, we check for
   * points in the "bad range", and set the appropriate flags in the (2D)
   * occupancy grid. The second time, we check each point against that grid,
   * and only keep those that we like. 
   */

  // This is our output; where we're done, std::swap takes care of things.
  PointCloud::Ptr output(new PointCloud());
  output->header = input->header;

  // Now, filter.
  boost::scoped_ptr<boost::mutex::scoped_lock> lock(OG.get_lock());
  foreach(const Point& p, *input)
  {
    if (OG.get(p.x, p.y) == MiniatureOccupancyGrid::EMPTY)
      output->points.push_back(p);
  }

  std::swap(input, output);
}

void blobs_cb(const Blobs::ConstPtr& blobs)
{
  // These days, you just get one.
  ROS_ASSERT(blobs->blobs.size() == 1);

  /* Convert point cloud message into pcl16 point cloud */
  PointCloud pcl16_cloud;
  pcl16::fromROSMsg(blobs->blobs[0], pcl16_cloud);

  /* We need our inputs to be in the canonical frame. */
  PointCloud::Ptr cluster(new PointCloud());
  pcl16_ros::transformPointCloud(config.canonical_frame, 
                                 pcl16_cloud, 
                                 *cluster, 
                                 *listener);

  /* 
   * Get rid of this that are unlikely to actually be objects, but instead are
   * probably parts of walls, using Bhaskara's occupancy grid trick.
   */
  occupancy_trick(cluster);
  if (cluster->points.size() == 0 || !ros::ok())
  {
    ROS_INFO_NAMED("object_tracker", "Occupancy trick erased everything.");
    return;
  }

  /* Now, cluster the points. */
  vector<PointCloud::Ptr> clusters;
  cluster_objects(cluster, clusters);
  if (clusters.size() == 0 || !ros::ok())
  {
    ROS_INFO_NAMED("object_tracker", "Clustering went nowhere.");
    return;
  }

  /* And throw away objects that are too planar. */
  deplane_objects(clusters);
  if (clusters.size() == 0 || !ros::ok())
  {
    ROS_INFO_NAMED("object_tracker", "Deplaning erased everything.");
    return;
  }

  foreach(PointCloud::Ptr cluster, clusters)
  {
    /* Build our new blob. */
    Blob* newblob = new Blob(cluster);

    push_one_frame_segment(newblob, 
                           blobs->rgb_image, 
                           blobs->rgb_info,
                           blobs->depth_image,
                           blobs->depth_info);

    /* See if it "belongs to" any of the other blobs. */
    Blob::MergeBlobWithSet(objects, newblob);
  }

  // Check after every long operation, of course.
  if (!ros::ok())
    return;

  //ROS_INFO("There are now %zu objects.", objects.size());
  
  /* Time to publish our data. */
  int id = 50;
  foreach(Blob* blob, objects)
  {
    Marker marker;
    marker.header.frame_id = blob->cloud->header.frame_id;
    marker.header.stamp = ros::Time::now();
    marker.id = id;
    marker.ns = "objects";
    marker.type = Marker::CUBE;
    marker.action = Marker::ADD;
    // We need the center of the object.
    Eigen::Vector4f center(semanticmodel::centroid(blob->cloud));
    marker.pose.position.x = center[0];
    marker.pose.position.y = center[1];
    marker.pose.position.z = center[2];
    marker.pose.orientation.x = 0.0;
    marker.pose.orientation.y = 0.0;
    marker.pose.orientation.z = 0.0;
    marker.pose.orientation.w = 1.0;
    blob->size(marker.scale.x, marker.scale.y, marker.scale.z);
    marker.color.a = 0.75;
    blob->RGB(marker.color.r, marker.color.g, marker.color.b);
    id++;
    if (config.send_boxes)
      viz_pub.publish(marker);
    // Either way, we do publish the tf frames.
    tf::Transform transform;
    transform.setOrigin(tf::Vector3(center[0], center[1], center[2]));
    transform.setRotation(tf::Quaternion(0, 0, 1.0, 0.0));
    std::string name = (boost::format("object_%02d") % blob->id).str();
    tf::StampedTransform trans(transform,
                         blob->cloud->header.stamp,
                         config.canonical_frame,
                         name);
    broadcaster->sendTransform(trans);
  }

  // Make sure we have enough colors for the point cloud.
  while (colors.size() < 3 * objects.size())
    colors.push_back(int(prng.uniform(0.0, 1.0) * 255));

  // Send 'em to the database.
  blob_store->update(objects, blobs->rgb_image, blobs->rgb_info);

  // Publish every blob at once, in one pointcloud.
  if (objects.size() == 0 || !ros::ok())
    return;

  PointCloud::Ptr all_objects(new PointCloud());
  all_objects->header.frame_id = objects.at(0)->cloud->header.frame_id;
  foreach(Blob* blob, objects)
  {
    (*all_objects) += *(blob->cloud);
  }
  all_objects_pub.publish(all_objects);
}

int main(int argc, char** argv)
{
  ros::init(argc, argv, "object_tracker");
  ros::NodeHandle node("~"); // Private nodehandle

  cluster_sub = node.subscribe<Blobs>("blobs", 100, blobs_cb);
  full_cloud_sub = node.subscribe<PointCloud>("/map_points", 100, fc_cb);
  all_objects_pub = node.advertise<PointCloud>("all_objects_cloud", 100, true);
  viz_pub = node.advertise<visualization_msgs::Marker>(
      "visualization_marker", 0);

  flat_frames_publisher = node.advertise<Segment>("/flat_frames", 10);
  og_pub = node.advertise<nav_msgs::OccupancyGrid>("/og", 100);
  std::string camera_frame;
  bool found = node.getParam("camera_frame", camera_frame);
  ROS_ASSERT(found);

  // Lifted straight from the dynamic_reconfigure docs.
  dynamic_reconfigure::Server<semanticmodel::ObjectTrackerConfig> srv;
  dynamic_reconfigure::Server<semanticmodel::ObjectTrackerConfig>::
    CallbackType f;
  f = boost::bind(&config_cb, _1, _2);
  srv.setCallback(f);

  ROS_INFO_STREAM("Waiting for service get_db_name");
  GetCollectionNamespace coll_ns_srv;
  ros::service::waitForService("get_collection_namespace");
  ros::service::call("get_collection_namespace", coll_ns_srv);
  ROS_INFO_STREAM ("Db name is " << coll_ns_srv.response.db_name << "/"
                                 << coll_ns_srv.response.collection_namespace);
  blob_store.reset(
      new semanticmodel::BlobStore(config.canonical_frame,
                                   camera_frame,
                                   coll_ns_srv.response.db_name,
                                   coll_ns_srv.response.collection_namespace));
  
  listener.reset(new tf::TransformListener(ros::Duration(600)));
  broadcaster.reset(new tf::TransformBroadcaster());
  ros::spin();

  // Just to be sure.
  foreach(Blob* blob, objects)
    delete blob;

  return 0;
}

void push_one_frame_segment(const Blob* blob, 
                            const sensor_msgs::Image& rgb_img,
                            const sensor_msgs::CameraInfo& rgb_info,
                            const sensor_msgs::Image& depth_img,
                            const sensor_msgs::CameraInfo& depth_info)
{
  // The easy parts.
  Segment seg;
  seg.rgb_image = rgb_img;
  seg.rgb_info = rgb_info;
  seg.depth_image = depth_img;
  seg.depth_info = depth_info;

  // Now I need project the blob into the image plane and push that.
  image_geometry::PinholeCameraModel pcm;
  pcm.fromCameraInfo(rgb_info);
  PointCloud::Ptr camcloud(new PointCloud());
  pcl16_ros::transformPointCloud(rgb_img.header.frame_id,
                                 rgb_img.header.stamp,
                                 *(blob->cloud),
                                 blob->cloud->header.frame_id,
                                 *camcloud,
                                 *listener);
  
  for (size_t i = 0 ; i < camcloud->size() ; ++i)
  {
    cv::Point3d cvpt(camcloud->points[i].x,
                     camcloud->points[i].y,
                     camcloud->points[i].z);
    cv::Point2d pt = pcm.project3dToPixel(cvpt);
    if (pt.x >= 0 && pt.x < rgb_img.width && 
        pt.y >= 0 && pt.y < rgb_img.height)
    {
      seg.coordinates.push_back(pt.y);
      seg.coordinates.push_back(pt.x);
    }
  }
  flat_frames_publisher.publish(seg);
}