openh264/formats/
yuv.rs

1use crate::formats::RGBSource;
2use crate::formats::rgb::RGB8Source;
3use crate::formats::rgb2yuv::{write_yuv_by_pixel, write_yuv_scalar};
4
5/// Allows the [Encoder](crate::encoder::Encoder) to be generic over a YUV source.
6pub trait YUVSource {
7    /// Size of the image as `(w, h)`.
8    #[must_use]
9    fn dimensions_i32(&self) -> (i32, i32) {
10        let (w, h) = self.dimensions();
11        (w as i32, h as i32)
12    }
13
14    /// Size of the image as `(w, h)`.
15    #[must_use]
16    fn dimensions(&self) -> (usize, usize);
17
18    /// YUV strides as `(y, u, v)`.
19    ///
20    /// For now you should make sure `u == v`.
21    #[must_use]
22    fn strides(&self) -> (usize, usize, usize);
23
24    /// YUV strides as `(y, u, v)`.
25    ///
26    /// For now you should make sure `u == v`.
27    #[must_use]
28    fn strides_i32(&self) -> (i32, i32, i32) {
29        let (y, u, v) = self.strides();
30        (y as i32, u as i32, v as i32)
31    }
32
33    /// Y buffer, should be of size `dimension.1 * strides.0`.
34    #[must_use]
35    fn y(&self) -> &[u8];
36
37    /// U buffer, should be of size `dimension.1 * strides.1`.
38    #[must_use]
39    fn u(&self) -> &[u8];
40
41    /// V buffer, should be of size `dimension.1 * strides.2`.
42    #[must_use]
43    fn v(&self) -> &[u8];
44
45    /// Estimates how many bytes you'll need to store this YUV in an `&[u8]` RGB array.
46    ///
47    /// This function should return `w * h * 3`.
48    #[must_use]
49    fn rgb8_len(&self) -> usize {
50        let (w, h) = self.dimensions();
51        w * h * 3
52    }
53
54    /// Estimates how many bytes you'll need to store this YUV in an `&[u8]` RGBA array.
55    ///
56    /// This function should return `w * h * 4`.
57    #[must_use]
58    fn rgba8_len(&self) -> usize {
59        let (w, h) = self.dimensions();
60        w * h * 4
61    }
62}
63
64/// Converts RGB to YUV data.
65#[must_use]
66pub struct YUVBuffer {
67    yuv: Vec<u8>,
68    width: usize,
69    height: usize,
70}
71
72impl YUVBuffer {
73    /// Creates a new YUV buffer from the given vec.
74    ///
75    /// The vec's length should be `3 * (width * height) / 2`.
76    ///
77    /// # Panics
78    ///
79    /// May panic if the given sizes are not multiples of 2, or the yuv buffer's size mismatches.
80    pub fn from_vec(yuv: Vec<u8>, width: usize, height: usize) -> Self {
81        assert_eq!(width % 2, 0, "width needs to be a multiple of 2");
82        assert_eq!(height % 2, 0, "height needs to be a multiple of 2");
83        assert_eq!(yuv.len(), (3 * (width * height)) / 2, "YUV buffer needs to be properly sized");
84
85        Self { yuv, width, height }
86    }
87
88    /// Allocates a new YUV buffer with the given width and height.
89    ///
90    /// Both dimensions must be even.
91    ///
92    /// # Panics
93    ///
94    /// May panic if the given sizes are not multiples of 2.
95    pub fn new(width: usize, height: usize) -> Self {
96        assert_eq!(width % 2, 0, "width needs to be a multiple of 2");
97        assert_eq!(height % 2, 0, "height needs to be a multiple of 2");
98
99        Self {
100            yuv: vec![0u8; (3 * (width * height)) / 2],
101            width,
102            height,
103        }
104    }
105
106    /// Allocates a new YUV buffer with the given width and height and data.
107    ///
108    /// # Panics
109    ///
110    /// May panic if invoked with an RGB source where the dimensions are not multiples of 2.
111    pub fn from_rgb_source(rgb: impl RGBSource) -> Self {
112        let mut rval = Self::new(rgb.dimensions().0, rgb.dimensions().1);
113        rval.read_rgb(rgb);
114        rval
115    }
116
117    /// Allocates a new YUV buffer with the given width and height and data.
118    ///
119    /// This is the faster version of [`Self::from_rgb_source`] and you should generally
120    /// use this one.
121    ///
122    /// # Panics
123    ///
124    /// May panic if invoked with an RGB source where the dimensions are not multiples of 2.
125    pub fn from_rgb8_source(rgb: impl RGB8Source) -> Self {
126        let mut rval = Self::new(rgb.dimensions().0, rgb.dimensions().1);
127        rval.read_rgb8(rgb);
128        rval
129    }
130
131    /// Reads an RGB buffer, converts it to YUV and stores it.
132    ///
133    /// # Panics
134    ///
135    /// May panic if the given `rgb` does not match the internal format.
136    #[allow(clippy::similar_names)]
137    pub fn read_rgb(&mut self, rgb: impl RGBSource) {
138        let dimensions = self.dimensions();
139        let u_base = self.width * self.height;
140        let v_base = u_base / 4;
141        let (y_buf, uv_buf) = self.yuv.split_at_mut(u_base);
142        let (u_buf, v_buf) = uv_buf.split_at_mut(v_base);
143        write_yuv_by_pixel(rgb, dimensions, y_buf, u_buf, v_buf);
144    }
145
146    /// Reads an RGB8 buffer, converts it to YUV and stores it.
147    ///
148    /// This is the faster version of [`Self::read_rgb`] and you should generally use this one.
149    ///
150    /// # Panics
151    ///
152    /// May panic if the given `rgb` does not match the internal format.
153    #[allow(clippy::similar_names)]
154    pub fn read_rgb8(&mut self, rgb: impl RGB8Source) {
155        let dimensions = self.dimensions();
156        let u_base = self.width * self.height;
157        let v_base = u_base / 4;
158        let (y_buf, uv_buf) = self.yuv.split_at_mut(u_base);
159        let (u_buf, v_buf) = uv_buf.split_at_mut(v_base);
160        write_yuv_scalar(rgb, dimensions, y_buf, u_buf, v_buf);
161    }
162}
163
164impl YUVSource for YUVBuffer {
165    fn dimensions(&self) -> (usize, usize) {
166        (self.width, self.height)
167    }
168
169    fn strides(&self) -> (usize, usize, usize) {
170        (self.width, self.width / 2, self.width / 2)
171    }
172
173    fn y(&self) -> &[u8] {
174        &self.yuv[0..self.width * self.height]
175    }
176
177    fn u(&self) -> &[u8] {
178        let base_u = self.width * self.height;
179        &self.yuv[base_u..base_u + base_u / 4]
180    }
181
182    fn v(&self) -> &[u8] {
183        let base_u = self.width * self.height;
184        let base_v = base_u + base_u / 4;
185        &self.yuv[base_v..]
186    }
187}
188
189/// Convenience wrapper if you already have YUV-sliced data from some other place.
190#[must_use]
191#[derive(Clone, Copy, Debug)]
192pub struct YUVSlices<'a> {
193    dimensions: (usize, usize),
194    yuv: (&'a [u8], &'a [u8], &'a [u8]),
195    strides: (usize, usize, usize),
196}
197
198impl<'a> YUVSlices<'a> {
199    /// Creates a new YUV slice in 4:2:0 format.
200    ///
201    /// Assume you have some dimension `(w, h)` that is your actual image size. In addition,
202    /// you will have strides `(sy, su, sv)` that specify how many pixels / bytes per row
203    /// are actually used be used. Strides must be larger or equal than `w` (y) or `w / 2` (uv)
204    /// respectively.
205    ///
206    /// # Panics
207    ///
208    /// This will panic if the given slices, strides or dimensions don't match.
209    pub fn new(yuv: (&'a [u8], &'a [u8], &'a [u8]), dimensions: (usize, usize), strides: (usize, usize, usize)) -> Self {
210        assert!(strides.0 >= dimensions.0);
211        assert!(strides.1 >= dimensions.0 / 2);
212        assert!(strides.2 >= dimensions.0 / 2);
213
214        assert_eq!(dimensions.1 * strides.0, yuv.0.len());
215        assert_eq!((dimensions.1 / 2) * strides.1, yuv.1.len());
216        assert_eq!((dimensions.1 / 2) * strides.2, yuv.2.len());
217
218        Self {
219            dimensions,
220            yuv,
221            strides,
222        }
223    }
224}
225
226impl YUVSource for YUVSlices<'_> {
227    fn dimensions(&self) -> (usize, usize) {
228        self.dimensions
229    }
230
231    fn strides(&self) -> (usize, usize, usize) {
232        self.strides
233    }
234
235    fn y(&self) -> &[u8] {
236        self.yuv.0
237    }
238
239    fn u(&self) -> &[u8] {
240        self.yuv.1
241    }
242
243    fn v(&self) -> &[u8] {
244        self.yuv.2
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::{YUVBuffer, YUVSlices};
251    use crate::formats::yuv2rgb::{write_rgb8_f32x8, write_rgb8_scalar};
252    use crate::formats::{RgbSliceU8, YUVSource};
253    use rand::prelude::IteratorRandom;
254    use rand::rngs::ThreadRng;
255
256    #[test]
257    fn rgb_to_yuv_conversion_black_2x2() {
258        let rgb_source = RgbSliceU8::new(&[0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8], (2, 2));
259        let yuv = YUVBuffer::from_rgb_source(rgb_source);
260        assert_eq!(yuv.y(), [16u8, 16u8, 16u8, 16u8]);
261        assert_eq!(yuv.u(), [128u8]);
262        assert_eq!(yuv.v(), [128u8]);
263        assert_eq!(yuv.strides_i32().0, 2);
264        assert_eq!(yuv.strides_i32().1, 1);
265        assert_eq!(yuv.strides_i32().2, 1);
266    }
267
268    #[test]
269    fn rgb_to_yuv_conversion_white_4x2() {
270        let data = &[
271            255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8,
272            255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8, 255u8,
273        ];
274        let rgb_source = RgbSliceU8::new(data, (4, 2));
275        let yuv = YUVBuffer::from_rgb_source(rgb_source);
276        assert_eq!(yuv.y(), [235u8, 235u8, 235u8, 235u8, 235u8, 235u8, 235u8, 235u8]);
277        assert_eq!(yuv.u(), [128u8, 128u8]);
278        assert_eq!(yuv.v(), [128u8, 128u8]);
279        assert_eq!(yuv.strides_i32().0, 4);
280        assert_eq!(yuv.strides_i32().1, 2);
281        assert_eq!(yuv.strides_i32().2, 2);
282    }
283
284    #[test]
285    fn rgb_to_yuv_conversion_red_2x4() {
286        let data = &[
287            255u8, 0u8, 0u8, 255u8, 0u8, 0u8, 255u8, 0u8, 0u8, 255u8, 0u8, 0u8, 255u8, 0u8, 0u8, 255u8, 0u8, 0u8, 255u8, 0u8,
288            0u8, 255u8, 0u8, 0u8,
289        ];
290        let rgb_source = RgbSliceU8::new(data, (4, 2));
291        let yuv = YUVBuffer::from_rgb_source(rgb_source);
292
293        assert_eq!(yuv.y(), [81u8, 81u8, 81u8, 81u8, 81u8, 81u8, 81u8, 81u8]);
294        assert_eq!(yuv.u(), [90u8, 90u8]);
295        assert_eq!(yuv.v(), [239u8, 239u8]);
296        assert_eq!(yuv.strides_i32().0, 4);
297        assert_eq!(yuv.strides_i32().1, 2);
298        assert_eq!(yuv.strides_i32().2, 2);
299    }
300
301    #[test]
302    #[should_panic = "strides.0 >= dimensions.0"]
303    fn test_new_stride_less_than_width() {
304        let y = vec![0u8; 10];
305        let u = vec![0u8; 5];
306        let v = vec![0u8; 5];
307        let _ = YUVSlices::new((&y, &u, &v), (10, 1), (9, 5, 5));
308    }
309
310    #[test]
311    #[should_panic = "strides.1 >= dimensions.0 / 2"]
312    fn test_new_u_stride_less_than_half_width() {
313        let y = vec![0u8; 20];
314        let u = vec![0u8; 5];
315        let v = vec![0u8; 5];
316        let _ = YUVSlices::new((&y, &u, &v), (10, 2), (10, 4, 5));
317    }
318
319    #[test]
320    #[should_panic = "strides.2 >= dimensions.0 / 2"]
321    fn test_new_v_stride_less_than_half_width() {
322        let y = vec![0u8; 20];
323        let u = vec![0u8; 5];
324        let v = vec![0u8; 5];
325        let _ = YUVSlices::new((&y, &u, &v), (10, 2), (10, 5, 4));
326    }
327
328    #[test]
329    #[should_panic = "assertion `left == right` failed"]
330    fn test_new_y_length_not_matching() {
331        let y = vec![0u8; 19];
332        let u = vec![0u8; 5];
333        let v = vec![0u8; 5];
334        let _ = YUVSlices::new((&y, &u, &v), (10, 2), (10, 5, 5));
335    }
336
337    #[test]
338    #[should_panic = "assertion `left == right` failed"]
339    fn test_new_u_length_not_matching() {
340        let y = vec![0u8; 20];
341        let u = vec![0u8; 4];
342        let v = vec![0u8; 5];
343        let _ = YUVSlices::new((&y, &u, &v), (10, 2), (10, 5, 5));
344    }
345
346    #[test]
347    #[should_panic = "assertion `left == right` failed"]
348    fn test_new_v_length_not_matching() {
349        let y = vec![0u8; 20];
350        let u = vec![0u8; 5];
351        let v = vec![0u8; 4];
352        let _ = YUVSlices::new((&y, &u, &v), (10, 2), (10, 5, 5));
353    }
354
355    #[test]
356    fn test_new_valid() {
357        let y = vec![0u8; 20];
358        let u = vec![0u8; 5];
359        let v = vec![0u8; 5];
360        let _ = YUVSlices::new((&y, &u, &v), (10, 2), (10, 5, 5));
361    }
362
363    /// Test every YUV value and see, if the SIMD version delivers a similar RGB value.
364    #[test]
365    fn test_write_rgb8_f32x8_spectrum() {
366        let mut rng = ThreadRng::default();
367        let dim = (8, 2);
368        let strides = (8, 4, 4);
369
370        // build artificial YUV planes containing the entire YUV spectrum
371        for y in (0..=255u8).choose_multiple(&mut rng, 10) {
372            // we sample probabilistically here, otherwise the test takes too long
373            for u in (0..=255u8).choose_multiple(&mut rng, 10) {
374                for v in (0..=255u8).choose_multiple(&mut rng, 10) {
375                    let (y_plane, u_plane, v_plane) = (vec![y; 16], vec![u; 4], vec![v; 4]);
376                    let mut target = vec![0; dim.0 * dim.1 * 3];
377                    write_rgb8_scalar(&y_plane, &u_plane, &v_plane, dim, strides, &mut target);
378
379                    let mut target2 = vec![0; dim.0 * dim.1 * 3];
380                    write_rgb8_f32x8(&y_plane, &u_plane, &v_plane, dim, strides, &mut target2);
381
382                    // compare first pixel
383                    for i in 0..3 {
384                        // Due to different CPU architectures the values may slightly change and may not be exactly equal.
385                        // allow difference of 1 / 255 (ca. 0.4%)
386                        let diff = (i32::from(target[i]) - i32::from(target2[i])).abs();
387                        assert!(diff <= 1, "YUV: {:?} yielded different results", (y, u, v));
388                    }
389                }
390            }
391        }
392    }
393}
openh264/formats/yuv.rs

openh264/formats/
yuv.rs
