From 9fa594045dd42245b7b2825e313e36ff2081aa71 Mon Sep 17 00:00:00 2001 From: Mark Hoekveen Date: Sun, 2 Jun 2024 14:58:14 +0200 Subject: [PATCH] Fixes compilation, horizontal flip for DCT. Adds PR calculation. DCT is now somewhat invariant to horizontal flips. First version of precision recall variable is calculated. --- src/bin/dctfilename.rs | 2 +- src/bin/dctquery.rs | 4 +- src/descriptors/dct.rs | 37 +++++++++------ src/main.rs | 9 +++- src/store.rs | 104 +++++++++++++++++++++++++++++++++++++++-- 5 files changed, 133 insertions(+), 23 deletions(-) diff --git a/src/bin/dctfilename.rs b/src/bin/dctfilename.rs index b81c50e..1d6cdec 100644 --- a/src/bin/dctfilename.rs +++ b/src/bin/dctfilename.rs @@ -15,7 +15,7 @@ fn main() { env_logger::init(); let cfg = Cfg::parse(); - let desc = DCT::new().with_quality(50); + let desc = DCT::new().with_quality(95); let img = image::open(cfg.path) .expect("Unable to open file"); diff --git a/src/bin/dctquery.rs b/src/bin/dctquery.rs index 5cd2a84..3db6f9d 100644 --- a/src/bin/dctquery.rs +++ b/src/bin/dctquery.rs @@ -18,7 +18,7 @@ struct Cfg { fn main() { env_logger::init(); - let desc = DCT::new().with_quality(50); + let desc = DCT::new().with_quality(95); let cfg = Cfg::parse(); @@ -32,7 +32,7 @@ fn main() { // viuer::print(&img, &conf).expect("Image printing failed."); let phash: u64 = desc.describe(&img); let store = - DescriptorStore::new(desc) + DescriptorStore::new(Box::new(desc)) .with_file("dct50.messagepack"); info!("Phash integer:\n{phash}"); info!("Phash binary:\n{phash:064b}"); diff --git a/src/descriptors/dct.rs b/src/descriptors/dct.rs index a903fe1..bc3169a 100644 --- a/src/descriptors/dct.rs +++ b/src/descriptors/dct.rs @@ -112,22 +112,26 @@ impl Descriptor for DCT { debug!("DCT-coefficients:\n {}", print_matrix(dct_values)); // Quantization: - //debug!("Using quantization matrix:\n{}", print_matrix(self.quantization_matrix)); - for i in 0..64 { - dct_values[i] = (dct_values[i] / self.quantization_matrix[i] as f64).round(); - } - debug!("DCT-coefficients, quantized:\n {}", print_matrix(dct_values)); + // debug!("Using quantization matrix:\n{}", print_matrix(self.quantization_matrix)); + // for i in 0..64 { + // dct_values[i] = (dct_values[i] / self.quantization_matrix[i] as f64).round(); + // if dct_values[i] == -0.0 && i % 2 != 0 { + // dct_values[i] = 0.0; + // } + // } + // debug!("DCT-coefficients, quantized:\n {}", print_matrix(dct_values)); if log_enabled!(Level::Debug) { resized.save("resize.png").expect("Error saving file"); // De-quantization: + let mut dequant = dct_values.clone(); for i in 0..64 { - dct_values[i] = dct_values[i] * self.quantization_matrix[i] as f64; + dequant[i] = dct_values[i] * self.quantization_matrix[i] as f64; } //debug!("DCT-coefficients, de-quantized:\n {}", print_matrix(dct_values)); // Reconstruction original pixel values: - let reconstructed = self.idct(dct_values); + let reconstructed = self.idct(dequant); save_buffer( "reconstructed.png", @@ -158,24 +162,27 @@ impl Descriptor for DCT { // If first horizontal AC coefficient is negative // This might account for horizontal flips when applied to all horizontal coefficients. let sign_mult = dct_values[1].signum(); + //debug!("Sign multiplier of first horizontal component: {}", sign_mult); // Mask that indicates if a coefficient is bigger or smaller than previous in order let mut pearson_mask: u64 = 0; - let mut prev = dct_values[0]; + let mut prev = dct_values[49].abs(); for i in zigzag { - let cur = dct_values[i]; + let cur = dct_values[i].abs(); if cur > prev { pearson_mask += 1; } prev = cur; - let signum = dct_values[i].signum(); + let mut signum = dct_values[i].signum(); + if i % 2 != 0 { + signum *= sign_mult; + } + //debug!("[{}]: {} has signum {}, and i%2 is {}", i, cur, signum, i%2); // Only multiply sign if dct-coefficient contains a horizontal component. if - i % 8 != 0 && sign_mult * signum < 0.0 - || signum < 0.0 { sign_mask += 1; @@ -184,9 +191,11 @@ impl Descriptor for DCT { // Shift masks sign_mask = sign_mask << 1; pearson_mask = pearson_mask << 1; + debug!("Sign mask: {:028b}", sign_mask); + debug!("Pearson mask: {:028b}", pearson_mask); } - debug!("Sign mask: {:028b}", sign_mask); - debug!("Pearson mask: {:028b}", pearson_mask); + // debug!("Sign mask: {:028b}", sign_mask); + // debug!("Pearson mask: {:028b}", pearson_mask); let mut mask = sign_mask; debug!("Mask: {:064b}", mask); diff --git a/src/main.rs b/src/main.rs index 54ac60c..33866f8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -26,7 +26,7 @@ fn main() { } for node in fs::read_dir(cfg.path).unwrap() { - println!("{:?}", node); + //println!("{:?}", node); let file = node.expect("Error walking directory"); let name = match file.file_name().into_string() { Ok(v) => v, @@ -56,8 +56,13 @@ fn main() { } } } - for store in &stores { println!("{}", store); } + + for store in &stores { + println!("Store for {}: ", store.descriptor.info()); + store.print_stats(); + //println!("{}", store); + } } diff --git a/src/store.rs b/src/store.rs index 62b818f..3756abb 100644 --- a/src/store.rs +++ b/src/store.rs @@ -3,6 +3,7 @@ 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; @@ -16,10 +17,19 @@ pub enum SaveError { File, } +/// Struct for calculating Recall-Precision per mutation +struct PRStats { + true_positives: u64, + false_positives: u64, + false_negatives: u64, + tag: String, + name: String +} + /// Uses a hashmap to map descriptors to buckets of files. /// Also keeps a BK-tree for quick distance ranking pub struct DescriptorStore { - descriptor: Box, + pub descriptor: Box, /// Main hashmap that maps descriptors to buckets of filenames map: HashMap>, @@ -71,6 +81,22 @@ impl DescriptorStore { 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 = HashSet::new(); + seen.insert(value); + self.seen = Some(seen); + } + } + } + } + pub fn save(&self) -> Result<(), SaveError> { let serialized: Vec = match rmp_serde::to_vec(&self.map) { Ok(value) => value, @@ -111,13 +137,14 @@ impl DescriptorStore { Some(b) => { let mut n = b.clone(); if !n.contains(&value) { - n.push(value); + n.push(value.clone()); } n }, - None => vec![value], + None => vec![value.clone()], }; self.map.insert(*key, bucket); + self.see(value); self.bktree.insert(*key); } @@ -152,7 +179,7 @@ impl DescriptorStore { self.save().expect("error"); } // We have done our bulk loading, so we can unset our hashset cache: - self.seen = None; + //self.seen = None; } pub fn store(&mut self, img: &DynamicImage, name: String) { @@ -230,6 +257,75 @@ impl DescriptorStore { } print!("DescriptorStore:\n{}Total: {}", output, self.map.len()) } + + pub fn print_stats(&self) { + let mut mutator_stats = Vec::new(); + for mutator in get_all_mutators() { + mutator_stats.push( + PRStats { + true_positives: 0, + false_positives: 0, + false_negatives: 0, + tag: "mut".to_string() + &mutator.tag(), + name: mutator.info(), + } + ); + } + + let set = self.seen.clone().unwrap(); + + for image in set.iter() { + debug!("Checking {}", image); + + // find the bucket this image is in + for (key, bucket) in self.map.iter() { + if bucket.contains(&image) { + debug!("{} found in bucket {}", image, key); + for value in bucket { + if value != image && value.ends_with(image) { + // Mutated version of image in same bucket + for mutator in &mut mutator_stats { + mutator.false_negatives += 1; + if value.starts_with(&mutator.tag) { + mutator.true_positives += 1; + mutator.false_negatives -= 1; + } + } + debug!("We matched {} to {}!", value, image); + } else if value != image { + // Other image in same bucket + //false_positives += 1; + for mutator in &mut mutator_stats { + if value.starts_with(&mutator.tag) { + mutator.false_positives += 1; + } + } + } + } + break; + } + } + } + // + for mutator in &mut mutator_stats { + // println!("\t{} true positives: {}", mutator.name, mutator.true_positives); + // println!("\t{} false positives: {}", mutator.name, mutator.false_positives); + println!( + "\t{} Precision: {:.3}", mutator.name, + mutator.true_positives as f64 / + (mutator.true_positives + mutator.false_positives) as f64 + ); + println!( + "\t{} Recall: {:.3}", mutator.name, + mutator.true_positives as f64 / + //(mutator.true_positives + mutator.false_negatives) as f64 + set.len() as f64 + ); + } + //println!("Images processed: {}", set.len()) + // println!("{} true positives", true_positives); + // println!("{} false positives", false_positives); + } } impl fmt::Display for DescriptorStore {