61 lines
1.8 KiB
Rust
61 lines
1.8 KiB
Rust
use std::collections::HashMap;
|
|
use crate::descriptors::Descriptor;
|
|
use std::fs;
|
|
use std::path::Path;
|
|
use std::fmt;
|
|
use log::info;
|
|
|
|
pub struct DescriptorStore {
|
|
map: HashMap<String, u64>
|
|
}
|
|
|
|
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<String, u64> = 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<T: Descriptor, U: AsRef<Path>>(&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()) {
|
|
info!("Processing {}", name);
|
|
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:064b}\n");
|
|
output.push_str(&entry);
|
|
}
|
|
write!(f, "DescriptorStore:\n{}", output)
|
|
}
|
|
} |