diff --git a/.drone.yml b/.drone.yml index 8551740..0cebe5e 100644 --- a/.drone.yml +++ b/.drone.yml @@ -6,6 +6,15 @@ steps: image: rust:latest commands: - cargo build --verbose --all +- name: test + image: rust:latest + commands: + - cargo test --all +- name: lint + image: rust:latest + commands: + - rustup component add clippy + - cargo clippy --all-targets -- -D warnings trigger: event: exclude: diff --git a/Cargo.lock b/Cargo.lock index d925023..369efd5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,6 +1,6 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. -version = 3 +version = 4 [[package]] name = "ab_glyph" @@ -767,7 +767,7 @@ dependencies = [ [[package]] name = "image_similarity" -version = "0.1.0" +version = "0.2.0" dependencies = [ "base64", "bktree", @@ -778,6 +778,7 @@ dependencies = [ "imageproc", "log", "plotters", + "rayon", "rmp-serde", "serde", "wasm-bindgen", diff --git a/Cargo.toml b/Cargo.toml index a9eed59..6a23d51 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,23 +1,50 @@ [package] name = "image_similarity" -version = "0.1.0" +version = "0.2.0" edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [dependencies] image = "0.25.1" -serde = "1.0.147" +imageproc = "0.25.0" +serde = { version = "1.0.147", features = ["derive"] } rmp-serde = "1.1.1" base64 = "0.22.1" log = "0.4.21" -env_logger = "0.11.3" bktree = "1.0.1" -clap = { version = "4.5.4", features = ["derive"] } -plotters = "0.3.6" wasm-bindgen = "0.2.92" console_error_panic_hook = "0.1.7" -imageproc = "0.25.0" + +# cli-only dependencies, skipped for wasm builds (--no-default-features) +clap = { version = "4.5.4", features = ["derive"], optional = true } +plotters = { version = "0.3.6", optional = true } +env_logger = { version = "0.11.3", optional = true } +rayon = { version = "1.10", optional = true } + +[features] +default = ["cli"] +cli = ["dep:clap", "dep:plotters", "dep:env_logger", "dep:rayon"] [lib] crate-type = ["cdylib", "rlib"] + +[[bin]] +name = "image_similarity" +path = "src/main.rs" +required-features = ["cli"] + +[[bin]] +name = "dctquery" +path = "src/bin/dctquery.rs" +required-features = ["cli"] + +[[bin]] +name = "dctfilename" +path = "src/bin/dctfilename.rs" +required-features = ["cli"] + +[[bin]] +name = "mutate" +path = "src/bin/mutate.rs" +required-features = ["cli"] diff --git a/pr_from_store.py b/pr_from_store.py new file mode 100644 index 0000000..4661b65 --- /dev/null +++ b/pr_from_store.py @@ -0,0 +1,120 @@ +"""Recompute PR curves from an image_similarity .store file (messagepack). + +Replicates DescriptorStore::get_stats semantics: +- micro-averaged over base-image queries +- TP_m(t): mutant m of query found within Hamming distance t +- FP: other base images count for every mutator; other images' mutants + count only for their own mutator +""" +import sys +import msgpack +import numpy as np + +TMAX = 33 # thresholds 0..32 + +def load(path): + with open(path, "rb") as f: + m = msgpack.unpack(f, strict_map_key=False) + if isinstance(m, (list, tuple)): + # versioned format: [format, descriptor, descriptor_version, map] + print(f"{path}: {m[1]} v{m[2]} (store format {m[0]})") + return m[3] + return m # legacy format: bare map + +def detect_mutators(buckets): + tags = set() + for bucket in buckets: + for name in bucket: + if name.startswith("mut."): + tags.add(name.split(".")[1]) + return sorted(tags) + +def main(path, label): + m = load(path) + global MUTATORS + MUTATORS = detect_mutators(m.values()) + print(f"mutators: {MUTATORS}") + keys = np.array(list(m.keys()), dtype=np.uint64) + buckets = list(m.values()) + K = len(keys) + + name2hash = {} + n_base = np.zeros(K, dtype=np.float64) + n_mut = {mut: np.zeros(K, dtype=np.float64) for mut in MUTATORS} + for ki, bucket in enumerate(buckets): + for name in bucket: + name2hash[name] = keys[ki] + if name.startswith("mut."): + for mut in MUTATORS: + if name.startswith(f"mut.{mut}."): + n_mut[mut][ki] += 1 + break + else: + n_base[ki] += 1 + + bases = [n for n in name2hash if not n.startswith("mut.")] + N = len(bases) + q = np.array([name2hash[b] for b in bases], dtype=np.uint64) + + # TP_m(t): distance from each base to its own mutant, cumulative over t + tp = {} + for mut in MUTATORS: + d = np.array( + [bin(int(name2hash[b]) ^ int(name2hash[f"mut.{mut}.{b}"])).count("1") + for b in bases]) + tp[mut] = np.cumsum(np.bincount(d, minlength=TMAX)[:TMAX]) + + # Histogram of (query, key) distances weighted by bucket composition + hist_base = np.zeros(TMAX) + hist_mut = {mut: np.zeros(TMAX) for mut in MUTATORS} + CHUNK = 512 + for i in range(0, N, CHUNK): + d = np.bitwise_count(q[i:i + CHUNK, None] ^ keys[None, :]).astype(np.uint8) + flat = d.ravel() + sel = flat < TMAX + flat = flat[sel] + rows = d.shape[0] + hist_base += np.bincount(flat, weights=np.broadcast_to(n_base, (rows, K)).ravel()[sel], minlength=TMAX)[:TMAX] + for mut in MUTATORS: + hist_mut[mut] += np.bincount(flat, weights=np.broadcast_to(n_mut[mut], (rows, K)).ravel()[sel], minlength=TMAX)[:TMAX] + + cum_base = np.cumsum(hist_base) - N # exclude self (d=0 always) + print(f"\n=== {label} ===") + print(f"{'t':>2} | " + " | ".join(f"{mut:>22}" for mut in MUTATORS)) + print(f"{'':>2} | " + " | ".join(f"{'recall':>10} {'precis':>11}" for _ in MUTATORS)) + curves = {} + for mut in MUTATORS: + fp = (np.cumsum(hist_mut[mut]) - tp[mut]) + cum_base + rec = tp[mut] / N + prec = np.divide(tp[mut], tp[mut] + fp, + out=np.zeros(TMAX), where=(tp[mut] + fp) > 0) + curves[mut] = (rec, prec) + for t in range(TMAX): + row = " | ".join(f"{curves[mut][0][t]:>10.4f} {curves[mut][1][t]:>11.6f}" for mut in MUTATORS) + print(f"{t:>2} | {row}") + return curves + +if __name__ == "__main__": + curves_by_store = {} + for path, label in [("dct.store", "DCT"), ("median.store", "Median")]: + full = f"/home/mark/workspace/repos/image-similarity/{path}" + curves_by_store[label] = main(full, label) + + try: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + fig, axes = plt.subplots(1, 2, figsize=(13, 5.5), sharey=True) + for ax, (label, curves) in zip(axes, curves_by_store.items()): + for mut, (rec, prec) in curves.items(): + ax.plot(rec, prec, marker=".", label=mut) + ax.set_title(f"{label} hash — 24,988 Flickr images, thresholds 0–32") + ax.set_xlabel("Recall") + ax.set_ylabel("Precision") + ax.grid(alpha=.3) + ax.legend() + fig.tight_layout() + fig.savefig("/tmp/imgsim/pr-full.png", dpi=110) + print("\nplot: /tmp/imgsim/pr-full.png") + except ImportError: + print("\nmatplotlib not available; table output only") diff --git a/src/bin/dctfilename.rs b/src/bin/dctfilename.rs index a76f56d..5b3c5cf 100644 --- a/src/bin/dctfilename.rs +++ b/src/bin/dctfilename.rs @@ -1,10 +1,10 @@ -//use image_similarity::descriptors::dct::DCT; use image_similarity::descriptors::{DCT, Descriptor}; use std::path::PathBuf; use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _}; use log::debug; use clap::Parser; +/// Print the DCT hash of an image as url-safe base64, e.g. for use in filenames #[derive(Parser)] #[command(version, about, long_about = None)] struct Cfg { @@ -15,13 +15,13 @@ fn main() { env_logger::init(); let cfg = Cfg::parse(); - let desc = DCT::new().with_quality(95); + let desc = DCT::new(); let img = image::open(cfg.path) .expect("Unable to open file"); - let phash: u64 = desc.describe(&img); - debug!("Phash integer:\n{phash}"); - debug!("Phash binary:\n{phash:064b}"); - let output = URL_SAFE_NO_PAD.encode(phash.to_be_bytes()); + let hash: u64 = desc.describe(&img); + debug!("Hash integer:\n{hash}"); + debug!("Hash binary:\n{hash:064b}"); + let output = URL_SAFE_NO_PAD.encode(hash.to_be_bytes()); println!("{output}"); } diff --git a/src/bin/dctquery.rs b/src/bin/dctquery.rs index ecac929..da86002 100644 --- a/src/bin/dctquery.rs +++ b/src/bin/dctquery.rs @@ -1,34 +1,60 @@ -//use image_similarity::descriptors::dct::DCT; -use image_similarity::descriptors::{DCT, Descriptor}; +use image_similarity::descriptors::DCT; use image_similarity::store::DescriptorStore; -use log::info; +use log::error; use clap::Parser; +/// Query a store for images similar to the given image #[derive(Parser)] #[command(version, about, long_about = None)] struct Cfg { - /// Path to the image + /// Path to the query image path: String, - /// Try to output the images to terminal with viuer - #[arg(short, long, default_value_t=false)] - show_images: bool, - distance: usize, + /// Store to query + #[arg(long, default_value = "dct.store")] + store: String, + /// Maximum hamming distance + #[arg(short, long, default_value_t = 10)] + distance: u64, + /// Maximum number of results + #[arg(short, long, default_value_t = 10)] + limit: usize, } fn main() { env_logger::init(); - - let desc = DCT::new().with_quality(95); - let cfg = Cfg::parse(); - let img = image::open(cfg.path) - .expect("Unable to open file"); + let desc = DCT::new(); + let store = match DescriptorStore::new(Box::new(desc)).with_file(&cfg.store) { + Ok(store) => store, + Err(e) => { + error!("{}: {e}", cfg.store); + std::process::exit(1); + } + }; + if !store.version_matches() { + error!( + "{} contains hashes from an older descriptor version, \ + they are not comparable to this query. Regenerate the store first.", + cfg.store + ); + std::process::exit(1); + } + if store.is_empty() { + error!("{} is empty", cfg.store); + std::process::exit(1); + } - let phash: u64 = desc.describe(&img); - let _store = - DescriptorStore::new(Box::new(desc)) - .with_file("dct50.messagepack"); - info!("Phash integer:\n{phash}"); - info!("Phash binary:\n{phash:064b}"); -} \ No newline at end of file + let img = image::open(&cfg.path).expect("Unable to open file"); + let hash = store.descriptor.describe(&img); + println!("query hash: {hash:016x}"); + + let results = store.query(hash, cfg.distance); + if results.is_empty() { + println!("no matches within distance {}", cfg.distance); + return; + } + for (name, distance) in results.iter().take(cfg.limit) { + println!("{distance:>3} {name}"); + } +} diff --git a/src/bin/mutate.rs b/src/bin/mutate.rs index 1774d4c..09b79a0 100644 --- a/src/bin/mutate.rs +++ b/src/bin/mutate.rs @@ -1,14 +1,12 @@ use image_similarity::mutators::get_all_mutators; use clap::Parser; +/// Write every mutated version of an image to the working directory #[derive(Parser)] #[command(version, about, long_about = None)] struct Cfg { /// Path to the image path: String, - /// Try to output the images to terminal with viuer - #[arg(short, long, default_value_t=false)] - show_images: bool, } fn main() { @@ -19,12 +17,10 @@ fn main() { let img = image::open(cfg.path) .expect("Unable to open file"); - let mutators = get_all_mutators(); - - for mutator in mutators { + for mutator in get_all_mutators() { let mutated = mutator.mutate(&img); - let filename = "mut".to_string() + &mutator.tag() + "png"; + let filename = format!("mut{}png", mutator.tag()); println!("Saving to {filename}"); mutated.save(filename).expect("Saving image failed"); } -} \ No newline at end of file +} diff --git a/src/descriptors/dct.rs b/src/descriptors/dct.rs index adbc140..44ce548 100644 --- a/src/descriptors/dct.rs +++ b/src/descriptors/dct.rs @@ -3,44 +3,59 @@ use std::f64::consts::{PI, SQRT_2}; use crate::descriptors::{Descriptor, DCT, print_matrix}; use log::{debug, log_enabled, Level}; +/// First 36 coefficients in zigzag order (everything with u+v <= 7) +/// for a row-major flattened 8x8 matrix +const ZIGZAG: [usize; 36] = [ + 0, + 1, 8, + 16, 9, 2, + 3, 10, 17, 24, + 32, 25, 18, 11, 4, + 5, 12, 19, 26, 33, 40, + 48, 41, 34, 27, 20, 13, 6, + 7, 14, 21, 28, 35, 42, 49, 56, +]; + +/// Number of sign bits taken from the zigzag +const SIGN_BITS: usize = 36; +/// Number of ordinal bits: comparisons between consecutive zigzag pairs +const ORDINAL_BITS: usize = 28; impl DCT { - /// Returns a new DCT instance with a base quality DCT matrix. + /// Returns a new DCT instance with the standard JPEG luminance + /// quantization matrix (only used for debug reconstruction). pub fn new() -> DCT { let quantization_matrix: [u8; 64] = [ 16, 11, 10, 16, 24, 40, 51, 61, 12, 12, 14, 19, 26, 58, 60, 55, 14, 13, 16, 24, 40, 57, 69, 56, 14, 17, 22, 29, 51, 87, 80, 62, - 18, 22, 37, 56, 6, 10, 103, 77, - 24, 35, 55, 64, 8, 10, 113, 92, - 49, 64, 78, 8, 10, 12, 12, 101, - 72, 92, 95, 9, 11, 10, 103, 99, + 18, 22, 37, 56, 68, 109, 103, 77, + 24, 35, 55, 64, 81, 104, 113, 92, + 49, 64, 78, 87, 103, 121, 120, 101, + 72, 92, 95, 98, 112, 100, 103, 99, ]; DCT { quantization_matrix } } - /// Builds DCT with given quality value - /// quality + /// Scales the quantization matrix for a given JPEG-style quality (1-100) pub fn with_quality(mut self, quality: u8) -> Self { - let mut quantization_matrix: [u8; 64] = [0; 64]; let scalar: f32 = match quality { - 1..=49 => 5000.0/quality as f32, - 50..=100 => 200.0 - 2.0*quality as f32, + 1..=49 => 5000.0 / quality as f32, + 50..=100 => 200.0 - 2.0 * quality as f32, _ => 100.0 // Invalid input: set to base quality }; - for (_, cell) in quantization_matrix.iter_mut().enumerate() { - *cell = ((scalar * *cell as f32 + 50.0) / 100.0).floor() as u8; - if *cell == 0 { - *cell = 1; - } + for cell in self.quantization_matrix.iter_mut() { + let scaled = ((scalar * *cell as f32 + 50.0) / 100.0).floor(); + *cell = scaled.max(1.0) as u8; } - self.quantization_matrix = quantization_matrix; self } - /// + /// 2D DCT-II of an 8x8 grayscale image, orthonormal scaling, + /// pixels centered around 0 pub fn dct(&self, img: &image::DynamicImage) -> [f64; 64] { + debug_assert_eq!((img.width(), img.height()), (8, 8)); let mut dct_values: [f64; 64] = [0.0; 64]; for u in 0..8 { for v in 0..8 { @@ -58,9 +73,9 @@ impl DCT { for (x, y, pix) in img.pixels() { let x: f64 = 1.0 + 2.0 * x as f64; let y: f64 = 1.0 + 2.0 * y as f64; - let pixel = (pix[0] as i16 - 127) as f64; + let pixel = (pix[0] as i16 - 128) as f64; sum += - pixel * + pixel * (x*u*PI/16.0).cos() * (y*v*PI/16.0).cos() } @@ -70,9 +85,8 @@ impl DCT { dct_values } - fn idct(&self, dct_values: [f64; 64]) -> [u8; 64] { + pub fn idct(&self, dct_values: [f64; 64]) -> [u8; 64] { let mut reconstructed: [u8; 64] = [0; 64]; - //for k in 0..64 { for (k, item) in reconstructed.iter_mut().enumerate() { let x = (k%8) as f64; let y = (k/8) as f64; @@ -96,11 +110,45 @@ impl DCT { ((2.0 * y + 1.0) * v * PI / 16.0).cos(); } } - sum = 127.0 + (0.25 * sum).round(); - *item = std::cmp::min(255_u8, sum as u8); + sum = 128.0 + (0.25 * sum).round(); + *item = sum.clamp(0.0, 255.0) as u8; } reconstructed } + + /// The two half-masks of the hash: 36 sign bits and 28 ordinal bits. + /// + /// Sign bits: sign of each zigzag coefficient. Coefficients with an odd + /// horizontal frequency are multiplied by the sign of the first horizontal + /// AC coefficient, which makes the mask invariant to horizontal flips. + /// + /// Ordinal bits: whether each zigzag coefficient is larger in magnitude + /// than its predecessor. Magnitudes are unaffected by flips. + pub fn masks(&self, dct_values: &[f64; 64]) -> (u64, u64) { + let sign_mult = dct_values[1].signum(); + let mut sign_mask: u64 = 0; + for &i in ZIGZAG.iter().take(SIGN_BITS) { + let mut signum = dct_values[i].signum(); + // Only flip sign if the coefficient has an odd horizontal component. + // i % 2 == u % 2 for a row-major index i = v*8+u. + if i % 2 != 0 { + signum *= sign_mult; + } + sign_mask <<= 1; + if signum < 0.0 { + sign_mask |= 1; + } + } + + let mut ordinal_mask: u64 = 0; + for pair in ZIGZAG[..=ORDINAL_BITS].windows(2) { + ordinal_mask <<= 1; + if dct_values[pair[1]].abs() > dct_values[pair[0]].abs() { + ordinal_mask |= 1; + } + } + (sign_mask, ordinal_mask) + } } impl Default for DCT { @@ -113,106 +161,143 @@ impl Descriptor for DCT { fn info(&self) -> String { "dct".to_string() } + + // v2: fixed mask packing (v1 lost the DC sign bit and wasted the low 9 bits), + // filled the previously unused 8 bits with the next zigzag diagonal, + // pixels now centered on 128 instead of 127 + fn version(&self) -> u32 { + 2 + } + fn describe(&self, img: &image::DynamicImage) -> u64 { let resized = self.resize(img); let dct_values = self.dct(&resized); 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(); - // 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; - for i in 0..64 { - dequant[i] = dct_values[i] * self.quantization_matrix[i] as f64; + // JPEG-style quantization roundtrip, to show what survives + let mut roundtrip = dct_values; + for (i, value) in roundtrip.iter_mut().enumerate() { + let q = self.quantization_matrix[i] as f64; + *value = (*value / q).round() * q; } - //debug!("DCT-coefficients, de-quantized:\n {}", print_matrix(dct_values)); - - // Reconstruction original pixel values: - let reconstructed = self.idct(dequant); - + let reconstructed = self.idct(roundtrip); save_buffer( "reconstructed.png", &reconstructed, - 8, + 8, 8, image::ColorType::L8 ).expect("Error saving buffer"); } - // Calculating descriptor from dct values: - - // Zigzag order for our flattened array, - // first 28 elements only - let zigzag: [usize; 28] = [ - 0, - 1, 8, - 16, 9, 2, - 3, 10, 17, 24, - 32, 25, 18, 11, 4, - 5, 12, 19, 26, 33, 40, - 48, 41, 34, 27, 20, 13, 6, - ]; - - // Mask that indicates if a dct coefficient is positive or negative - // By convention, when sign bit is 1, number is negative - let mut sign_mask: u64 = 0; - // 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[49].abs(); - - - for i in zigzag { - let cur = dct_values[i].abs(); - if cur > prev { - pearson_mask += 1; - } - prev = cur; - - 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 - signum < 0.0 - { - sign_mask += 1; - } - - // Shift masks - sign_mask <<= 1; - 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); - - let mut mask = sign_mask; - debug!("Mask: {:064b}", mask); - mask <<= 28; - debug!("Mask: {:064b}", mask); - mask += pearson_mask; - debug!("Mask: {:064b}", mask); - mask <<= 8; - debug!("Mask: {:064b}", mask); - // TODO: Do something with these last 8 bits. - mask + let (sign_mask, ordinal_mask) = self.masks(&dct_values); + debug!("Sign mask: {sign_mask:036b}"); + debug!("Ordinal mask: {ordinal_mask:028b}"); + (sign_mask << ORDINAL_BITS) | ordinal_mask + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::descriptors::testimg; + + #[test] + fn zigzag_is_valid() { + let mut seen = [false; 64]; + for &i in ZIGZAG.iter() { + assert!(!seen[i]); + seen[i] = true; + } + assert_eq!(SIGN_BITS + ORDINAL_BITS, 64); + } + + #[test] + fn deterministic() { + let img = testimg::gradient(64); + assert_eq!(DCT::new().describe(&img), DCT::new().describe(&img)); + } + + #[test] + fn invariant_to_horizontal_flip() { + // Exact invariance holds when downsampling is mirror-symmetric, + // i.e. image dimensions are a multiple of 8. Other sizes land close + // but not exactly equal because the sampling grid is asymmetric. + let dct = DCT::new(); + for img in [ + testimg::noise8x8(), + testimg::colorful(128), + testimg::checkerboard(256), + ] { + assert_eq!(dct.describe(&img), dct.describe(&img.fliph())); + } + } + + #[test] + fn invariant_to_horizontal_flip_real_photo() { + let img = image::open("img/meowl.jpg").unwrap(); + let dct = DCT::new(); + assert_eq!(dct.describe(&img), dct.describe(&img.fliph())); + } + + #[test] + fn distinguishes_images() { + let dct = DCT::new(); + assert_ne!( + dct.describe(&testimg::gradient(64)), + dct.describe(&testimg::checkerboard(64)) + ); + } + + #[test] + fn dc_sign_lands_on_top_bit() { + // regression: v1 packing shifted the DC sign bit out of the hash + use image::{DynamicImage, ImageBuffer, Luma}; + let shifted = |offset: u32| { + DynamicImage::ImageLuma8(ImageBuffer::from_fn(8, 8, move |x, y| { + Luma([(offset + x * 4 + y * 3) as u8]) + })) + }; + let dark = DCT::new().describe(&shifted(20)); // mean well below 128 + let bright = DCT::new().describe(&shifted(180)); // mean well above 128 + assert_eq!(dark >> 63, 1); + assert_eq!(bright >> 63, 0); + } + + #[test] + fn masks_pack_exactly() { + // regression: v1 packing left the lowest 9 bits always zero + let mut values = [1.0f64; 64]; + values[0] = -2.0; // DC sign, first sign bit + values[56] = -5.0; // zigzag[35], last sign bit + values[7] = 3.0; // zigzag[28], larger than predecessor: last ordinal bit + let (sign_mask, ordinal_mask) = DCT::new().masks(&values); + let hash = (sign_mask << ORDINAL_BITS) | ordinal_mask; + assert_eq!(hash, (1 << 63) | (1 << 28) | 1); + } + + #[test] + fn idct_roundtrip() { + let dct = DCT::new(); + let resized = dct.resize(&testimg::gradient(64)); + let values = dct.dct(&resized); + let reconstructed = dct.idct(values); + use image::GenericImageView; + for (i, (_, _, pix)) in resized.pixels().enumerate() { + let diff = (pix[0] as i16 - reconstructed[i] as i16).abs(); + assert!(diff <= 1, "pixel {i} off by {diff}"); + } + } + + #[test] + fn quality_scales_quantization() { + let base = DCT::new().quantization_matrix; + let low = DCT::new().with_quality(10).quantization_matrix; + let high = DCT::new().with_quality(95).quantization_matrix; + assert!(low[0] > base[0]); + assert!(high[0] < base[0]); + assert!(low.iter().all(|&c| c >= 1)); } } diff --git a/src/descriptors/mod.rs b/src/descriptors/mod.rs index 0fd53d9..cdea547 100644 --- a/src/descriptors/mod.rs +++ b/src/descriptors/mod.rs @@ -3,9 +3,11 @@ use image::GenericImageView; -pub trait Descriptor { - /// Print the name of the descriptor, +pub trait Descriptor: Send + Sync { + /// Name of the descriptor, also used as store filename fn info(&self) -> String; + /// Bump whenever the hash output changes, so stores can detect stale data + fn version(&self) -> u32; fn resize(&self, img: &image::DynamicImage) -> image::DynamicImage { img.grayscale().thumbnail_exact(8, 8) } @@ -17,7 +19,7 @@ pub trait Descriptor { /// Interprets a 64-element array as an 8x8 matrix /// returns a nicely printable string -fn print_matrix(array: [T; 64]) -> String { +pub fn print_matrix(array: [T; 64]) -> String { let mut output = String::new(); for x in 0..8 { for y in 0..8 { @@ -38,6 +40,10 @@ pub struct DCT { } pub mod dct; +/// The classic pHash DCT hash, as a baseline to compare against +pub struct PHash; +pub mod phash; + fn median64(values: &[T]) -> T { let mut sorted_values = values.to_vec(); sorted_values.sort(); @@ -52,6 +58,11 @@ impl Descriptor for Median { "median".to_string() } + // v2: pixel (0,0) used to be shifted out of the mask entirely + fn version(&self) -> u32 { + 2 + } + fn describe(&self, img: &image::DynamicImage) -> u64 { let img = self.resize(img); let mut values: [u8; 64] = [0; 64]; @@ -60,21 +71,103 @@ impl Descriptor for Median { } let median = median64(&values); let mut mask: u64 = 0; - img.save("debug.png").unwrap(); - for (_, _, pix) in img.pixels() { - if pix[0] > median { - mask += 1; - } + for value in values { mask <<= 1; + if value > median { + mask |= 1; + } } mask } } pub fn get_all_descriptors() -> Vec> { - let descriptors: Vec> = - vec![Box::new(DCT::new()), Box::new(Median)]; - // descriptors.push(Box::new(DCT::new())); - // descriptors.push(Box::new(Median)); - descriptors -} \ No newline at end of file + vec![Box::new(DCT::new()), Box::new(Median), Box::new(PHash)] +} + +#[cfg(test)] +pub(crate) mod testimg { + use image::{DynamicImage, ImageBuffer, Luma, Rgb}; + + /// Asymmetric test image: diagonal gradient with a bright blob off-center + pub fn gradient(size: u32) -> DynamicImage { + DynamicImage::ImageLuma8(ImageBuffer::from_fn(size, size, |x, y| { + let base = (x * 2 + y) * 255 / (size * 3); + let blob = if x < size / 4 && y > size / 2 { 80 } else { 0 }; + Luma([(base + blob).min(255) as u8]) + })) + } + + pub fn checkerboard(size: u32) -> DynamicImage { + DynamicImage::ImageLuma8(ImageBuffer::from_fn(size, size, |x, y| { + Luma([if (x / 8 + y / 8) % 2 == 0 { 30 } else { 220 }]) + })) + } + + pub fn colorful(size: u32) -> DynamicImage { + DynamicImage::ImageRgb8(ImageBuffer::from_fn(size, size, |x, y| { + Rgb([ + (x * 255 / size) as u8, + (y * 255 / size) as u8, + ((x + y) * 128 / size) as u8, + ]) + })) + } + + /// Fixed pseudo-random 8x8 image: decisive coefficients, no ties + pub fn noise8x8() -> DynamicImage { + let mut state: u32 = 0x2545f491; + DynamicImage::ImageLuma8(ImageBuffer::from_fn(8, 8, move |_, _| { + state = state.wrapping_mul(1664525).wrapping_add(1013904223); + Luma([(state >> 24) as u8]) + })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use image::{DynamicImage, ImageBuffer, Luma}; + + fn image_with_first_pixel(value: u8) -> DynamicImage { + // 30 pixels of 10 and 33 pixels of 250, so the median stays at 250 + // no matter which side the first pixel lands on + DynamicImage::ImageLuma8(ImageBuffer::from_fn(8, 8, move |x, y| { + let i = y * 8 + x; + match i { + 0 => Luma([value]), + 1..=30 => Luma([10]), + _ => Luma([250]), + } + })) + } + + #[test] + fn median_uses_pixel_zero() { + // regression: the first pixel used to be shifted out of the hash + let bright = Median.describe(&image_with_first_pixel(255)); + let dark = Median.describe(&image_with_first_pixel(0)); + assert_eq!(bright, 1 << 63); + assert_eq!(dark, 0); + } + + #[test] + fn median_deterministic() { + let img = testimg::gradient(64); + assert_eq!(Median.describe(&img), Median.describe(&img)); + } + + #[test] + fn median_distinguishes_images() { + assert_ne!( + Median.describe(&testimg::gradient(64)), + Median.describe(&testimg::checkerboard(64)) + ); + } + + #[test] + fn distance_is_popcount() { + assert_eq!(Median.distance(0b1011, 0b0010), 2); + assert_eq!(Median.distance(u64::MAX, 0), 64); + } +} diff --git a/src/descriptors/phash.rs b/src/descriptors/phash.rs new file mode 100644 index 0000000..0e0f4c6 --- /dev/null +++ b/src/descriptors/phash.rs @@ -0,0 +1,120 @@ +use std::f64::consts::PI; +use crate::descriptors::{Descriptor, PHash, median64}; + +/// Classic pHash DCT hash, following the widely used imagehash recipe: +/// grayscale, resize to 32x32, 2D DCT, keep the top-left 8x8 low-frequency +/// block, threshold each coefficient against the block median. +/// +/// Included as a baseline: same 64-bit budget, same Hamming distance, +/// different bit extraction than our DCT descriptor. +const SIZE: usize = 32; +const KEEP: usize = 8; + +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(); + + // Row pass: keep the first KEEP coefficients of every row + let mut rows = [[0.0f64; KEEP]; SIZE]; + for (y, row) in rows.iter_mut().enumerate() { + for (u, coeff) in row.iter_mut().enumerate() { + let mut sum = 0.0; + for x in 0..SIZE { + let pix = gray.get_pixel(x as u32, y as u32)[0] as f64; + sum += pix * ((2.0 * x as f64 + 1.0) * u as f64 * PI / (2.0 * SIZE as f64)).cos(); + } + *coeff = sum; + } + } + + // Column pass over the kept coefficients + let mut block = [0.0f64; KEEP * KEEP]; + for u in 0..KEEP { + for v in 0..KEEP { + let mut sum = 0.0; + for (y, row) in rows.iter().enumerate() { + sum += row[u] * ((2.0 * y as f64 + 1.0) * v as f64 * PI / (2.0 * SIZE as f64)).cos(); + } + block[v * KEEP + u] = sum; + } + } + block + } +} + +impl Descriptor for PHash { + fn info(&self) -> String { + "phash".to_string() + } + + fn version(&self) -> u32 { + 1 + } + + fn describe(&self, img: &image::DynamicImage) -> u64 { + let block = self.lowfreq(img); + let mut sortable = [0i64; 64]; + for (i, value) in block.iter().enumerate() { + sortable[i] = (value * 1024.0) as i64; + } + let median = median64(&sortable); + let mut mask: u64 = 0; + for value in sortable { + mask <<= 1; + if value > median { + mask |= 1; + } + } + mask + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::descriptors::testimg; + + #[test] + fn deterministic() { + let img = testimg::gradient(64); + assert_eq!(PHash.describe(&img), PHash.describe(&img)); + } + + #[test] + fn distinguishes_images() { + assert_ne!( + PHash.describe(&testimg::gradient(64)), + PHash.describe(&testimg::checkerboard(64)) + ); + } + + #[test] + fn survives_jpeg_compression() { + // Structured images are stable under compression. Smooth gradients + // are the worst case for median thresholding: most coefficients sit + // near the median, so those hashes move a lot more. + let img = testimg::gradient(200); + let mut buffer = std::io::Cursor::new(Vec::new()); + img.to_rgb8() + .write_with_encoder(image::codecs::jpeg::JpegEncoder::new_with_quality( + &mut buffer, + 60, + )) + .unwrap(); + let compressed = image::load_from_memory(buffer.get_ref()).unwrap(); + let distance = PHash.distance(PHash.describe(&img), PHash.describe(&compressed)); + assert!(distance <= 10, "jpeg roundtrip moved hash by {distance}"); + } + + #[test] + fn not_flip_invariant() { + // documents the difference with our DCT descriptor + let img = testimg::gradient(64); + assert_ne!(PHash.describe(&img), PHash.describe(&img.fliph())); + } +} diff --git a/src/lib.rs b/src/lib.rs index 69f07ea..4b3b786 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,8 +1,8 @@ //! Near-copy image similarity detection. - -//! Descriptors that can be used to describe an image. -//! If two images are (almost) the same, their descriptions will be the same. +//! +//! 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 descriptors; pub mod store; pub mod mutators; -pub mod wasm; \ No newline at end of file +pub mod wasm; diff --git a/src/main.rs b/src/main.rs index 0b2491a..8aecf59 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,93 +1,133 @@ use image_similarity::store::DescriptorStore; use image_similarity::descriptors::get_all_descriptors; use image_similarity::mutators::get_all_mutators; -use log::{error}; +use log::{error, info}; use std::path::PathBuf; -use std::{fs}; +use std::fs; use clap::Parser; use plotters::prelude::*; +use rayon::prelude::*; #[derive(Parser)] #[command(version, about, long_about = None)] struct Cfg { /// Path that contains the input images. Will not traverse directories. path: PathBuf, + /// 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)] + batch_size: usize, } fn main() { env_logger::init(); let cfg = Cfg::parse(); - let mutators = get_all_mutators(); - //let descriptors = get_all_descriptors(); - let mut stores: Vec = Vec::new(); + let mutators = get_all_mutators(); + let mut stores: Vec = Vec::new(); for descriptor in get_all_descriptors() { let filename = format!("{}.store", descriptor.info()); - let store = DescriptorStore::new(descriptor).with_file(filename); - stores.push(store); - } - - for node in fs::read_dir(cfg.path).unwrap() { - let file = node.expect("Error walking directory"); - let name = match file.file_name().into_string() { - Ok(v) => v, + let store = match DescriptorStore::new(descriptor).with_file(&filename) { + Ok(store) => store, Err(e) => { - error!("Error reading {}: {:?}:", file.path().to_string_lossy(), e); - continue + error!("{filename}: {e}"); + std::process::exit(1); } }; - - let mut get = true; - for store in &stores { - if store.has_value(&name) { - get = false; - break; - } - } - - // 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 - } - }.thumbnail_exact(8, 8); - - //Store the phashes of the base image - for store in &mut stores { - 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()); - } - } + if !store.version_matches() { + error!( + "{filename} contains hashes from an older descriptor version. \ + New hashes would not be comparable. Move it away or delete it first." + ); + std::process::exit(1); } + stores.push(store); } + + let mut entries: Vec<(String, PathBuf)> = fs::read_dir(&cfg.path) + .expect("Error reading directory") + .filter_map(|node| { + let file = node.expect("Error walking directory"); + match file.file_name().into_string() { + Ok(name) => Some((name, file.path())), + Err(e) => { + error!("Error reading {}: {:?}", file.path().to_string_lossy(), e); + None + } + } + }) + .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 + .into_iter() + .filter(|(name, _)| !stores.iter().any(|store| store.has_value(name))) + .collect(); + info!("{} new images to process", todo.len()); + + let mut done = 0; + 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. + let descriptors: Vec<&dyn image_similarity::descriptors::Descriptor> = + stores.iter().map(|store| store.descriptor.as_ref()).collect(); + 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; + } + }; + 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); + for (i, descriptor) in descriptors.iter().enumerate() { + out.push((i, mutated_name.clone(), descriptor.describe(&mutated))); + } + } + Some(out) + }) + .collect(); + drop(descriptors); + + for file_hashes in hashes { + for (i, name, hash) in file_hashes { + stores[i].insert(hash, name); + } + } + for store in &stores { + store.save().expect("Error saving store"); + } + done += batch.len(); + info!("{done}/{} images done", todo.len()); + } + for store in &stores { - store.save().expect("Error saving store"); println!("{store}"); } // Create PR-curve graphs from stores - let max_threshold = 8; + let max_threshold = cfg.max_threshold.clamp(1, 64); for store in &stores { println!("Stats for {}", store.descriptor.info()); - let stats = store.get_stats(max_threshold); + let stats = store.get_stats(&mutators, max_threshold); 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()) + .caption(format!("PR curves ({})", store.descriptor.info()), ("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) @@ -98,16 +138,13 @@ fn main() { .x_label_formatter(&|x| format!("{x:.3}")) .y_label_formatter(&|x| format!("{x:.3}")) .draw().unwrap(); - let colors = [&RED, &BLUE, &CYAN, &MAGENTA, &BLACK, &GREEN, &YELLOW]; - let n = colors.len(); for (i, mutator_stats) in stats.into_iter().enumerate() { - // And we can draw something in the drawing area let pr_curve = mutator_stats.pr_curve(max_threshold); println!("{}:", mutator_stats.name); for (x, y) in pr_curve.clone() { println!("{x}, {y}"); } - let color = colors[i % n]; + let color = Palette99::pick(i).to_rgba(); chart.draw_series(LineSeries::new( pr_curve, color.filled(), diff --git a/src/mutators/mod.rs b/src/mutators/mod.rs index 8c2f659..a532f30 100644 --- a/src/mutators/mod.rs +++ b/src/mutators/mod.rs @@ -1,18 +1,22 @@ //! Mutators that can be used to mutate an image. -//! These mutated images can be used as a "near copy" -use image::{Rgb, Rgba}; -use imageproc::definitions::Image; -use imageproc::geometric_transformations::*; +//! These mutated images can be used as a "near copy". +//! +//! The suite follows the transformation categories from +//! Thomee et al., "Large Scale Image Copy Detection Evaluation" (MIR '08): +//! recoding, resampling, content processing, framing and insertion of +//! elements. Flip and rotation are our own additions. +use image::Rgb; +use imageproc::drawing::draw_filled_rect_mut; +use imageproc::geometric_transformations::{rotate_about_center, Interpolation}; +use imageproc::rect::Rect; -pub trait Mutator { +pub trait Mutator: Send + Sync { /// Name/description of the mutator fn info(&self) -> String; - /// Short tag that represents the mutator. + /// Short tag that represents the mutator, including parameters. /// To be used in the filename such that it can be recognized as mutated image - fn tag(&self) -> String { - "MUT".to_string() - } + fn tag(&self) -> String; /// Returns the mutated form of the input image fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage; @@ -32,22 +36,27 @@ impl Mutator for Flip { } } -/// Hue shift -pub struct Hue; +/// Hue shift, degrees +pub struct Hue { + pub degrees: i32, +} impl Mutator for Hue { fn info(&self) -> String { - "Hue shift".to_string() + format!("Hue shift {}", self.degrees) } fn tag(&self) -> String { - ".hue.".to_string() + format!(".hue{}.", self.degrees) } fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { - img.huerotate(5) + img.huerotate(self.degrees) } } /// Unsharp mask -pub struct Sharp; +pub struct Sharp { + pub sigma: f32, + pub threshold: i32, +} impl Mutator for Sharp { fn info(&self) -> String { "Unsharp mask".to_string() @@ -56,58 +65,269 @@ impl Mutator for Sharp { ".sharp.".to_string() } fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { - img.unsharpen(1.5, 20) + img.unsharpen(self.sigma, self.threshold) } } -pub struct Blur; +/// Gaussian blur, sigma in tenths so the tag stays dot-free +pub struct Blur { + pub sigma: f32, +} impl Mutator for Blur { fn info(&self) -> String { - "Gaussian blur".to_string() + format!("Gaussian blur {:.1}", self.sigma) } fn tag(&self) -> String { - ".blur.".to_string() + format!(".blur{}.", (self.sigma * 10.0) as u32) } fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { - img.blur(1.5) + img.blur(self.sigma) } } -pub struct BlurSharp; -impl Mutator for BlurSharp { +/// JPEG encode/decode roundtrip at a given quality (recoding) +pub struct Jpeg { + pub quality: u8, +} +impl Mutator for Jpeg { fn info(&self) -> String { - "Blur and sharpen".to_string() + format!("JPEG quality {}", self.quality) } fn tag(&self) -> String { - ".blsh.".to_string() + format!(".jpeg{}.", self.quality) } fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { - img.blur(1.5).unsharpen(1.5, 20) + let mut buffer = std::io::Cursor::new(Vec::new()); + let encoder = image::codecs::jpeg::JpegEncoder::new_with_quality(&mut buffer, self.quality); + img.to_rgb8() + .write_with_encoder(encoder) + .expect("jpeg encoding failed"); + image::load_from_memory(buffer.get_ref()).expect("jpeg decoding failed") } } -pub struct RotateCrop; -impl Mutator for RotateCrop { +/// Downscale to a percentage of the original size (resampling) +pub struct Scale { + pub percent: u32, +} +impl Mutator for Scale { fn info(&self) -> String { - "Rotate and crop".to_string() + format!("Rescale {}%", self.percent) } fn tag(&self) -> String { - ".rcrop.".to_string() + format!(".scale{}.", self.percent) } fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { - let img: Image> = img.to_rgb8(); + let w = (img.width() * self.percent / 100).max(1); + let h = (img.height() * self.percent / 100).max(1); + img.resize_exact(w, h, image::imageops::FilterType::Triangle) + } +} + +/// Contrast adjustment (content processing) +pub struct Contrast { + pub amount: f32, +} +impl Mutator for Contrast { + fn info(&self) -> String { + format!("Contrast {:+}", self.amount) + } + fn tag(&self) -> String { + format!(".contr{}.", self.amount as i32) + } + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { + img.adjust_contrast(self.amount) + } +} + +/// Brightness adjustment (content processing) +pub struct Brightness { + pub delta: i32, +} +impl Mutator for Brightness { + fn info(&self) -> String { + format!("Brightness {:+}", self.delta) + } + fn tag(&self) -> String { + format!(".bright{}.", self.delta) + } + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { + img.brighten(self.delta) + } +} + +/// Central crop keeping a percentage of both dimensions (framing) +pub struct CropCenter { + pub keep_percent: u32, +} +impl Mutator for CropCenter { + fn info(&self) -> String { + format!("Crop to {}%", self.keep_percent) + } + fn tag(&self) -> String { + format!(".crop{}.", self.keep_percent) + } + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { + let w = (img.width() * self.keep_percent / 100).max(1); + let h = (img.height() * self.keep_percent / 100).max(1); + img.crop_imm((img.width() - w) / 2, (img.height() - h) / 2, w, h) + } +} + +/// Black bars on top and bottom (framing) +pub struct Letterbox { + pub bar_percent: u32, +} +impl Mutator for Letterbox { + fn info(&self) -> String { + format!("Letterbox {}%", self.bar_percent) + } + fn tag(&self) -> String { + format!(".lbox{}.", self.bar_percent) + } + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { + let mut img = img.to_rgb8(); + let bar = (img.height() * self.bar_percent / 100).max(1); + let (w, h) = (img.width(), img.height()); + draw_filled_rect_mut(&mut img, Rect::at(0, 0).of_size(w, bar), Rgb([0, 0, 0])); + draw_filled_rect_mut( + &mut img, + Rect::at(0, (h - bar) as i32).of_size(w, bar), + Rgb([0, 0, 0]), + ); + image::DynamicImage::ImageRgb8(img) + } +} + +/// White square in the bottom-right corner, mimics a logo (insertion of elements) +pub struct Logo { + pub size_percent: u32, +} +impl Mutator for Logo { + fn info(&self) -> String { + format!("Logo insert {}%", self.size_percent) + } + fn tag(&self) -> String { + format!(".logo{}.", self.size_percent) + } + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { + let mut img = img.to_rgb8(); + let side = (img.width().min(img.height()) * self.size_percent / 100).max(1); + let x = img.width() - side - side / 2; + let y = img.height() - side - side / 2; + draw_filled_rect_mut( + &mut img, + Rect::at(x as i32, y as i32).of_size(side, side), + Rgb([255, 255, 255]), + ); + image::DynamicImage::ImageRgb8(img) + } +} + +/// Small rotation around the center, edges filled with the corner pixel +pub struct Rotate { + pub degrees: f32, +} +impl Mutator for Rotate { + fn info(&self) -> String { + format!("Rotate {}", self.degrees) + } + fn tag(&self) -> String { + format!(".rot{}.", self.degrees as i32) + } + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { + let img = img.to_rgb8(); let default = *img.get_pixel(0, 0); - image::DynamicImage::ImageRgb8(rotate_about_center(&img, 0.1, Interpolation::Bicubic, default)) + image::DynamicImage::ImageRgb8(rotate_about_center( + &img, + self.degrees.to_radians(), + Interpolation::Bicubic, + default, + )) + } +} + +/// Additive gaussian noise, fixed seed for reproducibility +pub struct Noise { + pub stddev: f64, +} +impl Mutator for Noise { + fn info(&self) -> String { + format!("Gaussian noise {}", self.stddev) + } + fn tag(&self) -> String { + format!(".noise{}.", self.stddev as u32) + } + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { + let noisy = imageproc::noise::gaussian_noise(&img.to_rgb8(), 0.0, self.stddev, 42); + image::DynamicImage::ImageRgb8(noisy) } } pub fn get_all_mutators() -> Vec> { - let mutators: Vec> = //Vec::with_capacity(4); vec![ Box::new(Flip), - Box::new(Hue), - Box::new(Sharp), - Box::new(Blur), - ]; - mutators -} \ No newline at end of file + Box::new(Hue { degrees: 30 }), + Box::new(Blur { sigma: 1.5 }), + Box::new(Sharp { sigma: 1.5, threshold: 20 }), + Box::new(Jpeg { quality: 90 }), + Box::new(Jpeg { quality: 50 }), + Box::new(Jpeg { quality: 20 }), + Box::new(Scale { percent: 50 }), + Box::new(Contrast { amount: 25.0 }), + Box::new(Brightness { delta: 30 }), + Box::new(CropCenter { keep_percent: 90 }), + Box::new(CropCenter { keep_percent: 70 }), + Box::new(Letterbox { bar_percent: 10 }), + Box::new(Logo { size_percent: 10 }), + Box::new(Rotate { degrees: 2.0 }), + Box::new(Noise { stddev: 10.0 }), + ] +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::descriptors::testimg; + + #[test] + fn tags_are_unique_and_wrapped_in_dots() { + let mutators = get_all_mutators(); + let mut tags: Vec = mutators.iter().map(|m| m.tag()).collect(); + tags.sort(); + let len_before = tags.len(); + tags.dedup(); + assert_eq!(tags.len(), len_before); + for tag in tags { + assert!(tag.starts_with('.') && tag.ends_with('.'), "bad tag {tag}"); + } + } + + #[test] + fn mutants_stay_decodable_with_expected_dimensions() { + let img = testimg::colorful(120); + for mutator in get_all_mutators() { + let out = mutator.mutate(&img); + assert!(out.width() > 0 && out.height() > 0, "{}", mutator.info()); + match mutator.tag().as_str() { + ".crop90." => assert_eq!(out.width(), 108), + ".crop70." => assert_eq!(out.width(), 84), + ".scale50." => assert_eq!(out.width(), 60), + ".flip." | ".lbox10." | ".logo10." | ".rot2." => { + assert_eq!((out.width(), out.height()), (120, 120)) + } + _ => {} + } + } + } + + #[test] + fn noise_is_deterministic() { + let img = testimg::colorful(64); + let noise = Noise { stddev: 10.0 }; + assert_eq!( + noise.mutate(&img).to_rgb8().as_raw(), + noise.mutate(&img).to_rgb8().as_raw() + ); + } +} diff --git a/src/store.rs b/src/store.rs index ac469ed..e42af7a 100644 --- a/src/store.rs +++ b/src/store.rs @@ -1,20 +1,61 @@ //! Persistent store of all calculated image descriptions. //! Can be queried to find identical or similar images. -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use crate::descriptors::Descriptor; -use crate::mutators::get_all_mutators; +use crate::mutators::Mutator; use std::fs; use std::path::Path; use std::fmt; use image::DynamicImage; -use log::{info, debug, error}; +use log::{info, warn, error}; use bktree::*; +use serde::{Serialize, Deserialize}; + +/// Bump when the on-disk layout of StoreFile changes +pub const STORE_FORMAT: u32 = 1; #[derive(Debug)] -pub enum SaveError { - Serialization, - File, +pub enum StoreError { + Io(std::io::Error), + Serialization(String), + /// The store on disk was written by a different descriptor or version, + /// its hashes are not comparable to freshly calculated ones + DescriptorMismatch { found: String, expected: String }, +} + +impl fmt::Display for StoreError { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + match self { + StoreError::Io(e) => write!(f, "io error: {e}"), + StoreError::Serialization(e) => write!(f, "serialization error: {e}"), + StoreError::DescriptorMismatch { found, expected } => write!( + f, + "store was written by descriptor {found}, expected {expected}. \ + Move the old store away or regenerate it." + ), + } + } +} + +impl std::error::Error for StoreError {} + +/// On-disk representation, owned variant for loading +#[derive(Deserialize)] +struct StoreFile { + format: u32, + descriptor: String, + descriptor_version: u32, + map: HashMap>, +} + +/// On-disk representation, borrowing variant for saving +#[derive(Serialize)] +struct StoreFileRef<'a> { + format: u32, + descriptor: String, + descriptor_version: u32, + map: &'a HashMap>, } /// Struct for calculating Recall-Precision per mutation @@ -28,126 +69,141 @@ pub struct PRStats { 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 t = threshold.min(63); let tpos = self.true_positives[t] as f64; let fpos = self.false_positives[t] as f64; + if tpos + fpos == 0.0 { + // nothing retrieved, nothing wrong + return 1.0; + } 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 t = threshold.min(63); let tpos = self.true_positives[t] as f64; let fneg = self.false_negatives[t] as f64; + if tpos + fneg == 0.0 { + return 0.0; + } 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) + (self.recall(threshold), self.precision(threshold)) } - 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 + pub fn pr_curve(&self, max_threshold: usize) -> Vec<(f64, f64)> { + (0..max_threshold.min(64)).map(|t| self.pr(t)).collect() } - } -/// Uses a hashmap to map descriptors to buckets of files. +/// 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. pub struct DescriptorStore { pub descriptor: Box, - /// Main hashmap that maps descriptors to buckets of filenames - map: HashMap>, + /// Main hashmap that maps hashes to buckets of filenames + map: HashMap>, + + /// Inverted index: filename to hash + names: HashMap, /// Place to load and store the map on file - save_location: std::path::PathBuf, + save_location: std::path::PathBuf, /// BK-tree for fast nearest neighbour - bktree: BkTree, + bktree: BkTree, - /// Cache for seen files to skip on initial load - seen: Option>, + /// Descriptor version that produced the data currently in the map. + /// 0 means unknown: loaded from a legacy store without metadata. + data_version: u32, } impl DescriptorStore { /// Makes a new empty `DescriptorStore` with default settings pub fn new(descriptor: Box) -> Self { - let map: HashMap> = HashMap::new(); - //let seen: HashSet = HashSet::new(); - //let seen = None; + let data_version = descriptor.version(); DescriptorStore { - map, + map: HashMap::new(), + names: HashMap::new(), descriptor, save_location: std::path::PathBuf::from("store.messagepack"), bktree: BkTree::new(hamming_distance), - seen: None, + data_version, } } - /// Sets the file location and loads data from file (if available) - pub fn with_file>(mut self, path: P) -> Self { + /// Sets the file location and loads data from file (if available). + /// Errors when the file belongs to a different descriptor (version). + /// Stores without metadata (legacy format) load with data version 0; + /// check [`Self::version_matches`] before adding new hashes to those. + pub fn with_file>(mut self, path: P) -> Result { 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 bytes = match fs::read(&self.save_location) { + Ok(b) => b, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + info!("{} not found, starting from empty store.", self.save_location.display()); + return Ok(self); } - let mut seen: HashSet = 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()); + Err(e) => return Err(StoreError::Io(e)), }; - 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); + let map = match rmp_serde::from_slice::(&bytes) { + Ok(file) => { + if file.format > STORE_FORMAT { + warn!( + "{} uses store format {}, this build knows up to {}", + self.save_location.display(), file.format, STORE_FORMAT + ); } + let expected = self.descriptor.info(); + if file.descriptor != expected || file.descriptor_version != self.descriptor.version() { + return Err(StoreError::DescriptorMismatch { + found: format!("{} v{}", file.descriptor, file.descriptor_version), + expected: format!("{} v{}", expected, self.descriptor.version()), + }); + } + self.data_version = file.descriptor_version; + file.map + } + // Legacy format: a bare hashmap without metadata + Err(_) => match rmp_serde::from_slice::>>(&bytes) { + Ok(map) => { + warn!( + "{} has no descriptor metadata, treating as legacy data (version 0)", + self.save_location.display() + ); + self.data_version = 0; + map + } + Err(e) => return Err(StoreError::Serialization(e.to_string())), + }, + }; + + for (key, bucket) in &map { + self.bktree.insert(*key); + for name in bucket { + self.names.insert(name.clone(), *key); } } + self.map = map; + Ok(self) } - pub fn save(&self) -> Result<(), SaveError> { - let serialized: Vec = match rmp_serde::to_vec(&self.map) { - Ok(value) => value, - Err(_e) => return Err(SaveError::Serialization), + /// True when the loaded data was produced by the current descriptor version, + /// i.e. new hashes are comparable to stored ones + pub fn version_matches(&self) -> bool { + self.data_version == self.descriptor.version() + } + + pub fn save(&self) -> Result<(), StoreError> { + let file = StoreFileRef { + format: STORE_FORMAT, + descriptor: self.descriptor.info(), + descriptor_version: self.data_version, + map: &self.map, }; - match fs::write(&self.save_location, serialized) { - Ok(()) => Ok(()), - Err(_e) => Err(SaveError::File), - } + 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) } /// Returns true iff the store already contains the key @@ -155,54 +211,46 @@ impl DescriptorStore { 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 - } - } + /// Returns true iff the store already contains the filename + pub fn has_value(&self, value: &str) -> bool { + self.names.contains_key(value) } - pub fn get(&self, value: &String) -> Option { - for (key, bucket) in self.map.iter() { - if bucket.contains(value) { - return Some(*key); - } - } - None + pub fn get(&self, value: &str) -> Option { + self.names.get(value).copied() + } + + pub fn len(&self) -> usize { + self.names.len() + } + + pub fn is_empty(&self) -> bool { + self.names.is_empty() } /// 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); + pub fn insert(&mut self, key: u64, value: String) { + let bucket = self.map.entry(key).or_default(); + if !bucket.contains(&value) { + bucket.push(value.clone()); + } + self.names.insert(value, key); + // the bktree ignores duplicate keys + self.bktree.insert(key); } - /// Calculates all descriptions with a given descriptor - pub fn insert_directory>(&mut self, dir: U) { - for node in fs::read_dir(dir).unwrap() { - let file = node.expect("Error walking directory"); + /// Describes and stores an image under the given name + pub fn store(&mut self, img: &DynamicImage, name: String) { + if !self.has_value(&name) { + let hash = self.descriptor.describe(img); + self.insert(hash, name); + } + } + + /// Calculates and stores descriptions for every image in a directory + pub fn insert_directory>(&mut self, dir: U) -> Result<(), StoreError> { + for node in fs::read_dir(dir).map_err(StoreError::Io)? { + let file = node.map_err(StoreError::Io)?; let name = match file.file_name().into_string() { Ok(v) => v, Err(e) => { @@ -210,172 +258,240 @@ impl DescriptorStore { 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); + if self.has_value(&name) { + continue; } - self.save().expect("error"); + info!("Processing {}", name); + let img = match image::open(file.path()) { + Ok(v) => v, + Err(e) => { + error!("Failed to process {}: {}", name, e); + continue + } + }; + let hash = self.descriptor.describe(&img); + self.insert(hash, name); } - // We have done our bulk loading, so we can unset our hashset cache: - //self.seen = None; + self.save() } - 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 { - let nn = self.nn(from, max_distance); - let mut results: Vec = 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); + /// All filenames within a given hamming distance of the given hash, + /// 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)); + } } + results.sort_by(|a, b| a.1.cmp(&b.1).then_with(|| a.0.cmp(&b.0))); 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; - } - } + /// Recall/precision statistics per mutator, for thresholds 0..max_threshold. + /// + /// Every unmutated image queries the store once. A found image counts as: + /// - true positive for mutator m when it is exactly the m-mutant of the query + /// - false positive for mutator m when it is an m-mutant of another image + /// - false positive for every mutator when it is a different unmutated image + /// + /// Mutants that are not found within the threshold are false negatives. + 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 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()) - } + 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]; - pub fn get_stats(&self, mut max_threshold: usize) -> Vec { - 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) { + let bases: Vec<&String> = self.names.keys().filter(|n| !n.starts_with("mut.")).collect(); + for base in &bases { + 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) { + let d = distance as usize; + for name in &self.map[key] { + if name == *base { 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; - } + } + 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; } + } else { + // Unmutated misses count for every mutator + fp_base_at[d] += 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) + let cumulative = |at: &[u64; 64]| { + let mut cum = [0u64; 64]; + let mut sum = 0; + for (t, count) in at.iter().enumerate() { + sum += count; + cum[t] = sum; } - } + cum + }; + + let fp_base = cumulative(&fp_base_at); + 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 mut false_positives = [0u64; 64]; + let mut false_negatives = [0u64; 64]; + for t in 0..64 { + false_positives[t] = fp_mut[t] + fp_base[t]; + false_negatives[t] = total - true_positives[t]; + } + PRStats { + true_positives, + false_positives, + false_negatives, + tag: format!("mut{}", mutator.tag()), + name: mutator.info(), + } + }).collect() } } 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()) + write!( + f, + "DescriptorStore {} v{}: {} images, {} distinct hashes", + self.descriptor.info(), + self.data_version, + self.names.len(), + self.map.len() + ) } -} \ No newline at end of file +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::descriptors::Median; + use crate::mutators::Mutator; + + struct TestMut; + impl Mutator for TestMut { + fn info(&self) -> String { + "X".to_string() + } + fn tag(&self) -> String { + ".x.".to_string() + } + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { + img.clone() + } + } + + fn temp_store_path(name: &str) -> std::path::PathBuf { + std::env::temp_dir().join(format!("imgsim-{}-{}.store", name, std::process::id())) + } + + #[test] + fn roundtrip_keeps_data_and_version() { + let path = temp_store_path("roundtrip"); + let mut store = DescriptorStore::new(Box::new(Median)) + .with_file(&path).unwrap(); + store.insert(42, "a.jpg".to_string()); + store.insert(42, "b.jpg".to_string()); + store.insert(7, "c.jpg".to_string()); + store.save().unwrap(); + + let loaded = DescriptorStore::new(Box::new(Median)).with_file(&path).unwrap(); + assert!(loaded.version_matches()); + assert_eq!(loaded.get("a.jpg"), Some(42)); + assert_eq!(loaded.get("c.jpg"), Some(7)); + assert_eq!(loaded.len(), 3); + std::fs::remove_file(path).unwrap(); + } + + #[test] + fn legacy_store_loads_with_version_zero() { + let path = temp_store_path("legacy"); + let mut legacy: HashMap> = HashMap::new(); + legacy.insert(3, vec!["old.jpg".to_string()]); + std::fs::write(&path, rmp_serde::to_vec(&legacy).unwrap()).unwrap(); + + let store = DescriptorStore::new(Box::new(Median)).with_file(&path).unwrap(); + assert!(!store.version_matches()); + assert_eq!(store.get("old.jpg"), Some(3)); + std::fs::remove_file(path).unwrap(); + } + + #[test] + fn descriptor_mismatch_is_an_error() { + let path = temp_store_path("mismatch"); + let file = StoreFileRef { + format: STORE_FORMAT, + descriptor: "median".to_string(), + descriptor_version: 999, + map: &HashMap::new(), + }; + std::fs::write(&path, rmp_serde::to_vec(&file).unwrap()).unwrap(); + + let result = DescriptorStore::new(Box::new(Median)).with_file(&path); + assert!(matches!(result, Err(StoreError::DescriptorMismatch { .. }))); + std::fs::remove_file(path).unwrap(); + } + + #[test] + fn query_is_sorted_by_distance() { + let mut store = DescriptorStore::new(Box::new(Median)); + store.insert(0b0000, "exact.jpg".to_string()); + store.insert(0b0001, "close.jpg".to_string()); + store.insert(0b0111, "far.jpg".to_string()); + store.insert(u64::MAX, "unrelated.jpg".to_string()); + let results = store.query(0, 3); + let names: Vec<&str> = results.iter().map(|(n, _)| n.as_str()).collect(); + assert_eq!(names, ["exact.jpg", "close.jpg", "far.jpg"]); + assert_eq!(results[2].1, 3); + } + + #[test] + fn stats_match_hand_computed_scenario() { + let mut store = DescriptorStore::new(Box::new(Median)); + store.insert(0b0000, "a".to_string()); + store.insert(0b0011, "mut.x.a".to_string()); + store.insert(0b0111, "b".to_string()); + store.insert(0b0111, "mut.x.b".to_string()); + + let mutators: Vec> = vec![Box::new(TestMut)]; + let stats = &store.get_stats(&mutators, 4)[0]; + + // query a (0b0000): finds mut.x.a at d=2 (TP), b at d=3 (base FP), + // mut.x.b at d=3 (mutant FP) + // query b (0b0111): finds mut.x.b at d=0 (TP), mut.x.a at d=1 + // (mutant FP), a at d=3 (base FP) + let expected = [ + (0.5, 1.0), + (0.5, 0.5), + (1.0, 2.0 / 3.0), + (1.0, 2.0 / 6.0), + ]; + for (t, (recall, precision)) in expected.into_iter().enumerate() { + assert_eq!(stats.pr(t), (recall, precision), "threshold {t}"); + } + } + + #[test] + fn precision_is_one_when_nothing_retrieved() { + let stats = PRStats { + true_positives: [0; 64], + false_positives: [0; 64], + false_negatives: [5; 64], + tag: "mut.x.".to_string(), + name: "X".to_string(), + }; + assert_eq!(stats.pr(70), (0.0, 1.0)); + } +} diff --git a/src/wasm/mod.rs b/src/wasm/mod.rs index 6e06b40..f74fa52 100644 --- a/src/wasm/mod.rs +++ b/src/wasm/mod.rs @@ -1,51 +1,96 @@ -use crate::descriptors::DCT; -extern crate wasm_bindgen; +//! Wasm bindings for the browser demos. +//! Hashes come back as BigInt, images go in and out as encoded bytes. +use crate::descriptors::{Descriptor, Median, DCT, PHash}; +use crate::mutators::{self, Mutator}; use wasm_bindgen::prelude::*; use std::io::Cursor; -use console_error_panic_hook; -#[wasm_bindgen] -extern { - fn alert(s: &str); -} - -#[wasm_bindgen] -pub fn greet() { - alert("Hello world!"); -} - -#[wasm_bindgen] -pub fn dctify(bytes: Vec) -> Vec { +#[wasm_bindgen(start)] +pub fn start() { console_error_panic_hook::set_once(); - let reader = - image::io::Reader - ::new(Cursor::new(bytes)).with_guessed_format().unwrap(); - let img: image::DynamicImage = - reader - .decode() - .unwrap(); - let dct = DCT::new(); - let values = dct.dct(&img); - Vec::from(values) } -#[wasm_bindgen] -pub fn resize(bytes: Vec) -> Vec { - console_error_panic_hook::set_once(); - let reader = - image::io::Reader - ::new(Cursor::new(bytes)).with_guessed_format().unwrap(); +fn load(bytes: &[u8]) -> Result { + image::load_from_memory(bytes).map_err(|e| JsError::new(&e.to_string())) +} - let img: image::DynamicImage = - reader - .decode() - .unwrap(); - - let resized = img.grayscale().thumbnail_exact(8, 8); +fn to_png(img: &image::DynamicImage) -> Result, JsError> { let mut buffer: Vec = Vec::new(); - let mut writer = std::io::Cursor::new(&mut buffer); + img.write_to(&mut Cursor::new(&mut buffer), image::ImageFormat::Png) + .map_err(|e| JsError::new(&e.to_string()))?; + Ok(buffer) +} - resized.write_to(&mut writer, image::ImageFormat::Png).unwrap(); +/// All three hashes of an image: [dct, median, phash] +#[wasm_bindgen] +pub fn hash_all(bytes: &[u8]) -> Result, JsError> { + let img = load(bytes)?; + Ok(vec![ + DCT::new().describe(&img), + Median.describe(&img), + PHash.describe(&img), + ]) +} - buffer -} \ No newline at end of file +#[wasm_bindgen] +pub fn hamming(a: u64, b: u64) -> u32 { + (a ^ b).count_ones() +} + +/// Grayscale thumbnail of the image as PNG, for pipeline visualisation +#[wasm_bindgen] +pub fn resize_preview(bytes: &[u8], size: u32) -> Result, JsError> { + let img = load(bytes)?; + to_png(&img.grayscale().thumbnail_exact(size, size)) +} + +/// The 64 coefficients of the 8x8 DCT used by the dct descriptor +#[wasm_bindgen] +pub fn dct_coefficients(bytes: &[u8]) -> Result, JsError> { + let img = load(bytes)?; + let dct = DCT::new(); + let resized = dct.resize(&img); + Ok(dct.dct(&resized).to_vec()) +} + +/// The two half-masks of the dct hash: [sign_mask, ordinal_mask] +#[wasm_bindgen] +pub fn dct_masks(bytes: &[u8]) -> Result, JsError> { + let img = load(bytes)?; + let dct = DCT::new(); + let resized = dct.resize(&img); + let values = dct.dct(&resized); + let (sign, ordinal) = dct.masks(&values); + Ok(vec![sign, ordinal]) +} + +/// The 8x8 low-frequency block of the 32x32 DCT used by phash +#[wasm_bindgen] +pub fn phash_lowfreq(bytes: &[u8]) -> Result, JsError> { + let img = load(bytes)?; + Ok(PHash.lowfreq(&img).to_vec()) +} + +/// Apply a single mutator to an image, returns PNG bytes. +/// The meaning of `amount` depends on the kind. +#[wasm_bindgen] +pub fn mutate(bytes: &[u8], kind: &str, amount: f64) -> Result, JsError> { + let img = load(bytes)?; + let mutator: Box = match kind { + "flip" => Box::new(mutators::Flip), + "hue" => Box::new(mutators::Hue { degrees: amount as i32 }), + "blur" => Box::new(mutators::Blur { sigma: amount as f32 }), + "sharpen" => Box::new(mutators::Sharp { sigma: amount as f32, threshold: 20 }), + "jpeg" => Box::new(mutators::Jpeg { quality: (amount as u8).clamp(1, 100) }), + "scale" => Box::new(mutators::Scale { percent: (amount as u32).clamp(1, 100) }), + "contrast" => Box::new(mutators::Contrast { amount: amount as f32 }), + "brightness" => Box::new(mutators::Brightness { delta: amount as i32 }), + "crop" => Box::new(mutators::CropCenter { keep_percent: (amount as u32).clamp(1, 100) }), + "letterbox" => Box::new(mutators::Letterbox { bar_percent: (amount as u32).clamp(1, 45) }), + "logo" => Box::new(mutators::Logo { size_percent: (amount as u32).clamp(1, 90) }), + "rotate" => Box::new(mutators::Rotate { degrees: amount as f32 }), + "noise" => Box::new(mutators::Noise { stddev: amount }), + _ => return Err(JsError::new(&format!("unknown mutator: {kind}"))), + }; + to_png(&mutator.mutate(&img)) +} diff --git a/web/colors-default.css b/web/colors-default.css new file mode 100644 index 0000000..8f70566 --- /dev/null +++ b/web/colors-default.css @@ -0,0 +1,10 @@ +/* Fallback palette for when the wal-generated colors.css is absent. + colors.css loads after this file and overrides everything here. */ +:root { + --background: #0f1017; + --foreground: #e3e2da; + --color1: #c9a227; + --color2: #b0552f; + --color3: #97742d; + --color4: #85902c; +} diff --git a/web/index.html b/web/index.html index 32005d1..99786f6 100644 --- a/web/index.html +++ b/web/index.html @@ -2,6 +2,7 @@ + @@ -12,18 +13,19 @@

Fucking around with perceptual hashes

Introduction

+

TODO: why this project exists, the master's thesis that wasn't, discovering pHash after the fact.

+

TODO: the rules of the game: 64 bits per image, hamming distance, nothing else.

+

Perceptual hashing

-

Locality-sensitive hashing

-

Subjective image similarity

-
-
-

Experimental setup

-

Mutators

-

Descriptor

-

Experiments

+

TODO: cryptographic hash vs perceptual hash. One flipped pixel: md5 avalanches, a perceptual hash shrugs.

+

Terminology

+

TODO: descriptor / fingerprint / perceptual hash / signature all mean roughly the same thing depending on which corner of the literature you are standing in. Copy detection vs near-duplicate detection vs LSH.

+

Subjective image similarity

+

TODO: what "the same image" even means. Same pixels? Same scene? Same vibe?

+

Demo

Select an image to run this demo with. Don't worry, nothing will be sent to any server! All calculations are done in the browser.

@@ -32,12 +34,15 @@ type="file" id="dctimage" accept="image/*" /> + or pick a sample: +
-
-

Resize

+ +
+

Step 1: Resize

The first step is to size the image down, and remove all color information. Converting the image to grayscale is done by simply averaging the pixels. @@ -58,22 +63,135 @@

-
-

DCT

-

What is DCT?

-

Frequency domain. Plaatje. Bla bla.

-
-
-

Mauris semper ipsum libero. Praesent laoreet massa sagittis enim consequat malesuada. Nullam viverra nibh sit amet lacus volutpat sollicitudin. Nullam lacus sem, commodo ut vulputate nec, consectetur in erat. Quisque eget nunc ac felis condimentum ultrices eu sit amet arcu. Praesent imperdiet faucibus aliquam. Curabitur sit amet faucibus erat.

-

Nam pulvinar, mi id sagittis laoreet, risus dolor pellentesque metus, sed mattis nunc erat sed urna. Ut facilisis, velit vel condimentum euismod, ex ante dictum leo, lobortis feugiat eros sem nec velit. Aliquam quis eros lacus. Duis venenatis purus at luctus tempor. Praesent gravida euismod ante, tristique placerat tortor mattis molestie. Duis viverra ex eget lectus tempus consectetur. Aliquam semper, ligula at molestie dapibus, sapien turpis rutrum massa, ac tristique lorem ex eget neque. Donec vel turpis odio.

-

Mauris semper ipsum libero. Praesent laoreet massa sagittis enim consequat malesuada. Nullam viverra nibh sit amet lacus volutpat sollicitudin. Nullam lacus sem, commodo ut vulputate nec, consectetur in erat. Quisque eget nunc ac felis condimentum ultrices eu sit amet arcu. Praesent imperdiet faucibus aliquam. Curabitur sit amet faucibus erat.

-

Mauris semper ipsum libero. Praesent laoreet massa sagittis enim consequat malesuada. Nullam viverra nibh sit amet lacus volutpat sollicitudin. Nullam lacus sem, commodo ut vulputate nec, consectetur in erat. Quisque eget nunc ac felis condimentum ultrices eu sit amet arcu. Praesent imperdiet faucibus aliquam. Curabitur sit amet faucibus erat.

-

Mauris semper ipsum libero. Praesent laoreet massa sagittis enim consequat malesuada. Nullam viverra nibh sit amet lacus volutpat sollicitudin. Nullam lacus sem, commodo ut vulputate nec, consectetur in erat. Quisque eget nunc ac felis condimentum ultrices eu sit amet arcu. Praesent imperdiet faucibus aliquam. Curabitur sit amet faucibus erat.

-

Mauris semper ipsum libero. Praesent laoreet massa sagittis enim consequat malesuada. Nullam viverra nibh sit amet lacus volutpat sollicitudin. Nullam lacus sem, commodo ut vulputate nec, consectetur in erat. Quisque eget nunc ac felis condimentum ultrices eu sit amet arcu. Praesent imperdiet faucibus aliquam. Curabitur sit amet faucibus erat.

-

Nam pulvinar, mi id sagittis laoreet, risus dolor pellentesque metus, sed mattis nunc erat sed urna. Ut facilisis, velit vel condimentum euismod, ex ante dictum leo, lobortis feugiat eros sem nec velit. Aliquam quis eros lacus. Duis venenatis purus at luctus tempor. Praesent gravida euismod ante, tristique placerat tortor mattis molestie. Duis viverra ex eget lectus tempus consectetur. Aliquam semper, ligula at molestie dapibus, sapien turpis rutrum massa, ac tristique lorem ex eget neque. Donec vel turpis odio.

-

Nam pulvinar, mi id sagittis laoreet, risus dolor pellentesque metus, sed mattis nunc erat sed urna. Ut facilisis, velit vel condimentum euismod, ex ante dictum leo, lobortis feugiat eros sem nec velit. Aliquam quis eros lacus. Duis venenatis purus at luctus tempor. Praesent gravida euismod ante, tristique placerat tortor mattis molestie. Duis viverra ex eget lectus tempus consectetur. Aliquam semper, ligula at molestie dapibus, sapien turpis rutrum massa, ac tristique lorem ex eget neque. Donec vel turpis odio.

-

Nam pulvinar, mi id sagittis laoreet, risus dolor pellentesque metus, sed mattis nunc erat sed urna. Ut facilisis, velit vel condimentum euismod, ex ante dictum leo, lobortis feugiat eros sem nec velit. Aliquam quis eros lacus. Duis venenatis purus at luctus tempor. Praesent gravida euismod ante, tristique placerat tortor mattis molestie. Duis viverra ex eget lectus tempus consectetur. Aliquam semper, ligula at molestie dapibus, sapien turpis rutrum massa, ac tristique lorem ex eget neque. Donec vel turpis odio.

-

Nam pulvinar, mi id sagittis laoreet, risus dolor pellentesque metus, sed mattis nunc erat sed urna. Ut facilisis, velit vel condimentum euismod, ex ante dictum leo, lobortis feugiat eros sem nec velit. Aliquam quis eros lacus. Duis venenatis purus at luctus tempor. Praesent gravida euismod ante, tristique placerat tortor mattis molestie. Duis viverra ex eget lectus tempus consectetur. Aliquam semper, ligula at molestie dapibus, sapien turpis rutrum massa, ac tristique lorem ex eget neque. Donec vel turpis odio.

+ +
+

Step 2: The simplest hash that could possibly work

+

TODO: median hash: one bit per pixel, brighter than the median or not. This is the "Median" baseline from Thomee et al.

+
+
+

Bits

+
+
+
+

Hash

+ +
+
+

TODO: why this breaks: global median shifts, flips scramble everything. Foreshadow the moon/plate anecdote.

+
+ +
+

Step 3: Frequency space

+

TODO: what the DCT does. JPEG uses the same trick. Low frequencies = global shape, high frequencies = detail we already threw away.

+
+
+

DCT coefficients

+
+
+
+

TODO: reading the coefficient grid: top-left is the average, first row is horizontal waves, first column vertical waves.

+

From coefficients to bits

+

TODO: sign bits of the first 36 zigzag coefficients + 28 ordinal comparisons between neighbours in zigzag order. The odd-column signs get multiplied by the sign of coefficient (1,0), which buys horizontal flip invariance for one bit.

+
+
+

Sign mask (36)

+
+
+
+

Ordinal mask (28)

+
+
+
+

Hash

+ +
+
+
+ +
+

Step 4: How pHash does it

+

TODO: same idea, different route: resize to 32x32 instead of 8x8, DCT, keep the top-left 8x8 block of low frequencies, threshold against the median coefficient.

+
+
+

32×32

+
+
+
+

Low-frequency block

+
+
+
+

Bits

+
+
+
+

Hash

+ +
+
+

TODO: what's genuinely different between our dct hash and phash (bit extraction, flip invariance) and what isn't (everything else).

+
+ +
+

Mutations

+

TODO: a copy is rarely byte-identical. Recoding, resampling, content processing, framing, inserted logos (the Thomee et al. taxonomy). Sweep the slider and watch which hash survives what.

+

+ + + +

+
+ original + mutated +
+

+ +

+ + + + + + + + +
dctmedianphash
hamming distance
+

TODO: what counts as "the same" now? Thresholds. pHash uses 22 of 64 bits, we will measure our own.

+
+ +
+

Find the copy

+

TODO: the fun part: a tiny search engine. Add a pile of images, click one, get the nearest neighbours per algorithm. Mention the white plate that matched the moon.

+

+ + +

+
+

Results

+
+
+ +
+

At scale

+

TODO: the browser demo is anecdote, this section is data. 25k images, 16 mutations each, precision-recall over the hamming threshold.

+
+ PR curves for the dct hash +
TODO: dct-pr.png from the experiment run
+
+
+ PR curves for median and phash +
TODO: median-pr.png and phash-pr.png
+
+

TODO: how the PR curves are computed, what a false positive means here, where the thresholds land per algorithm.

+
+ +
+

Loose ends

+

TODO: vertical flips and 180 rotations (same trick, one more bit). Coefficient stability near zero. Video: keyframes vs temporally averaged frames vs 3D-DCT. The seen-images daemon idea. Sorting a folder by visual similarity.

diff --git a/web/resize.js b/web/resize.js deleted file mode 100644 index 4991609..0000000 --- a/web/resize.js +++ /dev/null @@ -1,11 +0,0 @@ -import init, { dctify, resize } from './pkg/image_similarity.js'; - -async function run() { - await init(); -} -run(); - -onmessage = (e) => { - let resizedBuffer = resize(e.data); - postMessage(resizedBuffer); -}; \ No newline at end of file diff --git a/web/script.js b/web/script.js index 971caf4..c836061 100644 --- a/web/script.js +++ b/web/script.js @@ -1,80 +1,300 @@ -import init, { dctify, resize } from './pkg/image_similarity.js'; +// Plumbing for the demos. All hashing happens in the worker (wasm). -async function run() { - await init(); +const worker = new Worker('worker.js', { type: 'module' }); +let nextId = 0; +const pending = new Map(); + +function rpc(op, args) { + return new Promise((resolve, reject) => { + const id = nextId++; + pending.set(id, { resolve, reject }); + worker.postMessage({ id, op, args }); + }); } -run(); -const resizeWorker = new Worker("resize.js", { type: 'module' }); -let imageContainers = document.getElementsByClassName("image-original"); -let resizedImageContainer = document.getElementById("image-resize"); -let dctCoefficientContainer = document.getElementById("image-dct"); -let formImage = document.getElementById("dctimage"); +worker.onmessage = (e) => { + const { id, ok, result, error } = e.data; + const promise = pending.get(id); + pending.delete(id); + if (!promise) return; + if (ok) { + promise.resolve(result); + } else { + promise.reject(new Error(error)); + } +}; -let resizedBuffer = null; +// Helpers + +function pngUrl(bytes) { + return URL.createObjectURL(new Blob([bytes], { type: 'image/png' })); +} + +function hamming(a, b) { + let x = a ^ b; + let count = 0; + while (x) { + x &= x - 1n; + count++; + } + return count; +} + +function hex64(hash) { + return hash.toString(16).padStart(16, '0'); +} + +// Renders the top `bits` bits of a hash, msb first +function bitGrid(container, hash, bits = 64) { + container.replaceChildren(); + for (let i = bits - 1; i >= 0; i--) { + const cell = document.createElement('div'); + cell.className = (hash >> BigInt(i)) & 1n ? 'bit on' : 'bit off'; + container.appendChild(cell); + } +} + +// 8x8 grid of coefficient magnitudes, log scale, red negative / blue positive +function heatGrid(container, values) { + container.replaceChildren(); + const max = Math.max(...values.map(Math.abs), 1e-9); + for (const value of values) { + const cell = document.createElement('div'); + const strength = Math.log1p(Math.abs(value)) / Math.log1p(max); + const hue = value < 0 ? 4 : 215; + cell.style.background = `hsl(${hue} 70% ${15 + strength * 45}%)`; + cell.title = value.toFixed(1); + container.appendChild(cell); + } +} + +function show(section) { + document.querySelectorAll('.needs-image').forEach((el) => el.classList.add('visible')); + if (section) document.getElementById(section).scrollIntoView({ behavior: 'smooth' }); +} + +// Demo state +let baseBuffer = null; // ArrayBuffer of the selected image +let baseUrl = null; +let baseHashes = null; // BigUint64Array [dct, median, phash] + +// Pipeline demo + +const imageContainers = document.getElementsByClassName('image-original'); +const formImage = document.getElementById('dctimage'); -// Resets all containers, deletes images, etc -// Redefines the event handlers for the base image function reset() { - //imageContainers.forEach((imageContainer) => { - for (let imageContainer of imageContainers) { - imageContainer.replaceChildren(); + for (const container of imageContainers) { + container.replaceChildren(); } - resizedImageContainer.replaceChildren(); - dctCoefficientContainer.replaceChildren(); - document.getElementById("resize").classList.remove("resize"); + document.getElementById('image-resize').replaceChildren(); + document.getElementById('resize').classList.remove('resize'); } -// Uses the resized buffer to get DCT coefficients -function getDCT() { - let buf = new Uint8Array(resizedBuffer) - let dct = dctify(buf); - - for (const c of dct) { - let cdiv = document.createElement("div"); - cdiv.innerHTML = c.toFixed(1); - dctCoefficientContainer.appendChild(cdiv); +async function loadImage(buffer, url) { + reset(); + baseBuffer = buffer; + baseUrl = url; + for (const container of imageContainers) { + const img = document.createElement('img'); + img.src = url; + container.appendChild(img); } + + const data = await rpc('pipeline', { buffer }); + baseHashes = data.hashes; + const [dct, median, phash] = data.hashes; + + // resize step + const resized = document.createElement('img'); + resized.src = pngUrl(data.resize8); + document.getElementById('image-resize').replaceChildren(resized); + const original = document.getElementById('resize-original').children[0]; + if (original && original.width) { + original.width = original.width; + original.height = original.height; + } + + // median + bitGrid(document.getElementById('median-bits'), median); + document.getElementById('median-hash').innerText = hex64(median); + + // dct + heatGrid(document.getElementById('image-dct'), data.coefficients); + bitGrid(document.getElementById('dct-sign-bits'), data.masks[0], 36); + bitGrid(document.getElementById('dct-ordinal-bits'), data.masks[1], 28); + document.getElementById('dct-hash').innerText = hex64(dct); + + // phash + const small = document.createElement('img'); + small.src = pngUrl(data.resize32); + small.className = 'pixelated med'; + document.getElementById('phash-resize').replaceChildren(small); + heatGrid(document.getElementById('phash-lowfreq'), data.lowfreq); + bitGrid(document.getElementById('phash-bits'), phash); + document.getElementById('phash-hash').innerText = hex64(phash); + + show(); + resetMutation(); + runMutation(); } -// Result from the resize worker, means we get resized buffer -resizeWorker.onmessage = function (e) { - console.log(e.data); - resizedBuffer = e.data.buffer; - const resizedImage = document.createElement("img"); - resizedImage.src = URL.createObjectURL( - new Blob([resizedBuffer], { type: 'image/png' }) - ); - resizedImageContainer.appendChild(resizedImage); - - // Set the width of the image explictly to help with the animation - let image = document.getElementById("resize-original").children[0]; - image.width = image.width; - image.height = image.height; -} - -document.getElementById("resize-button").addEventListener("click", function() { - document.getElementById("resize").classList.add("resize"); +formImage.addEventListener('change', () => { + const file = formImage.files[0]; + if (!file) return; + file.arrayBuffer().then((buffer) => loadImage(buffer, URL.createObjectURL(file))); }); -// User selected an image from disk. Start the demo. -formImage.addEventListener("change", function() { - reset(); - // Should only get one file from picker - for (const file of formImage.files) { - let originalUrl = URL.createObjectURL(file); - for (let imageContainer of imageContainers) { - let originalImage = document.createElement("img"); - originalImage.src = originalUrl; - imageContainer.appendChild(originalImage); - } +document.getElementById('resize-button').addEventListener('click', () => { + document.getElementById('resize').classList.add('resize'); +}); - file.arrayBuffer().then((buf) => { - buf = new Uint8Array(buf); - resizeWorker.postMessage(buf); - }); +// Sample images, also used to seed the ranking pool +const samples = ['img/moon1.jpg', 'img/moon2.jpg', 'img/sunflower1.jpg', 'img/sunflower2.jpg']; +const sampleContainer = document.getElementById('sample-images'); +for (const src of samples) { + fetch(src) + .then((response) => (response.ok ? response.arrayBuffer() : Promise.reject(response.status))) + .then((buffer) => { + const button = document.createElement('img'); + button.src = src; + button.className = 'sample'; + button.addEventListener('click', () => loadImage(buffer, src)); + sampleContainer.appendChild(button); + addToPool(src.split('/').pop(), buffer, src); + }) + .catch(() => {}); +} + +// Mutation demo + +const mutators = { + flip: { label: 'Horizontal flip' }, + jpeg: { label: 'JPEG quality', min: 1, max: 100, step: 1, value: 50 }, + blur: { label: 'Gaussian blur', min: 0, max: 10, step: 0.1, value: 1.5 }, + sharpen: { label: 'Unsharp mask', min: 0, max: 10, step: 0.1, value: 1.5 }, + scale: { label: 'Rescale %', min: 5, max: 100, step: 1, value: 50 }, + crop: { label: 'Crop, keep %', min: 50, max: 100, step: 1, value: 90 }, + rotate: { label: 'Rotate deg', min: -45, max: 45, step: 0.5, value: 2 }, + hue: { label: 'Hue shift deg', min: 0, max: 180, step: 1, value: 30 }, + brightness: { label: 'Brightness', min: -100, max: 100, step: 1, value: 30 }, + contrast: { label: 'Contrast', min: -100, max: 100, step: 1, value: 25 }, + letterbox: { label: 'Letterbox bar %', min: 1, max: 40, step: 1, value: 10 }, + logo: { label: 'Logo size %', min: 1, max: 60, step: 1, value: 10 }, + noise: { label: 'Gaussian noise', min: 0, max: 50, step: 1, value: 10 }, +}; + +const kindSelect = document.getElementById('mutator-kind'); +const amountSlider = document.getElementById('mutator-amount'); +const amountValue = document.getElementById('mutator-amount-value'); + +for (const [kind, config] of Object.entries(mutators)) { + const option = document.createElement('option'); + option.value = kind; + option.innerText = config.label; + kindSelect.appendChild(option); +} + +function resetMutation() { + const config = mutators[kindSelect.value]; + if (config.min === undefined) { + amountSlider.disabled = true; + amountValue.innerText = ''; + } else { + amountSlider.disabled = false; + amountSlider.min = config.min; + amountSlider.max = config.max; + amountSlider.step = config.step; + amountSlider.value = config.value; + amountValue.innerText = config.value; } +} - // Show the rest of the demo. - document.getElementById("demo-resize").classList.add("visible"); -}); \ No newline at end of file +let mutationTimer = null; +async function runMutation() { + if (!baseBuffer) return; + const kind = kindSelect.value; + const amount = amountSlider.disabled ? 0 : Number(amountSlider.value); + amountValue.innerText = amountSlider.disabled ? '' : amount; + const data = await rpc('mutate', { buffer: baseBuffer, kind, amount }); + + document.getElementById('compare-a').src = baseUrl; + document.getElementById('compare-b').src = pngUrl(data.png); + + const names = ['dct', 'median', 'phash']; + names.forEach((name, i) => { + document.getElementById(`dist-${name}`).innerText = hamming(baseHashes[i], data.hashes[i]); + }); +} + +function scheduleMutation() { + clearTimeout(mutationTimer); + mutationTimer = setTimeout(runMutation, 120); +} + +kindSelect.addEventListener('change', () => { + resetMutation(); + scheduleMutation(); +}); +amountSlider.addEventListener('input', scheduleMutation); +resetMutation(); + +document.getElementById('compare-slider').addEventListener('input', (e) => { + document.getElementById('compare-b').style.clipPath = `inset(0 0 0 ${e.target.value}%)`; +}); + +// Ranking demo + +const pool = []; // {name, url, hashes} +const poolContainer = document.getElementById('ranking-pool'); +const resultsContainer = document.getElementById('ranking-results'); +const algoSelect = document.getElementById('ranking-algo'); +let queryIndex = null; + +async function addToPool(name, buffer, url) { + const { hashes } = await rpc('hashes', { buffer }); + const index = pool.length; + pool.push({ name, url, hashes }); + + const thumb = document.createElement('figure'); + const img = document.createElement('img'); + img.src = url; + const caption = document.createElement('figcaption'); + caption.innerText = name; + thumb.append(img, caption); + thumb.addEventListener('click', () => { + queryIndex = index; + rank(); + }); + poolContainer.appendChild(thumb); +} + +function rank() { + if (queryIndex === null) return; + const algo = Number(algoSelect.value); + const query = pool[queryIndex]; + const ranked = pool + .map((entry) => ({ entry, distance: hamming(query.hashes[algo], entry.hashes[algo]) })) + .sort((a, b) => a.distance - b.distance) + .slice(0, 10); + + resultsContainer.replaceChildren(); + for (const { entry, distance } of ranked) { + const thumb = document.createElement('figure'); + const img = document.createElement('img'); + img.src = entry.url; + const caption = document.createElement('figcaption'); + caption.innerText = `${distance} ${entry.name}`; + if (entry === query) thumb.className = 'query'; + thumb.append(img, caption); + resultsContainer.appendChild(thumb); + } +} + +algoSelect.addEventListener('change', rank); + +document.getElementById('ranking-files').addEventListener('change', (e) => { + for (const file of e.target.files) { + file.arrayBuffer().then((buffer) => addToPool(file.name, buffer, URL.createObjectURL(file))); + } +}); diff --git a/web/style.css b/web/style.css index 4dab206..d437cb4 100644 --- a/web/style.css +++ b/web/style.css @@ -2,7 +2,7 @@ body { margin: 0; padding: 0; color: var(--foreground); - background: url('bg.jpeg'); + background: var(--background) url('bg.jpeg'); background-size: cover; font-family: sans-serif; height: 100%; @@ -75,12 +75,12 @@ img { /* Demo */ -#demo-resize, #demo-dct { - visibility: hidden; +.needs-image { + display: none; } -#demo-resize.visible, #demo-dct.visible { - visibility: visible; +.needs-image.visible { + display: block; } #base .image-original img { @@ -209,4 +209,138 @@ figure { 100% { transform: rotate(360deg); } -} \ No newline at end of file +} +/* Skeleton */ + +.todo { + opacity: 0.55; + font-style: italic; + border-left: 3px solid var(--color3); + padding-left: 8px; +} + +/* Hash visualisation */ + +.hash-panel { + display: flex; + flex-flow: row wrap; + gap: 24px; + align-items: flex-start; +} + +.hash-value { + font-size: 18px; + letter-spacing: 2px; +} + +.bit-grid { + display: grid; + grid-template-columns: repeat(8, 20px); +} + +.bit-grid.wide { + grid-template-columns: repeat(12, 20px); +} + +.bit-grid .bit { + height: 20px; + box-sizing: border-box; + border: 1px solid var(--background); +} + +.bit-grid .bit.on { + background: var(--foreground); +} + +.bit-grid .bit.off { + background: color-mix(in srgb, var(--foreground) 12%, transparent); +} + +.heat-grid { + display: grid; + grid-template-columns: repeat(8, 32px); +} + +.heat-grid > div { + height: 32px; + box-sizing: border-box; +} + +.pixelated, #image-resize img, #phash-resize img { + image-rendering: pixelated; + image-rendering: -moz-crisp-edges; +} + +#phash-resize img { + width: 128px; + height: 128px; +} + +.sample { + height: 48px; + margin: 0 4px; + cursor: pointer; + vertical-align: middle; +} + +/* Mutation compare */ + +.compare { + position: relative; + max-width: 640px; +} + +.compare img { + display: block; + width: 100%; +} + +.compare #compare-b { + position: absolute; + top: 0; + left: 0; + clip-path: inset(0 0 0 50%); +} + +#compare-slider, #mutator-amount { + width: 320px; +} + +.distances td { + text-align: center; + font-size: 18px; + min-width: 64px; + color: var(--color1); +} + +/* Ranking */ + +.thumb-grid { + display: flex; + flex-flow: row wrap; + gap: 8px; +} + +.thumb-grid figure { + width: 96px; + margin: 0; + cursor: pointer; + text-align: center; +} + +.thumb-grid figure.query { + outline: 2px solid var(--color1); +} + +.thumb-grid img { + width: 96px; + height: 96px; + object-fit: cover; +} + +.thumb-grid figcaption { + font-size: 11px; + overflow: hidden; + text-overflow: ellipsis; + white-space: nowrap; +} diff --git a/web/worker.js b/web/worker.js new file mode 100644 index 0000000..178aceb --- /dev/null +++ b/web/worker.js @@ -0,0 +1,37 @@ +import init, { hash_all, resize_preview, dct_coefficients, dct_masks, phash_lowfreq, mutate } from './pkg/image_similarity.js'; + +const ready = init(); + +onmessage = async (e) => { + const { id, op, args } = e.data; + await ready; + try { + const bytes = new Uint8Array(args.buffer); + let result; + switch (op) { + case 'pipeline': + result = { + resize8: resize_preview(bytes, 8), + resize32: resize_preview(bytes, 32), + coefficients: dct_coefficients(bytes), + masks: dct_masks(bytes), + lowfreq: phash_lowfreq(bytes), + hashes: hash_all(bytes), + }; + break; + case 'hashes': + result = { hashes: hash_all(bytes) }; + break; + case 'mutate': { + const png = mutate(bytes, args.kind, args.amount); + result = { png, hashes: hash_all(png) }; + break; + } + default: + throw new Error(`unknown op ${op}`); + } + postMessage({ id, ok: true, result }); + } catch (err) { + postMessage({ id, ok: false, error: String(err) }); + } +};