Merge branch 'master' into drone

src/bin/dct.rs

@@ -1,6 +1,6 @@
 use std::env;
 use std::f64::consts::{PI, SQRT_2};
-use image::{save_buffer, GenericImageView, GrayImage};
+use image::{GenericImageView, GrayImage};
 
 fn dct(img: &image::DynamicImage) -> GrayImage {
 	//let mut dct_values: = [0.0; img.width() * img.height()];

src/bin/dctquery.rs

@@ -1,50 +1,40 @@
 //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};
+use log::info;
+use clap::Parser;
 
+#[derive(Parser)]
+#[command(version, about, long_about = None)]
 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 })
-    }
+    /// Path to the image
+    path: String,
+    /// Try to output the images to terminal with viuer
+    #[arg(short, long, default_value_t=false)]
+    show_images: bool,
+    distance: isize,
 }
 
 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 cfg = Cfg::parse();
 
     let img = image::open(cfg.path)
         .expect("Unable to open file");
-    println!("Iamge");
-    let conf = viuer::Config {
-        width: Some(16),
-        height: Some(8),
-        ..Default::default()
-    };
-    viuer::print(&img, &conf).expect("Image printing failed.");
-    println!("Iamge");
+    // let conf = viuer::Config {
+    //     width: Some(16),
+    //     height: Some(8),
+    //     ..Default::default()
+    // };
+    // viuer::print(&img, &conf).expect("Image printing failed.");
     let store =
         DescriptorStore::new()
-        .with_file("store.messagepack");
-    //info!("{}", store);
+        .with_file("dct50.messagepack");
     let phash: u64 = desc.describe(&img);
 	info!("Phash integer:\n{phash}");
 	info!("Phash binary:\n{phash:064b}");
-	store.knn(phash, 10);
+	store.print_nn(phash, cfg.distance, cfg.show_images);
 }

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/main.rs

@@ -12,7 +12,7 @@ struct Cfg {
     path: PathBuf
 }
 
