static void vl_kdtree_calc_bounds_recursively (VlKDTree * tree, vl_uindex nodeIndex, double * searchBounds) { VlKDTreeNode * node = tree->nodes + nodeIndex ; vl_uindex i = node->splitDimension ; double t = node->splitThreshold ; node->lowerBound = searchBounds [2 * i + 0] ; node->upperBound = searchBounds [2 * i + 1] ; if (node->lowerChild > 0) { searchBounds [2 * i + 1] = t ; vl_kdtree_calc_bounds_recursively (tree, node->lowerChild, searchBounds) ; searchBounds [2 * i + 1] = node->upperBound ; } if (node->upperChild > 0) { searchBounds [2 * i + 0] = t ; vl_kdtree_calc_bounds_recursively (tree, node->upperChild, searchBounds) ; searchBounds [2 * i + 0] = node->lowerBound ; } }
void vl_kdforest_build (VlKDForest * self, vl_size numData, void const * data) { vl_uindex di, ti ; vl_size maxNumNodes ; double * searchBounds; /* need to check: if alredy built, clean first */ self->data = data ; self->numData = numData ; self->trees = vl_malloc (sizeof(VlKDTree*) * self->numTrees) ; maxNumNodes = 0 ; for (ti = 0 ; ti < self->numTrees ; ++ ti) { self->trees[ti] = vl_malloc (sizeof(VlKDTree)) ; self->trees[ti]->dataIndex = vl_malloc (sizeof(VlKDTreeDataIndexEntry) * self->numData) ; for (di = 0 ; di < self->numData ; ++ di) { self->trees[ti]->dataIndex[di].index = di ; } self->trees[ti]->numUsedNodes = 0 ; /* num. nodes of a complete binary tree with numData leaves */ self->trees[ti]->numAllocatedNodes = 2 * self->numData - 1 ; self->trees[ti]->nodes = vl_malloc (sizeof(VlKDTreeNode) * self->trees[ti]->numAllocatedNodes) ; self->trees[ti]->depth = 0 ; vl_kdtree_build_recursively (self, self->trees[ti], vl_kdtree_node_new(self->trees[ti], 0), 0, self->numData, 0) ; maxNumNodes += self->trees[ti]->numUsedNodes ; } searchBounds = vl_malloc(sizeof(double) * 2 * self->dimension); for (ti = 0 ; ti < self->numTrees ; ++ ti) { double * iter = searchBounds ; double * end = iter + 2 * self->dimension ; while (iter < end) { *iter++ = - VL_INFINITY_F ; *iter++ = + VL_INFINITY_F ; } vl_kdtree_calc_bounds_recursively (self->trees[ti], 0, searchBounds) ; } vl_free(searchBounds); self -> maxNumNodes = maxNumNodes; }
VL_EXPORT vl_size vl_kdforest_query (VlKDForest * self, VlKDForestNeighbor * neighbors, vl_size numNeighbors, void const * query) { vl_uindex i, ti ; vl_bool exactSearch = (self->searchMaxNumComparisons == 0) ; VlKDForestSearchState * searchState ; vl_size numAddedNeighbors = 0 ; assert (neighbors) ; assert (numNeighbors > 0) ; assert (query) ; /* this number is used to differentiate a query from the next */ self -> searchId += 1 ; self -> searchNumRecursions = 0 ; if (! self -> searchHeapArray) { /* count number of tree nodes */ /* add support structures */ vl_size maxNumNodes = 0 ; for (ti = 0 ; ti < self->numTrees ; ++ti) { maxNumNodes += self->trees[ti]->numUsedNodes ; } self -> searchHeapArray = vl_malloc (sizeof(VlKDForestSearchState) * maxNumNodes) ; self -> searchIdBook = vl_calloc (sizeof(vl_uindex), self->numData) ; for (ti = 0 ; ti < self->numTrees ; ++ti) { double * searchBounds = vl_malloc(sizeof(double) * 2 * self->dimension) ; double * iter = searchBounds ; double * end = iter + 2 * self->dimension ; while (iter < end) { *iter++ = - VL_INFINITY_F ; *iter++ = + VL_INFINITY_F ; } vl_kdtree_calc_bounds_recursively (self->trees[ti], 0, searchBounds) ; vl_free (searchBounds) ; } } self->searchNumComparisons = 0 ; self->searchNumSimplifications = 0 ; /* put the root node into the search heap */ self->searchHeapNumNodes = 0 ; for (ti = 0 ; ti < self->numTrees ; ++ ti) { searchState = self->searchHeapArray + self->searchHeapNumNodes ; searchState -> tree = self->trees[ti] ; searchState -> nodeIndex = 0 ; searchState -> distanceLowerBound = 0 ; vl_kdforest_search_heap_push (self->searchHeapArray, &self->searchHeapNumNodes) ; } /* branch and bound */ while (exactSearch || self->searchNumComparisons < self->searchMaxNumComparisons) { /* pop the next optimal search node */ VlKDForestSearchState * searchState ; /* break if search space completed */ if (self->searchHeapNumNodes == 0) { break ; } searchState = self->searchHeapArray + vl_kdforest_search_heap_pop (self->searchHeapArray, &self->searchHeapNumNodes) ; /* break if no better solution may exist */ if (numAddedNeighbors == numNeighbors && neighbors[0].distance < searchState->distanceLowerBound) { self->searchNumSimplifications ++ ; break ; } vl_kdforest_query_recursively (self, searchState->tree, searchState->nodeIndex, neighbors, numNeighbors, &numAddedNeighbors, searchState->distanceLowerBound, query) ; } /* sort neighbors by increasing distance */ for (i = numAddedNeighbors ; i < numNeighbors ; ++ i) { neighbors[i].index = -1 ; neighbors[i].distance = VL_NAN_F ; } while (numAddedNeighbors) { vl_kdforest_neighbor_heap_pop (neighbors, &numAddedNeighbors) ; } return self->searchNumComparisons ; }