diff --git a/.gitignore b/.gitignore index 13ced13..24748fc 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,8 @@ web/colors.css smolflickr data spook.jpg +run*.log +thumbcache/ +legacy/ +*.pdf +run-*.sh diff --git a/Cargo.lock b/Cargo.lock index 369efd5..fbc61f1 100644 --- a/Cargo.lock +++ b/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", diff --git a/Cargo.toml b/Cargo.toml index 6a23d51..9e66993 100644 --- a/Cargo.toml +++ b/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"] diff --git a/src/bin/bench.rs b/src/bin/bench.rs new file mode 100644 index 0000000..1753b6e --- /dev/null +++ b/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, + /// Skip the image benchmarks (decode/describe/mutate) + #[arg(long, default_value_t = false)] + trees_only: bool, +} + +fn ms(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 { + (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 = cfg.sizes.split(',').map(|s| s.trim().parse().expect("bad size")).collect(); + let radii: Vec = cfg.radii.split(',').map(|s| s.trim().parse().expect("bad radius")).collect(); + + if !cfg.trees_only { + let mut paths: Vec = 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 = 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::() / 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 = 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> = rmp_serde::from_slice(&bytes) + .or_else(|_| { + rmp_serde::from_slice::<(u32, String, u32, HashMap>)>(&bytes).map(|f| f.3) + }) + .expect("unrecognized store format"); + let keys: Vec = 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 = 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); +} diff --git a/src/bin/dctfilename.rs b/src/bin/dctfilename.rs index 5b3c5cf..54e0b47 100644 --- a/src/bin/dctfilename.rs +++ b/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}"); diff --git a/src/bin/dctquery.rs b/src/bin/dctquery.rs index da86002..8399c24 100644 --- a/src/bin/dctquery.rs +++ b/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}"); diff --git a/src/bin/mutate.rs b/src/bin/mutate.rs index 09b79a0..08a46ad 100644 --- a/src/bin/mutate.rs +++ b/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() { diff --git a/src/bk.rs b/src/bk.rs new file mode 100644 index 0000000..6530551 --- /dev/null +++ b/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::() + } +} + +#[derive(Default)] +pub struct BkTree { + root: Option, +} + +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 { + 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()); + } +} diff --git a/src/cache.rs b/src/cache.rs new file mode 100644 index 0000000..454c5bf --- /dev/null +++ b/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, Vec); + +#[derive(Serialize)] +struct CacheFileRef<'a> { + format: u32, + thumbs: &'a HashMap, +} + +#[derive(Deserialize)] +struct CacheFile { + format: u32, + thumbs: HashMap, +} + +pub struct ThumbCache { + dir: PathBuf, +} + +impl ThumbCache { + pub fn new>(dir: P) -> std::io::Result { + 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 { + let Ok(bytes) = std::fs::read(self.file(name)) else { + return HashMap::new(); + }; + match rmp_serde::from_slice::(&bytes) { + Ok(file) if file.format == THUMB_FORMAT => file.thumbs, + _ => HashMap::new(), // unknown format: recompute + } + } + + pub fn save(&self, name: &str, thumbs: &HashMap) -> 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 { + 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()); + } +} diff --git a/src/descriptors/mod.rs b/src/descriptors/mod.rs index cdea547..ce54cbb 100644 --- a/src/descriptors/mod.rs +++ b/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 { diff --git a/src/descriptors/phash.rs b/src/descriptors/phash.rs index 0e0f4c6..93016cb 100644 --- a/src/descriptors/phash.rs +++ b/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 - .grayscale() - .resize_exact(SIZE as u32, SIZE as u32, image::imageops::FilterType::Lanczos3) - .into_luma8(); + // area-averaged downscale, close to the ANTIALIAS resize imagehash uses + // but a single pass over the source instead of a full Lanczos kernel + let gray = img.thumbnail_exact(SIZE as u32, SIZE as u32).grayscale().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]; diff --git a/src/lib.rs b/src/lib.rs index 4b3b786..137a88f 100644 --- a/src/lib.rs +++ b/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>(path: P) -> image::ImageResult { + image::io::Reader::open(path)? + .with_guessed_format() + .map_err(image::ImageError::IoError)? + .decode() +} diff --git a/src/main.rs b/src/main.rs index 8aecf59..1b05554 100644 --- a/src/main.rs +++ b/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 - #[arg(long, default_value_t = 32)] + /// Images per parallel batch + #[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 - let todo: Vec<(String, PathBuf)> = entries + // Resume: skip an image only when every store already has it. A run that + // 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> = batch .par_iter() .filter_map(|(name, path)| { - let img = match image::open(path) { - Ok(v) => v, - Err(e) => { - error!("Failed to process {}: {}", name, e); - return None; + // (version name, mutator index) for the base image and every mutant + let versions: Vec<(String, Option)> = std::iter::once((name.clone(), None)) + .chain(mutators.iter().enumerate().map(|(m, mutator)| { + (format!("mut{}{}", mutator.tag(), name), Some(m)) + })) + .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)> = 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 mutated_name = format!("mut{}{}", mutator.tag(), name); + + let mut out = Vec::new(); + 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 { - store.save().expect("Error saving store"); + if last_save.elapsed().as_secs() >= cfg.save_every { + 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(); diff --git a/src/store.rs b/src/store.rs index e42af7a..7b53314 100644 --- a/src/store.rs +++ b/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, + 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 name in &self.map[key] { - results.push((name.clone(), distance as u64)); + for (key, distance) in self.bktree.find(from, max_distance) { + for name in &self.map[&key] { + 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], max_threshold: usize) -> Vec { let max_t = max_threshold.clamp(1, 64); let prefixes: Vec = 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 = prefixes.iter().map(|p| format!("{p}{base}")).collect(); - for (key, distance) in self.bktree.find(hash, (max_t - 1) as isize) { + for (key, distance) in self.bktree.find(hash, (max_t - 1) as u64) { 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) { - tp_at[m][d] += 1; - } else if name.starts_with("mut") { - // Mutated misses only count for their own mutator - if let Some(m) = prefixes.iter().position(|p| name.starts_with(p.as_str())) { - fp_at[m][d] += 1; + if let Some(m) = prefixes.iter().position(|p| name.starts_with(p.as_str())) { + if name[prefixes[m].len()..] == ***base { + counts.tp[m][d] += 1; + } else { + // Mutated misses only count for their own mutator + 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 fp_mut = cumulative(&fp_at[m]); + let true_positives = cumulative(&counts.tp[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 {