CRForestEstimator.cpp

#ifdef DWIN
#include "stdafx.h"
#include <windows.h>
#endif

#include "CRForestEstimator.h"
#include <vector>
#include <iostream>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>

using namespace std;
using namespace cv;

bool CRForestEstimator::loadForest(const char* treespath, int ntrees){

	// Init forest with number of trees
	crForest = new CRForest( ntrees );
	// Load forest
	if( !crForest->loadForest( treespath ) )
		return false;

	return true;
}

Rect CRForestEstimator::getBoundingBox(const Mat& im3D){

	Rect bbox;
	int min_x = im3D.cols;
	int min_y = im3D.rows;
	int max_x = 0;
	int max_y = 0;
	int p_width = crForest->getPatchWidth();
	int p_height = crForest->getPatchHeight();

	for(int y = 0; y < im3D.rows; y++)
	{
	    const Vec3f* Mi = im3D.ptr<Vec3f>(y);
	    for(int x = 0; x < im3D.cols; x++){

	    	if( Mi[x][2] > 0) {

				min_x = MIN(min_x,x); min_y = MIN(min_y,y);
				max_x = MAX(max_x,x); max_y = MAX(max_y,y);
			}

	    }
	}

	int new_w = max_x - min_x + p_width;
	int new_h = max_y - min_y + p_height;

	bbox.x = MIN( im3D.cols-1, MAX( 0 , min_x - p_width/2 ) );
	bbox.y = MIN( im3D.rows-1, MAX( 0 , min_y - p_height/2) );

	bbox.width  = MAX(0, MIN( new_w, im3D.cols-bbox.x));
	bbox.height = MAX(0, MIN( new_h, im3D.rows-bbox.y));

	return bbox;

}

