visioncortex/image/
embedding.rs

1use crate::{Color, ColorImage, SpiralWalker};
2
3pub struct ImageEmbedding {
4    pub values: Vec<u8>,
5}
6
7impl ImageEmbedding {
8    pub fn ideal_image_size(_w: usize, _h: usize) -> (usize, usize) {
9        (128, 128)
10    }
11
12    /// Compute image embedding.
13    ///
14    /// Only accepts image of ideal_image_size.
15    pub fn encode(mut image: ColorImage) -> Self {
16        assert_eq!(image.width, 128);
17        assert_eq!(image.height, 128);
18
19        // 0| 1| 2| 3|  4|  5|  6|   7
20        // 1| 2| 4| 8| 16| 32| 64| 128
21
22        // image pyramid
23        let mut images = Vec::new();
24        let (mut w, mut h) = (image.width, image.height);
25        while w != 1 && h != 1 {
26            images.push(std::mem::take(&mut image));
27            let parent = images.last().unwrap();
28            w /= 2; h /= 2;
29            image = ColorImage::new_w_h(w, h);
30            for y in 0..h {
31                for x in 0..w {
32                    let xx = x * 2; let yy = y * 2;
33                    image.set_pixel(x, y, &Color::ave_4(
34                        parent.get_pixel(xx  , yy  ),
35                        parent.get_pixel(xx  , yy+1),
36                        parent.get_pixel(xx+1, yy  ),
37                        parent.get_pixel(xx+1, yy+1),
38                    ));
39                }
40            }
41            #[cfg(test)]
42            println!("\n{}", image.to_binary_image(|c| c.r != 0));
43        }
44        images.push(std::mem::take(&mut image));
45
46        let mut values = Vec::new();
47
48        for (i, image) in images.into_iter().rev().enumerate() {
49            debug_assert_eq!(i, power_of_two(image.width));
50            for (x, y) in SpiralWalker::new(image.width) {
51                values.extend(&image.get_pixel(x as usize, y as usize).rgb_u8());
52            }
53        }
54
55        debug_assert_eq!(
56            values.len(), 3 * (1*1 + 2*2 + 4*4 + 8*8 + 16*16 + 32*32 + 64*64 + 128*128)
57        );
58
59        Self { values }
60    }
61
62    /// Reconstruct original input image
63    pub fn decode(&self) -> ColorImage {
64        let (w, h) = ImageEmbedding::ideal_image_size(0, 0);
65        let idx = self.values.len() - w * h * 3;
66        let mut image = ColorImage::new_w_h(w, h);
67        for (i, (x, y)) in SpiralWalker::new(w).enumerate() {
68            image.set_pixel(x as usize, y as usize, &Color::new(
69                self.values[idx + i * 3 + 0],
70                self.values[idx + i * 3 + 1],
71                self.values[idx + i * 3 + 2],
72            ));
73        }
74        image
75    }
76
77    /// Get the 0th layer as color
78    pub fn value_0th(&self) -> Color {
79        Color::new(
80            self.values[0],
81            self.values[1],
82            self.values[2],
83        )
84    }
85
86    pub fn values_1st(&self) -> [u8; 12] {
87        let idx = 3;
88        let mut values = [0; 12];
89        for i in 0..4 {
90            values[i * 3 + 0] = self.values[idx + i * 3 + 0];
91            values[i * 3 + 1] = self.values[idx + i * 3 + 1];
92            values[i * 3 + 2] = self.values[idx + i * 3 + 2];
93        }
94        values
95    }
96
97    pub fn values_2nd(&self) -> &'_[u8] {
98        &self.values[15 .. (3 * (1*1 + 2*2 + 4*4))]
99    }
100
101    pub fn load(values: Vec<u8>) -> ImageEmbedding {
102        ImageEmbedding { values }
103    }
104}
105
106fn power_of_two(n: usize) -> usize {
107    // n must be power of 2
108    n.trailing_zeros() as usize
109    // slightly better: (usize::BITS - 1) - n.leading_zeros()
110}
111
112#[cfg(test)]
113mod test {
114    use super::*;
115
116    #[test]
117    fn test_embedding_1() {
118        let (w, h) = ImageEmbedding::ideal_image_size(0, 0);
119        let mut image = ColorImage::new_w_h(w, h);
120        let red = Color::new(255, 0, 0);
121        image.set_pixel(55, 55, &red);
122        for x in 56..64 {
123            for y in 56..64 {
124                image.set_pixel(x, y, &red);
125            }
126        }
127        let embedding = ImageEmbedding::encode(image).values;
128        // panic!("{embedding:?}");
129
130        let red_fraction = 255u32 * 65 / 128 / 128; // = 1
131        assert_eq!(embedding[0] as u32, red_fraction);
132        assert_eq!(embedding[1] as u32, 0);
133        assert_eq!(embedding[2] as u32, 0);
134
135        // layer 1
136        let red_fraction = 255u32 * 65 / 64 / 64; // = 4
137        assert_eq!(embedding[3] as u32, red_fraction);
138        assert_eq!(embedding[4] as u32, 0);
139        assert_eq!(embedding[5] as u32, 0);
140        assert_eq!(embedding[6] as u32, 0);
141        assert_eq!(embedding[7] as u32, 0);
142        assert_eq!(embedding[8] as u32, 0);
143
144        // layer 2
145        let red_fraction = 255u32 * 65 / 32 / 32; // = 16
146        assert_eq!(embedding[15] as u32, red_fraction);
147        assert_eq!(embedding[16] as u32, 0);
148        assert_eq!(embedding[17] as u32, 0);
149        assert_eq!(embedding[18] as u32, 0);
150        assert_eq!(embedding[19] as u32, 0);
151        assert_eq!(embedding[20] as u32, 0);
152    }
153}
visioncortex/image/embedding.rs

visioncortex/image/
embedding.rs
