Feature extractor (#48)
* Create FeatureTest.rs * Create mod.rs * Delete FeatureTest.rs * Added FeatureExtr FeatureExtractor provided by Servostar * Added AverageBrightness Feature * Added Dimension Compare Feature * Update mod.rs * added feature module --------- Co-authored-by: SirTalksalot75 <132705706+SirTalksalot75@users.noreply.github.com>
This commit is contained in:
parent
7928bf0725
commit
f913855fb3
|
@ -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<FeatureResult>),
|
||||||
|
/// Standard RGBA color
|
||||||
|
RGBA(f32, f32, f32, f32),
|
||||||
|
/// Indices intended for the usage in historgrams
|
||||||
|
Indices(Vec<u64>)
|
||||||
|
}
|
||||||
|
|
||||||
|
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<dyn Fn(crate::Arc<Image<f32>>) -> (String, FeatureResult)>;
|
||||||
|
|
||||||
|
#[derive(Serialize, Deserialize, Default)]
|
||||||
|
struct Database {
|
||||||
|
images: HashMap<String, HashMap<String, FeatureResult>>,
|
||||||
|
|
||||||
|
/// 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<Vec<FeatureGenerator>>
|
||||||
|
}
|
||||||
|
|
||||||
|
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<Image<f32>>) -> (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<Image<f32>>, image1: Arc<Image<f32>>) -> (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)
|
||||||
|
}
|
|
@ -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<Image<f32>>) -> (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<Image<f32>>) -> (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<Image<f32>>) -> (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<Image<f32>>) -> (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::path::PathBuf> = std::fs::read_dir("res/integration/")
|
||||||
|
.unwrap()
|
||||||
|
.map(|f| f.unwrap().path())
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let feats: Vec<FeatureGenerator> = 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::path::PathBuf> = std::fs::read_dir("res/integration/")
|
||||||
|
.unwrap()
|
||||||
|
.map(|f| f.unwrap().path())
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let feats: Vec<FeatureGenerator> = 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::path::PathBuf> = std::fs::read_dir("res/integration/")
|
||||||
|
.unwrap()
|
||||||
|
.map(|f| f.unwrap().path())
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
let feats: Vec<FeatureGenerator> = 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);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
19
src/lib.rs
19
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<PathBuf> = std::fs::read_dir("image/folder/")
|
||||||
|
//! .unwrap()
|
||||||
|
//! .map(|f| f.unwrap().path())
|
||||||
|
//! .collect();
|
||||||
|
//!
|
||||||
|
//! let feats: Vec<FeatureGenerator> = vec![average_rgb_value];
|
||||||
|
//!
|
||||||
|
//! let db = Database::new(&files, feats).unwrap();
|
||||||
|
//!
|
||||||
|
//! db.write_to_file(json);
|
||||||
|
//! ```
|
||||||
|
|
||||||
extern crate core;
|
extern crate core;
|
||||||
|
|
||||||
pub mod image;
|
pub mod image;
|
||||||
pub mod image_loader;
|
pub mod image_loader;
|
||||||
pub mod multithreading;
|
pub mod multithreading;
|
||||||
pub mod search_index;
|
pub mod search_index;
|
||||||
|
pub mod feature;
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue