vfh_nn_classifier.h

Go to the documentation of this file.

/*
 * Software License Agreement (BSD License)
 *
 *  Copyright (c) 2010, Willow Garage, Inc.
 *  All rights reserved.
 *
 *  Redistribution and use in source and binary forms, with or without
 *  modification, are permitted provided that the following conditions
 *  are met:
 *
 *   * Redistributions of source code must retain the above copyright
 *     notice, this list of conditions and the following disclaimer.
 *   * Redistributions in binary form must reproduce the above
 *     copyright notice, this list of conditions and the following
 *     disclaimer in the documentation and/or other materials provided
 *     with the distribution.
 *   * Neither the name of Willow Garage, Inc. nor the names of its
 *     contributors may be used to endorse or promote products derived
 *     from this software without specific prior written permission.
 *
 *  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 *  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 *  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 *  FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 *  COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 *  INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 *  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 *  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 *  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 *  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
 *  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 *  POSSIBILITY OF SUCH DAMAGE.
 *
 * $Id$
 *
 */

#ifndef VFHCLASSIFICATION_H_
#define VFHCLASSIFICATION_H_

#include <fstream>
#include <pcl/point_types.h>
#include <pcl/io/pcd_io.h>
#include <pcl/features/vfh.h>
#include <pcl/features/normal_3d.h>
#include <pcl/apps/nn_classification.h>
#include <pcl/PCLPointCloud2.h>

namespace pcl
{
  template <typename PointT> pcl::PointCloud<pcl::VFHSignature308>::Ptr
  computeVFH (typename PointCloud<PointT>::ConstPtr cloud, double radius)
  {
    using namespace pcl;

    // Create an empty kdtree representation, and pass it to the objects.
    // Its content will be filled inside the object, based on the given input dataset (as no other search surface is given).
    typename pcl::search::KdTree<PointT>::Ptr tree (new pcl::search::KdTree<PointT> ());

    // Create the normal estimation class, and pass the input dataset to it
    NormalEstimation<PointT, Normal> ne;
    ne.setInputCloud (cloud);
    ne.setSearchMethod (tree);

    // Use all neighbors in a sphere of given radius to compute the normals
    PointCloud<Normal>::Ptr normals (new PointCloud<Normal> ());
    ne.setRadiusSearch (radius);
    ne.compute (*normals);

    // Create the VFH estimation class, and pass the input dataset+normals to it
    VFHEstimation<PointT, Normal, VFHSignature308> vfh;
    vfh.setInputCloud (cloud);
    vfh.setInputNormals (normals);
    vfh.setSearchMethod (tree);

    // Output datasets
    PointCloud<VFHSignature308>::Ptr vfhs (new PointCloud<VFHSignature308>);

    // Compute the features and return
    vfh.compute (*vfhs);
    return vfhs;
  }

  class VFHClassifierNN
  {
    public:

      typedef pcl::PointCloud<pcl::VFHSignature308> FeatureCloud;
      typedef pcl::PointCloud<pcl::VFHSignature308>::Ptr FeatureCloudPtr;
      typedef pcl::PointCloud<pcl::VFHSignature308>::ConstPtr FeatureCloudConstPtr;
      typedef NNClassification<pcl::VFHSignature308>::Result Result;
      typedef NNClassification<pcl::VFHSignature308>::ResultPtr ResultPtr;

    private:

      FeatureCloudPtr training_features_;
      std::vector<std::string> labels_;
      NNClassification<pcl::VFHSignature308> classifier_;

    public:

      VFHClassifierNN () : training_features_ (), labels_ (), classifier_ ()
      {
        reset ();
      }

      void reset ()
      {
        training_features_.reset (new FeatureCloud);
        labels_.clear ();
        classifier_ = NNClassification<pcl::VFHSignature308> ();
      }

      void finalizeTraining ()
      {
        finalizeTree ();
        finalizeLabels ();
      }

      void finalizeTree ()
      {
        classifier_.setTrainingFeatures(training_features_);
      }

      void finalizeLabels ()
      {
        classifier_.setTrainingLabels(labels_);
      }

      bool saveTrainingFeatures(std::string file_name, std::string labels_file_name)
      {
        if (labels_.size () == training_features_->points.size ())
        {
          if (pcl::io::savePCDFile (file_name.c_str (), *training_features_) != 0)
            return false;
          std::ofstream f (labels_file_name.c_str ());
          BOOST_FOREACH (std::string s, labels_)
            f << s << "\n";
          return true;
        }
        return false;
      }

      bool addTrainingFeatures (const FeatureCloudPtr training_features, const std::vector<std::string> &labels)
      {
        if (labels.size () == training_features->points.size ())
        {
          labels_.insert (labels_.end (), labels.begin (), labels.end ());
          training_features_->points.insert (training_features_->points.end (), training_features->points.begin (), training_features->points.end ());
          training_features_->header = training_features->header;
          training_features_->height = 1;
          training_features_->width  = static_cast<uint32_t> (training_features_->points.size ());
          training_features_->is_dense &= training_features->is_dense;
          training_features_->sensor_origin_ = training_features->sensor_origin_;
          training_features_->sensor_orientation_ = training_features->sensor_orientation_;
          return true;
        }
        return false;
      }

      bool loadTrainingFeatures(std::string file_name, std::string labels_file_name)
      {
        FeatureCloudPtr cloud (new FeatureCloud);
        if (pcl::io::loadPCDFile (file_name.c_str (), *cloud) != 0)
          return false;
        std::vector<std::string> labels;
        std::ifstream f (labels_file_name.c_str ());
        std::string label;
        while (getline (f, label))
          if (label.size () > 0)
            labels.push_back(label);
        return addTrainingFeatures (cloud, labels);
      }

      bool loadTrainingData (std::string file_name, std::string label)
      {
        pcl::PCLPointCloud2 cloud_blob;
        if (pcl::io::loadPCDFile (file_name.c_str (), cloud_blob) != 0)
          return false;
        return addTrainingData (cloud_blob, label);
      }

      bool addTrainingData (const pcl::PCLPointCloud2 &training_data, std::string &label)
      {
        // Create label list containing the single label
        std::vector<std::string> labels;
        labels.push_back (label);

        // Compute the feature from the cloud and add it as a training example
        FeatureCloudPtr vfhs = computeFeature (training_data);
        return addTrainingFeatures(vfhs, labels);
      }

      ResultPtr classify (const pcl::PCLPointCloud2 &testing_data, double radius = 300, double min_score = 0.002)
      {
        // compute the VFH feature for this point cloud
        FeatureCloudPtr vfhs = computeFeature (testing_data);
        // compute gaussian parameter producing the desired minimum score (around 50 for the default values)
        float gaussian_param = - static_cast<float> (radius / log (min_score));
        // TODO accept result to be filled in by reference
        return classifier_.classify(vfhs->points.at (0), radius, gaussian_param);
      }

      FeatureCloudPtr computeFeature (const pcl::PCLPointCloud2 &points, double radius = 0.03)
      {
        pcl::PointCloud<pcl::PointXYZ>::Ptr cloud (new pcl::PointCloud<pcl::PointXYZ> ());
        pcl::fromPCLPointCloud2 (points, *cloud);
        return pcl::computeVFH<pcl::PointXYZ> (cloud, radius);
      }
  };
}

#endif /* VFHCLASSIFICATION_H_ */