381 lines
10 KiB
Rust
381 lines
10 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 crate::mutators::get_all_mutators;
|
|
use std::fs;
|
|
use std::path::Path;
|
|
use std::fmt;
|
|
use image::DynamicImage;
|
|
use log::{info, debug, error};
|
|
use bktree::*;
|
|
|
|
#[derive(Debug)]
|
|
pub enum SaveError {
|
|
Serialization,
|
|
File,
|
|
}
|
|
|
|
/// Struct for calculating Recall-Precision per mutation
|
|
pub struct PRStats {
|
|
true_positives: [u64; 64],
|
|
false_positives: [u64; 64],
|
|
false_negatives: [u64; 64],
|
|
pub tag: String,
|
|
pub name: String
|
|
}
|
|
|
|
impl PRStats {
|
|
pub fn precision(&self, threshold: usize) -> f64 {
|
|
let t = match threshold {
|
|
0..=64 => threshold,
|
|
_ => std::cmp::max(0, std::cmp::min(64, threshold))
|
|
};
|
|
let tpos = self.true_positives[t] as f64;
|
|
let fpos = self.false_positives[t] as f64;
|
|
tpos / (tpos + fpos)
|
|
}
|
|
pub fn recall(&self, threshold: usize) -> f64 {
|
|
let t = match threshold {
|
|
0..=64 => threshold,
|
|
_ => std::cmp::max(0, std::cmp::min(64, threshold))
|
|
};
|
|
let tpos = self.true_positives[t] as f64;
|
|
let fneg = self.false_negatives[t] as f64;
|
|
tpos / (tpos + fneg)
|
|
}
|
|
pub fn pr(&self, threshold: usize) -> (f64, f64) {
|
|
let t = match threshold {
|
|
0..=64 => threshold,
|
|
_ => std::cmp::max(0, std::cmp::min(64, threshold))
|
|
};
|
|
let tpos = self.true_positives[t] as f64;
|
|
let fpos = self.false_positives[t] as f64;
|
|
let fneg = self.false_negatives[t] as f64;
|
|
let p = tpos / (tpos + fpos);
|
|
let r = tpos / (tpos + fneg);
|
|
(r, p)
|
|
}
|
|
pub fn pr_curve(&self, threshold: usize) -> Vec<(f64, f64)> {
|
|
let mut curve = Vec::new();
|
|
for i in 0..threshold {
|
|
curve.push(self.pr(i));
|
|
}
|
|
curve
|
|
}
|
|
|
|
}
|
|
|
|
/// Uses a hashmap to map descriptors to buckets of files.
|
|
/// Also keeps a BK-tree for quick distance ranking
|
|
pub struct DescriptorStore {
|
|
pub descriptor: Box<dyn Descriptor>,
|
|
/// 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(descriptor: Box<dyn Descriptor>) -> 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);
|
|
if let Ok(f) = map_file {
|
|
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 {
|
|
if !element.starts_with("mut.") {
|
|
seen.insert(element.clone());
|
|
}
|
|
}
|
|
}
|
|
self.seen = Some(seen);
|
|
} else {
|
|
info!("{} not found, starting from empty store.", self.save_location.display());
|
|
};
|
|
self
|
|
}
|
|
|
|
/// Stores non-mutated images in a hashset for quick membership checks
|
|
pub fn see(&mut self, value: String) {
|
|
if !value.starts_with("mut.") {
|
|
match &mut self.seen {
|
|
Some(set) => {
|
|
set.insert(value);
|
|
},
|
|
None => {
|
|
let mut seen: HashSet<String> = HashSet::new();
|
|
seen.insert(value);
|
|
self.seen = Some(seen);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn get(&self, value: &String) -> Option<u64> {
|
|
for (key, bucket) in self.map.iter() {
|
|
if bucket.contains(value) {
|
|
return Some(*key);
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
/// 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.clone());
|
|
}
|
|
n
|
|
},
|
|
None => vec![value.clone()],
|
|
};
|
|
self.map.insert(*key, bucket);
|
|
self.see(value);
|
|
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) {
|
|
debug!("{} already known, skipping.", name);
|
|
} else {
|
|
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);
|
|
}
|
|
self.save().expect("error");
|
|
}
|
|
// We have done our bulk loading, so we can unset our hashset cache:
|
|
//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: usize) -> Vec<(&u64, isize)> {
|
|
let neighbours = self.bktree.find(from, max_distance.try_into().unwrap());
|
|
//neighbours.sort_by(|a, b| a.1.cmp(&b.1));
|
|
neighbours
|
|
}
|
|
|
|
|
|
/// Returns a flat vector with all the filenames that are within a given distance
|
|
pub fn nn_flat_results(&self, from: u64, max_distance: usize) -> Vec<String> {
|
|
let nn = self.nn(from, max_distance);
|
|
let mut results: Vec<String> = Vec::new();
|
|
for (key, _) in nn {
|
|
// We know the key exists, so getting the result should never be none
|
|
let mut bucket = self.map.get(key).unwrap().clone();
|
|
results.append(&mut bucket);
|
|
}
|
|
results
|
|
}
|
|
|
|
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('\n');
|
|
}
|
|
}
|
|
print!("DescriptorStore:\n{}Total: {}", output, self.map.len())
|
|
}
|
|
|
|
pub fn get_stats(&self, mut max_threshold: usize) -> Vec<PRStats> {
|
|
max_threshold = std::cmp::min(64, max_threshold);
|
|
let mut mutator_stats = Vec::new();
|
|
for mutator in get_all_mutators() {
|
|
mutator_stats.push(
|
|
PRStats {
|
|
true_positives: [0; 64],
|
|
false_positives: [0; 64],
|
|
false_negatives: [0; 64],
|
|
tag: "mut".to_string() + &mutator.tag(),
|
|
name: mutator.info(),
|
|
}
|
|
);
|
|
}
|
|
|
|
// Seen should only contain a list of base images
|
|
let set = self.seen.clone().unwrap();
|
|
for image in set.iter() {
|
|
let phash = self.get(image).unwrap();
|
|
debug!("Checking {}, with phash: {}", image, phash);
|
|
for threshold in 0..max_threshold {
|
|
debug!("Threshold: {}", threshold);
|
|
//Assume miss, therefore a false negative
|
|
//Undo the miss when there is a true positive
|
|
for mutator in &mut mutator_stats {
|
|
mutator.false_negatives[threshold] += 1;
|
|
}
|
|
|
|
for found in self.nn_flat_results(phash, threshold) {
|
|
if found.ends_with(&format!(".{}", image)) { // True positive
|
|
for mutator in &mut mutator_stats {
|
|
if found.starts_with(&mutator.tag) {
|
|
debug!("{} is hit for {}", found, mutator.name);
|
|
mutator.true_positives[threshold] += 1;
|
|
mutator.false_negatives[threshold] -= 1;
|
|
}
|
|
}
|
|
} else if found.eq(image) {
|
|
continue;
|
|
} else { // False positive!
|
|
// Mutated misses only count for the mutator
|
|
// Unmutated misses count for everyone
|
|
if found.starts_with("mut") {
|
|
for mutator in &mut mutator_stats {
|
|
if found.starts_with(&mutator.tag) {
|
|
debug!("{} is false positive for {}", found, mutator.name);
|
|
mutator.false_positives[threshold] += 1;
|
|
}
|
|
}
|
|
} else {
|
|
for mutator in &mut mutator_stats {
|
|
mutator.false_positives[threshold] += 1;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
mutator_stats
|
|
}
|
|
|
|
pub fn print_stats(&self) {
|
|
let mut mutator_stats = self.get_stats(64);
|
|
for mutator in &mut mutator_stats {
|
|
println!("{}:", mutator.name);
|
|
for threshold in 0..64 {
|
|
let tpos = mutator.true_positives[threshold] as f64;
|
|
let fpos = mutator.false_positives[threshold] as f64;
|
|
let fneg = mutator.false_negatives[threshold] as f64;
|
|
let p = tpos / (tpos + fpos);
|
|
let r = tpos / (tpos + fneg);
|
|
let f1 = 2.0*tpos / (2.0*tpos + fpos + fneg);
|
|
debug!("Precision: {}, Recall: {}, F1 score: {}", p, r, f1);
|
|
debug!("Tp: {}, Fp: {}, Fn: {}", tpos, fpos, fneg);
|
|
println!("{}, {}" , p, r)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
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() {
|
|
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('\n');
|
|
}
|
|
write!(f, "DescriptorStore:\n{}Total: {}", output, self.map.len())
|
|
}
|
|
} |