3D Object tracking with PCL

CMakeLists.txt

cmake_minimum_required(VERSION 2.6 FATAL_ERROR)

project(correspondence_grouping_openni)

set(CMAKE_RUNTIME_OUTPUT_DIRECTORY "../bin/")
set(PCL_DIR "/Users/kikko/Dev/kinect/pcl/pcl-trunk/build/")

find_package(PCL 1.5 REQUIRED)

include_directories(${PCL_INCLUDE_DIRS})
link_directories(${PCL_LIBRARY_DIRS})
add_definitions(${PCL_DEFINITIONS})

add_executable (correspondence_grouping_openni correspondence_grouping_openni.cpp)
target_link_libraries (correspondence_grouping_openni ${PCL_LIBRARIES})

correspondence_grouping_openni.cpp

#define PCD_SCENE "bottle_scene.pcd"

#ifndef PCD_SCENE
    #include <pcl/io/openni_grabber.h>
#endif

#include <pcl/console/parse.h>
#include <pcl/io/pcd_io.h>
#include <pcl/point_cloud.h>
#include <pcl/correspondence.h>
#include <pcl/features/normal_3d_omp.h>
#include <pcl/features/shot_omp.h>
#include <pcl/features/board.h>
#include <pcl/keypoints/uniform_sampling.h>
#include <pcl/recognition/cg/hough_3d.h>
#include <pcl/recognition/cg/geometric_consistency.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <pcl/kdtree/impl/kdtree_flann.hpp>
#include <pcl/common/transforms.h>
#include <pcl/filters/passthrough.h>

typedef pcl::PointXYZRGBA PointType;
typedef pcl::Normal NormalType;
typedef pcl::ReferenceFrame RFType;
typedef pcl::SHOT352 DescriptorType;

class CGCloud {
    
public:
    
    CGCloud():
    cloud(new pcl::PointCloud<PointType> ()),
    keypoints (new pcl::PointCloud<PointType> ()),
    normals (new pcl::PointCloud<NormalType> ()),
    descriptors (new pcl::PointCloud<DescriptorType> ()),
    refFrames (new pcl::PointCloud<RFType> ())
    {}
    
    virtual ~CGCloud() {}
    
    pcl::PointCloud<PointType>::Ptr         cloud;
    pcl::PointCloud<PointType>::Ptr         keypoints;
    pcl::PointCloud<NormalType>::Ptr        normals;
    pcl::PointCloud<DescriptorType>::Ptr    descriptors;
    
    pcl::PointCloud<RFType>::Ptr refFrames;
};


//Algorithm params
bool use_cloud_resolution_ (false);

float model_ss_ (0.01f);
float scene_ss_ (0.03f);
float rf_rad_ (0.015f);
float descr_rad_ (0.02f);
float cg_size_ (0.01f);
float cg_thresh_ (5.0f);

bool bCompute = false;
bool bSaveScene = false;

std::string dataPath = "/Users/kikko/Dev/kinect/pcl/tests/correspondance_grouping_openni/bin/Debug/";

void keyboardEventOccurred (const pcl::visualization::KeyboardEvent &event,
                            void* viewer_void)
{
    std::cout << "key occured" << std::endl;
    boost::shared_ptr<pcl::visualization::PCLVisualizer> viewer = *static_cast<boost::shared_ptr<pcl::visualization::PCLVisualizer> *> (viewer_void);
    if (event.getKeySym () == " " && event.keyDown ())
    {
        bCompute = true;
    }
}

void mouseEventOccurred (const pcl::visualization::MouseEvent &event,
                         void* viewer_void)
{
    boost::shared_ptr<pcl::visualization::PCLVisualizer> viewer = *static_cast<boost::shared_ptr<pcl::visualization::PCLVisualizer> *> (viewer_void);
    if (event.getButton () == pcl::visualization::MouseEvent::LeftButton &&
        event.getType () == pcl::visualization::MouseEvent::MouseButtonRelease)
    {
        if(event.getX() < 20 && event.getY() < 20) {
            bCompute = true;
            bSaveScene = true;
        }
    }
}


class ObjectDetection {
    
public:
    
    CGCloud model, scene;
    
    // openNI
    boost::mutex cloud_mutex;
    
    // visualization
    pcl::visualization::PCLVisualizer viewer;
    pcl::PointCloud<PointType>::Ptr off_scene_model;
    pcl::PointCloud<PointType>::Ptr off_scene_model_keypoints;
    
    // correspondences
    pcl::CorrespondencesPtr model_scene_corrs;
    std::vector<pcl::Correspondences> clustered_corrs;
    
