kdtree.c

Rob Hess Linux下的SIFT提取源码

/*
  Functions and structures for maintaining a k-d tree database of image
  features.
  
  For more information, refer to:
  
  Beis, J. S. and Lowe, D. G.  Shape indexing using approximate
  nearest-neighbor search in high-dimensional spaces.  In <EM>Conference
  on Computer Vision and Pattern Recognition (CVPR)</EM> (2003),
  pp. 1000--1006.
  
  Copyright (C) 2006-2007  Rob Hess <[e-mail]>

  @version 1.1.1-20070913
*/

#include "kdtree.h"
#include "minpq.h"
#include "imgfeatures.h"
#include "utils.h"

#include <cxcore.h>

#include <stdio.h>

struct bbf_data
{
  double d;
  void* old_data;
};

/************************* Local Function Prototypes *************************/

struct kd_node* kd_node_init( struct feature*, int );
void expand_kd_node_subtree( struct kd_node* );
void assign_part_key( struct kd_node* );
double median_select( double*, int );
double rank_select( double*, int, int );
void insertion_sort( double*, int );
int partition_array( double*, int, double );
void partition_features( struct kd_node* );
struct kd_node* explore_to_leaf( struct kd_node*, struct feature*,
				 struct min_pq* );
int insert_into_nbr_array( struct feature*, struct feature**, int, int );
int within_rect( CvPoint2D64f, CvRect );


/******************** Functions prototyped in keyptdb.h **********************/


/*
  A function to build a k-d tree database from keypoints in an array.
  
  @param features an array of features
  @param n the number of features in features
  
  @return Returns the root of a kd tree built from features or NULL on
    error.
*/
struct kd_node* kdtree_build( struct feature* features, int n )
{
  struct kd_node* kd_root;

  if( ! features  ||  n <= 0 )
    {
      fprintf( stderr, "Warning: kdtree_build(): no features, %s, line %d\n",
	       __FILE__, __LINE__ );
      return NULL;
    }

  kd_root = kd_node_init( features, n );
  expand_kd_node_subtree( kd_root );

  return kd_root;
}



/*
  Finds an image feature's approximate k nearest neighbors in a kd tree using
  Best Bin First search.
  
  @param kd_root root of an image feature kd tree
  @param feat image feature for whose neighbors to search
  @param k number of neighbors to find
   @param nbrs pointer to an array in which to store pointers to neighbors
     in order of increasing descriptor distance
  @param max_nn_chks search is cut off after examining this many tree entries
  
  @return Returns the number of neighbors found and stored in nbrs, or
    -1 on error.
*/
int kdtree_bbf_knn( struct kd_node* kd_root, struct feature* feat, int k,
		    struct feature*** nbrs, int max_nn_chks )
{
  struct kd_node* expl;
  struct min_pq* min_pq;
  struct feature* tree_feat, ** _nbrs;
  struct bbf_data* bbf_data;
  int i, t = 0, n = 0;

  if( ! nbrs  ||  ! feat  ||  ! kd_root )
    {
      fprintf( stderr, "Warning: NULL pointer error, %s, line %d\n",
	       __FILE__, __LINE__ );
      return -1;
    }

  _nbrs = calloc( k, sizeof( struct feature* ) );
  min_pq = minpq_init();
  minpq_insert( min_pq, kd_root, 0 );
  while( min_pq->n > 0  &&  t < max_nn_chks )
    {
      expl = (struct kd_node*)minpq_extract_min( min_pq );
      if( ! expl )
	{
	  fprintf( stderr, "Warning: PQ unexpectedly empty, %s line %d\n",
		   __FILE__, __LINE__ );
	  goto fail;
	}

      expl = explore_to_leaf( expl, feat, min_pq );
      if( ! expl )
	{
	  fprintf( stderr, "Warning: PQ unexpectedly empty, %s line %d\n",
		   __FILE__, __LINE__ );
	  goto fail;
	}

      for( i = 0; i < expl->n; i++ )
	{
	  tree_feat = &expl->features[i];
	  bbf_data = malloc( sizeof( struct bbf_data ) );
	  if( ! bbf_data )
	    {
	      fprintf( stderr, "Warning: unable to allocate memory,"
		       " %s line %d\n", __FILE__, __LINE__ );
	      goto fail;
	    }
	  bbf_data->old_data = tree_feat->feature_data;
	  bbf_data->d = descr_dist_sq(feat, tree_feat);
	  tree_feat->feature_data = bbf_data;
	  n += insert_into_nbr_array( tree_feat, _nbrs, n, k );
	}
      t++;
    }