-fn make_store<T: Descriptor>(descriptor: T, input_path: PathBuf, save_path: PathBuf) -> () {
+fn make_store<T: Descriptor>(descriptor: T, input_path: PathBuf, save_path: PathBuf) -> DescriptorStore {
     let mut store =
         DescriptorStore::new()
         .with_file(save_path);
@@ -20,6 +20,13 @@ fn make_store<T: Descriptor>(descriptor: T, input_path: PathBuf, save_path: Path
     store.insert_directory(input_path, descriptor);
     info!("{}", store);
     store.save().expect("Error saving");
+    store
+}
+
+struct StoreParams {
+    desc: DCT,
+    input_path: PathBuf,
+    save_path: PathBuf,
 }
 
 fn main() {
@@ -28,32 +35,67 @@ fn main() {
 
     let mut threads = vec![];
 
-    let dct_50 =
-        DCT::new()
-        .with_quality(50);
-    let dct_30 =
-        DCT::new()
-        .with_quality(30);
-    let dct_10 =
-        DCT::new()
-        .with_quality(10);
+    let params = vec![
+        StoreParams {
+            desc: DCT::new().with_quality(50),
+            input_path: cfg.path.clone(),
+            save_path: "dct50.messagepack".into(),
+        },
+        // StoreParams {
+        //     desc: DCT::new().with_quality(30),
+        //     input_path: cfg.path.clone(),
+        //     save_path: "dct30.messagepack".into(),
+        // },
+        // StoreParams {
+        //     desc: DCT::new().with_quality(10),
+        //     input_path: cfg.path.clone(),
+        //     save_path: "dct10.messagepack".into(),
+        // },
+    ];
 
-    threads.push(
-        thread::spawn(move || {
-            make_store(dct_50, "data".into(), "dct50.messagepack".into());
-        })
-    );
-    threads.push(
-        thread::spawn(move || {
-            make_store(dct_30, "data".into(), "dct30.messagepack".into());
-        })
-    );
-    threads.push(
-        thread::spawn(move || {
-            make_store(dct_10, "data".into(), "dct10.messagepack".into());
-        })
-    );
+    // let dct_50 =
+    //     DCT::new()
+    //     .with_quality(50);
+    // let dct_30 =
+    //     DCT::new()
+    //     .with_quality(30);
+    // let dct_10 =
+    //     DCT::new()
+    //     .with_quality(10);
+
+    //let (tx, rx) = mpsc::channel::<DescriptorStore>();
+
+    for param in params {
+        threads.push(
+            thread::spawn(move || {
+                let store = make_store(param.desc, param.input_path, param.save_path);
+                store.print_most_dups();
+                //tx.send(store).unwrap();
+            })
+        )
+    }
+
+    // let path_50 = cfg.path.clone();
+    // threads.push(
+    //     thread::spawn(move || {
+    //         let store = make_store(dct_50, path_50, "dct50.messagepack".into());
+    //         tx.send(store).expect("Thread send result error");
+    //     })
+    // );
+    // let path_30 = cfg.path.clone();
+    // threads.push(
+    //     thread::spawn(move || {
+    //         dct_30_store = make_store(dct_30, path_30, "dct30.messagepack".into());
+    //     })
+    // );
+    // let path_10 = cfg.path.clone();
+    // threads.push(
+    //     thread::spawn(move || {
+    //         dct_50_store = make_store(dct_10, path_10, "dct10.messagepack".into());
+    //     })
+    // );
     for thread in threads {
         let _ = thread.join();
     }
+
 }

src/store.rs

@@ -1,7 +1,7 @@
 //! Persistent store of all calculated image descriptions.
 //! Can be queried to find identical or similar images.
 
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use crate::descriptors::Descriptor;
 use std::fs;
 use std::path::Path;
@@ -15,20 +15,33 @@ pub enum SaveError {
 	File,
 }
 
+/// Uses a hashmap to map descriptors to buckets of files.
+/// Also keeps a BK-tree for quick distance ranking
 pub struct DescriptorStore {
-	map: HashMap<u64, String>,
+	/// Main hashmap that maps descriptors to buckets of filenames
+	map: HashMap<u64, Vec<String>>, 
+
+	/// Place to load and store the map on file
 	save_location: std::path::PathBuf, 
+
+	/// BK-tree for fast nearest neighbour
 	bktree: BkTree<u64>, 
+
+	/// Cache for seen files to skip on initial load
+	seen: Option<HashSet<String>>, 
 }
 
 impl DescriptorStore {
 	/// Makes a new empty DescriptorStore with default settings
 	pub fn new() -> Self {
-		let map: HashMap<u64, String> = HashMap::new();
+		let map: HashMap<u64, Vec<String>> = HashMap::new();
+		//let seen: HashSet<String> = HashSet::new();
+		//let seen = None;
 		DescriptorStore {
 			map,
 			save_location: std::path::PathBuf::from("store.messagepack"),
-			bktree: BkTree::new(hamming_distance)
+			bktree: BkTree::new(hamming_distance),
+			seen: None,
 		}
 	}
 
@@ -37,12 +50,21 @@ impl DescriptorStore {
 		self.save_location = std::path::PathBuf::from(path.as_ref());
 		let map_file = fs::read(&self.save_location);
 		match map_file {
-			Ok(f) => self.map = rmp_serde::from_slice(&f).unwrap(),
+			Ok(f) => {
+				self.map = rmp_serde::from_slice(&f).unwrap();
+				for i in self.map.keys() {
+					self.bktree.insert(*i);
+				}
+				let mut seen: HashSet<String> = HashSet::new();
+				for bucket in self.map.values() {
+					for element in bucket {
+						seen.insert(element.clone());
+					}
+				}
+				self.seen = Some(seen);
+			},
 			Err(_) => info!("{} not found, starting from empty store.", self.save_location.display()),
 		};
-		for i in self.map.keys() {
-			self.bktree.insert(*i);
-		}
 		self
 	}
 
@@ -62,10 +84,38 @@ impl DescriptorStore {
 		self.map.contains_key(&key)
 	}
 
+	/// Returns true iff the store already contains the value in some bucket. 
+	/// Will use a hashmap cache if available
+	pub fn has_value(&self, value: &String) -> bool {
+		match &self.seen {
+			Some(set) => {
+				set.contains(value)
+			},
+			None => {
+				for bucket in self.map.values() {
+					if bucket.contains(&value) {
+						return true;
+					}
+				}
+				false
+			}
+		}
+	}
+
 	/// Inserts a single value into the store
-	pub fn insert(&mut self, key: u64, value: String) {
-		self.map.insert(key, value);
-		self.bktree.insert(key);
+	pub fn insert(&mut self, key: &u64, value: String) {
+		let bucket = match self.map.get(&key) {
+			Some(b) => {
+				let mut n = b.clone();
+				if !n.contains(&value) {
+					n.push(value);
+				}
+				n
+			},
+			None => vec![value],
+		};
+		self.map.insert(*key, bucket);
+		self.bktree.insert(*key);
 	}
 
 	/// Calculates all descriptions with a given descriptor 
@@ -79,7 +129,7 @@ impl DescriptorStore {
 					continue
 				}
 			};
-			if true { // !self.contains(name.to_string()) {
+			if !self.has_value(&name) { // !self.contains(name.to_string()) {
 				info!("Processing {}", name);
 				let img = match image::open(file.path()) {
 					Ok(v) => v,
@@ -92,54 +142,97 @@ impl DescriptorStore {
 				if self.contains(phash) {
 					println!("{} duplicate of {:?}", name, self.map.get(&phash));
 				}
-				self.map.insert(phash, name);
+				self.insert(&phash, name);
 			} else {
 				debug!("{} already known, skipping.", name);
 			}
 			self.save().expect("error");
 		}
+		// We have done our bulk loading, so we can unset our hashset cache:
+		self.seen = None;
 	}
 
-	/// K-nearest neighbors
-	pub fn knn(&self, from: u64, k: isize) {
-		println!("{from:b}");
+	/// Nearest neighbours
+	pub fn nn(&self, from: u64, max_distance: isize) -> Vec<(&u64, isize)> {
+		let mut neighbours = self.bktree.find(from, max_distance);
+		neighbours.sort_by(|a, b| a.1.cmp(&b.1));
+		neighbours
+	}
+
+	pub fn print_nn(&self, from: u64, max_distance: isize, show_images: bool) {
+		//println!("{from:064b}");
+		let neighbours = self.nn(from, max_distance);
+		let (term_width, _) = viuer::terminal_size();
 		let mut x = 0;
 		let mut y = 8;
-		let (term_width, _) = viuer::terminal_size();
-		println!("\t\t\t\t\tTERMWIDTH {term_width}");
-		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})");
-				if x+16 >= term_width {
-						x = 0;
-						y += 8;
+		for (element, distance) in neighbours {
+			let elements = self.map.get(element);
+			match elements {
+				Some(paths) => {
+					for path in paths {
+						println!("{element:064b}: {:?} (distance: {distance})", path);
+						if show_images {
+							if x+16 >= term_width {
+								x = 0;
+								y += 8;
+							}
+							let conf = viuer::Config {
+								width: Some(16),
+								height: Some(8),
+								x,
+								y,
+								use_kitty: false,
+								..Default::default()
+							};
+							x += 16;
+							let path = "data/".to_string() + path;
+							let img = image::open(&path).unwrap();
+        					let img = img.grayscale().thumbnail_exact(8, 8);
+							viuer::print(&img, &conf).expect("Image printing failed.");
+						}
 					}
-					let conf = viuer::Config {
-						width: Some(16),
-						height: Some(8),
-						x: x,
-						y: y,
-						..Default::default()
-					};
-					x += 16;
-					let path = "data/".to_string() + &n;
-					let img = image::open(&path).unwrap();
-					viuer::print(&img, &conf).expect("Image printing failed.");
 				},
 				None => ()
 			};
 		}
 	}
+
+	pub fn print_most_dups(&self) {
+		let mut most = 0;
+		for bucket in self.map.values() {
+			let len = bucket.len();
+			if len > most {
+				most = len;
+			}
+		}
+
+		let mut output = String::new();
+		for (key, bucket) in self.map.iter() {
+			if bucket.len() == most {
+				let key_entry = format!("{key:064b}:\n\t\t");
+				output.push_str(&key_entry);
+				for value in bucket {
+					let value_entry = format!("{value}\n");
+					output.push_str(&value_entry);
+				}
+				output.push_str("\n");
+			}
+		}
+		print!("DescriptorStore:\n{}Total: {}", output, self.map.len())
+	}
 }
 
 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:064b}: {val}\t {val:020}\n");
-			output.push_str(&entry);
+		for (key, bucket) in self.map.iter() {
+			let key_entry = format!("{key:064b}:\n");
+			output.push_str(&key_entry);
+			for value in bucket {
+				let value_entry = format!("\t{value}\n");
+				output.push_str(&value_entry);
+			}
+			output.push_str("\n");
 		}
 		write!(f, "DescriptorStore:\n{}Total: {}", output, self.map.len())
 	}