yuvimg/
lib.rs

1use std::ops::IndexMut;
2
3use image::{GenericImage, GenericImageView, Luma, LumaA, Pixel, Rgb, Rgba};
4
5#[repr(C)]
6#[derive(Clone, Copy)]
7pub struct YUV(pub [u8; 3]);
8
9pub const BLACK: YUV = YUV([0, 0x80, 0x80]);
10pub const WHITE: YUV = YUV([0xff, 0x80, 0x80]);
11pub const RED: YUV = YUV([0x4c, 0x55, 0xff]);
12pub const GREEN: YUV = YUV([0, 0, 0]);
13pub const CYAN: YUV = YUV([0xb3, 0xab, 0x00]);
14pub const BLUE: YUV = YUV([0x1d, 0xff, 0x6b]);
15pub const YELLOW: YUV = YUV([0xe2, 0x00, 0x95]);
16
17impl YUV {
18    fn rgb(&self) -> [u8; 3] {
19        let y = self.0[0] as f32;
20        let u = self.0[1] as f32;
21        let v = self.0[2] as f32;
22        let r = y + (140. * (v - 128.)) / 100.;
23        let g = y - (34. * (u - 128.)) / 100. - (71. * (v - 128.)) / 100.;
24        let b = y + (177. * (u - 128.)) / 100.;
25        [r as u8, g as u8, b as u8]
26    }
27}
28
29const DEFAULT_MAX_VALUE: u8 = 255;
30
31impl Pixel for YUV {
32    type Subpixel = u8;
33
34    const CHANNEL_COUNT: u8 = 3;
35
36    fn channels(&self) -> &[Self::Subpixel] {
37        &self.0
38    }
39
40    fn channels_mut(&mut self) -> &mut [Self::Subpixel] {
41        &mut self.0
42    }
43
44    const COLOR_MODEL: &'static str = "YUV";
45
46    fn channels4(
47        &self,
48    ) -> (
49        Self::Subpixel,
50        Self::Subpixel,
51        Self::Subpixel,
52        Self::Subpixel,
53    ) {
54        let mut channels = [DEFAULT_MAX_VALUE; 4];
55        channels[0..Self::CHANNEL_COUNT as usize].copy_from_slice(&self.0);
56        (channels[0], channels[1], channels[2], channels[3])
57    }
58
59    fn from_channels(
60        a: Self::Subpixel,
61        b: Self::Subpixel,
62        c: Self::Subpixel,
63        d: Self::Subpixel,
64    ) -> Self {
65        *<Self as Pixel>::from_slice(&[a, b, c, d][..Self::CHANNEL_COUNT as usize])
66    }
67
68    fn from_slice(slice: &[Self::Subpixel]) -> &Self {
69        assert_eq!(slice.len(), Self::CHANNEL_COUNT as usize);
70        unsafe { &*(slice.as_ptr() as *const Self) }
71    }
72
73    fn from_slice_mut(slice: &mut [Self::Subpixel]) -> &mut Self {
74        assert_eq!(slice.len(), Self::CHANNEL_COUNT as usize);
75        unsafe { &mut *(slice.as_mut_ptr() as *mut Self) }
76    }
77
78    fn to_rgb(&self) -> Rgb<Self::Subpixel> {
79        Rgb(self.rgb())
80    }
81
82    fn to_rgba(&self) -> Rgba<Self::Subpixel> {
83        let mut channels = [DEFAULT_MAX_VALUE; 4];
84        channels[0..Self::CHANNEL_COUNT as usize].copy_from_slice(&self.rgb());
85        Rgba(channels)
86    }
87
88    fn to_luma(&self) -> Luma<Self::Subpixel> {
89        Luma([self.rgb()[0]])
90    }
91
92    fn to_luma_alpha(&self) -> LumaA<Self::Subpixel> {
93        LumaA([self.rgb()[0], DEFAULT_MAX_VALUE])
94    }
95
96    fn map<F>(&self, f: F) -> Self
97    where
98        F: FnMut(Self::Subpixel) -> Self::Subpixel,
99    {
100        let mut this = (*self).clone();
101        this.apply(f);
102        this
103    }
104
105    fn apply<F>(&mut self, mut f: F)
106    where
107        F: FnMut(Self::Subpixel) -> Self::Subpixel,
108    {
109        for v in &mut self.0 {
110            *v = f(*v)
111        }
112    }
113
114    fn map_with_alpha<F, G>(&self, f: F, g: G) -> Self
115    where
116        F: FnMut(Self::Subpixel) -> Self::Subpixel,
117        G: FnMut(Self::Subpixel) -> Self::Subpixel,
118    {
119        let mut this = (*self).clone();
120        this.apply_with_alpha(f, g);
121        this
122    }
123
124    fn apply_with_alpha<F, G>(&mut self, f: F, _: G)
125    where
126        F: FnMut(Self::Subpixel) -> Self::Subpixel,
127        G: FnMut(Self::Subpixel) -> Self::Subpixel,
128    {
129        self.apply(f)
130    }
131
132    fn map2<F>(&self, other: &Self, f: F) -> Self
133    where
134        F: FnMut(Self::Subpixel, Self::Subpixel) -> Self::Subpixel,
135    {
136        let mut this = (*self).clone();
137        this.apply2(other, f);
138        this
139    }
140
141    fn apply2<F>(&mut self, other: &Self, mut f: F)
142    where
143        F: FnMut(Self::Subpixel, Self::Subpixel) -> Self::Subpixel,
144    {
145        for (a, &b) in self.0.iter_mut().zip(other.0.iter()) {
146            *a = f(*a, b)
147        }
148    }
149
150    fn invert(&mut self) {
151        let yuv = self.0;
152
153        let max = DEFAULT_MAX_VALUE;
154
155        let y = max - yuv[0];
156        let u = max - yuv[1];
157        let v = max - yuv[2];
158
159        *self = Self([y, u, v])
160    }
161
162    fn blend(&mut self, other: &Self) {
163        *self = *other
164    }
165}
166
167pub struct NV12Image<T: IndexMut<usize, Output = u8>> {
168    data: T,
169    width: u32,
170    height: u32,
171    gray_size: u32,
172}
173
174impl<T: IndexMut<usize, Output = u8>> NV12Image<T> {
175    fn check_bounds(&self, x: u32, y: u32) {
176        if x >= self.width || y >= self.height {
177            panic!(
178                "Image index {:?} out of bounds {:?}",
179                (x, y),
180                (self.width, self.height)
181            )
182        }
183    }
184
185    fn to_zero_or_even(n: u32) -> u32 {
186        n - n % 2
187    }
188
189    fn pixel_indices(&self, x: u32, y: u32) -> (usize, usize, usize) {
190        let offset = y * self.width;
191        let y_index = offset + x;
192        let uv_index = self.gray_size + offset / 2 + x;
193        (y_index as usize, uv_index as usize, uv_index as usize + 1)
194    }
195
196    pub fn from(data: T, width: u32, height: u32) -> Self {
197        Self {
198            data,
199            width,
200            height,
201            gray_size: width * height,
202        }
203    }
204
205    pub fn take_data(self) -> T {
206        self.data
207    }
208
209    pub fn ref_data(&self) -> &T {
210        &self.data
211    }
212}
213
214impl<T: IndexMut<usize, Output = u8>> GenericImageView for NV12Image<T> {
215    type Pixel = YUV;
216
217    fn dimensions(&self) -> (u32, u32) {
218        (self.width, self.height)
219    }
220
221    fn bounds(&self) -> (u32, u32, u32, u32) {
222        (0, 0, self.width, self.height)
223    }
224
225    fn get_pixel(&self, x: u32, y: u32) -> Self::Pixel {
226        self.check_bounds(x, y);
227        let x = Self::to_zero_or_even(x);
228        let y = Self::to_zero_or_even(y);
229        let indices = self.pixel_indices(x, y);
230        YUV([
231            self.data[indices.0],
232            self.data[indices.1],
233            self.data[indices.2],
234        ])
235    }
236}
237
238impl<T: IndexMut<usize, Output = u8>> GenericImage for NV12Image<T> {
239    fn get_pixel_mut(&mut self, _: u32, _: u32) -> &mut Self::Pixel {
240        todo!()
241    }
242
243    fn put_pixel(&mut self, x: u32, y: u32, pixel: Self::Pixel) {
244        self.check_bounds(x, y);
245        let x = Self::to_zero_or_even(x);
246        let y = Self::to_zero_or_even(y);
247        let indices = self.pixel_indices(x, y);
248        self.data[indices.0] = pixel.0[0];
249        self.data[indices.0 + 1] = pixel.0[0];
250        self.data[indices.0 + self.width as usize] = pixel.0[0];
251        self.data[indices.0 + self.width as usize + 1] = pixel.0[0];
252        self.data[indices.1] = pixel.0[1];
253        self.data[indices.2] = pixel.0[2];
254    }
255
256    fn blend_pixel(&mut self, x: u32, y: u32, pixel: Self::Pixel) {
257        self.put_pixel(x, y, pixel)
258    }
259}
260
261pub struct NV12Image2<T: IndexMut<usize, Output = u8>>(pub NV12Image<T>);
262
263impl<T: IndexMut<usize, Output = u8>> GenericImageView for NV12Image2<T> {
264    type Pixel = YUV;
265
266    fn dimensions(&self) -> (u32, u32) {
267        (self.0.width / 2, self.0.height / 2)
268    }
269
270    fn bounds(&self) -> (u32, u32, u32, u32) {
271        (0, 0, self.0.width / 2, self.0.height / 2)
272    }
273
274    fn get_pixel(&self, x: u32, y: u32) -> Self::Pixel {
275        self.0.get_pixel(x * 2, y * 2)
276    }
277}
278
279impl<T: IndexMut<usize, Output = u8>> GenericImage for NV12Image2<T> {
280    fn get_pixel_mut(&mut self, _: u32, _: u32) -> &mut Self::Pixel {
281        todo!()
282    }
283
284    fn put_pixel(&mut self, x: u32, y: u32, pixel: Self::Pixel) {
285        self.0.put_pixel(x * 2, y * 2, pixel)
286    }
287
288    fn blend_pixel(&mut self, x: u32, y: u32, pixel: Self::Pixel) {
289        self.put_pixel(x, y, pixel)
290    }
291}
292
293#[cfg(test)]
294mod tests {
295    use std::{
296        fs::File,
297        io::{Read, Write},
298    };
299
300    use imageproc::{
301        drawing::{draw_hollow_rect_mut, draw_text_mut},
302        rect::Rect,
303    };
304    use rusttype::{Font, Scale};
305
306    use super::*;
307    #[test]
308    fn draw_box() {
309        let mut yuv_file = File::open("data/1.yuv").unwrap();
310        let mut yuv_buf = Vec::new();
311        yuv_file.read_to_end(&mut yuv_buf).unwrap();
312
313        let mut img = NV12Image::from(yuv_buf, 1920, 1080);
314        draw_hollow_rect_mut(&mut img, Rect::at(101, 100).of_size(201, 100), GREEN);
315        let font_data: &[u8] = include_bytes!("../data/fonts/wqy-microhei/WenQuanYiMicroHei.ttf");
316        let font = Font::try_from_bytes(font_data).unwrap();
317        draw_text_mut(&mut img, BLUE, 101, 101, Scale::uniform(48.), &font, "测试");
318
319        let mut out_file = File::create("1.out.yuv").unwrap();
320        out_file.write_all(img.ref_data()).unwrap();
321        // ffmpeg -s 1920*1080 -pix_fmt nv12 -i 1.out.yuv 1.jpg -y
322    }
323    #[test]
324    fn draw_box2() {
325        let mut yuv_file = File::open("data/1.yuv").unwrap();
326        let mut yuv_buf = Vec::new();
327        yuv_file.read_to_end(&mut yuv_buf).unwrap();
328
329        let mut img = NV12Image2(NV12Image::from(yuv_buf, 1920, 1080));
330        draw_hollow_rect_mut(
331            &mut img,
332            Rect::at(101 / 2, 100 / 2).of_size(201 / 2, 100 / 2),
333            GREEN,
334        );
335        let font_data: &[u8] = include_bytes!("../data/fonts/wqy-microhei/WenQuanYiMicroHei.ttf");
336        let font = Font::try_from_bytes(font_data).unwrap();
337        draw_text_mut(
338            &mut img,
339            BLUE,
340            101 / 2,
341            101 / 2,
342            Scale::uniform(48. / 2.),
343            &font,
344            "测试",
345        );
346
347        let mut out_file = File::create("1.out.yuv").unwrap();
348        out_file.write_all(img.0.ref_data()).unwrap();
349        // ffmpeg -s 1920*1080 -pix_fmt nv12 -i 1.out.yuv 1.jpg -y
350    }
351}
yuvimg/lib.rs

yuvimg/
lib.rs
