aboutsummaryrefslogtreecommitdiff
path: root/src
diff options
context:
space:
mode:
Diffstat (limited to 'src')
-rw-r--r--src/embedders.rs25
-rw-r--r--src/main.rs39
2 files changed, 36 insertions, 28 deletions
diff --git a/src/embedders.rs b/src/embedders.rs
index 0693b5e..8911e95 100644
--- a/src/embedders.rs
+++ b/src/embedders.rs
@@ -1,5 +1,6 @@
+use anyhow::{bail, Result};
use rayon::prelude::*;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
use serde::{Deserialize, Serialize};
pub trait MetricElem: Send + Sync + 'static + Serialize + for<'a> Deserialize<'a> {
@@ -16,21 +17,21 @@ pub trait EmbedderT: Send + Sync {
type Embedding: MetricElem;
const NAME: &'static str;
- fn embed(&self, _: &PathBuf) -> Result<Self::Embedding, String>;
+ fn embed(&self, _: &Path) -> Result<Self::Embedding>;
}
pub trait BatchEmbedder: Send + Sync {
type Embedding: MetricElem;
const NAME: &'static str;
- fn embeds(&mut self, _: &[PathBuf]) -> Result<Vec<Self::Embedding>, String>;
+ fn embeds(&mut self, _: &[PathBuf]) -> Result<Vec<Self::Embedding>>;
}
impl<T: EmbedderT> BatchEmbedder for T {
type Embedding = T::Embedding;
const NAME: &'static str = T::NAME;
- fn embeds(&mut self, paths: &[PathBuf]) -> Result<Vec<Self::Embedding>, String> {
+ fn embeds(&mut self, paths: &[PathBuf]) -> Result<Vec<Self::Embedding>> {
paths.par_iter()
.map(|p| self.embed(p))
.collect::<Vec<_>>()
@@ -44,12 +45,12 @@ impl EmbedderT for BrightnessEmbedder {
type Embedding = f64;
const NAME: &'static str = "Brightness";
- fn embed(&self, path: &PathBuf) -> Result<f64, String> {
- let im = image::open(path).map_err(|e| e.to_string())?;
+ fn embed(&self, path: &Path) -> Result<f64> {
+ let im = image::open(path)?;
let num_bytes = 3 * (im.height() * im.width());
if num_bytes == 0 {
- return Err("Encountered NaN brightness, due to an empty image".to_string());
+ bail!("Encountered NaN brightness, due to an empty image");
}
Ok(im.to_rgb8()
@@ -74,8 +75,8 @@ impl EmbedderT for HueEmbedder {
type Embedding = Hue;
const NAME: &'static str = "Hue";
- fn embed(&self, path: &PathBuf) -> Result<Hue, String> {
- let im = image::open(path).map_err(|e| e.to_string())?;
+ fn embed(&self, path: &Path) -> Result<Hue> {
+ let im = image::open(path)?;
let num_pixels = im.height() * im.width();
let [sr, sg, sb] = im
.to_rgb8()
@@ -98,7 +99,7 @@ impl EmbedderT for HueEmbedder {
};
if hue.is_nan() {
- return Err("Encountered NaN hue, possibly because of a colorless or empty image".to_string());
+ bail!("Encountered NaN hue, possibly because of a colorless or empty image");
}
Ok(Hue(hue))
@@ -118,8 +119,8 @@ impl EmbedderT for ColorEmbedder {
type Embedding = (f64, f64, f64);
const NAME: &'static str = "Color";
- fn embed(&self, path: &PathBuf) -> Result<(f64, f64, f64), String> {
- let im = image::open(path).map_err(|e| e.to_string())?;
+ fn embed(&self, path: &Path) -> Result<(f64, f64, f64)> {
+ let im = image::open(path)?;
let num_pixels = im.height() * im.width();
let [sr, sg, sb] = im
.to_rgb8()
diff --git a/src/main.rs b/src/main.rs
index c4a0c26..5caa4ce 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,10 +1,10 @@
#![feature(iterator_try_collect)]
+use anyhow::Result;
use clap::Parser;
use priority_queue::PriorityQueue;
-use std::cmp::Ordering;
-use std::collections::HashMap;
-use std::path::PathBuf;
+use sha2::{Sha512_256, Digest};
+use std::{cmp::Ordering, collections::HashMap, fs, io, path::PathBuf};
use embedders::*;
@@ -30,9 +30,8 @@ struct Config {
base_dirs: xdg::BaseDirectories,
}
-fn get_config() -> Result<Config, String> {
- let dirs = xdg::BaseDirectories::with_prefix("embeddings-sort")
- .map_err(|_| "oh no")?;
+fn get_config() -> Result<Config> {
+ let dirs = xdg::BaseDirectories::with_prefix("embeddings-sort")?;
Ok(Config{base_dirs: dirs})
}
@@ -113,21 +112,30 @@ fn tsp_from_mst(mst: HashMap<usize, Vec<usize>>) -> Vec<usize> {
tsp_path
}
-fn process_embedder<E>(mut e: E, args: Args, cfg: Config) -> Result<Vec<PathBuf>, String>
+fn hash_file(p: &PathBuf) -> Result<[u8; 32]> {
+ let mut f = fs::File::open(p)?;
+ let mut hasher = Sha512_256::new();
+ io::copy(&mut f, &mut hasher)?;
+ Ok(hasher.finalize().into_iter().collect::<Vec<u8>>().try_into().unwrap())
+}
+
+fn process_embedder<E>(mut e: E, args: Args, cfg: Config) -> Result<Vec<PathBuf>>
where E: EmbedderT
{
- if args.images.len() == 0 {
+ if args.images.is_empty() {
return Ok(Vec::new());
}
- let db = sled::open(cfg.base_dirs.place_cache_file("embeddings.db")
- .map_err(|e| e.to_string())?).map_err(|e| e.to_string())?;
- let tree = typed_sled::Tree::<PathBuf, E::Embedding>::open(&db, E::NAME);
+ let db = sled::open(cfg.base_dirs.place_cache_file("embeddings.db")?)?;
+ let tree = typed_sled::Tree::<[u8; 32], E::Embedding>::open(&db, E::NAME);
- // TODO nicht pfad, sondern hash vom bild als key nehmen
let mut embeds: Vec<Option<_>> = args.images
.iter()
- .map(|p| tree.get(p).map_err(|e| e.to_string()))
+ .map(|p| {
+ let h = hash_file(p)?;
+ let r: Result<Option<E::Embedding>> = tree.get(&h).map_err(|e| e.into());
+ r
+ })
.try_collect()?;
let missing_embeds_indices: Vec<_> = embeds
@@ -145,8 +153,7 @@ fn process_embedder<E>(mut e: E, args: Args, cfg: Config) -> Result<Vec<PathBuf>
for (idx, emb) in missing_embeds_indices
.into_iter().zip(missing_embeds.into_iter())
{
- // TODO hier auch hash statt pfad
- tree.insert(&args.images[idx], &emb).map_err(|e| e.to_string())?;
+ tree.insert(&hash_file(&args.images[idx])?, &emb)?;
embeds[idx] = Some(emb);
}
@@ -156,7 +163,7 @@ fn process_embedder<E>(mut e: E, args: Args, cfg: Config) -> Result<Vec<PathBuf>
Ok(tsp_path.iter().map(|i| args.images[*i].clone()).collect())
}
-fn main() -> Result<(), String> {
+fn main() -> Result<()> {
let cfg = get_config()?;
let args = Args::parse();