mark/image-similarity
1
0
Fork 0
Code Pull Requests Packages Projects Releases Activity

Various updates

Browse Source 
too lazy to split commits. some refactoring, clap, more binaries.
This commit is contained in:
Mark Hoekveen
2024-05-24 09:59:49 +02:00
parent 1b2536f1e0
commit 46216c230e
5 changed files with 224 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/bin/dct.rs
+1 -1
View File
@@ -1,6 +1,6 @@
use std::env;
use std::f64::consts::{PI, SQRT_2};
use image::{save_buffer, GenericImageView, GrayImage};
use image::{GenericImageView, GrayImage};
fn dct(img: &image::DynamicImage) -> GrayImage {
//let mut dct_values: = [0.0; img.width() * img.height()];
src/bin/dctquery.rs
+19 -29
View File
@@ -1,50 +1,40 @@
//use image_similarity::descriptors::dct::DCT;
use image_similarity::descriptors::{DCT, Descriptor};
use image_similarity::store::DescriptorStore;
use std::env;
use log::{info, debug, error};
use log::info;
use clap::Parser;
#[derive(Parser)]
#[command(version, about, long_about = None)]
struct Cfg {
path: String
}
impl Cfg{
fn load(args: &[String]) -> Result<Cfg, &'static str> {
if args.len() < 2 {
return Err("Filename argument required");
}
let path = args[1].clone();
Ok(Cfg { path })
}
/// 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,
distance: isize,
}
fn main() {
env_logger::init();
//Init all descriptors:
let desc = DCT::new().with_quality(50);
let args: Vec<String> = env::args().collect();
let cfg = Cfg::load(&args)
.expect("Error loading config");
let cfg = Cfg::parse();
let img = image::open(cfg.path)
.expect("Unable to open file");
println!("Iamge");
let conf = viuer::Config {
width: Some(16),
height: Some(8),
..Default::default()
};
viuer::print(&img, &conf).expect("Image printing failed.");
println!("Iamge");
// let conf = viuer::Config {
// width: Some(16),
// height: Some(8),
// ..Default::default()
// };
// viuer::print(&img, &conf).expect("Image printing failed.");
let store =
DescriptorStore::new()
.with_file("store.messagepack");
//info!("{}", store);
.with_file("dct50.messagepack");
let phash: u64 = desc.describe(&img);
info!("Phash integer:\n{phash}");
info!("Phash binary:\n{phash:064b}");
store.knn(phash, 10);
store.print_nn(phash, cfg.distance, cfg.show_images);
}
src/descriptors/dct.rs
+2 -2
View File
@@ -174,8 +174,8 @@ impl Descriptor for DCT {
let signum = dct_values[i].signum();
// Only multiply sign if dct-coefficient contains a horizontal component.
if
//i % 8 != 0 && sign_mult * signum < 0.0
// ||
i % 8 != 0 && sign_mult * signum < 0.0
||
signum < 0.0
{
sign_mask += 1;
src/main.rs
+64 -22
View File
@@ -12,7 +12,7 @@ struct Cfg {
path: PathBuf
}
fn make_store<T: Descriptor>(descriptor: T, input_path: PathBuf, save_path: PathBuf) -> () {
fn make_store<T: Descriptor>(descriptor: T, input_path: PathBuf, save_path: PathBuf) -> DescriptorStore {
let mut store =
DescriptorStore::new()
.with_file(save_path);
@@ -20,6 +20,13 @@ fn make_store<T: Descriptor>(descriptor: T, input_path: PathBuf, save_path: Path
store.insert_directory(input_path, descriptor);
info!("{}", store);
store.save().expect("Error saving");
store
}
struct StoreParams {
desc: DCT,
input_path: PathBuf,
save_path: PathBuf,
}
fn main() {
@@ -28,32 +35,67 @@ fn main() {
let mut threads = vec![];
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 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(),
// },
];
// 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::<DescriptorStore>();
for param in params {
threads.push(
thread::spawn(move || {
make_store(dct_50, "data".into(), "dct50.messagepack".into());
let store = make_store(param.desc, param.input_path, param.save_path);
store.print_most_dups();
//tx.send(store).unwrap();
})
);
threads.push(
thread::spawn(move || {
make_store(dct_30, "data".into(), "dct30.messagepack".into());
})
);
threads.push(
thread::spawn(move || {
make_store(dct_10, "data".into(), "dct10.messagepack".into());
})
);
)
}
// 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();
}
}
src/store.rs
+121 -28
View File
@@ -1,7 +1,7 @@
//! Persistent store of all calculated image descriptions.
//! Can be queried to find identical or similar images.
use std::collections::HashMap;
use std::collections::{HashMap, HashSet};
use crate::descriptors::Descriptor;
use std::fs;
use std::path::Path;
@@ -15,20 +15,33 @@ pub enum SaveError {
File,
}
/// Uses a hashmap to map descriptors to buckets of files.
/// Also keeps a BK-tree for quick distance ranking
pub struct DescriptorStore {
map: HashMap<u64, String>,
/// Main hashmap that maps descriptors to buckets of filenames
map: HashMap<u64, Vec<String>>,
/// Place to load and store the map on file
save_location: std::path::PathBuf,
/// BK-tree for fast nearest neighbour
bktree: BkTree<u64>,
/// Cache for seen files to skip on initial load
seen: Option<HashSet<String>>,
}
impl DescriptorStore {
/// Makes a new empty DescriptorStore with default settings
pub fn new() -> Self {
let map: HashMap<u64, String> = HashMap::new();
let map: HashMap<u64, Vec<String>> = HashMap::new();
//let seen: HashSet<String> = HashSet::new();
//let seen = None;
DescriptorStore {
map,
save_location: std::path::PathBuf::from("store.messagepack"),
bktree: BkTree::new(hamming_distance)
bktree: BkTree::new(hamming_distance),
seen: None,
}
}
@@ -37,12 +50,21 @@ impl DescriptorStore {
self.save_location = std::path::PathBuf::from(path.as_ref());
let map_file = fs::read(&self.save_location);
match map_file {
Ok(f) => self.map = rmp_serde::from_slice(&f).unwrap(),
Err(_) => info!("{} not found, starting from empty store.", self.save_location.display()),
};
Ok(f) => {
self.map = rmp_serde::from_slice(&f).unwrap();
for i in self.map.keys() {
self.bktree.insert(*i);
}
let mut seen: HashSet<String> = HashSet::new();
for bucket in self.map.values() {
for element in bucket {
seen.insert(element.clone());
}
}
self.seen = Some(seen);
},
Err(_) => info!("{} not found, starting from empty store.", self.save_location.display()),
};
self
}
@@ -62,10 +84,38 @@ 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
}
}
}
/// Inserts a single value into the store
pub fn insert(&mut self, key: u64, value: String) {
self.map.insert(key, value);
self.bktree.insert(key);
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);
}
n
},
None => vec![value],
};
self.map.insert(*key, bucket);
self.bktree.insert(*key);
}
/// Calculates all descriptions with a given descriptor
@@ -79,7 +129,7 @@ impl DescriptorStore {
continue
}
};
if true { // !self.contains(name.to_string()) {
if !self.has_value(&name) { // !self.contains(name.to_string()) {
info!("Processing {}", name);
let img = match image::open(file.path()) {
Ok(v) => v,
@@ -92,26 +142,36 @@ impl DescriptorStore {
if self.contains(phash) {
println!("{} duplicate of {:?}", name, self.map.get(&phash));
}
self.map.insert(phash, name);
self.insert(&phash, name);
} else {
debug!("{} already known, skipping.", name);
}
self.save().expect("error");
}
// We have done our bulk loading, so we can unset our hashset cache:
self.seen = None;
}
/// K-nearest neighbors
pub fn knn(&self, from: u64, k: isize) {
println!("{from:b}");
/// Nearest neighbours
pub fn nn(&self, from: u64, max_distance: isize) -> Vec<(&u64, isize)> {
let mut neighbours = self.bktree.find(from, max_distance);
neighbours.sort_by(|a, b| a.1.cmp(&b.1));
neighbours
}
pub fn print_nn(&self, from: u64, max_distance: isize, show_images: bool) {
//println!("{from:064b}");
let neighbours = self.nn(from, max_distance);
let (term_width, _) = viuer::terminal_size();
let mut x = 0;
let mut y = 8;
let (term_width, _) = viuer::terminal_size();
println!("\t\t\t\t\tTERMWIDTH {term_width}");
for (element, distance) in self.bktree.find(from, k) {
let name = self.map.get(element);
match name {
Some(n) => {
println!("{element:b}: {n} (distance: {distance})");
for (element, distance) in neighbours {
let elements = self.map.get(element);
match elements {
Some(paths) => {
for path in paths {
println!("{element:064b}: {:?} (distance: {distance})", path);
if show_images {
if x+16 >= term_width {
x = 0;
y += 8;
@@ -119,27 +179,60 @@ impl DescriptorStore {
let conf = viuer::Config {
width: Some(16),
height: Some(8),
x: x,
y: y,
x,
y,
use_kitty: false,
..Default::default()
};
x += 16;
let path = "data/".to_string() + &n;
let path = "data/".to_string() + path;
let img = image::open(&path).unwrap();
let img = img.grayscale().thumbnail_exact(8, 8);
viuer::print(&img, &conf).expect("Image printing failed.");
}
}
},
None => ()
};
}
}
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;
}
}
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_str("\n");
}
}
print!("DescriptorStore:\n{}Total: {}", output, self.map.len())
}
}
impl fmt::Display for DescriptorStore {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let mut output = String::new();
for (key, val) in self.map.iter() {
let entry = format!("{key:064b}: {val}\t {val:020}\n");
output.push_str(&entry);
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_str("\n");
}
write!(f, "DescriptorStore:\n{}Total: {}", output, self.map.len())
}