better pr curve, plotter, store fixes

src/bin/dctquery.rs

@@ -12,7 +12,7 @@ struct Cfg {
 	/// Try to output the images to terminal with viuer
 	#[arg(short, long, default_value_t=false)]
 	show_images: bool,
-	distance: isize,
+	distance: usize,
 }
 
 fn main() {

src/descriptors/dct.rs

@@ -108,7 +108,7 @@ impl Descriptor for DCT {
     }
     fn describe(&self, img: &image::DynamicImage) -> u64 {
         let resized = self.resize(img);
-        let mut dct_values = self.dct(&resized);
+        let dct_values = self.dct(&resized);
         debug!("DCT-coefficients:\n {}", print_matrix(dct_values));
 
         // Quantization:

src/descriptors/mod.rs

@@ -76,6 +76,6 @@ impl Descriptor for Median {
 pub fn get_all_descriptors() -> Vec<Box<dyn Descriptor>> {
 	let mut descriptors: Vec<Box<dyn Descriptor>> = Vec::with_capacity(2);
 	descriptors.push(Box::new(DCT::new()));
-	//descriptors.push(Box::new(Median));
+	descriptors.push(Box::new(Median));
 	descriptors
 }

src/main.rs

@@ -1,10 +1,11 @@
 use image_similarity::store::DescriptorStore;
 use image_similarity::descriptors::get_all_descriptors;
 use image_similarity::mutators::get_all_mutators;
-use log::{error, info};
+use log::{error};
 use std::path::PathBuf;
-use std::{fs, thread};
+use std::{fs};
 use clap::Parser;
+use plotters::prelude::*;
 
 #[derive(Parser)]
 #[command(version, about, long_about = None)]
@@ -21,7 +22,8 @@ fn main() {
     let mut stores: Vec<DescriptorStore> = Vec::new();
 
     for descriptor in get_all_descriptors() {
-        let store = DescriptorStore::new(descriptor);
+        let filename = format!("{}.store", descriptor.info());
+        let store = DescriptorStore::new(descriptor).with_file(filename);
         stores.push(store);
     }
     
@@ -35,34 +37,84 @@ fn main() {
                 continue
             }
         };
-        let img = match image::open(file.path()) {
-            Ok(v) => v,
-            Err(e) => {
-                error!("Failed to process {}: {}", name, e);
-                continue
-            }
-        };
 
-        //Store the phashes of the base image
-        for store in &mut stores {
-            store.store(&img, name.clone());
+        let mut get = true;
+        for store in &stores {
+            if store.has_value(&name) {
+                get = false;
+                break;
+            }
         }
 
-        for mutator in &mutators {
-            let mutated = mutator.mutate(&img);
-            let mutated_name = "mut".to_string() + &mutator.tag() + &name.clone();
+        // We assume that if a base image exists in the store, the mutated images also exist
+        if get {
+            let img = match image::open(file.path()) {
+                Ok(v) => v,
+                Err(e) => {
+                    error!("Failed to process {}: {}", name, e);
+                    continue
+                }
+            };
+
+            //Store the phashes of the base image
             for store in &mut stores {
-                store.store(&mutated, mutated_name.clone());
+                store.store(&img, name.clone());
+            }
+
+            for mutator in &mutators {
+                let mutated = mutator.mutate(&img);
+                let mutated_name = "mut".to_string() + &mutator.tag() + &name.clone();
+                for store in &mut stores {
+                    store.store(&mutated, mutated_name.clone());
+                }
             }
         }
     }
     for store in &stores {
+        store.save().expect("Error saving store");
         println!("{}", store);
     }
 
+    // Create PR-curve graphs from stores
     for store in &stores {
         println!("Store for {}: ", store.descriptor.info());
-        store.print_stats();
+        //store.print_stats();
+        let stats = store.get_stats(16);
+        let filename = format!("{}-pr.png", store.descriptor.info());
+        let root = BitMapBackend::new(&filename, (1024, 768)).into_drawing_area();
+        root.fill(&WHITE).unwrap();
+        let mut chart = ChartBuilder::on(&root)
+            // Set the caption of the chart
+            .caption("PR curves", ("sans-serif", 40).into_font())
+            .margin(25)
+            .set_all_label_area_size(50)
+            // Finally attach a coordinate on the drawing area and make a chart context
+            .build_cartesian_2d(0f64..1f64, 0f64..1f64).unwrap();
+        chart.configure_mesh()
+            .x_labels(10)
+            .y_labels(10)
+            .disable_mesh()
+            .x_desc("Recall")
+            .y_desc("Precision")
+            .x_label_formatter(&|x| format!("{:.3}", x))
+            .y_label_formatter(&|x| format!("{:.3}", x))
+            .draw().unwrap();
+        let colors = [&RED, &BLUE, &CYAN, &MAGENTA, &BLACK, &GREEN, &YELLOW];
+        let n = colors.len();
+        let mut i = 0;
+        for mutator_stats in stats {
+            // And we can draw something in the drawing area
+            let pr_curve = mutator_stats.pr_curve();
+            let color = colors[i % n];
+            chart.draw_series(LineSeries::new(
+                pr_curve,
+                color,
+            )).unwrap()
+                .label(mutator_stats.name)
+                .legend(move |(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], color));
+            i += 1;
+        }
+        chart.configure_series_labels().border_style(BLACK).draw().unwrap();
         //println!("{}", store);
     }
 }

src/store.rs

@@ -18,12 +18,55 @@ pub enum SaveError {
 }
 
 /// Struct for calculating Recall-Precision per mutation
-struct PRStats {
+pub struct PRStats {
 	true_positives: [u64; 64],
 	false_positives: [u64; 64],
 	false_negatives: [u64; 64],
-	tag: String,
-	name: String
+	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;
+		let p = tpos / (tpos + fpos);
+		p
+	}
+	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;
+		let r = tpos / (tpos + fneg);
+		r
+	}
+	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) -> Vec<(f64, f64)> {
+		let mut curve = Vec::new();
+		for i in 0..64 {
+			curve.push(self.pr(i));
+		}
+		curve
+	}
+
 }
 
 /// Uses a hashmap to map descriptors to buckets of files.
