update with continues and bk tree

.gitignore

@@ -10,3 +10,8 @@ web/colors.css
 smolflickr
 data
 spook.jpg
+run*.log
+thumbcache/
+legacy/
+*.pdf
+run-*.sh

Cargo.lock

@@ -200,15 +200,6 @@ version = "2.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "415f8399438eb5e4b2f73ed3152a3448b98149dda642a957ee704e1daa5cf1d8"
 
-[[package]]
-name = "bktree"
-version = "1.0.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "0bb1e744816f6a3b9e962186091867f3e5959d4dac995777ec254631cb00b21c"
-dependencies = [
- "num",
-]
-
 [[package]]
 name = "built"
 version = "0.7.3"
@@ -770,7 +761,6 @@ name = "image_similarity"
 version = "0.2.0"
 dependencies = [
  "base64",
- "bktree",
  "clap",
  "console_error_panic_hook",
  "env_logger",

Cargo.toml

@@ -12,7 +12,6 @@ serde = { version = "1.0.147", features = ["derive"] }
 rmp-serde = "1.1.1"
 base64 = "0.22.1"
 log = "0.4.21"
-bktree = "1.0.1"
 wasm-bindgen = "0.2.92"
 console_error_panic_hook = "0.1.7"
 
@@ -48,3 +47,8 @@ required-features = ["cli"]
 name = "mutate"
 path = "src/bin/mutate.rs"
 required-features = ["cli"]
+
+[[bin]]
+name = "bench"
+path = "src/bin/bench.rs"
+required-features = ["cli"]

src/bin/bench.rs

@@ -0,0 +1,280 @@
+use clap::Parser;
+use image_similarity::bk::BkTree;
+use image_similarity::descriptors::get_all_descriptors;
+use image_similarity::mutators::get_all_mutators;
+use image_similarity::store::DescriptorStore;
+use std::collections::HashMap;
+use std::path::PathBuf;
+use std::time::Instant;
+
+/// Benchmarks for the whole pipeline: decoding, hashing, mutating,
+/// BK-tree scaling and store overhead
+#[derive(Parser)]
+#[command(version, about, long_about = None)]
+struct Cfg {
+	/// Directory with sample images
+	#[arg(long, default_value = "img")]
+	images: PathBuf,
+	/// How many sample images to use
+	#[arg(long, default_value_t = 6)]
+	samples: usize,
+	/// BK-tree sizes to test, comma separated
+	#[arg(long, default_value = "10000,100000,1000000")]
+	sizes: String,
+	/// Query radii to test, comma separated
+	#[arg(long, default_value = "4,8,12,16,24")]
+	radii: String,
+	/// Queries per radius
+	#[arg(long, default_value_t = 200)]
+	queries: usize,
+	/// Also query a tree built from the keys of an existing store (legacy or current format)
+	#[arg(long)]
+	real_store: Option<PathBuf>,
+	/// Skip the image benchmarks (decode/describe/mutate)
+	#[arg(long, default_value_t = false)]
+	trees_only: bool,
+}
+
+fn ms<F: FnMut()>(iters: usize, mut f: F) -> f64 {
+	f(); // warmup
+	let start = Instant::now();
+	for _ in 0..iters {
+		f();
+	}
+	start.elapsed().as_secs_f64() * 1000.0 / iters as f64
+}
+
+fn lcg_keys(n: usize, mut state: u64) -> Vec<u64> {
+	(0..n)
+		.map(|_| {
+			state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
+			state
+		})
+		.collect()
+}
+
+/// Reference BK-tree with the same nested layout as the bktree crate,
+/// to measure what the arena layout buys us
+struct RefNode {
+	key: u64,
+	children: Vec<(u64, RefNode)>,
+}
+
+impl RefNode {
+	fn insert(&mut self, key: u64) {
+		let d = (self.key ^ key).count_ones() as u64;
+		if d == 0 {
+			return;
+		}
+		match self.children.iter_mut().find(|(cd, _)| *cd == d) {
+			Some((_, child)) => child.insert(key),
+			None => self.children.push((d, RefNode { key, children: Vec::new() })),
+		}
+	}
+
+	fn find(&self, key: u64, radius: u64, out: &mut Vec<(u64, u64)>) {
+		let d = (self.key ^ key).count_ones() as u64;
+		if d <= radius {
+			out.push((self.key, d));
+		}
+		for (cd, child) in &self.children {
+			if cd + radius >= d && *cd <= d + radius {
+				child.find(key, radius, out);
+			}
+		}
+	}
+}
+
+fn bench_tree_queries(label: &str, find: &dyn Fn(u64, u64) -> usize, queries: &[u64], radii: &[u64]) {
+	for &radius in radii {
+		let mut results = 0usize;
+		let start = Instant::now();
+		for &q in queries {
+			results += find(q, radius);
+		}
+		let per_query = start.elapsed().as_secs_f64() * 1000.0 / queries.len() as f64;
+		println!(
+			"  {label} radius {radius:>2}: {per_query:>9.3} ms/query, avg {} results",
+			results / queries.len()
+		);
+	}
+}
+
+fn main() {
+	env_logger::init();
+	let cfg = Cfg::parse();
+	let sizes: Vec<usize> = cfg.sizes.split(',').map(|s| s.trim().parse().expect("bad size")).collect();
+	let radii: Vec<u64> = cfg.radii.split(',').map(|s| s.trim().parse().expect("bad radius")).collect();
+
+	if !cfg.trees_only {
+		let mut paths: Vec<PathBuf> = std::fs::read_dir(&cfg.images)
+			.expect("cannot read image dir")
+			.filter_map(|e| e.ok())
+			.map(|e| e.path())
+			.filter(|p| p.is_file())
+			.collect();
+		paths.sort();
+		paths.truncate(cfg.samples);
+		assert!(!paths.is_empty(), "no images in {}", cfg.images.display());
+
+		println!("== decode ({} images from {})", paths.len(), cfg.images.display());
+		let start = Instant::now();
+		let images: Vec<image::DynamicImage> = paths.iter().map(|p| image_similarity::open_image(p).expect("decode failed")).collect();
+		let decode_ms = start.elapsed().as_secs_f64() * 1000.0 / images.len() as f64;
+		let avg_pixels: u64 = images.iter().map(|i| (i.width() * i.height()) as u64).sum::<u64>() / images.len() as u64;
+		println!("  {decode_ms:>8.2} ms/image (avg {avg_pixels} pixels)");
+
+		println!("== describe (full-size input, cycling {} images)", images.len());
+		let mut describe_total = 0.0;
+		for descriptor in get_all_descriptors() {
+			let mut i = 0;
+			let t = ms(12, || {
+				std::hint::black_box(descriptor.describe(&images[i % images.len()]));
+				i += 1;
+			});
+			describe_total += t;
+			println!("  {:<8} {t:>8.2} ms/image", descriptor.info());
+		}
+
+		println!("== mutate (full-size input)");
+		let mut mutate_total = 0.0;
+		for mutator in get_all_mutators() {
+			let mut i = 0;
+			let t = ms(4, || {
+				std::hint::black_box(mutator.mutate(&images[i % images.len()]));
+				i += 1;
+			});
+			mutate_total += t;
+			println!("  {:<18} {t:>8.2} ms/image", mutator.info());
+		}
+
+		// The real per-image cost: mutate everything, describe every version
+		println!("== pipeline (per image: {} mutators, {} descriptors on every version)", get_all_mutators().len(), get_all_descriptors().len());
+		let descriptors = get_all_descriptors();
+		let mutators = get_all_mutators();
+		let mut i = 0;
+		let pipeline_ms = ms(3, || {
+			let img = &images[i % images.len()];
+			for d in &descriptors {
+				std::hint::black_box(d.describe(img));
+			}
+			for m in &mutators {
+				let mutated = m.mutate(img);
+				for d in &descriptors {
+					std::hint::black_box(d.describe(&mutated));
+				}
+			}
+			i += 1;
+		});
+		println!("  hash work:   {pipeline_ms:>8.1} ms/image (single thread)");
+		println!("  with decode: {:>8.1} ms/image (single thread)", pipeline_ms + decode_ms);
+
+		// Same work served from the thumbnail cache: no decode, no mutation
+		use image_similarity::cache::{self, ThumbCache};
+		let cache_dir = std::env::temp_dir().join(format!("imgsim-bench-cache-{}", std::process::id()));
+		let cache = ThumbCache::new(&cache_dir).expect("cannot create cache dir");
+		let img = &images[0];
+		let mut thumbs = std::collections::HashMap::new();
+		thumbs.insert("bench.jpg".to_string(), cache::thumbs_of(img));
+		for m in &mutators {
+			thumbs.insert(format!("mut{}bench.jpg", m.tag()), cache::thumbs_of(&m.mutate(img)));
+		}
+		cache.save("bench.jpg", &thumbs).unwrap();
+		let warm_ms = ms(50, || {
+			let thumbs = cache.load("bench.jpg");
+			for (small, large) in thumbs.values() {
+				for d in &descriptors {
+					let (pixels, size) = match d.input_size() {
+						cache::SMALL => (small, cache::SMALL),
+						_ => (large, cache::LARGE),
+					};
+					std::hint::black_box(d.describe(&cache::to_image(pixels, size).unwrap()));
+				}
+			}
+		});
+		println!("  warm cache:  {warm_ms:>8.2} ms/image (single thread, {} versions from disk)", thumbs.len());
+		let _ = std::fs::remove_dir_all(cache_dir);
+		println!("SUMMARY per_image_ms={:.1} warm_cache_ms={warm_ms:.2} decode_ms={decode_ms:.1} mutate_ms={mutate_total:.1} describe_ms={describe_total:.1}", pipeline_ms + decode_ms);
+	}
+
+	println!("== bk-tree (uniform random keys)");
+	for &n in &sizes {
+		let keys = lcg_keys(n, 0x2545f4914f6cdd1d);
+		// half the queries are existing keys, half perturbed versions
+		let queries: Vec<u64> = keys
+			.iter()
+			.take(cfg.queries)
+			.enumerate()
+			.map(|(i, &k)| if i % 2 == 0 { k } else { k ^ 0b1011 })
+			.collect();
+
+		let start = Instant::now();
+		let mut tree = BkTree::new();
+		for &k in &keys {
+			tree.insert(k);
+		}
+		let build = start.elapsed().as_secs_f64();
+
+		let start = Instant::now();
+		let mut reference = RefNode { key: keys[0], children: Vec::new() };
+		for &k in &keys[1..] {
+			reference.insert(k);
+		}
+		let build_ref = start.elapsed().as_secs_f64();
+
+		println!("  n={n}: bk build {build:.2}s, boxed build {build_ref:.2}s");
+		bench_tree_queries("bk   ", &|q, r| tree.find(q, r).len(), &queries, &radii);
+		bench_tree_queries(
+			"boxed",
+			&|q, r| {
+				let mut out = Vec::new();
+				reference.find(q, r, &mut out);
+				out.len()
+			},
+			&queries,
+			&radii,
+		);
+	}
+
+	if let Some(path) = &cfg.real_store {
+		println!("== bk-tree (real keys from {})", path.display());
+		let bytes = std::fs::read(path).expect("cannot read store");
+		let map: HashMap<u64, Vec<String>> = rmp_serde::from_slice(&bytes)
+			.or_else(|_| {
+				rmp_serde::from_slice::<(u32, String, u32, HashMap<u64, Vec<String>>)>(&bytes).map(|f| f.3)
+			})
+			.expect("unrecognized store format");
+		let keys: Vec<u64> = map.keys().copied().collect();
+		let mut tree = BkTree::new();
+		let start = Instant::now();
+		for &k in &keys {
+			tree.insert(k);
+		}
+		println!("  n={}: build {:.2}s", keys.len(), start.elapsed().as_secs_f64());
+		let queries: Vec<u64> = keys.iter().take(cfg.queries).copied().collect();
+		bench_tree_queries("real ", &|q, r| tree.find(q, r).len(), &queries, &radii);
+	}
+
+	println!("== store (insert + checkpoint save)");
+	let n = 200_000;
+	let keys = lcg_keys(n, 0x9e3779b97f4a7c15);
+	let tmp = std::env::temp_dir().join(format!("imgsim-bench-{}.store", std::process::id()));
+	let _ = std::fs::remove_file(&tmp);
+	let mut store = DescriptorStore::new(Box::new(image_similarity::descriptors::Median))
+		.with_file(&tmp)
+		.unwrap();
+	let start = Instant::now();
+	for (i, &k) in keys.iter().enumerate() {
+		// shifted keys collide, so buckets hold several names like real stores
+		store.insert(k >> 3, format!("mut.jpeg50.im{i}.jpg"));
+	}
+	let insert = start.elapsed().as_secs_f64();
+	println!("  insert {n}: {insert:.2}s ({:.0} items/s)", n as f64 / insert);
+	let start = Instant::now();
+	store.save().unwrap();
+	let save = start.elapsed().as_secs_f64() * 1000.0;
+	let bytes = std::fs::metadata(&tmp).map(|m| m.len()).unwrap_or(0);
+	println!("  save {n}: {save:.0} ms ({:.1} MB)", bytes as f64 / 1e6);
+	println!("SUMMARY store_insert_per_s={:.0} store_save_ms_at_200k={save:.0}", n as f64 / insert);
+	let _ = std::fs::remove_file(tmp);
+}

src/bin/dctfilename.rs

@@ -17,7 +17,7 @@ fn main() {
 
     let desc = DCT::new();
 
-    let img = image::open(cfg.path)
+    let img = image_similarity::open_image(cfg.path)
         .expect("Unable to open file");
     let hash: u64 = desc.describe(&img);
     debug!("Hash integer:\n{hash}");

src/bin/dctquery.rs

@@ -45,7 +45,7 @@ fn main() {
 		std::process::exit(1);
 	}
 
-	let img = image::open(&cfg.path).expect("Unable to open file");
+	let img = image_similarity::open_image(&cfg.path).expect("Unable to open file");
 	let hash = store.descriptor.describe(&img);
 	println!("query hash: {hash:016x}");
 

src/bin/mutate.rs

@@ -14,7 +14,7 @@ fn main() {
 
 	let cfg = Cfg::parse();
 
-	let img = image::open(cfg.path)
+	let img = image_similarity::open_image(cfg.path)
 		.expect("Unable to open file");
 
 	for mutator in get_all_mutators() {

src/bk.rs

@@ -0,0 +1,135 @@
+//! BK-tree over 64-bit hashes with hamming distance.
+//! Children are stored by value inside their parent, which keeps subtree
+//! traversal cache-friendly (benchmarked ~2x faster at a million keys than
+//! an index-based arena). Plain data, so it is Sync and queries can run from
+//! multiple threads. Rebuilt from the store map on load, never serialized.
+
+fn hamming(a: u64, b: u64) -> u64 {
+	(a ^ b).count_ones() as u64
+}
+
+struct Node {
+	key: u64,
+	/// (distance to child, child), distances are unique within one node
+	children: Vec<(u8, Node)>,
+}
+
+impl Node {
+	fn insert(&mut self, key: u64) {
+		let distance = hamming(self.key, key) as u8;
+		if distance == 0 {
+			return;
+		}
+		match self.children.iter_mut().find(|(d, _)| *d == distance) {
+			Some((_, child)) => child.insert(key),
+			None => self.children.push((distance, Node { key, children: Vec::new() })),
+		}
+	}
+
+	fn find(&self, key: u64, radius: u64, found: &mut Vec<(u64, u64)>) {
+		let distance = hamming(self.key, key);
+		if distance <= radius {
+			found.push((self.key, distance));
+		}
+		// triangle inequality: a child at distance d from this node can only
+		// contain matches when |d - distance| <= radius
+		for (child_distance, child) in &self.children {
+			let child_distance = *child_distance as u64;
+			if child_distance + radius >= distance && child_distance <= distance + radius {
+				child.find(key, radius, found);
+			}
+		}
+	}
+
+	fn count(&self) -> usize {
+		1 + self.children.iter().map(|(_, c)| c.count()).sum::<usize>()
+	}
+}
+
+#[derive(Default)]
+pub struct BkTree {
+	root: Option<Node>,
+}
+
+impl BkTree {
+	pub fn new() -> Self {
+		BkTree { root: None }
+	}
+
+	pub fn len(&self) -> usize {
+		self.root.as_ref().map_or(0, Node::count)
+	}
+
+	pub fn is_empty(&self) -> bool {
+		self.root.is_none()
+	}
+
+	/// Inserts a key. Duplicate keys are ignored.
+	pub fn insert(&mut self, key: u64) {
+		match &mut self.root {
+			Some(root) => root.insert(key),
+			None => self.root = Some(Node { key, children: Vec::new() }),
+		}
+	}
+
+	/// All keys within the given hamming radius, with their distances
+	pub fn find(&self, key: u64, radius: u64) -> Vec<(u64, u64)> {
+		let mut found = Vec::new();
+		if let Some(root) = &self.root {
+			root.find(key, radius, &mut found);
+		}
+		found
+	}
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+
+	fn lcg_keys(n: usize) -> Vec<u64> {
+		let mut state: u64 = 0x853c49e6748fea9b;
+		(0..n)
+			.map(|_| {
+				state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
+				state
+			})
+			.collect()
+	}
+
+	#[test]
+	fn matches_brute_force() {
+		let keys = lcg_keys(500);
+		let mut tree = BkTree::new();
+		for &key in &keys {
+			tree.insert(key);
+		}
+		for &query in keys.iter().step_by(37).chain(lcg_keys(520)[500..].iter()) {
+			for radius in [0, 4, 16, 32] {
+				let mut expected: Vec<(u64, u64)> = keys
+					.iter()
+					.map(|&k| (k, hamming(k, query)))
+					.filter(|&(_, d)| d <= radius)
+					.collect();
+				let mut got = tree.find(query, radius);
+				expected.sort();
+				got.sort();
+				assert_eq!(got, expected, "query {query} radius {radius}");
+			}
+		}
+	}
+
+	#[test]
+	fn duplicates_are_ignored() {
+		let mut tree = BkTree::new();
+		tree.insert(42);
+		tree.insert(42);
+		tree.insert(7);
+		assert_eq!(tree.len(), 2);
+		assert_eq!(tree.find(42, 0), vec![(42, 0)]);
+	}
+
+	#[test]
+	fn empty_tree_finds_nothing() {
+		assert!(BkTree::new().find(1, 64).is_empty());
+	}
+}

src/cache.rs

@@ -0,0 +1,135 @@
+//! Disk cache for descriptor-input thumbnails, keyed on filename.
+//!
+//! One cache file per base image holds the small grayscale versions of the
+//! image and every mutant. Re-hashing after a descriptor change then skips
+//! all decoding and mutating, which is where nearly all the time goes.
+//!
+//! Filename-keyed on purpose: keep paths non-colliding. Stale entries are
+//! not detected when file contents change under the same name.
+
+use image::DynamicImage;
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::path::PathBuf;
+
+/// Bump when the resize pipeline or the stored sizes change
+pub const THUMB_FORMAT: u32 = 1;
+
+/// The two descriptor input sizes we store per version
+pub const SMALL: u32 = 8;
+pub const LARGE: u32 = 32;
+
+/// Grayscale pixels of one image version: (8x8, 32x32)
+pub type CachedThumbs = (Vec<u8>, Vec<u8>);
+
+#[derive(Serialize)]
+struct CacheFileRef<'a> {
+	format: u32,
+	thumbs: &'a HashMap<String, CachedThumbs>,
+}
+
+#[derive(Deserialize)]
+struct CacheFile {
+	format: u32,
+	thumbs: HashMap<String, CachedThumbs>,
+}
+
+pub struct ThumbCache {
+	dir: PathBuf,
+}
+
+impl ThumbCache {
+	pub fn new<P: Into<PathBuf>>(dir: P) -> std::io::Result<Self> {
+		let dir = dir.into();
+		std::fs::create_dir_all(&dir)?;
+		Ok(ThumbCache { dir })
+	}
+
+	fn file(&self, name: &str) -> PathBuf {
+		self.dir.join(format!("{name}.thumbs"))
+	}
+
+	/// Cached thumbnails for a base image and its mutants, keyed by version name
+	pub fn load(&self, name: &str) -> HashMap<String, CachedThumbs> {
+		let Ok(bytes) = std::fs::read(self.file(name)) else {
+			return HashMap::new();
+		};
+		match rmp_serde::from_slice::<CacheFile>(&bytes) {
+			Ok(file) if file.format == THUMB_FORMAT => file.thumbs,
+			_ => HashMap::new(), // unknown format: recompute
+		}
+	}
+
+	pub fn save(&self, name: &str, thumbs: &HashMap<String, CachedThumbs>) -> std::io::Result<()> {
+		let file = CacheFileRef { format: THUMB_FORMAT, thumbs };
+		let bytes = rmp_serde::to_vec(&file)
+			.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
+		let tmp = self.file(name).with_extension("tmp");
+		std::fs::write(&tmp, bytes)?;
+		std::fs::rename(tmp, self.file(name))
+	}
+}
+
+/// The descriptor inputs of an image version
+pub fn thumbs_of(img: &DynamicImage) -> CachedThumbs {
+	let small = img.thumbnail_exact(SMALL, SMALL).grayscale().into_luma8();
+	let large = img.thumbnail_exact(LARGE, LARGE).grayscale().into_luma8();
+	(small.into_raw(), large.into_raw())
+}
+
+/// Rebuilds a grayscale image of the given size from cached pixels
+pub fn to_image(pixels: &[u8], size: u32) -> Option<DynamicImage> {
+	let buffer = image::ImageBuffer::from_raw(size, size, pixels.to_vec())?;
+	Some(DynamicImage::ImageLuma8(buffer))
+}
+
+#[cfg(test)]
+mod tests {
+	use super::*;
+	use crate::descriptors::{get_all_descriptors, testimg};
+
+	#[test]
+	fn roundtrip() {
+		let dir = std::env::temp_dir().join(format!("imgsim-cache-{}", std::process::id()));
+		let cache = ThumbCache::new(&dir).unwrap();
+		let mut thumbs = HashMap::new();
+		thumbs.insert("foo.jpg".to_string(), thumbs_of(&testimg::gradient(64)));
+		thumbs.insert("mut.flip.foo.jpg".to_string(), thumbs_of(&testimg::checkerboard(64)));
+		cache.save("foo.jpg", &thumbs).unwrap();
+		assert_eq!(cache.load("foo.jpg"), thumbs);
+		assert!(cache.load("other.jpg").is_empty());
+		std::fs::remove_dir_all(dir).unwrap();
+	}
+
+	#[test]
+	fn cached_thumbs_hash_identically() {
+		// the whole point of the cache: describing a cached thumbnail must
+		// give the same hash as describing the original image
+		for img in [
+			testimg::gradient(64),
+			testimg::colorful(100),
+			crate::open_image("img/meowl.jpg").unwrap(),
+		] {
+			let (small, large) = thumbs_of(&img);
+			for descriptor in get_all_descriptors() {
+				let (pixels, size) = match descriptor.input_size() {
+					SMALL => (&small, SMALL),
+					_ => (&large, LARGE),
+				};
+				let thumb = to_image(pixels, size).unwrap();
+				assert_eq!(
+					descriptor.describe(&img),
+					descriptor.describe(&thumb),
+					"{} differs through cache",
+					descriptor.info()
+				);
+			}
+		}
+	}
+
+	#[test]
+	fn rejects_wrong_pixel_count() {
+		assert!(to_image(&[0u8; 63], 8).is_none());
+		assert!(to_image(&[0u8; 64], 8).is_some());
+	}
+}

src/descriptors/mod.rs

@@ -8,8 +8,15 @@ pub trait Descriptor: Send + Sync {
     fn info(&self) -> String;
     /// Bump whenever the hash output changes, so stores can detect stale data
     fn version(&self) -> u32;
+    /// Side length of the square grayscale image this descriptor works on,
+    /// also decides which cached thumbnail can stand in for the full image
+    fn input_size(&self) -> u32 {
+        8
+    }
+    // thumbnail before grayscale: same weighted sums, but the grayscale
+    // conversion then runs on 64 pixels instead of the full image
     fn resize(&self, img: &image::DynamicImage) -> image::DynamicImage {
-        img.grayscale().thumbnail_exact(8, 8)
+        img.thumbnail_exact(8, 8).grayscale()
     }
     fn describe(&self, img: &image::DynamicImage) -> u64;
     fn distance(&self, a: u64, b: u64) -> u64 {

src/descriptors/phash.rs

@@ -14,10 +14,9 @@ impl PHash {
     /// The 8x8 low-frequency block of the 32x32 DCT, unnormalized DCT-II,
     /// computed separably (rows then columns)
     pub fn lowfreq(&self, img: &image::DynamicImage) -> [f64; 64] {
-        let gray = img
+        // area-averaged downscale, close to the ANTIALIAS resize imagehash uses
-            .grayscale()
+        // but a single pass over the source instead of a full Lanczos kernel
-            .resize_exact(SIZE as u32, SIZE as u32, image::imageops::FilterType::Lanczos3)
+        let gray = img.thumbnail_exact(SIZE as u32, SIZE as u32).grayscale().into_luma8();
-            .into_luma8();
 
         // Row pass: keep the first KEEP coefficients of every row
         let mut rows = [[0.0f64; KEEP]; SIZE];
@@ -56,6 +55,10 @@ impl Descriptor for PHash {
         1
     }
 
+    fn input_size(&self) -> u32 {
+        SIZE as u32
+    }
+
     fn describe(&self, img: &image::DynamicImage) -> u64 {
         let block = self.lowfreq(img);
         let mut sortable = [0i64; 64];

src/lib.rs

@@ -2,7 +2,18 @@
 //!
 //! Descriptors map an image to a compact 64-bit hash.
 //! If two images are (almost) the same, their hashes will be close in Hamming distance.
+pub mod bk;
+pub mod cache;
 pub mod descriptors;
 pub mod store;
 pub mod mutators;
 pub mod wasm;
+
+/// Opens an image by sniffing its content instead of trusting the file
+/// extension. Mirflickr contains a handful of PNGs disguised as .jpg.
+pub fn open_image<P: AsRef<std::path::Path>>(path: P) -> image::ImageResult<image::DynamicImage> {
+    image::io::Reader::open(path)?
+        .with_guessed_format()
+        .map_err(image::ImageError::IoError)?
+        .decode()
+}

src/main.rs

@@ -1,7 +1,8 @@
+use image_similarity::cache::{self, ThumbCache};
 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, info, warn};
 use std::path::PathBuf;
 use std::fs;
 use clap::Parser;
@@ -16,9 +17,21 @@ struct Cfg {
     /// Highest hamming distance to sweep in the PR curves
     #[arg(long, default_value_t = 24)]
     max_threshold: usize,
-    /// Images per parallel batch between store saves
+    /// Images per parallel batch
-    #[arg(long, default_value_t = 32)]
+    #[arg(long, default_value_t = 64)]
     batch_size: usize,
+    /// Seconds between checkpoint saves of the stores
+    #[arg(long, default_value_t = 60)]
+    save_every: u64,
+    /// Only process the first N new images (0 = all)
+    #[arg(long, default_value_t = 0)]
+    limit: usize,
+    /// Directory for cached thumbnails of every image version
+    #[arg(long, default_value = "thumbcache")]
+    thumb_cache: PathBuf,
+    /// Skip the thumbnail cache entirely
+    #[arg(long, default_value_t = false)]
+    no_cache: bool,
 }
 
 fn main() {
@@ -50,6 +63,9 @@ fn main() {
         .expect("Error reading directory")
         .filter_map(|node| {
             let file = node.expect("Error walking directory");
+            if !file.path().is_file() {
+                return None;
+            }
             match file.file_name().into_string() {
                 Ok(name) => Some((name, file.path())),
                 Err(e) => {
@@ -61,14 +77,40 @@ fn main() {
         .collect();
     entries.sort();
 
-    // We assume that if a base image exists in the store, the mutated images also exist
+    // Resume: skip an image only when every store already has it. A run that
-    let todo: Vec<(String, PathBuf)> = entries
+    // died between checkpoint saves leaves the stores unequally far along,
+    // reprocessing those images is cheap and inserts deduplicate.
+    let mut todo: Vec<(String, PathBuf)> = entries
         .into_iter()
-        .filter(|(name, _)| !stores.iter().any(|store| store.has_value(name)))
+        .filter(|(name, _)| !stores.iter().all(|store| store.has_value(name)))
         .collect();
+    if cfg.limit > 0 {
+        todo.truncate(cfg.limit);
+    }
     info!("{} new images to process", todo.len());
 
+    // The cache serves the two sizes the current descriptors consume
+    let cache_supported = stores.iter().all(|store| {
+        [cache::SMALL, cache::LARGE].contains(&store.descriptor.input_size())
+    });
+    let thumb_cache = if cfg.no_cache {
+        None
+    } else if !cache_supported {
+        warn!("a descriptor needs an uncached input size, thumbnail cache disabled");
+        None
+    } else {
+        match ThumbCache::new(&cfg.thumb_cache) {
+            Ok(cache) => Some(cache),
+            Err(e) => {
+                warn!("cannot use thumbnail cache at {}: {e}", cfg.thumb_cache.display());
+                None
+            }
+        }
+    };
+
     let mut done = 0;
+    let phase_start = std::time::Instant::now();
+    let mut last_save = std::time::Instant::now();
     for batch in todo.chunks(cfg.batch_size.max(1)) {
         // Hash batches in parallel, insert on the main thread.
         // Only the descriptors cross threads, the stores themselves are not Sync.
@@ -77,22 +119,56 @@ fn main() {
         let hashes: Vec<Vec<(usize, String, u64)>> = batch
             .par_iter()
             .filter_map(|(name, path)| {
-                let img = match image::open(path) {
+                // (version name, mutator index) for the base image and every mutant
-                    Ok(v) => v,
+                let versions: Vec<(String, Option<usize>)> = std::iter::once((name.clone(), None))
-                    Err(e) => {
+                    .chain(mutators.iter().enumerate().map(|(m, mutator)| {
-                        error!("Failed to process {}: {}", name, e);
+                        (format!("mut{}{}", mutator.tag(), name), Some(m))
-                        return None;
+                    }))
+                    .collect();
+
+                let mut thumbs = thumb_cache.as_ref().map(|c| c.load(name)).unwrap_or_default();
+                // damaged entries count as missing and get recomputed
+                thumbs.retain(|_, (small, large)| {
+                    small.len() == (cache::SMALL * cache::SMALL) as usize
+                        && large.len() == (cache::LARGE * cache::LARGE) as usize
+                });
+                let missing: Vec<&(String, Option<usize>)> = versions
+                    .iter()
+                    .filter(|(version, _)| !thumbs.contains_key(version))
+                    .collect();
+
+                if !missing.is_empty() {
+                    let img = match image_similarity::open_image(path) {
+                        Ok(v) => v,
+                        Err(e) => {
+                            error!("Failed to process {}: {}", name, e);
+                            return None;
+                        }
+                    };
+                    for (version, mutator) in missing {
+                        let cached = match mutator {
+                            Some(m) => cache::thumbs_of(&mutators[*m].mutate(&img)),
+                            None => cache::thumbs_of(&img),
+                        };
+                        thumbs.insert(version.clone(), cached);
+                    }
+                    if let Some(cache) = &thumb_cache {
+                        if let Err(e) = cache.save(name, &thumbs) {
+                            warn!("cannot cache thumbnails for {name}: {e}");
+                        }
                     }
-                };
-                let mut out = Vec::new();
-                for (i, descriptor) in descriptors.iter().enumerate() {
-                    out.push((i, name.clone(), descriptor.describe(&img)));
                 }
-                for mutator in &mutators {
+
-                    let mutated = mutator.mutate(&img);
+                let mut out = Vec::new();
-                    let mutated_name = format!("mut{}{}", mutator.tag(), name);
+                for (version, _) in &versions {
+                    let (small, large) = &thumbs[version];
                     for (i, descriptor) in descriptors.iter().enumerate() {
-                        out.push((i, mutated_name.clone(), descriptor.describe(&mutated)));
+                        let (pixels, size) = match descriptor.input_size() {
+                            cache::SMALL => (small, cache::SMALL),
+                            _ => (large, cache::LARGE),
+                        };
+                        let thumb = cache::to_image(pixels, size).expect("corrupt cached thumbnail");
+                        out.push((i, version.clone(), descriptor.describe(&thumb)));
                     }
                 }
                 Some(out)
@@ -105,12 +181,19 @@ fn main() {
                 stores[i].insert(hash, name);
             }
         }
-        for store in &stores {
+        if last_save.elapsed().as_secs() >= cfg.save_every {
-            store.save().expect("Error saving store");
+            for store in &stores {
+                store.save().expect("Error saving store");
+            }
+            last_save = std::time::Instant::now();
         }
         done += batch.len();
         info!("{done}/{} images done", todo.len());
     }
+    for store in &stores {
+        store.save().expect("Error saving store");
+    }
+    info!("hashing phase took {:.1?}", phase_start.elapsed());
 
     for store in &stores {
         println!("{store}");
@@ -120,7 +203,9 @@ fn main() {
     let max_threshold = cfg.max_threshold.clamp(1, 64);
     for store in &stores {
         println!("Stats for {}", store.descriptor.info());
+        let stats_start = std::time::Instant::now();
         let stats = store.get_stats(&mutators, max_threshold);
+        info!("stats for {} took {:.1?}", store.descriptor.info(), stats_start.elapsed());
         let filename = format!("{}-pr.png", store.descriptor.info());
         let root = BitMapBackend::new(&filename, (1024, 768)).into_drawing_area();
         root.fill(&WHITE).unwrap();

src/store.rs

@@ -2,6 +2,7 @@
 //! Can be queried to find identical or similar images.
 
 use std::collections::HashMap;
+use crate::bk::BkTree;
 use crate::descriptors::Descriptor;
 use crate::mutators::Mutator;
 use std::fs;
@@ -9,7 +10,6 @@ use std::path::Path;
 use std::fmt;
 use image::DynamicImage;
 use log::{info, warn, error};
-use bktree::*;
 use serde::{Serialize, Deserialize};
 
 /// Bump when the on-disk layout of StoreFile changes
@@ -95,6 +95,41 @@ impl PRStats {
 	}
 }
 
+/// Per-thread accumulator for get_stats: hit counts per exact distance,
+/// per mutator, plus the shared unmutated false positives
+struct StatCounts {
+	tp: Vec<[u64; 64]>,
+	fp: Vec<[u64; 64]>,
+	fp_base: [u64; 64],
+}
+
+impl StatCounts {
+	fn new(mutators: usize) -> Self {
+		StatCounts {
+			tp: vec![[0; 64]; mutators],
+			fp: vec![[0; 64]; mutators],
+			fp_base: [0; 64],
+		}
+	}
+
+	fn merge(mut self, other: Self) -> Self {
+		for (mine, theirs) in self.tp.iter_mut().zip(&other.tp) {
+			for (a, b) in mine.iter_mut().zip(theirs) {
+				*a += b;
+			}
+		}
+		for (mine, theirs) in self.fp.iter_mut().zip(&other.fp) {
+			for (a, b) in mine.iter_mut().zip(theirs) {
+				*a += b;
+			}
+		}
+		for (a, b) in self.fp_base.iter_mut().zip(&other.fp_base) {
+			*a += b;
+		}
+		self
+	}
+}
+
 /// Uses a hashmap to map hashes to buckets of files.
 /// Also keeps a BK-tree for quick distance ranking
 /// and an inverted index from filename to hash.
@@ -110,7 +145,7 @@ pub struct DescriptorStore {
 	save_location: std::path::PathBuf,
 
 	/// BK-tree for fast nearest neighbour
-	bktree: BkTree<u64>,
+	bktree: BkTree,
 
 	/// Descriptor version that produced the data currently in the map.
 	/// 0 means unknown: loaded from a legacy store without metadata.
@@ -126,7 +161,7 @@ impl DescriptorStore {
 			names: HashMap::new(),
 			descriptor,
 			save_location: std::path::PathBuf::from("store.messagepack"),
-			bktree: BkTree::new(hamming_distance),
+			bktree: BkTree::new(),
 			data_version,
 		}
 	}
@@ -203,7 +238,10 @@ impl DescriptorStore {
 		};
 		let serialized = rmp_serde::to_vec(&file)
 			.map_err(|e| StoreError::Serialization(e.to_string()))?;
-		fs::write(&self.save_location, serialized).map_err(StoreError::Io)
+		// write-then-rename so a crash mid-save cannot corrupt the store
+		let tmp = self.save_location.with_extension("tmp");
+		fs::write(&tmp, serialized).map_err(StoreError::Io)?;
+		fs::rename(&tmp, &self.save_location).map_err(StoreError::Io)
 	}
 
 	/// Returns true iff the store already contains the key
@@ -262,7 +300,7 @@ impl DescriptorStore {
 				continue;
 			}
 			info!("Processing {}", name);
-			let img = match image::open(file.path()) {
+			let img = match crate::open_image(file.path()) {
 				Ok(v) => v,
 				Err(e) => {
 					error!("Failed to process {}: {}", name, e);
@@ -279,9 +317,9 @@ impl DescriptorStore {
 	/// sorted nearest first
 	pub fn query(&self, from: u64, max_distance: u64) -> Vec<(String, u64)> {
 		let mut results = Vec::new();
-		for (key, distance) in self.bktree.find(from, max_distance as isize) {
+		for (key, distance) in self.bktree.find(from, max_distance) {
-			for name in &self.map[key] {
+			for name in &self.map[&key] {
-				results.push((name.clone(), distance as u64));
+				results.push((name.clone(), distance));
 			}
 		}
 		results.sort_by(|a, b| a.1.cmp(&b.1).then_with(|| a.0.cmp(&b.0)));
@@ -299,35 +337,42 @@ impl DescriptorStore {
 	pub fn get_stats(&self, mutators: &[Box<dyn Mutator>], max_threshold: usize) -> Vec<PRStats> {
 		let max_t = max_threshold.clamp(1, 64);
 		let prefixes: Vec<String> = mutators.iter().map(|m| format!("mut{}", m.tag())).collect();
-
-		let mut tp_at = vec![[0u64; 64]; mutators.len()];
-		let mut fp_at = vec![[0u64; 64]; mutators.len()];
-		let mut fp_base_at = [0u64; 64];
-
 		let bases: Vec<&String> = self.names.keys().filter(|n| !n.starts_with("mut.")).collect();
-		for base in &bases {
+
+		let tally = |mut counts: StatCounts, base: &&String| {
 			let hash = self.names[*base];
-			let expected: Vec<String> = prefixes.iter().map(|p| format!("{p}{base}")).collect();
+			for (key, distance) in self.bktree.find(hash, (max_t - 1) as u64) {
-			for (key, distance) in self.bktree.find(hash, (max_t - 1) as isize) {
 				let d = distance as usize;
-				for name in &self.map[key] {
+				for name in &self.map[&key] {
 					if name == *base {
 						continue;
 					}
-					if let Some(m) = expected.iter().position(|e| e == name) {
+					if let Some(m) = prefixes.iter().position(|p| name.starts_with(p.as_str())) {
-						tp_at[m][d] += 1;
+						if name[prefixes[m].len()..] == ***base {
-					} else if name.starts_with("mut") {
+							counts.tp[m][d] += 1;
-						// Mutated misses only count for their own mutator
+						} else {
-						if let Some(m) = prefixes.iter().position(|p| name.starts_with(p.as_str())) {
+							// Mutated misses only count for their own mutator
-							fp_at[m][d] += 1;
+							counts.fp[m][d] += 1;
 						}
-					} else {
+					} else if !name.starts_with("mut") {
 						// Unmutated misses count for every mutator
-						fp_base_at[d] += 1;
+						counts.fp_base[d] += 1;
 					}
 				}
 			}
-		}
+			counts
+		};
+
+		#[cfg(feature = "cli")]
+		let counts = {
+			use rayon::prelude::*;
+			bases
+				.par_iter()
+				.fold(|| StatCounts::new(mutators.len()), tally)
+				.reduce(|| StatCounts::new(mutators.len()), StatCounts::merge)
+		};
+		#[cfg(not(feature = "cli"))]
+		let counts = bases.iter().fold(StatCounts::new(mutators.len()), tally);
 
 		let cumulative = |at: &[u64; 64]| {
 			let mut cum = [0u64; 64];
@@ -339,11 +384,11 @@ impl DescriptorStore {
 			cum
 		};
 
-		let fp_base = cumulative(&fp_base_at);
+		let fp_base = cumulative(&counts.fp_base);
 		let total = bases.len() as u64;
 		mutators.iter().enumerate().map(|(m, mutator)| {
-			let true_positives = cumulative(&tp_at[m]);
+			let true_positives = cumulative(&counts.tp[m]);
-			let fp_mut = cumulative(&fp_at[m]);
+			let fp_mut = cumulative(&counts.fp[m]);
 			let mut false_positives = [0u64; 64];
 			let mut false_negatives = [0u64; 64];
 			for t in 0..64 {