Calculates basic phashes for images and their mutations. More stats are needed, but it's a start
This commit is contained in:
+13
-5
@@ -6,6 +6,7 @@ use crate::descriptors::Descriptor;
|
||||
use std::fs;
|
||||
use std::path::Path;
|
||||
use std::fmt;
|
||||
use image::DynamicImage;
|
||||
use log::{info, debug, error};
|
||||
use bktree::*;
|
||||
|
||||
@@ -17,8 +18,8 @@ pub enum SaveError {
|
||||
|
||||
/// 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,
|
||||
pub struct DescriptorStore {
|
||||
descriptor: Box<dyn Descriptor>,
|
||||
/// Main hashmap that maps descriptors to buckets of filenames
|
||||
map: HashMap<u64, Vec<String>>,
|
||||
|
||||
@@ -32,9 +33,9 @@ pub struct DescriptorStore<T: Descriptor> {
|
||||
seen: Option<HashSet<String>>,
|
||||
}
|
||||
|
||||
impl<T: Descriptor> DescriptorStore<T> {
|
||||
impl DescriptorStore {
|
||||
/// Makes a new empty DescriptorStore with default settings
|
||||
pub fn new(descriptor: T) -> Self {
|
||||
pub fn new(descriptor: Box<dyn Descriptor>) -> Self {
|
||||
let map: HashMap<u64, Vec<String>> = HashMap::new();
|
||||
//let seen: HashSet<String> = HashSet::new();
|
||||
//let seen = None;
|
||||
@@ -154,6 +155,13 @@ impl<T: Descriptor> DescriptorStore<T> {
|
||||
self.seen = None;
|
||||
}
|
||||
|
||||
pub fn store(&mut self, img: &DynamicImage, name: String) {
|
||||
if !self.has_value(&name) {
|
||||
let phash = self.descriptor.describe(&img);
|
||||
self.insert(&phash, name);
|
||||
}
|
||||
}
|
||||
|
||||
/// Nearest neighbours
|
||||
pub fn nn(&self, from: u64, max_distance: isize) -> Vec<(&u64, isize)> {
|
||||
let mut neighbours = self.bktree.find(from, max_distance);
|
||||
@@ -224,7 +232,7 @@ impl<T: Descriptor> DescriptorStore<T> {
|
||||
}
|
||||
}
|
||||
|
||||
impl<T: Descriptor> fmt::Display for DescriptorStore<T> {
|
||||
impl fmt::Display for DescriptorStore {
|
||||
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
|
||||
let mut output = String::new();
|
||||
for (key, bucket) in self.map.iter() {
|
||||
|
||||
Reference in New Issue
Block a user