diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/embedders.rs | 94 | ||||
| -rw-r--r-- | src/main.rs | 121 |
2 files changed, 215 insertions, 0 deletions
diff --git a/src/embedders.rs b/src/embedders.rs new file mode 100644 index 0000000..a8d34b4 --- /dev/null +++ b/src/embedders.rs @@ -0,0 +1,94 @@ +use std::path::PathBuf; + +pub trait MetricElem { + fn dist(&self, _: &Self) -> f64; +} + +impl MetricElem for f64 { + fn dist(&self, b: &f64) -> f64 { + (self - b).abs() + } +} + +pub trait EmbedderT { + type Embedding: MetricElem; + + fn embed(&mut self, _: &[PathBuf]) -> Result<Vec<Self::Embedding>, String>; +} + +pub struct BrightnessEmbedder; +impl EmbedderT for BrightnessEmbedder { + type Embedding = f64; + + fn embed(&mut self, paths: &[PathBuf]) -> Result<Vec<f64>, String> { + paths + .iter() + .map(|p| { + println!("{:?}", p); + let im = image::open(p).map_err(|e| e.to_string())?; + let num_bytes = 3 * (im.height() * im.width()); + + if num_bytes == 0 { + Err("Encountered NaN brightness, due to an empty image")?; + } + + Ok(im.to_rgb8() + .iter() + .map(|e| *e as u64) + .sum::<u64>() as f64 / num_bytes as f64) + }) + .try_collect() + } +} + +struct Hue (f64); +impl MetricElem for Hue { + fn dist(&self, b: &Hue) -> f64 { + let d = self.0.dist(&b.0); + d.min(6.-d) + } +} + +pub struct HueEmbedder; +impl EmbedderT for HueEmbedder { + type Embedding = f64; // TODO anderes Ding was Winkel vergleicht + + fn embed(&mut self, paths: &[PathBuf]) -> Result<Vec<f64>, String> { + paths + .iter() + .map(|p| { + let im = image::open(p).map_err(|e| e.to_string())?; + let num_pixels = im.height() * im.width(); + let [sr, sg, sb] = im + .to_rgb8() + .pixels() + .fold([0, 0, 0], |[or, og, ob], n| { + let [nr, ng, nb] = n.0; + [or + nr as u64, og + ng as u64, ob + nb as u64] + }) + .map(|e| e as f64 / 255. / num_pixels as f64); + + let mut hue = + if sr >= sg && sr >= sb { + (sg - sb) / (sr - sg.min(sb)) + } + else if sg >= sb { + 2. + (sb - sr) / (sg - sr.min(sb)) + } + else { + 4. + (sr - sg) / (sb - sr.min(sg)) + }; + + if hue < 0. { + hue += 6.; + } + + if hue != hue { + Err("Encountered NaN hue, possibly because of a colorless or empty image")?; + } + + Ok(hue) + }) + .try_collect() + } +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..5cfc10c --- /dev/null +++ b/src/main.rs @@ -0,0 +1,121 @@ +#![feature(iterator_try_collect)] + +use clap::Parser; +use priority_queue::PriorityQueue; +use std::cmp::Ordering; +use std::collections::HashSet; +use std::path::PathBuf; + +use embedders::*; + +mod embedders; + +#[derive(Debug, Clone, Copy, clap::ValueEnum)] +enum Embedder { + Brightness, + Hue, +} + +#[derive(Debug, Parser)] +struct Args { + #[arg(short, long, default_value = "hue")] + embedder: Embedder, + + images: Vec<PathBuf>, +} + +//#[derive(Debug)] +//struct Config { +// base_dir: xdg::BaseDirectories, +//} +// +//fn get_config() -> Result<Config, String> { +// let dirs = xdg::BaseDirectories::with_prefix("embeddings-sort") +// .map_err(|_| "oh no")?; +// +// Ok(Config{base_dir: dirs}) +//} + +fn process_embedder<E>(mut e: E, args: Args) -> Result<(), String> + where E: EmbedderT +{ + // wrapper struct to + // - reverse the ordering + // - implement Ord, even though the type is backed by an f64 + #[repr(transparent)] + #[derive(Debug, PartialEq)] + struct DownOrd (f64); + impl Eq for DownOrd {} + impl PartialOrd for DownOrd { + fn partial_cmp(&self, other: &Self) -> Option<Ordering> { + self.0.partial_cmp(&other.0).map(|e| e.reverse()) + } + } + impl Ord for DownOrd { + fn cmp(&self, other: &Self) -> Ordering {self.partial_cmp(other).unwrap()} + } + + let num_ims = args.images.len(); + if num_ims == 0 { + return Ok(()); + } + + let embeds: Vec<_> = e + .embed(&args.images)? + .into_iter() + .collect(); + + let mut possibleEdges = + PriorityQueue::with_capacity((num_ims * num_ims - num_ims) / 2); + let mut visitedImages = HashSet::with_capacity(num_ims); + let mut takenEdges = HashSet::with_capacity(num_ims - 1); // TODO das muss anders + + // here, we start at 0. + // we might get a better result in the end if we started with a vertex next + // to the lowest-cost edge, but we don't know which one that is (though we + // could compute that without changing our asymptotic complexity) + visitedImages.insert(0); + for i in 1..num_ims { + possibleEdges.push((0, i), DownOrd(embeds[0].dist(&embeds[i]))); + } + + // prims algorithm or something like that + while visitedImages.len() < num_ims { + // find the edge with the least cost that connects us to a new vertex + let (new, old) = loop { + let ((a, b), _) = possibleEdges.pop().unwrap(); + if !visitedImages.contains(&a) { + break (a, b); + } + else if !visitedImages.contains(&b) { + break (b, a); + } + }; + visitedImages.insert(new); + + println!("next edge: {:?} -> {:?}", old, new); + + // insert all the new edges we could take + takenEdges.insert((new, old)); + for i in 0..num_ims { + // don't consider edges taking us to nodes we already visited + if visitedImages.contains(&i) { + continue; + } + + possibleEdges.push((new, i), DownOrd(embeds[new].dist(&embeds[i]))); + } + } + + Ok(()) +} + +fn main() -> Result<(), String> { + //let cfg = get_config()?; + let args = Args::parse(); + + match args.embedder { + Embedder::Brightness => process_embedder(BrightnessEmbedder, args), + Embedder::Hue => process_embedder(HueEmbedder, args), + } +} |