diff options
Diffstat (limited to 'src/embedders/pure.rs')
-rw-r--r-- | src/embedders/pure.rs | 91 |
1 files changed, 91 insertions, 0 deletions
diff --git a/src/embedders/pure.rs b/src/embedders/pure.rs new file mode 100644 index 0000000..09c8321 --- /dev/null +++ b/src/embedders/pure.rs @@ -0,0 +1,91 @@ +use anyhow::{bail, Result}; +use serde::{Deserialize, Serialize}; +use std::path::Path; + +use crate::{EmbedderT, MetricElem}; + +pub(crate) struct BrightnessEmbedder; +impl EmbedderT for BrightnessEmbedder { + type Embedding = f64; + const NAME: &'static str = "Brightness"; + + fn embed(&self, path: &Path) -> Result<f64> { + let im = image::open(path)?; + let num_bytes = 3 * (im.height() * im.width()); + + if num_bytes == 0 { + bail!("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) + } +} + +#[repr(transparent)] +#[derive(Serialize, Deserialize)] +pub(crate) 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(crate) struct HueEmbedder; +impl EmbedderT for HueEmbedder { + type Embedding = Hue; + const NAME: &'static str = "Hue"; + + 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() + .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 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.is_nan() { + bail!("Encountered NaN hue, possibly because of a colorless or empty image"); + } + + Ok(Hue(hue)) + } +} + +impl MetricElem for (f64, f64, f64) { + fn dist(&self, o: &(f64, f64, f64)) -> f64 { + let (dr, dg, db) = ((self.0 - o.0), (self.1 - o.1), (self.2 - o.2)); + (dr * dr + dg * dg + db * db).sqrt() + } +} +pub(crate) struct ColorEmbedder; +impl EmbedderT for ColorEmbedder { + type Embedding = (f64, f64, f64); + const NAME: &'static str = "Color"; + + 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() + .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 / num_pixels as f64); + + Ok((sr, sg, sb)) + } +} |