inverted hashmap, added bktree

src/store.rs

@@ -7,6 +7,7 @@ use std::fs;
 use std::path::Path;
 use std::fmt;
 use log::{info, debug, error};
+use bktree::*;
 
 #[derive(Debug)]
 pub enum SaveError {
@@ -15,15 +16,20 @@ pub enum SaveError {
 }
 
 pub struct DescriptorStore {
-	map: HashMap<String, u64>,
+	map: HashMap<u64, String>,
 	save_location: std::path::PathBuf,
+	bktree: BkTree<u64>,
 }
 
 impl DescriptorStore {
 	/// Makes a new empty DescriptorStore with default settings
 	pub fn new() -> Self {
-		let map: HashMap<String, u64> = HashMap::new();
-		DescriptorStore { map, save_location: std::path::PathBuf::from("store.messagepack")}
+		let map: HashMap<u64, String> = HashMap::new();
+		DescriptorStore {
+			map,
+			save_location: std::path::PathBuf::from("store.messagepack"),
+			bktree: BkTree::new(hamming_distance)
+		}
 	}
 
 	/// Sets the file location and loads data from file (if available)
@@ -34,6 +40,9 @@ impl DescriptorStore {
 			Ok(f) => self.map = rmp_serde::from_slice(&f).unwrap(),
 			Err(_) => info!("{} not found, starting from empty store.", self.save_location.display()),
 		};
+		for i in self.map.keys() {
+			self.bktree.insert(*i);
+		}
 		self
 	}
 
@@ -49,12 +58,12 @@ impl DescriptorStore {
 	}
 
 	/// Returns true iff the store already contains the key
-	pub fn contains(&self, key: String) -> bool {
+	pub fn contains(&self, key: u64) -> bool {
 		self.map.contains_key(&key)
 	}
 
 	/// Inserts a single value into the store
-	pub fn insert(&mut self, key: String, value: u64) {
+	pub fn insert(&mut self, key: u64, value: String) {
 		self.map.insert(key, value);
 	}
 
@@ -69,7 +78,7 @@ impl DescriptorStore {
 					continue
 				}
 			};
-			if !self.contains(name.to_string()) {
+			if true { // !self.contains(name.to_string()) {
 				info!("Processing {}", name);
 				let img = match image::open(file.path()) {
 					Ok(v) => v,
@@ -79,20 +88,35 @@ impl DescriptorStore {
 					}
 				};
 				let phash = desc.describe(&img);
-				self.map.insert(name, phash);
+				if self.contains(phash) {
+					println!("{} duplicate of {:?}", name, self.map.get(&phash));
+				}
+				self.map.insert(phash, name);
 			} else {
 				debug!("{} already known, skipping.", name);
 			}
 			self.save().expect("error");
 		}
 	}
+
+	/// K-nearest neighbors
+	pub fn knn(&self, from: u64, k: isize) {
+		println!("{from:b}");
+		for (element, distance) in self.bktree.find(from, k) {
+			let name = self.map.get(element);
+			match name {
+				Some(n) => println!("{element:b}: {n} (distance: {distance})"),
+				None => ()
+			};
+		}
+	}
 }
 
 impl fmt::Display for DescriptorStore {
 	fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 		let mut output = String::new();
 		for (key, val) in self.map.iter() {
-			let entry = format!("{key:020}: {val}\t {val:064b}\n");
+			let entry = format!("{key:064b}: {val}\t {val:020}\n");
 			output.push_str(&entry);
 		}
 		write!(f, "DescriptorStore:\n{}Total: {}", output, self.map.len())