aboutsummaryrefslogtreecommitdiff
path: root/src/embedders/pure.rs
blob: 09c8321c916ac7c44637495d01274f4fdad5270d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
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))
    }
}