diff --git a/src/bin/dct.rs b/src/bin/dct.rs deleted file mode 100644 index 29881d4..0000000 --- a/src/bin/dct.rs +++ /dev/null @@ -1,66 +0,0 @@ -use std::env; -use std::f64::consts::{PI, SQRT_2}; -use image::{GenericImageView, GrayImage}; - -fn dct(img: &image::DynamicImage) -> GrayImage { - //let mut dct_values: = [0.0; img.width() * img.height()]; - let w = img.width(); - let h = img.height(); - let mut dct_img = GrayImage::new(w, h); - - //let mut dct_values: Vec = Vec::new(); - for u in 0..w { - for v in 0..h { - let w = w as f64; - let h = h as f64; - //let k = (v*img.width())+u; - let mut alpha = 0.25; - if u == 0 { - alpha = alpha / SQRT_2 - } - if v == 0 { - alpha = alpha / SQRT_2 - } - let v: f64 = v as f64; - let u: f64 = u as f64; - let mut sum: f64 = 0.0; - for (x, y, pix) in img.pixels() { - let x: f64 = x as f64; - let y: f64 = y as f64; - let pixel = pix[0] as f64; - // sum += - // pixel * - // (x*u*PI/16.0).cos() * - // (y*v*PI/16.0).cos() - sum += - pixel * - ( - PI/w* - (x+0.5)*u - ).cos() - * - ( - PI/h* - (y+0.5)*v - ).cos() - } - //dct_values[k] = alpha * sum; - //dct_values.push(alpha * sum); - print!("{sum}\t"); - //sum = (sum + 200000.0)/784.0; - print!("{}\n", sum * alpha); - let pixel = (sum).abs() as u8; - println!("{u}, {v}:\t {sum}"); - dct_img.put_pixel(u as u32, v as u32, image::Luma([pixel])); - } - } - dct_img -} - -fn main() { - let args: Vec = env::args().collect(); - let path = args[1].clone(); - let img = image::open(&path).unwrap(); - let dct = dct(&img); - dct.save("dct.png").unwrap(); -} \ No newline at end of file diff --git a/src/bin/mutate.rs b/src/bin/mutate.rs new file mode 100644 index 0000000..650e47a --- /dev/null +++ b/src/bin/mutate.rs @@ -0,0 +1,38 @@ +use image_similarity::mutators::get_all_mutators; +use clap::Parser; + +#[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() { + env_logger::init(); + + let cfg = Cfg::parse(); + + let img = image::open(&cfg.path) + .expect("Unable to open file"); + + let mutators = get_all_mutators(); + + let conf = viuer::Config { + width: Some(16), + height: Some(8), + ..Default::default() + }; + for mutator in mutators { + let mutated = mutator.mutate(&img); + let filename = "mut".to_string() + &mutator.tag() + &"jpg".to_string(); + println!("Saving to {}", filename); + mutated.save(filename).expect("Saving image failed"); + if cfg.show_images { + viuer::print(&mutated, &conf).expect("Image printing failed."); + } + } +} \ No newline at end of file diff --git a/src/descriptors/dct.rs b/src/descriptors/dct.rs index 52cfc5b..a903fe1 100644 --- a/src/descriptors/dct.rs +++ b/src/descriptors/dct.rs @@ -104,7 +104,7 @@ impl DCT { impl Descriptor for DCT { fn info(&self) -> String { - "DCT".to_string() + "dct".to_string() } fn describe(&self, img: &image::DynamicImage) -> u64 { let resized = self.resize(img); diff --git a/src/descriptors/mod.rs b/src/descriptors/mod.rs index 8d5806c..5e0e747 100644 --- a/src/descriptors/mod.rs +++ b/src/descriptors/mod.rs @@ -1,6 +1,8 @@ //! Descriptors that can be used to describe an image. //! If two images are (almost) the same, their descriptions will be the same. +use image::GenericImageView; + pub trait Descriptor { /// Print the name of the descriptor, fn info(&self) -> String; @@ -36,34 +38,44 @@ pub struct DCT { } pub mod dct; -// fn median64(values: &[T]) -> T { -// let mut sorted_values = values.to_vec(); -// sorted_values.sort(); -// let len = sorted_values.len(); -// sorted_values[len/2] -// } +fn median64(values: &[T]) -> T { + let mut sorted_values = values.to_vec(); + sorted_values.sort(); + let len = sorted_values.len(); + sorted_values[len/2] +} +pub struct Median; -// pub struct Median; +impl Descriptor for Median { + fn info(&self) -> String { + "median".to_string() + } -// impl Descriptor for Median { -// fn describe(&self, img: image::DynamicImage) -> u64 { -// let img = self.resize(img); -// let mut values: [u8; 64] = [0; 64]; -// let mut i: usize = 0; -// for (_, _, pix) in img.pixels() { -// values[i] = pix[0]; -// i = i+1; -// } -// 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; -// } -// mask = mask << 1; -// } -// mask -// } -// } + fn describe(&self, img: &image::DynamicImage) -> u64 { + let img = self.resize(img); + let mut values: [u8; 64] = [0; 64]; + let mut i: usize = 0; + for (_, _, pix) in img.pixels() { + values[i] = pix[0]; + i = i+1; + } + 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; + } + mask = mask << 1; + } + mask + } +} + +pub fn get_all_descriptors() -> Vec> { + let mut descriptors: Vec> = Vec::with_capacity(2); + descriptors.push(Box::new(DCT::new())); + descriptors.push(Box::new(Median)); + descriptors +} \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index ea0f982..dbce76e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,4 +3,5 @@ //! Descriptors that can be used to describe an image. //! If two images are (almost) the same, their descriptions will be the same. pub mod descriptors; -pub mod store; \ No newline at end of file +pub mod store; +pub mod mutators; \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index c1317ff..54ac60c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,101 +1,63 @@ use image_similarity::store::DescriptorStore; -use image_similarity::descriptors::{Descriptor, DCT}; -use log::info; +use image_similarity::descriptors::get_all_descriptors; +use image_similarity::mutators::get_all_mutators; +use log::{error, info}; use std::path::PathBuf; -use std::thread; +use std::{fs, thread}; use clap::Parser; #[derive(Parser)] #[command(version, about, long_about = None)] struct Cfg { /// Path that contains the input images. Will not traverse directories. - path: PathBuf -} - -fn make_store(descriptor: T, input_path: PathBuf, save_path: PathBuf) -> DescriptorStore { - let mut store = - DescriptorStore::new(descriptor) - .with_file(save_path); - info!("{}", store); - store.insert_directory(input_path); - info!("{}", store); - store.save().expect("Error saving"); - store -} - -struct StoreParams { - desc: DCT, - input_path: PathBuf, - save_path: PathBuf, + path: PathBuf, } 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 mut threads = vec![]; + for descriptor in get_all_descriptors() { + let store = DescriptorStore::new(descriptor); + stores.push(store); + } + + for node in fs::read_dir(cfg.path).unwrap() { + println!("{:?}", node); + let file = node.expect("Error walking directory"); + let name = match file.file_name().into_string() { + Ok(v) => v, + Err(e) => { + error!("Error reading {}: {:?}:", file.path().to_string_lossy(), e); + continue + } + }; + let img = match image::open(file.path()) { + Ok(v) => v, + Err(e) => { + error!("Failed to process {}: {}", name, e); + continue + } + }; - let params = vec![ - StoreParams { - desc: DCT::new().with_quality(50), - input_path: cfg.path.clone(), - save_path: "dct50.messagepack".into(), - }, - // StoreParams { - // desc: DCT::new().with_quality(30), - // input_path: cfg.path.clone(), - // save_path: "dct30.messagepack".into(), - // }, - // StoreParams { - // desc: DCT::new().with_quality(10), - // input_path: cfg.path.clone(), - // save_path: "dct10.messagepack".into(), - // }, - ]; + //Store the phashes of the base image + for store in &mut stores { + store.store(&img, name.clone()); + } - // let dct_50 = - // DCT::new() - // .with_quality(50); - // let dct_30 = - // DCT::new() - // .with_quality(30); - // let dct_10 = - // DCT::new() - // .with_quality(10); - - //let (tx, rx) = mpsc::channel::(); - - for param in params { - threads.push( - thread::spawn(move || { - let store = make_store(param.desc, param.input_path, param.save_path); - store.print_most_dups(); - //tx.send(store).unwrap(); - }) - ) + 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()); + } + } } - // let path_50 = cfg.path.clone(); - // threads.push( - // thread::spawn(move || { - // let store = make_store(dct_50, path_50, "dct50.messagepack".into()); - // tx.send(store).expect("Thread send result error"); - // }) - // ); - // let path_30 = cfg.path.clone(); - // threads.push( - // thread::spawn(move || { - // dct_30_store = make_store(dct_30, path_30, "dct30.messagepack".into()); - // }) - // ); - // let path_10 = cfg.path.clone(); - // threads.push( - // thread::spawn(move || { - // dct_50_store = make_store(dct_10, path_10, "dct10.messagepack".into()); - // }) - // ); - for thread in threads { - let _ = thread.join(); + for store in &stores { + println!("{}", store); } - } diff --git a/src/mutators/mod.rs b/src/mutators/mod.rs new file mode 100644 index 0000000..0aebd3c --- /dev/null +++ b/src/mutators/mod.rs @@ -0,0 +1,81 @@ +//! Mutators that can be used to mutate an image. +//! These mutated images can be used as a "near copy" + +pub trait Mutator { + /// Name/description of the mutator + fn info(&self) -> String; + + /// Short tag that represents the mutator. + /// To be used in the filename such that it can be recognized as mutated image + fn tag(&self) -> String { + "MUT".to_string() + } + + /// Returns the mutated form of the input image + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage; +} + +/// Horizontal flip/mirror +pub struct Flip; +impl Mutator for Flip { + fn info(&self) -> String { + "Horizontal flip".to_string() + } + fn tag(&self) -> String { + ".flip.".to_string() + } + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { + img.fliph() + } +} + +/// Slight hue-shift +pub struct Hue; +impl Mutator for Hue { + fn info(&self) -> String { + "Hue shift".to_string() + } + fn tag(&self) -> String { + ".hue.".to_string() + } + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { + img.huerotate(5) + } +} + +/// Slight sharpen (unsharp mask) +pub struct Sharp; +impl Mutator for Sharp { + fn info(&self) -> String { + "Unsharp mask".to_string() + } + fn tag(&self) -> String { + ".sharp.".to_string() + } + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { + img.unsharpen(1.2, 50) + } +} + +/// Slight blur +pub struct Blur; +impl Mutator for Blur { + fn info(&self) -> String { + "Gaussian blur".to_string() + } + fn tag(&self) -> String { + ".blur.".to_string() + } + fn mutate(&self, img: &image::DynamicImage) -> image::DynamicImage { + img.blur(1.5) + } +} + +pub fn get_all_mutators() -> Vec> { + let mut mutators: Vec> = Vec::with_capacity(4); + mutators.push(Box::new(Flip)); + mutators.push(Box::new(Hue)); + mutators.push(Box::new(Sharp)); + mutators.push(Box::new(Blur)); + mutators +} \ No newline at end of file diff --git a/src/store.rs b/src/store.rs index 231c7b3..62b818f 100644 --- a/src/store.rs +++ b/src/store.rs @@ -6,6 +6,7 @@ use crate::descriptors::Descriptor; use std::fs; use std::path::Path; use std::fmt; +use image::DynamicImage; use log::{info, debug, error}; use bktree::*; @@ -17,8 +18,8 @@ pub enum SaveError { /// Uses a hashmap to map descriptors to buckets of files. /// Also keeps a BK-tree for quick distance ranking -pub struct DescriptorStore { - descriptor: T, +pub struct DescriptorStore { + descriptor: Box, /// Main hashmap that maps descriptors to buckets of filenames map: HashMap>, @@ -32,9 +33,9 @@ pub struct DescriptorStore { seen: Option>, } -impl DescriptorStore { +impl DescriptorStore { /// Makes a new empty DescriptorStore with default settings - pub fn new(descriptor: T) -> Self { + pub fn new(descriptor: Box) -> Self { let map: HashMap> = HashMap::new(); //let seen: HashSet = HashSet::new(); //let seen = None; @@ -154,6 +155,13 @@ impl DescriptorStore { self.seen = None; } + 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: isize) -> Vec<(&u64, isize)> { let mut neighbours = self.bktree.find(from, max_distance); @@ -224,7 +232,7 @@ impl DescriptorStore { } } -impl fmt::Display for DescriptorStore { +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() {