void CRForestEstimator::estimate( const Mat & im3D,
								   std::vector< cv::Vec<float,6> >& means, //output
                                   std::vector< std::vector< Vote > >& clusters, //all clusters
                                   std::vector< Vote >& votes, //all votes
                                   int stride,
                                   float max_variance,
                                   float prob_th,
                                   float larger_radius_ratio,
                                   float smaller_radius_ratio,
                                   bool verbose , int threshold ){

    unsigned int max_clusters = 20;
    int max_ms_iterations = 10;

	int p_width = crForest->getPatchWidth();
	int p_height = crForest->getPatchHeight();

	//get bounding box around the data in the image
	Rect bbox = getBoundingBox(im3D);

	vector<Mat> channels;
	split(im3D, channels);

	int rows = im3D.rows;
	int cols = im3D.cols;

	//integral image of the depth channel
	Mat depthInt( rows+1, cols+1, CV_64FC1);
	integral( channels[2], depthInt );

	//feature channels vector, in our case it contains only the depth integral image, but it could have other channels (e.g., greyscale)
	std::vector< cv::Mat > featureChans;
	featureChans.push_back(depthInt);

	//mask
    Mat mask( rows, cols, CV_32FC1); mask.setTo(0);
	cv::threshold( channels[2], mask, 10, 1, THRESH_BINARY);

	//integral image of the mask
	Mat maskIntegral( rows+1, cols+1, CV_64FC1); maskIntegral.setTo(0);
	integral( mask, maskIntegral );

	//defines the test patch
    Rect roi = Rect(0,0,p_width,p_height);

    int min_no_pixels = p_width*p_height/10;
    int half_w = roi.width/2;
    int half_h = roi.height/2;

    //process each patch
    for(roi.y=bbox.y; roi.y<bbox.y+bbox.height-p_height; roi.y+=stride) {

    	float* rowX = channels[0].ptr<float>(roi.y + half_h);
    	float* rowY = channels[1].ptr<float>(roi.y + half_h);
    	float* rowZ = channels[2].ptr<float>(roi.y + half_h);

    	double* maskIntY1 = maskIntegral.ptr<double>(roi.y);
    	double* maskIntY2 = maskIntegral.ptr<double>(roi.y + roi.height);

        for(roi.x=bbox.x; roi.x<bbox.x+bbox.width-p_width; roi.x+=stride) {

        	//discard if the middle of the patch does not have depth data
			if( rowZ[roi.x + half_w] <= 0.f )
				continue;

			//discard if the patch is filled with data for less than 10%
            if( (maskIntY1[roi.x] + maskIntY2[roi.x + roi.width] - maskIntY1[roi.x + roi.width] - maskIntY2[roi.x]) <= min_no_pixels )
               continue;

            //send the patch down the trees and retrieve leaves
            std::vector< const LeafNode* > leaves = crForest->regressionIntegral( featureChans, maskIntegral, roi );

            //go through the results
            for(unsigned int t=0;t<leaves.size();++t){

            	//discard bad votes
                if ( leaves[t]->pfg < prob_th || leaves[t]->trace > max_variance )
                    continue;

                Vote v;

                //add the 3D location under the patch center to the vote for the head center
                v.vote[0] = leaves[t]->mean.at<float>(0) + rowX[roi.x + half_w];
                v.vote[1] = leaves[t]->mean.at<float>(1) + rowY[roi.x + half_w];
                v.vote[2] = leaves[t]->mean.at<float>(2) + rowZ[roi.x + half_w];

                //angles, leave as in the leaf
                v.vote[3] = leaves[t]->mean.at<float>(3);
                v.vote[4] = leaves[t]->mean.at<float>(4);
                v.vote[5] = leaves[t]->mean.at<float>(5);

                v.trace = &(leaves[t]->trace);
                v.conf = &(leaves[t]->pfg);

                votes.push_back(v);
            }

        } // end for x

    } // end for y

    if(verbose)
        cout << endl << "votes : " << votes.size() << endl;


    Vec<float,POSE_SIZE> temp_mean;
    vector< vector< Vote* > > temp_clusters;
    vector< Vec<float,POSE_SIZE> > cluster_means;

    //radius for clustering votes
    float large_radius = AVG_FACE_DIAMETER2/(larger_radius_ratio*larger_radius_ratio);

    //cluster using the head centers
    for(unsigned int l=0;l<votes.size();++l){

        bool found = false;
        float best_dist = FLT_MAX;
        unsigned int best_cluster = 0;

		//for each cluster
		for(unsigned int c=0; ( c<cluster_means.size() && found==false ); ++c){

			float norm = 0;
			for(int n=0;n<3;++n)
			    norm += (votes[l].vote[n]-cluster_means[c][n])*(votes[l].vote[n]-cluster_means[c][n]);

			//is the offset smaller than the radius?
			if( norm < large_radius ){

				best_cluster = c;
				found = true;

				 //add (pointer to) vote to the closest cluster (well, actually, the first cluster found within the distance)
				temp_clusters[best_cluster].push_back( &(votes[l]) );

				//update cluster's mean
				((cluster_means[best_cluster])) = 0;
				if ( temp_clusters[best_cluster].size() > 0 ){

					for(unsigned int i=0;i < temp_clusters[best_cluster].size(); ++i)
						((cluster_means[best_cluster])) = ((cluster_means[best_cluster])) + temp_clusters[best_cluster][i]->vote;

					float div = float(MAX(1,temp_clusters[best_cluster].size()));
					for(int n=0;n<POSE_SIZE;++n)
						cluster_means[best_cluster][n] /= div;
				}

			}

		}

		//create a new cluster
        if( !found && temp_clusters.size() < max_clusters ){

            vector< Vote* > new_cluster;
            new_cluster.push_back( &(votes[l]) );
            temp_clusters.push_back( new_cluster );

            Vec<float,POSE_SIZE> vote( votes[l].vote );
            cluster_means.push_back( vote );

        }

    }

    if(verbose){
        cout << cluster_means.size() << " CLUSTERS ";
        for(unsigned int c = 0 ; c<cluster_means.size(); ++c)
            cout << temp_clusters[c].size() << " ";
        cout << endl;
    }

    std::vector< std::vector< Vote* > > new_clusters; //after MS
    vector< Vec<float,POSE_SIZE> > new_means;

    int count = 0;
    float ms_radius2 = AVG_FACE_DIAMETER*AVG_FACE_DIAMETER/(smaller_radius_ratio*smaller_radius_ratio);


    //threshold defining if the cluster belongs to a head: it depends on the stride and on the number of trees
    int th = cvRound((double)threshold*crForest->getSize()/(double)(stride*stride));

    //do MS for each cluster
    for(unsigned int c=0;c<cluster_means.size();++c){

        if(verbose){
            cout << endl << "MS cluster " << c << " : ";
            for(int n=0;n<3;++n)
            	cout << cluster_means[c][n] << " ";
            cout << endl;
        }

        if ( temp_clusters[c].size() <= th ){
            if(verbose)
                cout << "skipping cluster " << endl;
            continue;
        }


        vector< Vote* > new_cluster;

        for(unsigned int it=0; it<max_ms_iterations; ++it){

            count = 0;
            temp_mean = 0;
            new_cluster.clear();

            //for each vote in the cluster
            for(unsigned int idx=0; idx < temp_clusters[c].size() ;++idx){

            	float norm = 0;
            	for(int n=0;n<3;++n)
            		norm += (temp_clusters[c][idx]->vote[n]-cluster_means[c][n])*(temp_clusters[c][idx]->vote[n]-cluster_means[c][n]);

                if( norm < ms_radius2 ){

                	temp_mean = temp_mean + temp_clusters[c][idx]->vote;
                	new_cluster.push_back( temp_clusters[c][idx] );
                    count++;

                }

            }

            for(int n=0;n<POSE_SIZE;++n)
            	temp_mean[n] /= (float)MAX(1,count);

            float distance_to_previous_mean2 = 0;
            for(int n=0;n<3;++n){
            	distance_to_previous_mean2 += (temp_mean[n]-cluster_means[c][n])*(temp_mean[n]-cluster_means[c][n]);
            }

            //cout << it << " " << distance_to_previous_mean2 << endl;

            //update the mean of the cluster
			cluster_means[c] = temp_mean;

            if( distance_to_previous_mean2 < 1 )
            	break;

        }

		new_clusters.push_back( new_cluster );
		new_means.push_back( cluster_means[c] );

    }

    for(unsigned int c=0; c < new_clusters.size() ;++c){

        if( new_clusters[c].size() < th ) //discard clusters with not enough votes
            continue;

        vector< Vote > cluster;
        cluster_means[c] = 0;

        //for each vote in the cluster
        for(unsigned int k=0; k < new_clusters[c].size(); k++ ){

        	cluster_means[c] = cluster_means[c] + new_clusters[c][k]->vote;
        	cluster.push_back( *(new_clusters[c][k]) );

        }

        float div = (float)MAX(1,new_clusters[c].size());
        for(int n=0;n<POSE_SIZE;++n)
        	cluster_means[c][n] /= div;

        means.push_back( cluster_means[c] );
        clusters.push_back( cluster );

    }

}