83 lines
1.9 KiB
Rust
83 lines
1.9 KiB
Rust
|
use std::{path::Path, sync::Arc};
|
||
|
|
|
||
|
|
use imsearch::{
|
||
|
|
image::Image,
|
||
|
|
search_index::{Database, FeatureGenerator, FeatureResult},
|
||
|
|
};
|
||
|
|
|
||
|
|
fn average_brightness(image: Arc<Image<f32>>) -> (String, FeatureResult) {
|
||
|
|
let bright = image
|
||
|
|
.pixels()
|
||
|
|
.iter()
|
||
|
|
.map(|(r, g, b, _)| (r + g + b) / 3.0 / 255.0)
|
||
|
|
.sum::<f32>()
|
||
|
|
/ image.pixels().len() as f32;
|
||
|
|
|
||
|
|
(
|
||
|
|
String::from("average_brightness"),
|
||
|
|
FeatureResult::Percent(bright),
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_database_basic() {
|
||
|
|
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_brightness];
|
||
|
|
|
||
|
|
let db = Database::new(&files, feats).unwrap();
|
||
|
|
|
||
|
|
for results in db
|
||
|
|
.search(
|
||
|
|
std::path::Path::new("res/integration/gray_image.png"),
|
||
|
|
average_brightness,
|
||
|
|
)
|
||
|
|
.unwrap()
|
||
|
|
{
|
||
|
|
println!(
|
||
|
|
"path: {} similarity: {}",
|
||
|
|
results.0.as_os_str().to_str().unwrap(),
|
||
|
|
results.1
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_database_files() {
|
||
|
|
let json = Path::new("db.json");
|
||
|
|
|
||
|
|
{
|
||
|
|
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_brightness];
|
||
|
|
|
||
|
|
let db = Database::new(&files, feats).unwrap();
|
||
|
|
|
||
|
|
db.write_to_file(json);
|
||
|
|
}
|
||
|
|
|
||
|
|
{
|
||
|
|
let db = Database::from_file(json);
|
||
|
|
|
||
|
|
for results in db
|
||
|
|
.search(
|
||
|
|
std::path::Path::new("res/integration/gray_image.png"),
|
||
|
|
average_brightness,
|
||
|
|
)
|
||
|
|
.unwrap()
|
||
|
|
{
|
||
|
|
println!(
|
||
|
|
"path: {} similarity: {}",
|
||
|
|
results.0.as_os_str().to_str().unwrap(),
|
||
|
|
results.1
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|