From 46216c230e0e4e64d545fbb71aba3cc3b7298a68 Mon Sep 17 00:00:00 2001 From: Mark Hoekveen Date: Fri, 24 May 2024 09:59:49 +0200 Subject: [PATCH] Various updates too lazy to split commits. some refactoring, clap, more binaries. --- src/bin/dct.rs | 2 +- src/bin/dctquery.rs | 48 +++++------ src/descriptors/dct.rs | 4 +- src/main.rs | 92 +++++++++++++++------ src/store.rs | 177 +++++++++++++++++++++++++++++++---------- 5 files changed, 224 insertions(+), 99 deletions(-) diff --git a/src/bin/dct.rs b/src/bin/dct.rs index 266f324..29881d4 100644 --- a/src/bin/dct.rs +++ b/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()]; diff --git a/src/bin/dctquery.rs b/src/bin/dctquery.rs index d8a74f8..924b6d0 100644 --- a/src/bin/dctquery.rs +++ b/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 { - 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 = 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); } \ No newline at end of file diff --git a/src/descriptors/dct.rs b/src/descriptors/dct.rs index 62df718..52cfc5b 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/main.rs b/src/main.rs index 5b6f311..e8fc6e9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -12,7 +12,7 @@ struct Cfg { path: PathBuf } -fn make_store(descriptor: T, input_path: PathBuf, save_path: PathBuf) -> () { +fn make_store(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(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::(); + + 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(); } + } diff --git a/src/store.rs b/src/store.rs index c5c7ef9..c31a90d 100644 --- a/src/store.rs +++ b/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, - save_location: std::path::PathBuf, - bktree: BkTree, + /// Main hashmap that maps descriptors to buckets of filenames + map: HashMap>, + + /// Place to load and store the map on file + save_location: std::path::PathBuf, + + /// BK-tree for fast nearest neighbour + bktree: BkTree, + + /// Cache for seen files to skip on initial load + seen: Option>, } impl DescriptorStore { /// Makes a new empty DescriptorStore with default settings pub fn new() -> Self { - let map: HashMap = HashMap::new(); + let map: HashMap> = HashMap::new(); + //let seen: HashSet = 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 = 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()) }