  minpq_release( &min_pq );
  for( i = 0; i < n; i++ )
    {
      bbf_data = _nbrs[i]->feature_data;
      _nbrs[i]->feature_data = bbf_data->old_data;
      free( bbf_data );
    }
  *nbrs = _nbrs;
  return n;

 fail:
  minpq_release( &min_pq );
  for( i = 0; i < n; i++ )
    {
      bbf_data = _nbrs[i]->feature_data;
      _nbrs[i]->feature_data = bbf_data->old_data;
      free( bbf_data );
    }
  free( _nbrs );
  *nbrs = NULL;
  return -1;
}



/*
  Finds an image feature's approximate k nearest neighbors within a specified
  spatial region in a kd tree using Best Bin First search.
  
  @param kd_root root of an image feature kd tree
  @param feat image feature for whose neighbors to search
  @param k number of neighbors to find
   @param nbrs pointer to an array in which to store pointers to neighbors
     in order of increasing descriptor distance
   @param max_nn_chks search is cut off after examining this many tree entries
   @param rect rectangular region in which to search for neighbors
   @param model if true, spatial search is based on kdtree features' model
     locations; otherwise it is based on their image locations
   
   @return Returns the number of neighbors found and stored in \a nbrs
     (in case \a k neighbors could not be found before examining
     \a max_nn_checks keypoint entries).
*/
int kdtree_bbf_spatial_knn( struct kd_node* kd_root, struct feature* feat,
			    int k, struct feature*** nbrs, int max_nn_chks,
			    CvRect rect, int model )
{
  struct feature** all_nbrs, ** sp_nbrs;
  CvPoint2D64f pt;
  int i, n, t = 0;

  n = kdtree_bbf_knn( kd_root, feat, max_nn_chks, &all_nbrs, max_nn_chks );
  sp_nbrs = calloc( k, sizeof( struct feature* ) );
  for( i = 0; i < n; i++ )
    {
      if( model )
	pt = all_nbrs[i]->mdl_pt;
      else
	pt = all_nbrs[i]->img_pt;

      if( within_rect( pt, rect ) )
	{
	  sp_nbrs[t++] = all_nbrs[i];
	  if( t == k )
	    goto end;
	}
    }
 end:
  free( all_nbrs );
  *nbrs = sp_nbrs;
  return t;
}



/*
  De-allocates memory held by a kd tree
  
  @param kd_root pointer to the root of a kd tree
*/
void kdtree_release( struct kd_node* kd_root )
{
  if( ! kd_root )
    return;
  kdtree_release( kd_root->kd_left );
  kdtree_release( kd_root->kd_right );
  free( kd_root );
}


/************************ Functions prototyped here **************************/


/*
  Initializes a kd tree node with a set of features.  The node is not
  expanded, and no ordering is imposed on the features.
  
  @param features an array of image features
  @param n number of features

  @return Returns an unexpanded kd-tree node.
*/
struct kd_node* kd_node_init( struct feature* features, int n )
{
  struct kd_node* kd_node;

  kd_node = malloc( sizeof( struct kd_node ) );
  memset( kd_node, 0, sizeof( struct kd_node ) );
  kd_node->ki = -1;
  kd_node->features = features;
  kd_node->n = n;

  return kd_node;
}



/*
  Recursively expands a specified kd tree node into a tree whose leaves
  contain one entry each.

  @param kd_node an unexpanded node in a kd tree
*/
void expand_kd_node_subtree( struct kd_node* kd_node )
{
  /* base case: leaf node */
  if( kd_node->n == 1  ||  kd_node->n == 0 )
    {
      kd_node->leaf = 1;
      return;
    }

  assign_part_key( kd_node );
  partition_features( kd_node );

  if( kd_node->kd_left )
    expand_kd_node_subtree( kd_node->kd_left );
  if( kd_node->kd_right )
    expand_kd_node_subtree( kd_node->kd_right );
}



/*
  Determines the descriptor index at which and the value with which to
  partition a kd tree node's features.

  @param kd_node a kd tree node
*/
void assign_part_key( struct kd_node* kd_node )
{
  struct feature* features;
  double kv, x, mean, var, var_max = 0;
  double* tmp;
  int d, n, i, j, ki = 0;

  features = kd_node->features;
  n = kd_node->n;
  d = features[0].d;

  /* partition key index is that along which descriptors have most variance */
  for( j = 0; j < d; j++ )
    {
      mean = var = 0;
      for( i = 0; i < n; i++ )
	mean += features[i].descr[j];
      mean /= n;
      for( i = 0; i < n; i++ )
	{
	  x = features[i].descr[j] - mean;
	  var += x * x;
	}
      var /= n;

      if( var > var_max )
	{
	  ki = j;
	  var_max = var;
	}
    }

  /* partition key value is median of descriptor values at ki */
  tmp = calloc( n, sizeof( double ) );