kdtree_index.h 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626
  1. /***********************************************************************
  2. * Software License Agreement (BSD License)
  3. *
  4. * Copyright 2008-2009 Marius Muja (mariusm@cs.ubc.ca). All rights reserved.
  5. * Copyright 2008-2009 David G. Lowe (lowe@cs.ubc.ca). All rights reserved.
  6. *
  7. * THE BSD LICENSE
  8. *
  9. * Redistribution and use in source and binary forms, with or without
  10. * modification, are permitted provided that the following conditions
  11. * are met:
  12. *
  13. * 1. Redistributions of source code must retain the above copyright
  14. * notice, this list of conditions and the following disclaimer.
  15. * 2. Redistributions in binary form must reproduce the above copyright
  16. * notice, this list of conditions and the following disclaimer in the
  17. * documentation and/or other materials provided with the distribution.
  18. *
  19. * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
  20. * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
  21. * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
  22. * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
  23. * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
  24. * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  25. * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  26. * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  27. * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
  28. * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  29. *************************************************************************/
  30. #ifndef OPENCV_FLANN_KDTREE_INDEX_H_
  31. #define OPENCV_FLANN_KDTREE_INDEX_H_
  32. #include <algorithm>
  33. #include <map>
  34. #include <cassert>
  35. #include <cstring>
  36. #include "general.h"
  37. #include "nn_index.h"
  38. #include "dynamic_bitset.h"
  39. #include "matrix.h"
  40. #include "result_set.h"
  41. #include "heap.h"
  42. #include "allocator.h"
  43. #include "random.h"
  44. #include "saving.h"
  45. namespace cvflann
  46. {
  47. struct KDTreeIndexParams : public IndexParams
  48. {
  49. KDTreeIndexParams(int trees = 4)
  50. {
  51. (*this)["algorithm"] = FLANN_INDEX_KDTREE;
  52. (*this)["trees"] = trees;
  53. }
  54. };
  55. /**
  56. * Randomized kd-tree index
  57. *
  58. * Contains the k-d trees and other information for indexing a set of points
  59. * for nearest-neighbor matching.
  60. */
  61. template <typename Distance>
  62. class KDTreeIndex : public NNIndex<Distance>
  63. {
  64. public:
  65. typedef typename Distance::ElementType ElementType;
  66. typedef typename Distance::ResultType DistanceType;
  67. /**
  68. * KDTree constructor
  69. *
  70. * Params:
  71. * inputData = dataset with the input features
  72. * params = parameters passed to the kdtree algorithm
  73. */
  74. KDTreeIndex(const Matrix<ElementType>& inputData, const IndexParams& params = KDTreeIndexParams(),
  75. Distance d = Distance() ) :
  76. dataset_(inputData), index_params_(params), distance_(d)
  77. {
  78. size_ = dataset_.rows;
  79. veclen_ = dataset_.cols;
  80. trees_ = get_param(index_params_,"trees",4);
  81. tree_roots_ = new NodePtr[trees_];
  82. // Create a permutable array of indices to the input vectors.
  83. vind_.resize(size_);
  84. for (size_t i = 0; i < size_; ++i) {
  85. vind_[i] = int(i);
  86. }
  87. mean_ = new DistanceType[veclen_];
  88. var_ = new DistanceType[veclen_];
  89. }
  90. KDTreeIndex(const KDTreeIndex&);
  91. KDTreeIndex& operator=(const KDTreeIndex&);
  92. /**
  93. * Standard destructor
  94. */
  95. ~KDTreeIndex()
  96. {
  97. if (tree_roots_!=NULL) {
  98. delete[] tree_roots_;
  99. }
  100. delete[] mean_;
  101. delete[] var_;
  102. }
  103. /**
  104. * Builds the index
  105. */
  106. void buildIndex() CV_OVERRIDE
  107. {
  108. /* Construct the randomized trees. */
  109. for (int i = 0; i < trees_; i++) {
  110. /* Randomize the order of vectors to allow for unbiased sampling. */
  111. #ifndef OPENCV_FLANN_USE_STD_RAND
  112. cv::randShuffle(vind_);
  113. #else
  114. std::random_shuffle(vind_.begin(), vind_.end());
  115. #endif
  116. tree_roots_[i] = divideTree(&vind_[0], int(size_) );
  117. }
  118. }
  119. flann_algorithm_t getType() const CV_OVERRIDE
  120. {
  121. return FLANN_INDEX_KDTREE;
  122. }
  123. void saveIndex(FILE* stream) CV_OVERRIDE
  124. {
  125. save_value(stream, trees_);
  126. for (int i=0; i<trees_; ++i) {
  127. save_tree(stream, tree_roots_[i]);
  128. }
  129. }
  130. void loadIndex(FILE* stream) CV_OVERRIDE
  131. {
  132. load_value(stream, trees_);
  133. if (tree_roots_!=NULL) {
  134. delete[] tree_roots_;
  135. }
  136. tree_roots_ = new NodePtr[trees_];
  137. for (int i=0; i<trees_; ++i) {
  138. load_tree(stream,tree_roots_[i]);
  139. }
  140. index_params_["algorithm"] = getType();
  141. index_params_["trees"] = tree_roots_;
  142. }
  143. /**
  144. * Returns size of index.
  145. */
  146. size_t size() const CV_OVERRIDE
  147. {
  148. return size_;
  149. }
  150. /**
  151. * Returns the length of an index feature.
  152. */
  153. size_t veclen() const CV_OVERRIDE
  154. {
  155. return veclen_;
  156. }
  157. /**
  158. * Computes the inde memory usage
  159. * Returns: memory used by the index
  160. */
  161. int usedMemory() const CV_OVERRIDE
  162. {
  163. return int(pool_.usedMemory+pool_.wastedMemory+dataset_.rows*sizeof(int)); // pool memory and vind array memory
  164. }
  165. /**
  166. * Find set of nearest neighbors to vec. Their indices are stored inside
  167. * the result object.
  168. *
  169. * Params:
  170. * result = the result object in which the indices of the nearest-neighbors are stored
  171. * vec = the vector for which to search the nearest neighbors
  172. * maxCheck = the maximum number of restarts (in a best-bin-first manner)
  173. */
  174. void findNeighbors(ResultSet<DistanceType>& result, const ElementType* vec, const SearchParams& searchParams) CV_OVERRIDE
  175. {
  176. int maxChecks = get_param(searchParams,"checks", 32);
  177. float epsError = 1+get_param(searchParams,"eps",0.0f);
  178. if (maxChecks==FLANN_CHECKS_UNLIMITED) {
  179. getExactNeighbors(result, vec, epsError);
  180. }
  181. else {
  182. getNeighbors(result, vec, maxChecks, epsError);
  183. }
  184. }
  185. IndexParams getParameters() const CV_OVERRIDE
  186. {
  187. return index_params_;
  188. }
  189. private:
  190. /*--------------------- Internal Data Structures --------------------------*/
  191. struct Node
  192. {
  193. /**
  194. * Dimension used for subdivision.
  195. */
  196. int divfeat;
  197. /**
  198. * The values used for subdivision.
  199. */
  200. DistanceType divval;
  201. /**
  202. * The child nodes.
  203. */
  204. Node* child1, * child2;
  205. };
  206. typedef Node* NodePtr;
  207. typedef BranchStruct<NodePtr, DistanceType> BranchSt;
  208. typedef BranchSt* Branch;
  209. void save_tree(FILE* stream, NodePtr tree)
  210. {
  211. save_value(stream, *tree);
  212. if (tree->child1!=NULL) {
  213. save_tree(stream, tree->child1);
  214. }
  215. if (tree->child2!=NULL) {
  216. save_tree(stream, tree->child2);
  217. }
  218. }
  219. void load_tree(FILE* stream, NodePtr& tree)
  220. {
  221. tree = pool_.allocate<Node>();
  222. load_value(stream, *tree);
  223. if (tree->child1!=NULL) {
  224. load_tree(stream, tree->child1);
  225. }
  226. if (tree->child2!=NULL) {
  227. load_tree(stream, tree->child2);
  228. }
  229. }
  230. /**
  231. * Create a tree node that subdivides the list of vecs from vind[first]
  232. * to vind[last]. The routine is called recursively on each sublist.
  233. * Place a pointer to this new tree node in the location pTree.
  234. *
  235. * Params: pTree = the new node to create
  236. * first = index of the first vector
  237. * last = index of the last vector
  238. */
  239. NodePtr divideTree(int* ind, int count)
  240. {
  241. NodePtr node = pool_.allocate<Node>(); // allocate memory
  242. /* If too few exemplars remain, then make this a leaf node. */
  243. if ( count == 1) {
  244. node->child1 = node->child2 = NULL; /* Mark as leaf node. */
  245. node->divfeat = *ind; /* Store index of this vec. */
  246. }
  247. else {
  248. int idx;
  249. int cutfeat;
  250. DistanceType cutval;
  251. meanSplit(ind, count, idx, cutfeat, cutval);
  252. node->divfeat = cutfeat;
  253. node->divval = cutval;
  254. node->child1 = divideTree(ind, idx);
  255. node->child2 = divideTree(ind+idx, count-idx);
  256. }
  257. return node;
  258. }
  259. /**
  260. * Choose which feature to use in order to subdivide this set of vectors.
  261. * Make a random choice among those with the highest variance, and use
  262. * its variance as the threshold value.
  263. */
  264. void meanSplit(int* ind, int count, int& index, int& cutfeat, DistanceType& cutval)
  265. {
  266. memset(mean_,0,veclen_*sizeof(DistanceType));
  267. memset(var_,0,veclen_*sizeof(DistanceType));
  268. /* Compute mean values. Only the first SAMPLE_MEAN values need to be
  269. sampled to get a good estimate.
  270. */
  271. int cnt = std::min((int)SAMPLE_MEAN+1, count);
  272. for (int j = 0; j < cnt; ++j) {
  273. ElementType* v = dataset_[ind[j]];
  274. for (size_t k=0; k<veclen_; ++k) {
  275. mean_[k] += v[k];
  276. }
  277. }
  278. for (size_t k=0; k<veclen_; ++k) {
  279. mean_[k] /= cnt;
  280. }
  281. /* Compute variances (no need to divide by count). */
  282. for (int j = 0; j < cnt; ++j) {
  283. ElementType* v = dataset_[ind[j]];
  284. for (size_t k=0; k<veclen_; ++k) {
  285. DistanceType dist = v[k] - mean_[k];
  286. var_[k] += dist * dist;
  287. }
  288. }
  289. /* Select one of the highest variance indices at random. */
  290. cutfeat = selectDivision(var_);
  291. cutval = mean_[cutfeat];
  292. int lim1, lim2;
  293. planeSplit(ind, count, cutfeat, cutval, lim1, lim2);
  294. if (lim1>count/2) index = lim1;
  295. else if (lim2<count/2) index = lim2;
  296. else index = count/2;
  297. /* If either list is empty, it means that all remaining features
  298. * are identical. Split in the middle to maintain a balanced tree.
  299. */
  300. if ((lim1==count)||(lim2==0)) index = count/2;
  301. }
  302. /**
  303. * Select the top RAND_DIM largest values from v and return the index of
  304. * one of these selected at random.
  305. */
  306. int selectDivision(DistanceType* v)
  307. {
  308. int num = 0;
  309. size_t topind[RAND_DIM];
  310. /* Create a list of the indices of the top RAND_DIM values. */
  311. for (size_t i = 0; i < veclen_; ++i) {
  312. if ((num < RAND_DIM)||(v[i] > v[topind[num-1]])) {
  313. /* Put this element at end of topind. */
  314. if (num < RAND_DIM) {
  315. topind[num++] = i; /* Add to list. */
  316. }
  317. else {
  318. topind[num-1] = i; /* Replace last element. */
  319. }
  320. /* Bubble end value down to right location by repeated swapping. */
  321. int j = num - 1;
  322. while (j > 0 && v[topind[j]] > v[topind[j-1]]) {
  323. std::swap(topind[j], topind[j-1]);
  324. --j;
  325. }
  326. }
  327. }
  328. /* Select a random integer in range [0,num-1], and return that index. */
  329. int rnd = rand_int(num);
  330. return (int)topind[rnd];
  331. }
  332. /**
  333. * Subdivide the list of points by a plane perpendicular on axe corresponding
  334. * to the 'cutfeat' dimension at 'cutval' position.
  335. *
  336. * On return:
  337. * dataset[ind[0..lim1-1]][cutfeat]<cutval
  338. * dataset[ind[lim1..lim2-1]][cutfeat]==cutval
  339. * dataset[ind[lim2..count]][cutfeat]>cutval
  340. */
  341. void planeSplit(int* ind, int count, int cutfeat, DistanceType cutval, int& lim1, int& lim2)
  342. {
  343. /* Move vector indices for left subtree to front of list. */
  344. int left = 0;
  345. int right = count-1;
  346. for (;; ) {
  347. while (left<=right && dataset_[ind[left]][cutfeat]<cutval) ++left;
  348. while (left<=right && dataset_[ind[right]][cutfeat]>=cutval) --right;
  349. if (left>right) break;
  350. std::swap(ind[left], ind[right]); ++left; --right;
  351. }
  352. lim1 = left;
  353. right = count-1;
  354. for (;; ) {
  355. while (left<=right && dataset_[ind[left]][cutfeat]<=cutval) ++left;
  356. while (left<=right && dataset_[ind[right]][cutfeat]>cutval) --right;
  357. if (left>right) break;
  358. std::swap(ind[left], ind[right]); ++left; --right;
  359. }
  360. lim2 = left;
  361. }
  362. /**
  363. * Performs an exact nearest neighbor search. The exact search performs a full
  364. * traversal of the tree.
  365. */
  366. void getExactNeighbors(ResultSet<DistanceType>& result, const ElementType* vec, float epsError)
  367. {
  368. // checkID -= 1; /* Set a different unique ID for each search. */
  369. if (trees_ > 1) {
  370. fprintf(stderr,"It doesn't make any sense to use more than one tree for exact search");
  371. }
  372. if (trees_>0) {
  373. searchLevelExact(result, vec, tree_roots_[0], 0.0, epsError);
  374. }
  375. assert(result.full());
  376. }
  377. /**
  378. * Performs the approximate nearest-neighbor search. The search is approximate
  379. * because the tree traversal is abandoned after a given number of descends in
  380. * the tree.
  381. */
  382. void getNeighbors(ResultSet<DistanceType>& result, const ElementType* vec, int maxCheck, float epsError)
  383. {
  384. int i;
  385. BranchSt branch;
  386. int checkCount = 0;
  387. Heap<BranchSt>* heap = new Heap<BranchSt>((int)size_);
  388. DynamicBitset checked(size_);
  389. /* Search once through each tree down to root. */
  390. for (i = 0; i < trees_; ++i) {
  391. searchLevel(result, vec, tree_roots_[i], 0, checkCount, maxCheck, epsError, heap, checked);
  392. }
  393. /* Keep searching other branches from heap until finished. */
  394. while ( heap->popMin(branch) && (checkCount < maxCheck || !result.full() )) {
  395. searchLevel(result, vec, branch.node, branch.mindist, checkCount, maxCheck, epsError, heap, checked);
  396. }
  397. delete heap;
  398. assert(result.full());
  399. }
  400. /**
  401. * Search starting from a given node of the tree. Based on any mismatches at
  402. * higher levels, all exemplars below this level must have a distance of
  403. * at least "mindistsq".
  404. */
  405. void searchLevel(ResultSet<DistanceType>& result_set, const ElementType* vec, NodePtr node, DistanceType mindist, int& checkCount, int maxCheck,
  406. float epsError, Heap<BranchSt>* heap, DynamicBitset& checked)
  407. {
  408. if (result_set.worstDist()<mindist) {
  409. // printf("Ignoring branch, too far\n");
  410. return;
  411. }
  412. /* If this is a leaf node, then do check and return. */
  413. if ((node->child1 == NULL)&&(node->child2 == NULL)) {
  414. /* Do not check same node more than once when searching multiple trees.
  415. Once a vector is checked, we set its location in vind to the
  416. current checkID.
  417. */
  418. int index = node->divfeat;
  419. if ( checked.test(index) || ((checkCount>=maxCheck)&& result_set.full()) ) return;
  420. checked.set(index);
  421. checkCount++;
  422. DistanceType dist = distance_(dataset_[index], vec, veclen_);
  423. result_set.addPoint(dist,index);
  424. return;
  425. }
  426. /* Which child branch should be taken first? */
  427. ElementType val = vec[node->divfeat];
  428. DistanceType diff = val - node->divval;
  429. NodePtr bestChild = (diff < 0) ? node->child1 : node->child2;
  430. NodePtr otherChild = (diff < 0) ? node->child2 : node->child1;
  431. /* Create a branch record for the branch not taken. Add distance
  432. of this feature boundary (we don't attempt to correct for any
  433. use of this feature in a parent node, which is unlikely to
  434. happen and would have only a small effect). Don't bother
  435. adding more branches to heap after halfway point, as cost of
  436. adding exceeds their value.
  437. */
  438. DistanceType new_distsq = mindist + distance_.accum_dist(val, node->divval, node->divfeat);
  439. // if (2 * checkCount < maxCheck || !result.full()) {
  440. if ((new_distsq*epsError < result_set.worstDist())|| !result_set.full()) {
  441. heap->insert( BranchSt(otherChild, new_distsq) );
  442. }
  443. /* Call recursively to search next level down. */
  444. searchLevel(result_set, vec, bestChild, mindist, checkCount, maxCheck, epsError, heap, checked);
  445. }
  446. /**
  447. * Performs an exact search in the tree starting from a node.
  448. */
  449. void searchLevelExact(ResultSet<DistanceType>& result_set, const ElementType* vec, const NodePtr node, DistanceType mindist, const float epsError)
  450. {
  451. /* If this is a leaf node, then do check and return. */
  452. if ((node->child1 == NULL)&&(node->child2 == NULL)) {
  453. int index = node->divfeat;
  454. DistanceType dist = distance_(dataset_[index], vec, veclen_);
  455. result_set.addPoint(dist,index);
  456. return;
  457. }
  458. /* Which child branch should be taken first? */
  459. ElementType val = vec[node->divfeat];
  460. DistanceType diff = val - node->divval;
  461. NodePtr bestChild = (diff < 0) ? node->child1 : node->child2;
  462. NodePtr otherChild = (diff < 0) ? node->child2 : node->child1;
  463. /* Create a branch record for the branch not taken. Add distance
  464. of this feature boundary (we don't attempt to correct for any
  465. use of this feature in a parent node, which is unlikely to
  466. happen and would have only a small effect). Don't bother
  467. adding more branches to heap after halfway point, as cost of
  468. adding exceeds their value.
  469. */
  470. DistanceType new_distsq = mindist + distance_.accum_dist(val, node->divval, node->divfeat);
  471. /* Call recursively to search next level down. */
  472. searchLevelExact(result_set, vec, bestChild, mindist, epsError);
  473. if (new_distsq*epsError<=result_set.worstDist()) {
  474. searchLevelExact(result_set, vec, otherChild, new_distsq, epsError);
  475. }
  476. }
  477. private:
  478. enum
  479. {
  480. /**
  481. * To improve efficiency, only SAMPLE_MEAN random values are used to
  482. * compute the mean and variance at each level when building a tree.
  483. * A value of 100 seems to perform as well as using all values.
  484. */
  485. SAMPLE_MEAN = 100,
  486. /**
  487. * Top random dimensions to consider
  488. *
  489. * When creating random trees, the dimension on which to subdivide is
  490. * selected at random from among the top RAND_DIM dimensions with the
  491. * highest variance. A value of 5 works well.
  492. */
  493. RAND_DIM=5
  494. };
  495. /**
  496. * Number of randomized trees that are used
  497. */
  498. int trees_;
  499. /**
  500. * Array of indices to vectors in the dataset.
  501. */
  502. std::vector<int> vind_;
  503. /**
  504. * The dataset used by this index
  505. */
  506. const Matrix<ElementType> dataset_;
  507. IndexParams index_params_;
  508. size_t size_;
  509. size_t veclen_;
  510. DistanceType* mean_;
  511. DistanceType* var_;
  512. /**
  513. * Array of k-d trees used to find neighbours.
  514. */
  515. NodePtr* tree_roots_;
  516. /**
  517. * Pooled memory allocator.
  518. *
  519. * Using a pooled memory allocator is more efficient
  520. * than allocating memory directly when there is a large
  521. * number small of memory allocations.
  522. */
  523. PooledAllocator pool_;
  524. Distance distance_;
  525. }; // class KDTreeForest
  526. }
  527. #endif //OPENCV_FLANN_KDTREE_INDEX_H_