kdtree_index.h 19 KB

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