241 lines
6.3 KiB
Rust
241 lines
6.3 KiB
Rust
//! Persistent store of all calculated image descriptions.
|
|
//! Can be queried to find identical or similar images.
|
|
|
|
use std::collections::{HashMap, HashSet};
|
|
use crate::descriptors::Descriptor;
|
|
use std::fs;
|
|
use std::path::Path;
|
|
use std::fmt;
|
|
use log::{info, debug, error};
|
|
use bktree::*;
|
|
|
|
#[derive(Debug)]
|
|
pub enum SaveError {
|
|
Serialization,
|
|
File,
|
|
}
|
|
|
|
/// Uses a hashmap to map descriptors to buckets of files.
|
|
/// Also keeps a BK-tree for quick distance ranking
|
|
pub struct DescriptorStore<T: Descriptor> {
|
|
descriptor: T,
|
|
/// 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<T: Descriptor> DescriptorStore<T> {
|
|
/// Makes a new empty DescriptorStore with default settings
|
|
pub fn new(descriptor: T) -> Self {
|
|
let map: HashMap<u64, Vec<String>> = HashMap::new();
|
|
//let seen: HashSet<String> = HashSet::new();
|
|
//let seen = None;
|
|
DescriptorStore {
|
|
map,
|
|
descriptor,
|
|
save_location: std::path::PathBuf::from("store.messagepack"),
|
|
bktree: BkTree::new(hamming_distance),
|
|
seen: None,
|
|
}
|
|
}
|
|
|
|
/// Sets the file location and loads data from file (if available)
|
|
pub fn with_file<P: AsRef<Path>>(mut self, path: P) -> Self {
|
|
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();
|
|
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()),
|
|
};
|
|
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
|
|
pub fn contains(&self, key: u64) -> bool {
|
|
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) {
|
|
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
|
|
pub fn insert_directory<U: AsRef<Path>>(&mut self, dir: U) {
|
|
for node in fs::read_dir(dir).unwrap() {
|
|
let file = node.expect("Error walking directory");
|
|
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.has_value(&name) { // !self.contains(name.to_string()) {
|
|
info!("Processing {}", name);
|
|
let img = match image::open(file.path()) {
|
|
Ok(v) => v,
|
|
Err(e) => {
|
|
error!("Failed to process {}: {}", name, e);
|
|
continue
|
|
}
|
|
};
|
|
let phash = self.descriptor.describe(&img);
|
|
if self.contains(phash) {
|
|
println!("{} duplicate of {:?}", name, self.map.get(&phash));
|
|
}
|
|
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;
|
|
}
|
|
|
|
/// 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;
|
|
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.");
|
|
}
|
|
}
|
|
},
|
|
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<T: Descriptor> fmt::Display for DescriptorStore<T> {
|
|
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
|
let mut output = String::new();
|
|
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())
|
|
}
|
|
} |