aboutsummaryrefslogtreecommitdiff
path: root/src/embedders.rs
blob: 0caf2c8146bb84dedb0b7dfb69ecede821214abf (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
92
93
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| {
                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()
    }
}