diff --git a/src/bin/dctfilename.rs b/src/bin/dctfilename.rs index ef9f440..593f043 100644 --- a/src/bin/dctfilename.rs +++ b/src/bin/dctfilename.rs @@ -18,6 +18,7 @@ impl Cfg{ } fn main() { + //Init all descriptors: let desc = DCT::new().with_quality(50); @@ -32,5 +33,4 @@ fn main() { println!("{phash}\n{phash:b}"); let output = URL_SAFE_NO_PAD.encode(phash.to_be_bytes()); println!("{output}") - } diff --git a/src/lib.rs b/src/lib.rs index e27c649..c807575 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1 +1,2 @@ -pub mod descriptors; \ No newline at end of file +pub mod descriptors; +pub mod store; \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index f8ed123..ccfd01b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,9 @@ +use image_similarity::store::DescriptorStore; +use image_similarity::descriptors::{DCT, Descriptor}; + fn main() { - print!("Hello world!") + let mut store = DescriptorStore::new(); + let desc = DCT::new().with_quality(30); + store.insert_directory("img", desc); + print!("{}", store); } diff --git a/src/store.rs b/src/store.rs new file mode 100644 index 0000000..a793382 --- /dev/null +++ b/src/store.rs @@ -0,0 +1,59 @@ +use std::collections::HashMap; +use crate::descriptors::Descriptor; +use std::fs; +use std::path::Path; +use std::fmt; + +pub struct DescriptorStore { + map: HashMap +} + +impl DescriptorStore { + /// Makes a new DescriptorStore. + /// Reads from file if possible, else empty + pub fn new() -> DescriptorStore { + //TODO: Set directory and file location from global config + let count = fs::read_dir("img").unwrap().count(); + let map_file = fs::read("map.messagepack"); + let map: HashMap = match map_file { + Ok(f) => rmp_serde::from_slice(&f).unwrap(), + Err(_e) => HashMap::with_capacity(count), + }; + DescriptorStore { map: map } + } + + /// Returns true iff the store already contains the key + pub fn contains(&self, key: String) -> bool { + self.map.contains_key(&key) + } + + /// Inserts a single value into the store + pub fn insert(mut self, key: String, value: u64) { + self.map.insert(key, value); + } + + + /// Calculates all descriptions with a given descriptor for a folder + pub fn insert_directory>(&mut self, dir: U, desc: T) { + for node in fs::read_dir(dir).unwrap() { + let file = node.expect("Error walking directory"); + let name = file.file_name().into_string().expect("Issue with filename"); + if !self.contains(name.to_string()) { + let img = image::open(file.path()).expect("Unable to open file"); + let phash = desc.describe(img); + self.map.insert(name, phash); + } + } + } +} + +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: >15}: {val}\t {val:b}\n"); + output.push_str(&entry); + } + write!(f, "{}", output) + } +} \ No newline at end of file