1# Copyright 2014 John O. Woods (john.o.woods@gmail.com), West Virginia 2# University's Applied Space Exploration Lab, and West Virginia Robotic 3# Technology Center. All rights reserved. 4# 5# THE BSD LICENSE 6# 7# Redistribution and use in source and binary forms, with or without 8# modification, are permitted provided that the following conditions 9# are met: 10# 11# 1. Redistributions of source code must retain the above copyright 12# notice, this list of conditions and the following disclaimer. 13# 2. Redistributions in binary form must reproduce the above copyright 14# notice, this list of conditions and the following disclaimer in the 15# documentation and/or other materials provided with the distribution. 16# 17# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR 18# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES 19# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. 20# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, 21# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT 22# NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 23# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 24# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 25# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF 26# THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 27 28class FFI::Pointer 29 class << self 30 def new_from_nmatrix nm 31 raise(StorageError, "dense storage expected") unless nm.dense? 32 c_type = Flann::dtype_to_c(nm.dtype) 33 c_type = :uchar if c_type == :byte 34 ::FFI::Pointer.new(c_type, nm.data_pointer).tap { |p| p.autorelease = false } 35 end 36 end 37end 38 39module Flann 40 class Index 41 42 # Constructor takes a block where we set each of the parameters. We need to be careful to do this since 43 # we're using the C API and not C++; so everything important needs to be initialized or there could be 44 # a segfault. For reasonable default definitions, see: 45 # 46 # * https://github.com/mariusmuja/flann/tree/master/src/cpp/flann/algorithms 47 # 48 def initialize index_dataset = nil, dtype: :float64, parameters: Flann::Parameters::DEFAULT 49 @dataset = index_dataset 50 #require 'pry' 51 #binding.pry if @dataset.nil? 52 @dtype = (!index_dataset.nil? && index_dataset.is_a?(NMatrix)) ? index_dataset.dtype : dtype 53 @index_ptr = nil 54 55 @parameters_ptr, @parameters = Flann::handle_parameters(parameters) 56 57 yield @parameters if block_given? 58 end 59 attr_reader :dtype, :dataset, :parameters, :parameters_ptr, :index_ptr 60 61 # Assign a new dataset. Requires that the old index be freed. 62 def dataset= index_dataset 63 free! 64 @dataset = index_dataset 65 end 66 67 # Build an index 68 def build! 69 raise("no dataset specified") if dataset.nil? 70 c_type = Flann::dtype_to_c(dtype) 71 c_method = "flann_build_index_#{c_type}".to_sym 72 speedup_float_ptr = FFI::MemoryPointer.new(:float) 73 @index_ptr = Flann.send(c_method, FFI::Pointer.new_from_nmatrix(dataset), dataset.shape[0], dataset.shape[1], speedup_float_ptr, parameters_ptr) 74 if index_ptr.address == 0 75 require 'pry' 76 binding.pry 77 raise("failed to allocate index_ptr") 78 end 79 80 81 # Return the speedup 82 speedup_float_ptr.read_float 83 end 84 85 # Get the nearest neighbors based on this index. Forces a build of the index if one hasn't been done yet. 86 def nearest_neighbors testset, k, parameters = {} 87 parameters = Parameters.new(Flann::Parameters::DEFAULT.merge(parameters)) 88 89 self.build! if index_ptr.nil? 90 91 parameters_ptr, parameters = Flann::handle_parameters(parameters) 92 result_size = testset.shape[0] * k 93 94 c_type = Flann::dtype_to_c(dataset.dtype) 95 c_method = "flann_find_nearest_neighbors_index_#{c_type}".to_sym 96 indices_int_ptr, distances_t_ptr = Flann::allocate_results_space(result_size, c_type) 97 98 Flann.send c_method, index_ptr, 99 FFI::Pointer.new_from_nmatrix(testset), 100 testset.shape[0], 101 indices_int_ptr, distances_t_ptr, 102 k, 103 parameters_ptr 104 105 106 [indices_int_ptr.read_array_of_int(result_size), 107 c_type == :double ? distances_t_ptr.read_array_of_double(result_size) : distances_t_ptr.read_array_of_float(result_size)] 108 end 109 110 # Perform a radius search on a single query point 111 def radius_search query, radius, max_k=nil, parameters = {} 112 max_k ||= dataset.shape[0] 113 parameters = Parameters.new(Flann::Parameters::DEFAULT.merge(parameters)) 114 115 self.build! if index_ptr.nil? 116 parameters_ptr, parameters = Flann::handle_parameters(parameters) 117 118 c_type = Flann::dtype_to_c(dataset.dtype) 119 c_method = "flann_radius_search_#{c_type}".to_sym 120 indices_int_ptr, distances_t_ptr = Flann::allocate_results_space(max_k, c_type) 121 122 Flann.send(c_method, index_ptr, FFI::Pointer.new_from_nmatrix(query), indices_int_ptr, distances_t_ptr, max_k, radius, parameters_ptr) 123 124 # Return results: two arrays, one of indices and one of distances. 125 indices = indices_int_ptr.read_array_of_int(max_k) 126 distances = c_type == :double ? distances_t_ptr.read_array_of_double(max_k) : distances_t_ptr.read_array_of_float(max_k) 127 128 # Stop where indices == -1 129 cutoff = indices.find_index(-1) 130 cutoff.nil? ? [indices, distances] : [indices[0...cutoff], distances[0...cutoff]] 131 end 132 133 # Save an index to a file (without the dataset). 134 def save filename 135 raise(IOError, "Cannot write an unbuilt index") if index_ptr.nil? # FIXME: This should probably have its own exception type. 136 c_method = "flann_save_index_#{Flann::dtype_to_c(dtype)}".to_sym 137 Flann.send(c_method, index_ptr, filename) 138 self 139 end 140 141 # Load an index from a file (with the dataset already known!). 142 # 143 # FIXME: This needs to free the previous dataset first. 144 def load! filename 145 c_method = "flann_load_index_#{Flann::dtype_to_c(dtype)}".to_sym 146 147 @index_ptr = Flann.send(c_method, filename, FFI::Pointer.new_from_nmatrix(dataset), dataset.shape[0], dataset.shape[1]) 148 self 149 end 150 151 # Free an index 152 def free! parameters = {} 153 parameters = Parameters.new(Flann::Parameters::DEFAULT.merge(parameters)) 154 c_method = "flann_free_index_#{Flann::dtype_to_c(dtype)}".to_sym 155 parameters_ptr, parameters = Flann::handle_parameters(parameters) 156 Flann.send(c_method, index_ptr, parameters_ptr) 157 @index_ptr = nil 158 self 159 end 160 161 end 162end