MotldNode.cpp

Go to the documentation of this file.

#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <image_transport/image_transport.h>
#include <opencv/cv.h>
#include <opencv/highgui.h>
#include <cv_bridge/CvBridge.h>
#include <boost/thread/mutex.hpp>

// point cloud related includes
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl_ros/impl/transforms.hpp>
#include <sensor_msgs/PointCloud2.h>
#include <tf/tfMessage.h>
#include <tf/transform_broadcaster.h>
// subscriber filters
#include <image_transport/image_transport.h>
#include <message_filters/subscriber.h>
#include <message_filters/synchronizer.h>
#include <message_filters/sync_policies/approximate_time.h>
#include <image_transport/subscriber_filter.h>

// motld messages
#include <motld/BoundingBox.h>
#include <motld/Save.h>
#include <motld/Config.h>
#include <motld/TrackedObjects.h>

// other messages
#include <geometry_msgs/PoseArray.h>

// tld classes and functions
#include <motld/MultiObjectTLD.h>
#ifndef SQR
#define SQR(X) ((X)*(X))
#endif

using namespace motld;

// typedefs
typedef pcl::PointXYZ Point_t;
typedef pcl::PointCloud<Point_t> Cloud_t;
typedef message_filters::sync_policies::ApproximateTime<sensor_msgs::Image, sensor_msgs::PointCloud2> MySyncPolicy;

struct CloudProperties {
  pcl::PointXYZ mean;
  float width;
  float height;
};

bool debug = false;

// image parameters
const int CHANNELS = 3;
int width_ = 640;
int height_ = 480;
bool force_resize_ = true;
float resize_factor_ = 1.f;

// pointcloud processing related variables
bool use_cloud_ = true;
Cloud_t cloud_;
std::string frame_id_;

// working buffers for the image
unsigned char *img_;
IplImage *ipl_img_;
std::string image_frame_;
sensor_msgs::CvBridge bridge_;

// for publishing the poses
motld::TrackedObjects tracked_objects_;
geometry_msgs::PoseArray debug_objects_;
ros::Publisher poses_pub_;
ros::Publisher debug_poses_pub_;

// for sending a debug image
// RGB channels of debug image
Matrix maRed;
Matrix maGreen;
Matrix maBlue;
// the actual ipl image and message
IplImage *out_img_;
sensor_msgs::Image image_msg_;
image_transport::Publisher image_pub_;

// motld instance and settings
bool learningEnabled_ = true;
MOTLDSettings settings = MOTLDSettings(COLOR_MODE_RGB);
MultiObjectTLD tld_ = MultiObjectTLD(width_, height_, settings);
boost::mutex tld_mutex_;

// initialization
bool initialized_ = false;
bool have_initial_box_ = false;
bool started_ = false;
bool read_model_from_file_ = false;
std::string model_file_("");
std::string camera_topic_("/camera/rgb/image_color");
std::string depth_topic_("/camera/depth_registered/points");

