store uitgewerkt
This commit is contained in:
+9
-3
@@ -4,8 +4,14 @@ use log::info;
|
|||||||
|
|
||||||
fn main() {
|
fn main() {
|
||||||
env_logger::init();
|
env_logger::init();
|
||||||
let mut store = DescriptorStore::new();
|
let mut store =
|
||||||
let desc = DCT::new().with_quality(30);
|
DescriptorStore::new()
|
||||||
store.insert_directory("img", desc);
|
.with_file("store.messagepack");
|
||||||
info!("{}", store);
|
info!("{}", store);
|
||||||
|
let desc =
|
||||||
|
DCT::new()
|
||||||
|
.with_quality(30);
|
||||||
|
store.insert_directory("data", desc);
|
||||||
|
info!("{}", store);
|
||||||
|
store.save().expect("Error saving");
|
||||||
}
|
}
|
||||||
|
|||||||
+55
-18
@@ -3,24 +3,46 @@ use crate::descriptors::Descriptor;
|
|||||||
use std::fs;
|
use std::fs;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
use std::fmt;
|
use std::fmt;
|
||||||
use log::info;
|
use log::{info, debug, error};
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
pub enum SaveError {
|
||||||
|
Serialization,
|
||||||
|
File,
|
||||||
|
}
|
||||||
|
|
||||||
pub struct DescriptorStore {
|
pub struct DescriptorStore {
|
||||||
map: HashMap<String, u64>
|
map: HashMap<String, u64>,
|
||||||
|
save_location: std::path::PathBuf,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl DescriptorStore {
|
impl DescriptorStore {
|
||||||
/// Makes a new DescriptorStore.
|
/// Makes a new empty DescriptorStore with default settings
|
||||||
/// Reads from file if possible, else empty
|
pub fn new() -> Self {
|
||||||
pub fn new() -> DescriptorStore {
|
let map: HashMap<String, u64> = HashMap::new();
|
||||||
//TODO: Set directory and file location from global config
|
DescriptorStore { map, save_location: std::path::PathBuf::from("store.messagepack")}
|
||||||
let count = fs::read_dir("img").unwrap().count();
|
}
|
||||||
let map_file = fs::read("map.messagepack");
|
|
||||||
let map: HashMap<String, u64> = match map_file {
|
/// Sets the file location and loads data from file (if available)
|
||||||
Ok(f) => rmp_serde::from_slice(&f).unwrap(),
|
pub fn with_file<P: AsRef<Path>>(mut self, path: P) -> Self {
|
||||||
Err(_e) => HashMap::with_capacity(count),
|
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(),
|
||||||
|
Err(_) => info!("{} not found, starting from empty store.", self.save_location.display()),
|
||||||
};
|
};
|
||||||
DescriptorStore { map }
|
self
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn save(&self) -> Result<(), SaveError> {
|
||||||
|
let serialized: Vec<u8> = match rmp_serde::to_vec(&self.map) {
|
||||||
|
Ok(value) => value,
|
||||||
|
Err(_e) => return Err(SaveError::Serialization),
|
||||||
|
};
|
||||||
|
match fs::write(&self.save_location, &serialized) {
|
||||||
|
Ok(_) => Ok(()),
|
||||||
|
Err(_e) => Err(SaveError::File),
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Returns true iff the store already contains the key
|
/// Returns true iff the store already contains the key
|
||||||
@@ -29,22 +51,37 @@ impl DescriptorStore {
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// Inserts a single value into the store
|
/// Inserts a single value into the store
|
||||||
pub fn insert(mut self, key: String, value: u64) {
|
pub fn insert(&mut self, key: String, value: u64) {
|
||||||
self.map.insert(key, value);
|
self.map.insert(key, value);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
/// Calculates all descriptions with a given descriptor for a folder
|
/// Calculates all descriptions with a given descriptor
|
||||||
pub fn insert_directory<T: Descriptor, U: AsRef<Path>>(&mut self, dir: U, desc: T) {
|
pub fn insert_directory<T: Descriptor, U: AsRef<Path>>(&mut self, dir: U, desc: T) {
|
||||||
for node in fs::read_dir(dir).unwrap() {
|
for node in fs::read_dir(dir).unwrap() {
|
||||||
let file = node.expect("Error walking directory");
|
let file = node.expect("Error walking directory");
|
||||||
let name = file.file_name().into_string().expect("Issue with filename");
|
let name = match file.file_name().into_string() {
|
||||||
|
Ok(v) => v,
|
||||||
|
Err(e) => {
|
||||||
|
error!("Error reading {}: {:?}:", file.path().to_string_lossy(), e);
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
};
|
||||||
if !self.contains(name.to_string()) {
|
if !self.contains(name.to_string()) {
|
||||||
info!("Processing {}", name);
|
info!("Processing {}", name);
|
||||||
let img = image::open(file.path()).expect("Unable to open file");
|
let img = match image::open(file.path()) {
|
||||||
|
Ok(v) => v,
|
||||||
|
Err(e) => {
|
||||||
|
error!("Failed to process {}: {}", name, e);
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
};
|
||||||
let phash = desc.describe(img);
|
let phash = desc.describe(img);
|
||||||
self.map.insert(name, phash);
|
self.map.insert(name, phash);
|
||||||
|
} else {
|
||||||
|
debug!("{} already known, skipping.", name);
|
||||||
}
|
}
|
||||||
|
self.save().expect("error");
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -53,9 +90,9 @@ impl fmt::Display for DescriptorStore {
|
|||||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||||
let mut output = String::new();
|
let mut output = String::new();
|
||||||
for (key, val) in self.map.iter() {
|
for (key, val) in self.map.iter() {
|
||||||
let entry = format!("{key: >15}: {val}\t {val:064b}\n");
|
let entry = format!("{key:020}: {val}\t {val:064b}\n");
|
||||||
output.push_str(&entry);
|
output.push_str(&entry);
|
||||||
}
|
}
|
||||||
write!(f, "DescriptorStore:\n{}", output)
|
write!(f, "DescriptorStore:\n{}Total: {}", output, self.map.len())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Reference in New Issue
Block a user