From 4327fe8c7ca1905c9321ac88f12864d18d8ee535 Mon Sep 17 00:00:00 2001 From: Mark Hoekveen Date: Tue, 14 May 2024 14:37:33 +0200 Subject: [PATCH] inverted hashmap, added bktree --- Cargo.lock | 73 ++++++++++++++++++++++++++++++++++++------ Cargo.toml | 1 + src/bin/dctquery.rs | 42 ++++++++++++++++++++++++ src/descriptors/dct.rs | 4 +-- src/store.rs | 40 ++++++++++++++++++----- 5 files changed, 140 insertions(+), 20 deletions(-) create mode 100644 src/bin/dctquery.rs diff --git a/Cargo.lock b/Cargo.lock index 54af603..af7c90e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -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", ] diff --git a/Cargo.toml b/Cargo.toml index ca7e200..305ba96 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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" diff --git a/src/bin/dctquery.rs b/src/bin/dctquery.rs new file mode 100644 index 0000000..65bec74 --- /dev/null +++ b/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 { + 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 = 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); +} \ No newline at end of file diff --git a/src/descriptors/dct.rs b/src/descriptors/dct.rs index fbcc9d9..62df718 100644 --- a/src/descriptors/dct.rs +++ b/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; diff --git a/src/store.rs b/src/store.rs index 99f9940..4461f1f 100644 --- a/src/store.rs +++ b/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, + map: HashMap, save_location: std::path::PathBuf, + bktree: BkTree, } impl DescriptorStore { /// Makes a new empty DescriptorStore with default settings pub fn new() -> Self { - let map: HashMap = HashMap::new(); - DescriptorStore { map, save_location: std::path::PathBuf::from("store.messagepack")} + let map: HashMap = 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())