inverted hashmap, added bktree

src/bin/dctquery.rs

@@ -0,0 +1,42 @@
+//use image_similarity::descriptors::dct::DCT;
+use image_similarity::descriptors::{DCT, Descriptor};
+use image_similarity::store::DescriptorStore;
+use std::env;
+use log::{info, debug, error};
+
+struct Cfg {
+    path: String
+}
+
+impl Cfg{
+    fn load(args: &[String]) -> Result<Cfg, &'static str> {
+        if args.len() < 2 {
+            return Err("Filename argument required");
+        }
+        let path = args[1].clone();
+        Ok(Cfg { path })
+    }
+}
+
+fn main() {
+    env_logger::init();
+
+    //Init all descriptors:
+    let desc = DCT::new().with_quality(50);
+
+    let args: Vec<String> = env::args().collect();
+
+    let cfg = Cfg::load(&args)
+        .expect("Error loading config");
+
+    let img = image::open(cfg.path)
+        .expect("Unable to open file");
+    let store =
+        DescriptorStore::new()
+        .with_file("store.messagepack");
+    //info!("{}", store);
+    let phash: u64 = desc.describe(&img);
+	info!("Phash integer:\n{phash}");
+	info!("Phash binary:\n{phash:064b}");
+	store.knn(phash, 4);
+}

src/descriptors/dct.rs

@@ -174,8 +174,8 @@ impl Descriptor for DCT {
             let signum = dct_values[i].signum();
             // Only multiply sign if dct-coefficient contains a horizontal component.
             if 
-                i % 8 != 0 && sign_mult * signum < 0.0
-                ||
+                //i % 8 != 0 && sign_mult * signum < 0.0
+                // ||
                 signum < 0.0
             {
                 sign_mask += 1;

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())