    // cluster
    std::vector<Eigen::Matrix4f, Eigen::aligned_allocator<Eigen::Matrix4f> > rototranslations;
    
    
    ObjectDetection()
    :viewer ("Correspondence Grouping"),
    off_scene_model (new pcl::PointCloud<PointType> ()),
    off_scene_model_keypoints (new pcl::PointCloud<PointType> ()),
    model_scene_corrs (new pcl::Correspondences ())
    {
    }
    
    double computeCloudResolution (const pcl::PointCloud<PointType>::ConstPtr &cloud)
    {
      double res = 0.0;
      int n_points = 0;
      int nres;
      std::vector<int> indices (2);
      std::vector<float> sqr_distances (2);
      pcl::search::KdTree<PointType> tree;
      tree.setInputCloud (cloud);

      for (size_t i = 0; i < cloud->size (); ++i)
      {
        if (! pcl_isfinite ((*cloud)[i].x))
        {
          continue;
        }
        //Considering the second neighbor since the first is the point itself.
        nres = tree.nearestKSearch (i, 2, indices, sqr_distances);
        if (nres == 2)
        {
          res += sqrt (sqr_distances[1]);
          ++n_points;
        }
      }
      if (n_points != 0)
      {
        res /= n_points;
      }
      return res;
    }

    void setupResolutionInvariance(const pcl::PointCloud<PointType>::ConstPtr &cloud)
    {
        float resolution = static_cast<float> (computeCloudResolution (cloud));
        if (resolution != 0.0f)
        {
            model_ss_   *= resolution;
            scene_ss_   *= resolution;
            rf_rad_     *= resolution;
            descr_rad_  *= resolution;
            cg_size_    *= resolution;
        }
        
        std::cout << "Model resolution:       " << resolution << std::endl;
        std::cout << "Model sampling size:    " << model_ss_ << std::endl;
        std::cout << "Scene sampling size:    " << scene_ss_ << std::endl;
        std::cout << "LRF support radius:     " << rf_rad_ << std::endl;
        std::cout << "SHOT descriptor radius: " << descr_rad_ << std::endl;
        std::cout << "Clustering bin size:    " << cg_size_ << std::endl << std::endl;
    }


    void computeNormals(CGCloud& cloud)
    {
        pcl::NormalEstimationOMP<PointType, NormalType> norm_est;
        norm_est.setKSearch (10);
        norm_est.setInputCloud (cloud.cloud);
        norm_est.compute (*cloud.normals);
    }

    void downSample(CGCloud& cloud, float ss)
    {
        pcl::PointCloud<int> sampled_indices;
        pcl::UniformSampling<PointType> uniform_sampling;
        uniform_sampling.setInputCloud (cloud.cloud);
        uniform_sampling.setRadiusSearch (ss);
        uniform_sampling.compute (sampled_indices);
        pcl::copyPointCloud (*cloud.cloud, sampled_indices.points, *cloud.keypoints);
        std::cout << "Cloud total points: " << cloud.cloud->size () << "; Selected Keypoints: " << cloud.keypoints->size () << std::endl;
    }

    void computeKeypointsDescriptor(CGCloud& cloud)
    {
        pcl::SHOTEstimationOMP<PointType, NormalType, DescriptorType> descr_est;
        descr_est.setRadiusSearch (descr_rad_);
        
        descr_est.setInputCloud (cloud.keypoints);
        descr_est.setInputNormals (cloud.normals);
        descr_est.setSearchSurface (cloud.cloud);
        descr_est.compute (*cloud.descriptors);
    }

    void findCorrespondences()
    {
        pcl::KdTreeFLANN<DescriptorType> match_search;
        match_search.setInputCloud (model.descriptors);
        
        model_scene_corrs->clear();
        
        //  For each scene keypoint descriptor, find nearest neighbor into the model keypoints descriptor cloud and add it to the correspondences vector.
        for (size_t i = 0; i < scene.descriptors->size (); ++i)
        {
            std::vector<int> neigh_indices (1);
            std::vector<float> neigh_sqr_dists (1);
            if (!pcl_isfinite (scene.descriptors->at (i).descriptor[0])) //skipping NaNs
            {
                continue;
            }
            int found_neighs = match_search.nearestKSearch (scene.descriptors->at (i), 1, neigh_indices, neigh_sqr_dists);
            if(found_neighs == 1 && neigh_sqr_dists[0] < 0.25f) //  add match only if the squared descriptor distance is less than 0.25 (SHOT descriptor distances are between 0 and 1 by design)
            {
                pcl::Correspondence corr (neigh_indices[0], static_cast<int> (i), neigh_sqr_dists[0]);
                model_scene_corrs->push_back (corr);
            }
        }
        std::cout << "Correspondences found: " << model_scene_corrs->size () << std::endl;
    }