CloudProperties calculateCloudProperties(ObjectBox box) {
  CloudProperties cp;

  int startX, startY;
  int endX, endY;
  // security check
  if (!initialized_ || !started_) {
    ROS_WARN("calculateCloudProperties called but tracker is not yet started");
    return cp;
  }
  // rescale box appropriately
  if (force_resize_) {
    box.x *= resize_factor_;
    box.y *= resize_factor_;
    box.width *= resize_factor_;
    box.height *= resize_factor_;
  }
  // TODO: this is a crude hack as it only supports equal resolution and
  //       half image resolution clouds that are registered with the image
  if (int(cloud_.width) != ipl_img_->width * resize_factor_) {
    startX = box.x/2;
    startY = box.y/2;
    endX = startX + box.width / 2;
    endY = startY + box.height / 2;
  } else {
    startX = box.x;
    startY = box.y;
    endX = startX + box.width;
    endY = startY + box.height;
  }
  //ROS_INFO("startX %d", startX);
  // compute center of mass of bounding box
  int count = 0;
  double mean_z = 0.;
  for (int x = startX; x < endX; ++x) {
    for (int y = startY; y < endY; ++y) {
      const pcl::PointXYZ &p = cloud_(x, y);
      if (std::isnan(p.x) || std::isnan(p.y) || std::isnan(p.z))
        continue;
      mean_z += p.z;
      ++count;
    }
  }
  mean_z /=count;

  double var_z = 0.;
  for (int x = startX; x < endX; ++x) {
    for (int y = startY; y < endY; ++y) {
      const pcl::PointXYZ &p = cloud_(x, y);
      if (std::isnan(p.x) || std::isnan(p.y) || std::isnan(p.z))
        continue;
      var_z += SQR(p.z - mean_z);
    }
  }
  var_z /= count;
  double std_dev_z = std::sqrt(var_z);

  // remove outliers, compute new mean and width/height
  float min_x = 1e7, max_x = -1e7;
  float min_y = 1e7, max_y = -1e7;
  count = 0;
  for (int x = startX; x < endX; ++x) {
    for (int y = startY; y < endY; ++y) {
      const pcl::PointXYZ &p = cloud_(x, y);
      if (std::isnan(p.x) || std::isnan(p.y) || std::isnan(p.z))
        continue;
      if (std::fabs(p.z - mean_z) > 2. * std_dev_z) {
        //ROS_INFO_STREAM("Removing depth of: " << p.z << " first mean was: " << mean_z << " std_dev: " << std_dev_z);
        // remove this point as it is too far away from the mean
        continue;
      }
      cp.mean.x += p.x;
      cp.mean.y += p.y;
      cp.mean.z += p.z;
      if (p.x < min_x)
        min_x = p.x;
      if (p.x > max_x)
        max_x = p.x;
      if (p.y < min_y)
        min_y = p.y;
      if (p.y > max_y)
        max_y = p.y;
      ++count;
    }
  }
  if (count != 0) {
    cp.mean.x /= count;
    cp.mean.y /= count;
    cp.mean.z /= count;
    cp.width  = max_x - min_x;
    cp.height = max_y - min_y;
  }

  if (cp.height < 0)
    cp.height = 0. ;
  if (cp.width < 0)
    cp.width = 0. ;

  return cp;
}

void init() {
  // create tld instance
  if (read_model_from_file_) 
  {
    tld_ = MultiObjectTLD::loadClassifier(model_file_.c_str());
    have_initial_box_ = true;
  }

  // allocate array holding all pixel values for computation
  img_ = new unsigned char[width_ * height_ * CHANNELS];
  // also allocate an ipimage for displaying debug info and resizing
  CvSize wsize = {width_, height_};
  ipl_img_ = cvCreateImage(wsize, IPL_DEPTH_8U, CHANNELS);
  out_img_ = cvCreateImage(wsize, IPL_DEPTH_8U, CHANNELS);
  initialized_ = true;
  ROS_INFO("done initializing");
}

void prepareForProcessing(IplImage *image) {
  if (!initialized_)
  {
    ROS_WARN("motld not initialized when calling prepareForProcessing()");
    return;
  }
  IplImage *copy_src = image;
  int size = width_ * height_;
  if(force_resize_)
  {
    cvResize(image, ipl_img_);
    copy_src = ipl_img_;
  }
  // copy into working buffer
  for(int i = 0; i < size; i++){
    img_[i] = copy_src->imageData[i*3+2];
    img_[i+size] = copy_src->imageData[i*3+1];
    img_[i+2*size] = copy_src->imageData[i*3];
  }
}

void processAndPublish() {
  if (!initialized_)
  {
    ROS_WARN("motld not initialized when calling processAndPublish()");
    return;
  }
  double start = ros::Time::now().toSec();
  tld_.processFrame(img_);
  double elapsed_time = ros::Time::now().toSec() - start;
  double fps = 1. / elapsed_time;
  ROS_DEBUG("MOTLD fps: %.2f", fps);
 
  // publish debug image if subscribers present
  if (image_pub_.getNumSubscribers() > 0)
  {
    tld_.getDebugImage(img_, maRed, maGreen, maBlue, 0);    
    for(int i = 0; i < width_*height_; ++i){
      out_img_->imageData[3*i+2] = maRed.data()[i];
      out_img_->imageData[3*i+1] = maGreen.data()[i];
      out_img_->imageData[3*i+0] = maBlue.data()[i];
    }
    bridge_.fromIpltoRosImage(out_img_, image_msg_);
    image_msg_.header.stamp = ros::Time::now();
    image_msg_.header.frame_id = image_frame_;
    image_pub_.publish(image_msg_);
  }
}

