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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

use thiserror::Error;

const IMAGE_WIDTH: usize = 320;
const IMAGE_HEIGHT: usize = 240;
const IMAGE_SIZE: usize = IMAGE_WIDTH * IMAGE_HEIGHT;

#[derive(Error, Debug)]
pub enum ImageCreationError {
    #[error("Invalid image size {0}. Should be {1}")]
    VectorSize(usize, usize),
    #[error("Invalid width {0}. Should be {1}")]
    InvalidWidth(usize, usize),
    #[error("Invalid height {0}. Should be {1}")]
    InvalidHeight(usize, usize),
}

pub struct R(u8);
pub struct G(u8);
pub struct B(u8);
pub struct Y(u8);
pub struct U(u8);
pub struct V(u8);
/// Represents an image that will be encoded
pub struct Robot36Image {
    data: Vec<Vec<(Y, U, V)>>,
}

impl Robot36Image {
    /// Creates a Robot36Image from a RGB8 vec.
    pub fn from_rgb8_vec(vec: Vec<(R, G, B)>) -> Result<Self, ImageCreationError> {
        if vec.len() != IMAGE_SIZE {
            return Err(ImageCreationError::VectorSize(vec.len(), IMAGE_SIZE));
        }
        let mut data: Vec<Vec<(Y, U, V)>> = vec![];
        for (i, rgb) in vec.into_iter().enumerate() {
            match data.get_mut(i / IMAGE_WIDTH) {
                Some(line) => line.push(to_yuv(rgb)),
                None => data.push(vec![to_yuv(rgb)]),
            }
        }
        Ok(Self { data })
    }

    pub fn get_height(&self) -> usize {
        IMAGE_HEIGHT
    }

    pub fn get_width(&self) -> usize {
        IMAGE_WIDTH
    }

    pub fn get_y(&self, x: usize, y: usize) -> &Y {
        &self.data[y][x].0
    }

    pub fn get_u(&self, x: usize, y: usize) -> &U {
        &self.data[y][x].1
    }

    pub fn get_v(&self, x: usize, y: usize) -> &V {
        &self.data[y][x].2
    }
}

#[cfg(feature = "image")]
impl Robot36Image {
    pub fn from_image(image: image::DynamicImage) -> Result<Robot36Image, ImageCreationError> {
        if image.width() != IMAGE_WIDTH as u32 {
            return Err(ImageCreationError::InvalidWidth(
                image.width() as usize,
                IMAGE_WIDTH,
            ));
        }
        if image.height() != IMAGE_HEIGHT as u32 {
            return Err(ImageCreationError::InvalidHeight(
                image.height() as usize,
                IMAGE_HEIGHT,
            ));
        }
        let rgb_image = image.to_rgb8().to_vec();
        let mut i = 0;
        let mut result = vec![];
        while i < rgb_image.len() {
            result.push((
                rgb_image[i].into(),
                rgb_image[i + 1].into(),
                rgb_image[i + 2].into(),
            ));
            i += 3;
        }
        Robot36Image::from_rgb8_vec(result)
    }
}

fn to_yuv(rgb: (R, G, B)) -> (Y, U, V) {
    (
        y_rgb(rgb.0 .0, rgb.1 .0, rgb.2 .0).into(),
        u_rgb(rgb.0 .0, rgb.1 .0, rgb.2 .0).into(),
        v_rgb(rgb.0 .0, rgb.1 .0, rgb.2 .0).into(),
    )
}

fn y_rgb(r: u8, g: u8, b: u8) -> u8 {
    yuv_clamp(
        16.0 + (0.003906 * ((65.738 * r as f32) + (129.057 * g as f32) + (25.064 * b as f32))),
    )
}

fn v_rgb(r: u8, g: u8, b: u8) -> u8 {
    yuv_clamp(
        128.0 + (0.003906 * ((112.439 * r as f32) + (-94.154 * g as f32) + (-18.285 * b as f32))),
    )
}

fn u_rgb(r: u8, g: u8, b: u8) -> u8 {
    yuv_clamp(
        128.0 + (0.003906 * ((-37.945 * r as f32) + (-74.494 * g as f32) + (112.439 * b as f32))),
    )
}

fn yuv_clamp(x: f32) -> u8 {
    let tmp = if x < 0.0 { 0.0 } else { x };
    (if tmp > 255.0 { 255.0 } else { tmp }) as u8
}

macro_rules! impl_traits {
    ($i:ident, $($is:ident),+) => {
        impl_traits!($i);
        impl_traits!($($is),+);
    };
    ($i:ident) => {
        impl From<&$i> for u16 {
            fn from(value: &$i) -> u16 {
                value.0 as u16
            }
        }
        impl From<&$i> for f64 {
            fn from(value: &$i) -> f64 {
                value.0 as f64
            }
        }
        impl From<u8> for $i {
            fn from(value: u8) -> Self {
                Self(value)
            }
        }
    };
}

impl_traits!(R, G, B, Y, U, V);