robot36_encoder/
image.rs

1use thiserror::Error;
2
3const IMAGE_WIDTH: usize = 320;
4const IMAGE_HEIGHT: usize = 240;
5const IMAGE_SIZE: usize = IMAGE_WIDTH * IMAGE_HEIGHT;
6
7#[derive(Error, Debug)]
8pub enum ImageCreationError {
9    #[error("Invalid image size {0}. Should be {1}")]
10    VectorSize(usize, usize),
11    #[error("Invalid width {0}. Should be {1}")]
12    InvalidWidth(usize, usize),
13    #[error("Invalid height {0}. Should be {1}")]
14    InvalidHeight(usize, usize),
15}
16
17pub struct R(u8);
18pub struct G(u8);
19pub struct B(u8);
20pub struct Y(u8);
21pub struct U(u8);
22pub struct V(u8);
23/// Represents an image that will be encoded
24pub struct Robot36Image {
25    data: Vec<Vec<(Y, U, V)>>,
26}
27
28impl Robot36Image {
29    /// Creates a Robot36Image from a RGB8 vec.
30    pub fn from_rgb8_vec(vec: Vec<(R, G, B)>) -> Result<Self, ImageCreationError> {
31        if vec.len() != IMAGE_SIZE {
32            return Err(ImageCreationError::VectorSize(vec.len(), IMAGE_SIZE));
33        }
34        let mut data: Vec<Vec<(Y, U, V)>> = vec![];
35        for (i, rgb) in vec.into_iter().enumerate() {
36            match data.get_mut(i / IMAGE_WIDTH) {
37                Some(line) => line.push(to_yuv(rgb)),
38                None => data.push(vec![to_yuv(rgb)]),
39            }
40        }
41        Ok(Self { data })
42    }
43
44    pub fn get_height(&self) -> usize {
45        IMAGE_HEIGHT
46    }
47
48    pub fn get_width(&self) -> usize {
49        IMAGE_WIDTH
50    }
51
52    pub fn get_y(&self, x: usize, y: usize) -> &Y {
53        &self.data[y][x].0
54    }
55
56    pub fn get_u(&self, x: usize, y: usize) -> &U {
57        &self.data[y][x].1
58    }
59
60    pub fn get_v(&self, x: usize, y: usize) -> &V {
61        &self.data[y][x].2
62    }
63}
64
65#[cfg(feature = "image")]
66impl Robot36Image {
67    pub fn from_image(image: image::DynamicImage) -> Result<Robot36Image, ImageCreationError> {
68        if image.width() != IMAGE_WIDTH as u32 {
69            return Err(ImageCreationError::InvalidWidth(
70                image.width() as usize,
71                IMAGE_WIDTH,
72            ));
73        }
74        if image.height() != IMAGE_HEIGHT as u32 {
75            return Err(ImageCreationError::InvalidHeight(
76                image.height() as usize,
77                IMAGE_HEIGHT,
78            ));
79        }
80        let rgb_image = image.to_rgb8().to_vec();
81        let mut i = 0;
82        let mut result = vec![];
83        while i < rgb_image.len() {
84            result.push((
85                rgb_image[i].into(),
86                rgb_image[i + 1].into(),
87                rgb_image[i + 2].into(),
88            ));
89            i += 3;
90        }
91        Robot36Image::from_rgb8_vec(result)
92    }
93}
94
95fn to_yuv(rgb: (R, G, B)) -> (Y, U, V) {
96    (
97        y_rgb(rgb.0 .0, rgb.1 .0, rgb.2 .0).into(),
98        u_rgb(rgb.0 .0, rgb.1 .0, rgb.2 .0).into(),
99        v_rgb(rgb.0 .0, rgb.1 .0, rgb.2 .0).into(),
100    )
101}
102
103fn y_rgb(r: u8, g: u8, b: u8) -> u8 {
104    yuv_clamp(
105        16.0 + (0.003906 * ((65.738 * r as f32) + (129.057 * g as f32) + (25.064 * b as f32))),
106    )
107}
108
109fn v_rgb(r: u8, g: u8, b: u8) -> u8 {
110    yuv_clamp(
111        128.0 + (0.003906 * ((112.439 * r as f32) + (-94.154 * g as f32) + (-18.285 * b as f32))),
112    )
113}
114
115fn u_rgb(r: u8, g: u8, b: u8) -> u8 {
116    yuv_clamp(
117        128.0 + (0.003906 * ((-37.945 * r as f32) + (-74.494 * g as f32) + (112.439 * b as f32))),
118    )
119}
120
121fn yuv_clamp(x: f32) -> u8 {
122    let tmp = if x < 0.0 { 0.0 } else { x };
123    (if tmp > 255.0 { 255.0 } else { tmp }) as u8
124}
125
126macro_rules! impl_traits {
127    ($i:ident, $($is:ident),+) => {
128        impl_traits!($i);
129        impl_traits!($($is),+);
130    };
131    ($i:ident) => {
132        impl From<&$i> for u16 {
133            fn from(value: &$i) -> u16 {
134                value.0 as u16
135            }
136        }
137        impl From<&$i> for f64 {
138            fn from(value: &$i) -> f64 {
139                value.0 as f64
140            }
141        }
142        impl From<u8> for $i {
143            fn from(value: u8) -> Self {
144                Self(value)
145            }
146        }
147    };
148}
149
150impl_traits!(R, G, B, Y, U, V);
robot36_encoder/image.rs

robot36_encoder/
image.rs