void processCloud() {
  // publish poses
  bool poses_pub_subscribed = poses_pub_.getNumSubscribers() > 0;
  bool debug_pub_subscribed = debug_poses_pub_.getNumSubscribers() > 0;
  if (poses_pub_subscribed || debug_pub_subscribed)
  {
    std::vector<ObjectBox> boxes = tld_.getObjectBoxes();

    // pose array
    tracked_objects_.pose.header.stamp = ros::Time::now();
    tracked_objects_.pose.header.frame_id = frame_id_;
    tracked_objects_.height.resize(boxes.size());
    tracked_objects_.width.resize(boxes.size());
    tracked_objects_.status.resize(boxes.size());
    tracked_objects_.name.resize(boxes.size());
    tracked_objects_.pose.poses.resize(boxes.size());

    // debug array 
    debug_objects_.header.stamp = ros::Time::now();
    debug_objects_.header.frame_id = frame_id_;
    debug_objects_.poses.resize(boxes.size());

    for (size_t o = 0; o < boxes.size(); ++o) 
    {
      int status = tld_.getStatus(boxes[o].objectId);
      // set name of object
      tracked_objects_.name[o] = boxes[o].objectId;

      // calculate cloud properties from cloud and publish
      CloudProperties cp  = calculateCloudProperties(boxes[o]);
      tracked_objects_.height[o] = cp.height;
      tracked_objects_.width[o] = cp.width;
      tracked_objects_.status[o] = status;
      
      tracked_objects_.pose.poses[o].position.x = cp.mean.x;
      tracked_objects_.pose.poses[o].position.y = cp.mean.y;
      tracked_objects_.pose.poses[o].position.z = cp.mean.z;

      // TODO: compute reasonable orientation from cloud here
      tracked_objects_.pose.poses[o].orientation.x = 0.f;
      tracked_objects_.pose.poses[o].orientation.y = 0.f;
      tracked_objects_.pose.poses[o].orientation.z = 0.f;
      tracked_objects_.pose.poses[o].orientation.w = 1.f;

      if (debug_pub_subscribed)
        {
          debug_objects_.poses[o].position.x = cp.mean.x;
          debug_objects_.poses[o].position.y = cp.mean.y;
          debug_objects_.poses[o].position.z = cp.mean.z;

          // TODO: compute reasonable orientation from cloud here
          debug_objects_.poses[o].orientation.x = 0.f;
          debug_objects_.poses[o].orientation.y = 0.f;
          debug_objects_.poses[o].orientation.z = 0.f;
          debug_objects_.poses[o].orientation.w = 1.f;
        }
    }
    
    if (poses_pub_subscribed)
        poses_pub_.publish(tracked_objects_);

    if (debug_pub_subscribed)
        debug_poses_pub_.publish(debug_objects_);
  }
}

void trackerCallback(const sensor_msgs::ImageConstPtr& msg) {
  // grab a frame
  IplImage *color_image = bridge_.imgMsgToCv(msg, "bgr8");
  image_frame_ = msg->header.frame_id;

  // first time initialization stuff
  if (!started_)
  {
    if (!force_resize_) {
      width_ = color_image->width;
      height_ = color_image->height;
    }
    else {
      resize_factor_ = color_image->width / width_;
    }
    started_ = true;
  }

  // prepare input
  prepareForProcessing(color_image);
  // lock and process
  {
    boost::mutex::scoped_lock(tld_mutex_);
    processAndPublish();
  }
}

