aboutsummaryrefslogtreecommitdiff
path: root/src/pure_embedders.rs
diff options
context:
space:
mode:
Diffstat (limited to 'src/pure_embedders.rs')
-rw-r--r--src/pure_embedders.rs98
1 files changed, 98 insertions, 0 deletions
diff --git a/src/pure_embedders.rs b/src/pure_embedders.rs
new file mode 100644
index 0000000..0f0c3ab
--- /dev/null
+++ b/src/pure_embedders.rs
@@ -0,0 +1,98 @@
+use anyhow::{bail, Result};
+use serde::{Deserialize, Serialize};
+use std::path::Path;
+
+use crate::{MetricElem, EmbedderT};
+
+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))
+ }
+}