    void computeReferenceFrames(CGCloud& cloud)
    {
        // only if using hough3D
        pcl::BOARDLocalReferenceFrameEstimation<PointType, NormalType, RFType> rf_est;
        rf_est.setFindHoles (true);
        rf_est.setRadiusSearch (rf_rad_);
        
        rf_est.setInputCloud (cloud.keypoints);
        rf_est.setInputNormals (cloud.normals);
        rf_est.setSearchSurface (cloud.cloud);
        rf_est.compute (*cloud.refFrames);
    }

    void clusterResult()
    {
        pcl::Hough3DGrouping<PointType, PointType, RFType, RFType> clusterer;

        clusterer.setHoughBinSize (cg_size_);
        clusterer.setHoughThreshold (cg_thresh_);
        clusterer.setUseInterpolation (true);
        clusterer.setUseDistanceWeight (false);
        
        clusterer.setInputCloud (model.keypoints);
        clusterer.setInputRf (model.refFrames);
        clusterer.setSceneCloud (scene.keypoints);
        clusterer.setSceneRf (scene.refFrames);
        clusterer.setModelSceneCorrespondences (model_scene_corrs);
        
        clustered_corrs.clear();
        
        //clusterer.cluster (clustered_corrs);
        clusterer.recognize (rototranslations, clustered_corrs);
        
        std::cout << "Clustered correspondences : " << clustered_corrs.size() << std::endl;
    }

    void printResult()
    {
        std::cout << "Model instances found: " << rototranslations.size () << std::endl;
        for (size_t i = 0; i < rototranslations.size (); ++i)
        {
            std::cout << "\n    Instance " << i + 1 << ":" << std::endl;
            std::cout << "        Correspondences belonging to this instance: " << clustered_corrs[i].size () << std::endl;
            
            // Print the rotation matrix and translation vector
            Eigen::Matrix3f rotation = rototranslations[i].block<3,3>(0, 0);
            Eigen::Vector3f translation = rototranslations[i].block<3,1>(0, 3);
            
            printf ("\n");
            printf ("            | %6.3f %6.3f %6.3f | \n", rotation (0,0), rotation (0,1), rotation (0,2));
            printf ("        R = | %6.3f %6.3f %6.3f | \n", rotation (1,0), rotation (1,1), rotation (1,2));
            printf ("            | %6.3f %6.3f %6.3f | \n", rotation (2,0), rotation (2,1), rotation (2,2));
            printf ("\n");
            printf ("        t = < %0.3f, %0.3f, %0.3f >\n", translation (0), translation (1), translation (2));
        }
    }
    
    void addModelSceneCorrespondencesView()
    {
        for (size_t i = 0; i < model_scene_corrs->size (); ++i)
        {
            std::stringstream ss_line;
            ss_line << "correspondence_line" << i;
            PointType& model_point = off_scene_model_keypoints->at (model_scene_corrs->at(i).index_query);
            PointType& scene_point = scene.keypoints->at (model_scene_corrs->at(i).index_match);
            
            //  We are drawing a line for each pair of clustered correspondences found between the model and the scene
            viewer.addLine<PointType, PointType> (model_point, scene_point, 0, 255, 0, ss_line.str ());
        }
    }
    
    void addOffSceneModelView()
    {
        //  We are translating the model so that it doesn't end in the middle of the scene representation
        pcl::transformPointCloud (*model.cloud, *off_scene_model, Eigen::Vector3f (0,0.2,0.2), Eigen::Quaternionf (sqrt(0.5), sqrt(0.5), 0, 0));
        pcl::transformPointCloud (*model.keypoints, *off_scene_model_keypoints, Eigen::Vector3f (0,0.2,0.2), Eigen::Quaternionf (sqrt(0.5), sqrt(0.5), 0, 0));
        
        pcl::visualization::PointCloudColorHandlerCustom<PointType> off_scene_model_color_handler (off_scene_model, 255, 255, 128);
        viewer.addPointCloud (off_scene_model, off_scene_model_color_handler, "off_scene_model");
    }
    