void synchronizedCallback(const sensor_msgs::ImageConstPtr& msg, const sensor_msgs::PointCloud2ConstPtr& msg_cloud) {
  // grab the image and process
  trackerCallback(msg);
  // also grab the cloud
  pcl::fromROSMsg(*msg_cloud, cloud_);
  frame_id_ = cloud_.header.frame_id;
  // and process it
  processCloud();
}

void boundingBoxCallback(const motld::BoundingBoxConstPtr& msg) {
  if (!initialized_ || !started_) 
  {
    ROS_WARN("boundingBoxCallback called but tracker is not yet started");
    return;
  }
  ROS_INFO("boundingBoxCallback called");
  ObjectBox box;
  box.x = msg->x / resize_factor_;
  box.y = msg->y / resize_factor_;
  box.width = msg->width / resize_factor_;
  box.height = msg->height / resize_factor_;
  {
    boost::mutex::scoped_lock(tld_mutex_);
    tld_.addObject(box);
    have_initial_box_ = true;
  }
}

void saveModelCallback(const motld::SaveConstPtr& msg) {
  boost::unique_lock<boost::mutex> state_lock_(tld_mutex_);
  ROS_INFO("Saving model to %s", msg->filename.c_str());
  tld_.saveClassifier(msg->filename.c_str());
}

void configCallback(const motld::ConfigConstPtr& msg) {
  ROS_INFO("Got new config parameters");
  if (msg->learning != -1) {
    learningEnabled_ = msg->learning > 0;
  }
  
  // set the learning parameter in the tld instance
  {
    boost::mutex::scoped_lock tld_lock(tld_mutex_);
    tld_.enableLearning(learningEnabled_);
  }
}


int main(int argc, char **argv) {
  ros::init(argc, argv, "motld_node");
  ros::NodeHandle nh;
  ros::NodeHandle private_nh("~");
  image_transport::ImageTransport it(nh);
  // read parameters
  private_nh.param("width", width_, width_);
  private_nh.param("height", width_, width_);
  private_nh.param("force_resize", force_resize_, force_resize_);
  private_nh.param("learning_enabled", learningEnabled_, learningEnabled_);
  private_nh.param("model_file", model_file_, model_file_);
  private_nh.param("use_cloud", use_cloud_, use_cloud_);
  private_nh.param("camera_topic", camera_topic_, camera_topic_);
  private_nh.param("depth_topic", depth_topic_, depth_topic_);
  if (model_file_.compare("") != 0)
  {
    read_model_from_file_ = true;
  }
  // init subscribers
  image_transport::SubscriberFilter sub(it, camera_topic_, 1);
  image_transport::Subscriber sub_img; 
  message_filters::Subscriber< sensor_msgs::PointCloud2 > sub_cloud(nh, depth_topic_, 1);
  message_filters::Synchronizer< MySyncPolicy > sync(MySyncPolicy(30), sub, sub_cloud);
  if (!use_cloud_ || depth_topic_.compare("") == 0) 
  {
    ROS_INFO("not registering with cloud");
    ROS_INFO_STREAM("topic: " << camera_topic_);
    sub_img = it.subscribe(camera_topic_, 1, trackerCallback);
  }
  else
  {
    ROS_INFO("using kinect registered version");
    ROS_INFO_STREAM("topics: " << camera_topic_ << " " << depth_topic_);
    sync.registerCallback(synchronizedCallback);
  }
  
  ros::Subscriber sub_box = private_nh.subscribe<motld::BoundingBox>("tracker_start_box", 1, boundingBoxCallback);
  ros::Subscriber sub_config = private_nh.subscribe<motld::Config>("config_opentld", 1, configCallback);
  ros::Subscriber sub_save = private_nh.subscribe<motld::Save>("save_tld_model", 1, saveModelCallback);        

  // init publishers
  image_pub_ = it.advertise("tracking_image", 1);
  poses_pub_ = private_nh.advertise<motld::TrackedObjects>("tracked_objects", 1);
  debug_poses_pub_ = private_nh.advertise<geometry_msgs::PoseArray>("debug_poses", 1);

  // init motld
  init();
  // and run (everything is handled via the trackerCallback)
  while (ros::ok()) 
  {
    ros::spinOnce();
  }
}