@@ -71,7 +114,9 @@ impl DescriptorStore {
 				let mut seen: HashSet<String> = HashSet::new();
 				for bucket in self.map.values() {
 					for element in bucket {
-						seen.insert(element.clone());
+						if !element.starts_with("mut.") {
+							seen.insert(element.clone());
+						}
 					}
 				}
 				self.seen = Some(seen);
@@ -199,15 +244,15 @@ impl DescriptorStore {
 	}
 
 	/// Nearest neighbours
-	pub fn nn(&self, from: u64, max_distance: isize) -> Vec<(&u64, isize)> {
-		let mut neighbours = self.bktree.find(from, max_distance);
+	pub fn nn(&self, from: u64, max_distance: usize) -> Vec<(&u64, isize)> {
+		let mut 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: isize) -> Vec<String> {
+	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 {
@@ -218,7 +263,7 @@ impl DescriptorStore {
 		results
 	}
 
-	pub fn print_nn(&self, from: u64, max_distance: isize, show_images: bool) {
+	pub fn print_nn(&self, from: u64, max_distance: usize, show_images: bool) {
 		//println!("{from:064b}");
 		let neighbours = self.nn(from, max_distance);
 		let (term_width, _) = viuer::terminal_size();
@@ -280,7 +325,8 @@ impl DescriptorStore {
 		print!("DescriptorStore:\n{}Total: {}", output, self.map.len())
 	}
 
-	pub fn print_stats(&self) {
+	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(
@@ -298,7 +344,7 @@ impl DescriptorStore {
 		for image in set.iter() {
 			let phash = self.get(image).unwrap();
 			debug!("Checking {}, with phash: {}", image, phash);
-			for threshold in 0..64 {
+			for threshold in 0..max_threshold {
 				debug!("Threshold: {}", threshold);
 				//Assume miss, therefore a false negative
 				//Undo the miss when there is a true positive
@@ -328,8 +374,13 @@ impl DescriptorStore {
 				}
 			}
 		}
+		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;
@@ -337,27 +388,11 @@ impl DescriptorStore {
 				let p = tpos / (tpos + fpos);
 				let r = tpos / (tpos + fneg);
 				let f1 = 2.0*tpos / (2.0*tpos + fpos + fneg);
-				println!("{} / {}:", mutator.name, threshold);
-				println!("Precision: {}, Recall: {}, F1 score: {}", p, r, f1);
-				println!("Tp: {}, Fp: {}, Fn: {}", tpos, fpos, fneg);
+				debug!("Precision: {}, Recall: {}, F1 score: {}", p, r, f1);
+				debug!("Tp: {}, Fp: {}, Fn: {}", tpos, fpos, fneg);
+				println!("{}, {}" , p, r)
 			}
-			// 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);
 	}
 }