    void addKeypointsView(pcl::PointCloud<PointType>::Ptr &cloud, std::string name )
    {
        pcl::visualization::PointCloudColorHandlerCustom<PointType> color_handler (cloud, 0, 0, 255);
        viewer.addPointCloud (cloud, color_handler, name);
        viewer.setPointCloudRenderingProperties (pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 5, name);
    }
    
    void passThroughFilter(const pcl::PointCloud<PointType>::ConstPtr &cloud, pcl::PointCloud<PointType>::Ptr &filtered_dest_cloud)
    {
        pcl::PassThrough<PointType> pass;
        pass.setInputCloud (cloud);
        pass.setFilterFieldName ("z");
        pass.setFilterLimits (0.0, 1.0);
        //pass.setFilterLimitsNegative (true);
        pass.filter (*filtered_dest_cloud);
    }

    void openNIcallback (const pcl::PointCloud<PointType>::ConstPtr& cloud)
    {
        if (!viewer.wasStopped ())
        {
            cloud_mutex.lock ();
            
            passThroughFilter(cloud, scene.cloud);
            
            if(bSaveScene) {
                std::string filename = dataPath + "/bottle_scene.pcd";
                pcl::io::savePCDFile(filename, *cloud);
                bSaveScene = false;
            }
            
            cloud_mutex.unlock ();
        }
    }
    
    void update()
    {
        computeNormals(scene);
        downSample(scene, scene_ss_);
        computeKeypointsDescriptor(scene);
        computeReferenceFrames(scene);
        
        findCorrespondences();
        
        clusterResult();
        
        printResult();
        
        viewer.removeAllShapes();
        addModelSceneCorrespondencesView();
        
        bCompute = false;
    }
    
    void run()
    {
        std::string filepath = dataPath + "bottle.pcd";
        if (pcl::io::loadPCDFile (filepath, *model.cloud) < 0)
        {
            std::cout << "Error loading model cloud." << std::endl;
        }
        
        if (use_cloud_resolution_)
        {
            setupResolutionInvariance(model.cloud);
        }
        
        std::cout << "Initiating Model." << std::endl;
        
        computeNormals(model);
        downSample(model, model_ss_);
        computeKeypointsDescriptor(model);
        computeReferenceFrames(model);

        std::cout << "Model initialization done." << std::endl;
        
        viewer.addCoordinateSystem (1.0);
        viewer.registerKeyboardCallback (keyboardEventOccurred, (void*)&viewer);
        viewer.registerMouseCallback (mouseEventOccurred, (void*)&viewer);
        
        addOffSceneModelView();
        addKeypointsView(off_scene_model_keypoints, "off_scene_model_keypoints");
        
        //setupVisualization();
        
#ifdef PCD_SCENE
        pcl::PointCloud<PointType>::Ptr cloud(new pcl::PointCloud<PointType>());
        if (pcl::io::loadPCDFile (dataPath + PCD_SCENE, *cloud) < 0)
        {
            std::cout << "Error loading scene cloud." << std::endl;
        }
        else {
            passThroughFilter(cloud, scene.cloud);
            update();
            viewer.addPointCloud<PointType> (scene.cloud, "scene_cloud");
            addKeypointsView(scene.keypoints, "scene_keypoints");
        }
#else
        std::cout << "Starting OpenNI..." << std::endl;
        
        pcl::Grabber* interface = new pcl::OpenNIGrabber ();
        boost::function<void(const pcl::PointCloud<PointType>::ConstPtr&)> f = boost::bind (&ObjectDetection::openNIcallback, this, _1);
        boost::signals2::connection c = interface->registerCallback (f);
        
        interface->start ();
        
        std::cout << "OpenNI started." << std::endl;
#endif
        
        while (!viewer.wasStopped ())
        {
            viewer.spinOnce (1);

#ifdef PCD_SCENE
            if (scene.cloud && cloud_mutex.try_lock ())
            {
#endif
                if(bCompute)  
                {
                    update();
                }
                
                if (!viewer.updatePointCloud<PointType> (scene.cloud, "scene_cloud"))
                    viewer.addPointCloud<PointType> (scene.cloud, "scene_cloud");
                
                if (!viewer.updatePointCloud<PointType> (scene.keypoints, "scene_keypoints"))
                    addKeypointsView(scene.keypoints, "scene_keypoints");
                
#ifdef PCD_SCENE    
                cloud_mutex.unlock ();
            }
#endif
        }
        
#ifndef PCD_SCENE
        interface->stop();
#endif
    }
};


int main (int argc, char *argv[])
{
    ObjectDetection obj_detect_app;
    obj_detect_app.run();
    return (0);
}