diff --git a/src/Feature/mod.rs b/src/Feature/mod.rs new file mode 100644 index 0000000..9507bce --- /dev/null +++ b/src/Feature/mod.rs @@ -0,0 +1,106 @@ +use std::sync::Arc; + +#[derive(Debug, Clone, Serialize, Deserialize)] +enum FeatureResult { + /// A boolean. Just a boolean + Bool(bool), + /// Signed 32-bit integer + I32(i32), + /// 32-bit single precision floating point + /// can be used for aspect ratio or luminance + F32(f32), + /// Vector for nested multidimensional + Vec(Vec), + /// Standard RGBA color + RGBA(f32, f32, f32, f32), + /// Indices intended for the usage in historgrams + Indices(Vec) +} + +impl Default for FeatureResult { + fn default() -> Self { + FeatureResult::Bool(false) + } +} + +/// For some feature return type we want to implement a custom compare function +/// for example: historgrams are compared with cosine similarity +impl PartialEq for FeatureResult { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (Self::Bool(l0), Self::Bool(r0)) => l0 == r0, + (Self::I32(l0), Self::I32(r0)) => l0 == r0, + (Self::F32(l0), Self::F32(r0)) => l0 == r0, + (Self::Vec(l0), Self::Vec(r0)) => l0 == r0, + (Self::RGBA(l0, l1, l2, l3), Self::RGBA(r0, r1, r2, r3)) => l0 == r0 && l1 == r1 && l2 == r2 && l3 == r3, + (Self::Indices(_), Self::Indices(_)) => todo!("implement cosine similarity"), + _ => false, + } + } +} + +type FeatureGenerator = Box>) -> (String, FeatureResult)>; + +#[derive(Serialize, Deserialize, Default)] +struct Database { + images: HashMap>, + + /// keep feature generator for the case when we add a new image + /// this field is not serialized and needs to be wrapped in an option + #[serde(skip)] + generators: Option> +} + +impl Database { + + pub fn add_feature(&mut self, feature: FeatureGenerator) { + for (path, features) in self.images.iter_mut() { + // compute feature for every image + todo!("run this as a closure parallel with a thread pool"); + let (name, res) = feature(todo!("load image from disk")); + features.insert(name, res); + } + + if let Some(generators) = self.generators.as_mut() { + generators.push(feature); + } else { + self.generators = Some(vec![feature]) + } + } + + pub fn add_image(&mut self, path: String) { + let image = todo!("load image from disk"); + let mut features = HashMap::new(); + if let Some(generators) = self.generators { + for generator in generators.iter() { + let (name, res) = generator(image); + features.insert(name, res); + } + } + self.images.insert(path, features); + } +} + +fn average_luminance(image: Arc>) -> (String, FeatureResult) { + let num_pixels = image.pixels.len() as u32; + let total_brightness: f32 = image.pixels + .iter() + .map(|(r, g, b, _)| 0.299 * r + 0.587 * g + 0.114 * b) // Calculate Y for each pixel + .sum(); + let average_brightness = total_brightness / num_pixels as f32; + + let feature_name = String::from("average-brightness"); + let feature_result = FeatureResult::F32(average_brightness); + + (feature_name, feature_result) +} +fn compare_Dim(image0: Arc>, image1: Arc>) -> (String, FeatureResult) { + let a = image0.width as f32 / image0.height as f32; + let b = image1.width as f32 / image1.height as f32; + let equal = a == b; + + let feature_name = String::from("Dimension-comparison"); + let feature_result = FeatureResult::Bool(equal); + + (feature_name, feature_result) +} diff --git a/src/feature/mod.rs b/src/feature/mod.rs new file mode 100644 index 0000000..b5bd2c9 --- /dev/null +++ b/src/feature/mod.rs @@ -0,0 +1,144 @@ +//! # Prebuild features +//! This module provides a set of prebuild features ready to be used with a database +//! to index images. +//! Features include: +//! - distribution of colors (via histogram) +//! - distribution of luminance (via histogram) +//! - average luminance +//! - aspect ratio of images computed a width/height +//! All features are designed to used with sRGB color channels only. + +use std::sync::Arc; +use crate::{image::Image, search_index::FeatureResult}; + +#[allow(unused)] +// from https://github.com/programmieren-mit-rust/pr-ferrisgroup/issues/8 by @SirTalksalot75 +/// Compute a basic distribution of values from all color channels and count their apprearances in buckets. +/// This function will use 5 buckets per channel. +fn color_distribution(image: Arc>) -> (String, FeatureResult) { + const N: usize = 5; + + let mut histogram = vec![0u64; N * 3 + 1]; + + const INV_255: f32 = 1./255. * N as f32; + for (r, g, b, _) in image.iter() { + // map linear channel value to bin index + histogram[ (r * INV_255) as usize] += 1; + histogram[ (g * INV_255) as usize * 2 ] += 1; + histogram[ (b * INV_255) as usize * 3 ] += 1; + } + + (String::from("luminance-distribution"), FeatureResult::Indices(histogram)) +} + +#[allow(unused)] +// from https://github.com/programmieren-mit-rust/pr-ferrisgroup/issues/8 by @SirTalksalot75 +/// Compute a basic distribution of luminance values and count their apprearances in buckets. +/// Luminance is calculated via Digital ITU BT.601 and NOT the more common Photometric ITU BT.709 +fn luminance_distribution(image: Arc>) -> (String, FeatureResult) { + let mut histogram = vec![0u64; 256]; // Assuming 256 bins for the histogram + + for (r, g, b, _) in image.iter() { + // map luminance to bin index + // luminance is a value between 0 and 255. + let luminance = (0.299 * r + 0.587 * g + 0.114 * b) as usize; + histogram[luminance] += 1; + } + + (String::from("luminance-distribution"), FeatureResult::Indices(histogram)) +} + +#[allow(unused)] +// from https://github.com/programmieren-mit-rust/pr-ferrisgroup/issues/8 by @SirTalksalot75 +/// Compute the average luminance of all pixels in a given image. +/// Luminance is calculated via Digital ITU BT.601 and NOT the more common Photometric ITU BT.709 +fn average_luminance(image: Arc>) -> (String, FeatureResult) { + let num_pixels = image.pixels().len() as u32; + let total_brightness: f32 = image + .iter() + .map(|(r, g, b, _)| (0.299 * r + 0.587 * g + 0.114 * b) / 255.0) // Calculate Y for each pixel + .sum(); + let average_brightness = total_brightness / num_pixels as f32; + + let feature_name = String::from("average-brightness"); + let feature_result = FeatureResult::Percent(average_brightness); + + (feature_name, feature_result) +} + +#[allow(unused)] +// from https://github.com/programmieren-mit-rust/pr-ferrisgroup/issues/8 by @SirTalksalot75 +fn aspect_ratio(image: Arc>) -> (String, FeatureResult) { + let a = image.width() as f32 / image.height() as f32; + + (String::from("aspect-ratio"), FeatureResult::Percent(a)) +} + +#[cfg(test)] +mod test { + use std::path::Path; + + use crate::search_index::{Database, FeatureGenerator}; + + use super::*; + + #[test] + fn test_histogram() { + let files: Vec = std::fs::read_dir("res/integration/") + .unwrap() + .map(|f| f.unwrap().path()) + .collect(); + + let feats: Vec = vec![color_distribution]; + + let db = Database::new(&files, feats).unwrap(); + + for (path, sim) in db.search(Path::new("res/integration/gray_image.png"), color_distribution).unwrap() { + let file_name = path.file_name().unwrap().to_str().unwrap(); + if file_name.eq("gray_image.png") { + assert_eq!(sim, 1.); + } + println!("{} {}", file_name, sim); + } + } + + #[test] + fn test_average_luminance() { + let files: Vec = std::fs::read_dir("res/integration/") + .unwrap() + .map(|f| f.unwrap().path()) + .collect(); + + let feats: Vec = vec![average_luminance]; + + let db = Database::new(&files, feats).unwrap(); + + for (path, sim) in db.search(Path::new("res/integration/gray_image.png"), average_luminance).unwrap() { + let file_name = path.file_name().unwrap().to_str().unwrap(); + if file_name.eq("gray_image.png") { + assert_eq!(sim, 1.); + } + println!("{} {}", file_name, sim); + } + } + + #[test] + fn test_aspect_ratio() { + let files: Vec = std::fs::read_dir("res/integration/") + .unwrap() + .map(|f| f.unwrap().path()) + .collect(); + + let feats: Vec = vec![aspect_ratio]; + + let db = Database::new(&files, feats).unwrap(); + + for (path, sim) in db.search(Path::new("res/integration/gray_image.png"), aspect_ratio).unwrap() { + let file_name = path.file_name().unwrap().to_str().unwrap(); + if file_name.eq("gray_image.png") { + assert_eq!(sim, 1.); + } + println!("{} {}", file_name, sim); + } + } +} \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index 178f861..2532a15 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,25 @@ +//! # Imsearch +//! Extensible library for creating an image based search engine. +//! The library exposes the functionality to create databases which index various images stored as png files. +//! # Examples +//! ```ignore +//! let files: Vec = std::fs::read_dir("image/folder/") +//! .unwrap() +//! .map(|f| f.unwrap().path()) +//! .collect(); +//! +//! let feats: Vec = vec![average_rgb_value]; +//! +//! let db = Database::new(&files, feats).unwrap(); +//! +//! db.write_to_file(json); +//! ``` + extern crate core; pub mod image; pub mod image_loader; pub mod multithreading; pub mod search_index; +pub mod feature; +