inverted hashmap, added bktree

2024-05-14 14:37:33 +02:00
parent 335f7e7d8c
commit 4327fe8c7c
5 changed files with 140 additions and 20 deletions
@@ -90,6 +90,15 @@ version = "1.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a"

+[[package]]
+name = "bktree"
+version = "1.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0bb1e744816f6a3b9e962186091867f3e5959d4dac995777ec254631cb00b21c"
+dependencies = [
+ "num",
+]
+
 [[package]]
 name = "bumpalo"
 version = "3.11.1"
@@ -334,6 +343,7 @@ name = "image_similarity"
 version = "0.1.0"
 dependencies = [
 "base64",
+ "bktree",
 "env_logger",
 "image",
 "log",
@@ -436,20 +446,52 @@ dependencies = [
 ]

 [[package]]
-name = "num-integer"
-version = "0.1.45"
+name = "num"
+version = "0.4.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9"
+checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23"
 dependencies = [
- "autocfg",
+ "num-bigint",
+ "num-complex",
+ "num-integer",
+ "num-iter",
+ "num-rational",
 "num-traits",
 ]

 [[package]]
-name = "num-rational"
-version = "0.4.1"
+name = "num-bigint"
+version = "0.4.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0638a1c9d0a3c0914158145bc76cff373a75a627e6ecbfb71cbe6f453a5a19b0"
+checksum = "c165a9ab64cf766f73521c0dd2cfdff64f488b8f0b3e621face3462d3db536d7"
+dependencies = [
+ "num-integer",
+ "num-traits",
+]
+
+[[package]]
+name = "num-complex"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-integer"
+version = "0.1.46"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-iter"
+version = "0.1.45"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1429034a0490724d0075ebb2bc9e875d6503c3cf69e235a8941aa757d83ef5bf"
 dependencies = [
 "autocfg",
 "num-integer",
@@ -457,10 +499,21 @@ dependencies = [
 ]

 [[package]]
-name = "num-traits"
-version = "0.2.15"
+name = "num-rational"
+version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "578ede34cf02f8924ab9447f50c28075b4d3e5b269972345e7e0372b38c6cdcd"
+checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824"
+dependencies = [
+ "num-bigint",
+ "num-integer",
+ "num-traits",
+]
+
+[[package]]
+name = "num-traits"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
 dependencies = [
 "autocfg",
 ]
@@ -12,3 +12,4 @@ rmp-serde = "1.1.1"
 base64 = "0.22.1"
 log = "0.4.21"
 env_logger = "0.11.3"
+bktree = "1.0.1"
@@ -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);
+}
@@ -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;
@@ -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())