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 { 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::() 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 { 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)) } }