Various updates
too lazy to split commits. some refactoring, clap, more binaries.
This commit is contained in:
+135
-42
@@ -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<u64, String>,
|
||||
save_location: std::path::PathBuf,
|
||||
bktree: BkTree<u64>,
|
||||
/// 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 DescriptorStore {
|
||||
/// Makes a new empty DescriptorStore with default settings
|
||||
pub fn new() -> Self {
|
||||
let map: HashMap<u64, String> = HashMap::new();
|
||||
let map: HashMap<u64, Vec<String>> = HashMap::new();
|
||||
//let seen: HashSet<String> = 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<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()),
|
||||
};
|
||||
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())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user