Skip to main content

nokhwa_core/
pixel_format.rs

1/*
2 * Copyright 2022 l1npengtul <l1npengtul@protonmail.com> / The Nokhwa Contributors
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *     http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16use crate::error::NokhwaError;
17use crate::types::{
18    buf_bgr_to_rgb, buf_mjpeg_to_rgb, buf_nv12_to_rgb, buf_yuyv422_to_rgb, color_frame_formats,
19    frame_formats, mjpeg_to_rgb, nv12_to_rgb, yuyv422_to_rgb, FrameFormat, Resolution,
20};
21use image::{Luma, LumaA, Pixel, Rgb, Rgba};
22use std::fmt::Debug;
23
24/// Trait that has methods to convert raw data from the webcam to a proper raw image.
25pub trait FormatDecoder: Clone + Sized + Send + Sync {
26    type Output: Pixel<Subpixel = u8>;
27    const FORMATS: &'static [FrameFormat];
28
29    /// Allocates and returns a `Vec`
30    /// # Errors
31    /// If the data is malformed, or the source [`FrameFormat`] is incompatible, this will error.
32    fn write_output(
33        fcc: FrameFormat,
34        resolution: Resolution,
35        data: &[u8],
36    ) -> Result<Vec<u8>, NokhwaError>;
37
38    /// Writes to a user provided buffer.
39    /// # Errors
40    /// If the data is malformed, the source [`FrameFormat`] is incompatible, or the user-alloted buffer is not large enough, this will error.
41    fn write_output_buffer(
42        fcc: FrameFormat,
43        resolution: Resolution,
44        data: &[u8],
45        dest: &mut [u8],
46    ) -> Result<(), NokhwaError>;
47}
48
49/// A Zero-Size-Type that contains the definition to convert a given image stream to an RGB888 in the [`Buffer`](crate::buffer::Buffer)'s [`.decode_image()`](crate::buffer::Buffer::decode_image)
50///
51/// ```.ignore
52/// use image::{ImageBuffer, Rgb};
53/// let image: ImageBuffer<Rgb<u8>, Vec<u8>> = buffer.to_image::<RgbFormat>();
54/// ```
55#[derive(Copy, Clone, Debug, Default, Hash, Ord, PartialOrd, Eq, PartialEq)]
56pub struct RgbFormat;
57
58impl FormatDecoder for RgbFormat {
59    type Output = Rgb<u8>;
60    const FORMATS: &'static [FrameFormat] = color_frame_formats();
61
62    #[inline]
63    fn write_output(
64        fcc: FrameFormat,
65        resolution: Resolution,
66        data: &[u8],
67    ) -> Result<Vec<u8>, NokhwaError> {
68        match fcc {
69            FrameFormat::MJPEG => mjpeg_to_rgb(data, false),
70            FrameFormat::YUYV => yuyv422_to_rgb(data, false),
71            FrameFormat::GRAY => Ok(data
72                .iter()
73                .flat_map(|x| {
74                    let pxv = *x;
75                    [pxv, pxv, pxv]
76                })
77                .collect()),
78            FrameFormat::RAWRGB => Ok(data.to_vec()),
79            FrameFormat::RAWBGR => {
80                let mut rgb = vec![0u8; data.len()];
81                data.chunks_exact(3).enumerate().for_each(|(idx, px)| {
82                    let index = idx * 3;
83                    rgb[index] = px[2];
84                    rgb[index + 1] = px[1];
85                    rgb[index + 2] = px[0];
86                });
87                Ok(rgb)
88            }
89            FrameFormat::NV12 => nv12_to_rgb(resolution, data, false),
90        }
91    }
92
93    #[inline]
94    fn write_output_buffer(
95        fcc: FrameFormat,
96        resolution: Resolution,
97        data: &[u8],
98        dest: &mut [u8],
99    ) -> Result<(), NokhwaError> {
100        match fcc {
101            FrameFormat::MJPEG => buf_mjpeg_to_rgb(data, dest, false),
102            FrameFormat::YUYV => buf_yuyv422_to_rgb(data, dest, false),
103            FrameFormat::GRAY => {
104                if dest.len() != data.len() * 3 {
105                    return Err(NokhwaError::ProcessFrameError {
106                        src: fcc,
107                        destination: "Luma => RGB".to_string(),
108                        error: "Bad buffer length".to_string(),
109                    });
110                }
111
112                data.iter().enumerate().for_each(|(idx, pixel_value)| {
113                    let index = idx * 3;
114                    dest[index] = *pixel_value;
115                    dest[index + 1] = *pixel_value;
116                    dest[index + 2] = *pixel_value;
117                });
118                Ok(())
119            }
120            FrameFormat::RAWRGB => {
121                dest.copy_from_slice(data);
122                Ok(())
123            }
124            FrameFormat::RAWBGR => buf_bgr_to_rgb(resolution, data, dest),
125            FrameFormat::NV12 => buf_nv12_to_rgb(resolution, data, dest, false),
126        }
127    }
128}
129
130/// A Zero-Size-Type that contains the definition to convert a given image stream to an RGBA8888 in the [`Buffer`](crate::buffer::Buffer)'s [`.decode_image()`](crate::buffer::Buffer::decode_image)
131///
132/// ```.ignore
133/// use image::{ImageBuffer, Rgba};
134/// let image: ImageBuffer<Rgba<u8>, Vec<u8>> = buffer.to_image::<RgbAFormat>();
135/// ```
136#[derive(Copy, Clone, Debug, Default, Hash, Ord, PartialOrd, Eq, PartialEq)]
137pub struct RgbAFormat;
138
139impl FormatDecoder for RgbAFormat {
140    type Output = Rgba<u8>;
141
142    const FORMATS: &'static [FrameFormat] = color_frame_formats();
143
144    #[inline]
145    fn write_output(
146        fcc: FrameFormat,
147        resolution: Resolution,
148        data: &[u8],
149    ) -> Result<Vec<u8>, NokhwaError> {
150        match fcc {
151            FrameFormat::MJPEG => mjpeg_to_rgb(data, true),
152            FrameFormat::YUYV => yuyv422_to_rgb(data, true),
153            FrameFormat::GRAY => Ok(data
154                .iter()
155                .flat_map(|x| {
156                    let pxv = *x;
157                    [pxv, pxv, pxv, 255]
158                })
159                .collect()),
160            FrameFormat::RAWRGB => Ok(data
161                .chunks_exact(3)
162                .flat_map(|x| [x[0], x[1], x[2], 255])
163                .collect()),
164            FrameFormat::RAWBGR => Ok(data
165                .chunks_exact(3)
166                .flat_map(|x| [x[2], x[1], x[0], 255])
167                .collect()),
168            FrameFormat::NV12 => nv12_to_rgb(resolution, data, true),
169        }
170    }
171
172    #[inline]
173    fn write_output_buffer(
174        fcc: FrameFormat,
175        resolution: Resolution,
176
177        data: &[u8],
178        dest: &mut [u8],
179    ) -> Result<(), NokhwaError> {
180        match fcc {
181            FrameFormat::MJPEG => buf_mjpeg_to_rgb(data, dest, true),
182            FrameFormat::YUYV => buf_yuyv422_to_rgb(data, dest, true),
183            FrameFormat::GRAY => {
184                if dest.len() != data.len() * 4 {
185                    return Err(NokhwaError::ProcessFrameError {
186                        src: fcc,
187                        destination: "Luma => RGBA".to_string(),
188                        error: "Bad buffer length".to_string(),
189                    });
190                }
191
192                data.iter().enumerate().for_each(|(idx, pixel_value)| {
193                    let index = idx * 4;
194                    dest[index] = *pixel_value;
195                    dest[index + 1] = *pixel_value;
196                    dest[index + 2] = *pixel_value;
197                    dest[index + 3] = 255;
198                });
199                Ok(())
200            }
201            FrameFormat::RAWRGB => {
202                data.chunks_exact(3).enumerate().for_each(|(idx, px)| {
203                    let index = idx * 4;
204                    dest[index] = px[0];
205                    dest[index + 1] = px[1];
206                    dest[index + 2] = px[2];
207                    dest[index + 3] = 255;
208                });
209                Ok(())
210            }
211            FrameFormat::RAWBGR => {
212                data.chunks_exact(3).enumerate().for_each(|(idx, px)| {
213                    let index = idx * 4;
214                    dest[index] = px[2];
215                    dest[index + 1] = px[1];
216                    dest[index + 2] = px[0];
217                    dest[index + 3] = 255;
218                });
219                Ok(())
220            }
221            FrameFormat::NV12 => buf_nv12_to_rgb(resolution, data, dest, true),
222        }
223    }
224}
225
226/// A Zero-Size-Type that contains the definition to convert a given image stream to an Luma8(Grayscale 8-bit) in the [`Buffer`](crate::buffer::Buffer)'s [`.decode_image()`](crate::buffer::Buffer::decode_image)
227///
228/// ```.ignore
229/// use image::{ImageBuffer, Luma};
230/// let image: ImageBuffer<Luma<u8>, Vec<u8>> = buffer.to_image::<LumaFormat>();
231/// ```
232#[derive(Copy, Clone, Debug, Default, Hash, Ord, PartialOrd, Eq, PartialEq)]
233pub struct LumaFormat;
234
235impl FormatDecoder for LumaFormat {
236    type Output = Luma<u8>;
237
238    const FORMATS: &'static [FrameFormat] = frame_formats();
239
240    #[allow(clippy::cast_possible_truncation)]
241    #[allow(clippy::cast_sign_loss)]
242    #[inline]
243    fn write_output(
244        fcc: FrameFormat,
245        resolution: Resolution,
246        data: &[u8],
247    ) -> Result<Vec<u8>, NokhwaError> {
248        match fcc {
249            FrameFormat::MJPEG => Ok(mjpeg_to_rgb(data, false)?
250                .as_slice()
251                .chunks_exact(3)
252                .map(|x| {
253                    let mut avg = 0;
254                    x.iter().for_each(|v| avg += u16::from(*v));
255                    (avg / 3) as u8
256                })
257                .collect()),
258            FrameFormat::YUYV => Ok(yuyv422_to_rgb(data, false)?
259                .as_slice()
260                .chunks_exact(3)
261                .map(|x| {
262                    let mut avg = 0;
263                    x.iter().for_each(|v| avg += u16::from(*v));
264                    (avg / 3) as u8
265                })
266                .collect()),
267            FrameFormat::NV12 => Ok(nv12_to_rgb(resolution, data, false)?
268                .as_slice()
269                .chunks_exact(3)
270                .map(|x| {
271                    let mut avg = 0;
272                    x.iter().for_each(|v| avg += u16::from(*v));
273                    (avg / 3) as u8
274                })
275                .collect()),
276            FrameFormat::GRAY => Ok(data.to_vec()),
277            FrameFormat::RAWRGB => Ok(data
278                .chunks(3)
279                .map(|px| ((i32::from(px[0]) + i32::from(px[1]) + i32::from(px[2])) / 3) as u8)
280                .collect()),
281            FrameFormat::RAWBGR => Ok(data
282                .chunks(3)
283                .map(|px| ((i32::from(px[2]) + i32::from(px[1]) + i32::from(px[0])) / 3) as u8)
284                .collect()),
285        }
286    }
287
288    #[inline]
289    fn write_output_buffer(
290        fcc: FrameFormat,
291        _resolution: Resolution,
292        data: &[u8],
293        dest: &mut [u8],
294    ) -> Result<(), NokhwaError> {
295        match fcc {
296            // TODO: implement!
297            FrameFormat::MJPEG | FrameFormat::YUYV | FrameFormat::NV12 => {
298                Err(NokhwaError::ProcessFrameError {
299                    src: fcc,
300                    destination: "RGB => Luma".to_string(),
301                    error: "Conversion Error".to_string(),
302                })
303            }
304            FrameFormat::GRAY => {
305                data.iter().zip(dest.iter_mut()).for_each(|(pxv, d)| {
306                    *d = *pxv;
307                });
308                Ok(())
309            }
310            FrameFormat::RAWRGB => Err(NokhwaError::ProcessFrameError {
311                src: fcc,
312                destination: "RGB => Luma".to_string(),
313                error: "Conversion Error".to_string(),
314            }),
315            FrameFormat::RAWBGR => Err(NokhwaError::ProcessFrameError {
316                src: fcc,
317                destination: "BGR => Luma".to_string(),
318                error: "Conversion Error".to_string(),
319            }),
320        }
321    }
322}
323
324/// A Zero-Size-Type that contains the definition to convert a given image stream to an LumaA8(Grayscale 8-bit with 8-bit alpha) in the [`Buffer`](crate::buffer::Buffer)'s [`.decode_image()`](crate::buffer::Buffer::decode_image)
325///
326/// ```.ignore
327/// use image::{ImageBuffer, LumaA};
328/// let image: ImageBuffer<LumaA<u8>, Vec<u8>> = buffer.to_image::<LumaAFormat>();
329/// ```
330#[derive(Copy, Clone, Debug, Default, Hash, Ord, PartialOrd, Eq, PartialEq)]
331pub struct LumaAFormat;
332
333impl FormatDecoder for LumaAFormat {
334    type Output = LumaA<u8>;
335
336    const FORMATS: &'static [FrameFormat] = frame_formats();
337
338    #[allow(clippy::cast_possible_truncation)]
339    #[inline]
340    fn write_output(
341        fcc: FrameFormat,
342        resolution: Resolution,
343        data: &[u8],
344    ) -> Result<Vec<u8>, NokhwaError> {
345        match fcc {
346            FrameFormat::MJPEG => Ok(mjpeg_to_rgb(data, false)?
347                .as_slice()
348                .chunks_exact(3)
349                .flat_map(|x| {
350                    let mut avg = 0;
351                    x.iter().for_each(|v| avg += u16::from(*v));
352                    [(avg / 3) as u8, 255]
353                })
354                .collect()),
355            FrameFormat::YUYV => Ok(yuyv422_to_rgb(data, false)?
356                .as_slice()
357                .chunks_exact(3)
358                .flat_map(|x| {
359                    let mut avg = 0;
360                    x.iter().for_each(|v| avg += u16::from(*v));
361                    [(avg / 3) as u8, 255]
362                })
363                .collect()),
364            FrameFormat::NV12 => Ok(nv12_to_rgb(resolution, data, false)?
365                .as_slice()
366                .chunks_exact(3)
367                .flat_map(|x| {
368                    let mut avg = 0;
369                    x.iter().for_each(|v| avg += u16::from(*v));
370                    [(avg / 3) as u8, 255]
371                })
372                .collect()),
373            FrameFormat::GRAY => Ok(data.iter().flat_map(|x| [*x, 255]).collect()),
374            FrameFormat::RAWRGB => Err(NokhwaError::ProcessFrameError {
375                src: fcc,
376                destination: "RGB => LumaA".to_string(),
377                error: "Conversion Error".to_string(),
378            }),
379            FrameFormat::RAWBGR => Err(NokhwaError::ProcessFrameError {
380                src: fcc,
381                destination: "BGR => LumaA".to_string(),
382                error: "Conversion Error".to_string(),
383            }),
384        }
385    }
386
387    #[inline]
388    fn write_output_buffer(
389        fcc: FrameFormat,
390        _resolution: Resolution,
391        data: &[u8],
392        dest: &mut [u8],
393    ) -> Result<(), NokhwaError> {
394        match fcc {
395            FrameFormat::MJPEG => {
396                // FIXME: implement!
397                Err(NokhwaError::ProcessFrameError {
398                    src: fcc,
399                    destination: "MJPEG => LumaA".to_string(),
400                    error: "Conversion Error".to_string(),
401                })
402            }
403            FrameFormat::YUYV => Err(NokhwaError::ProcessFrameError {
404                src: fcc,
405                destination: "YUYV => LumaA".to_string(),
406                error: "Conversion Error".to_string(),
407            }),
408            FrameFormat::NV12 => Err(NokhwaError::ProcessFrameError {
409                src: fcc,
410                destination: "NV12 => LumaA".to_string(),
411                error: "Conversion Error".to_string(),
412            }),
413            FrameFormat::GRAY => {
414                if dest.len() != data.len() * 2 {
415                    return Err(NokhwaError::ProcessFrameError {
416                        src: fcc,
417                        destination: "GRAY8 => LumaA".to_string(),
418                        error: "Conversion Error".to_string(),
419                    });
420                }
421
422                data.iter()
423                    .zip(dest.chunks_exact_mut(2))
424                    .enumerate()
425                    .for_each(|(idx, (pxv, d))| {
426                        let index = idx * 2;
427                        d[index] = *pxv;
428                        d[index + 1] = 255;
429                    });
430                Ok(())
431            }
432            FrameFormat::RAWRGB => Err(NokhwaError::ProcessFrameError {
433                src: fcc,
434                destination: "RGB => LumaA".to_string(),
435                error: "Conversion Error".to_string(),
436            }),
437            FrameFormat::RAWBGR => Err(NokhwaError::ProcessFrameError {
438                src: fcc,
439                destination: "BGR => LumaA".to_string(),
440                error: "Conversion Error".to_string(),
441            }),
442        }
443    }
444}
445
446/// let image: ImageBuffer<Rgb<u8>, Vec<u8>> = buffer.to_image::<YuyvFormat>();
447/// ```
448#[derive(Copy, Clone, Debug, Default, Hash, Ord, PartialOrd, Eq, PartialEq)]
449pub struct YuyvFormat;
450
451impl FormatDecoder for YuyvFormat {
452    type Output = Rgb<u8>;
453    const FORMATS: &'static [FrameFormat] = color_frame_formats();
454
455    #[inline]
456    fn write_output(
457        fcc: FrameFormat,
458        resolution: Resolution,
459        data: &[u8],
460    ) -> Result<Vec<u8>, NokhwaError> {
461        match fcc {
462            FrameFormat::YUYV => {
463                let i420 = private_convert_yuyv_to_i420(
464                    data,
465                    resolution.width() as usize,
466                    resolution.height() as usize,
467                );
468                Ok(i420)
469            }
470            _ => Err(NokhwaError::GeneralError("Invalid FrameFormat".into())),
471        }
472    }
473
474    #[inline]
475    fn write_output_buffer(
476        fcc: FrameFormat,
477        resolution: Resolution,
478        data: &[u8],
479        dest: &mut [u8],
480    ) -> Result<(), NokhwaError> {
481        match fcc {
482            FrameFormat::YUYV => {
483                convert_yuyv_to_i420_direct(
484                    data,
485                    dest,
486                    resolution.width() as usize,
487                    resolution.height() as usize,
488                )?;
489                Ok(())
490            }
491            _ => Err(NokhwaError::GeneralError("Invalid FrameFormat".into())),
492        }
493    }
494}
495
496fn private_convert_yuyv_to_i420(yuyv: &[u8], width: usize, height: usize) -> Vec<u8> {
497    assert!(
498        width % 2 == 0 && height % 2 == 0,
499        "Width and height must be even numbers."
500    );
501
502    let mut i420 = vec![0u8; width * height + 2 * (width / 2) * (height / 2)];
503    let (y_plane, uv_plane) = i420.split_at_mut(width * height);
504    let (u_plane, v_plane) = uv_plane.split_at_mut(uv_plane.len() / 2);
505
506    for y in 0..height {
507        for x in (0..width).step_by(2) {
508            let base_index = (y * width + x) * 2;
509            let y0 = yuyv[base_index];
510            let u = yuyv[base_index + 1];
511            let y1 = yuyv[base_index + 2];
512            let v = yuyv[base_index + 3];
513
514            y_plane[y * width + x] = y0;
515            y_plane[y * width + x + 1] = y1;
516
517            if y % 2 == 0 {
518                u_plane[y / 2 * (width / 2) + x / 2] = u;
519                v_plane[y / 2 * (width / 2) + x / 2] = v;
520            }
521        }
522    }
523
524    i420
525}
526
527fn convert_yuyv_to_i420_direct(
528    yuyv: &[u8],
529    dest: &mut [u8],
530    width: usize,
531    height: usize,
532) -> Result<(), NokhwaError> {
533    // Ensure the destination buffer is large enough
534    if dest.len() < width * height + 2 * (width / 2) * (height / 2) {
535        return Err(NokhwaError::GeneralError(
536            "Destination buffer is too small".into(),
537        ));
538    }
539
540    // Split the destination buffer into Y, U, and V planes
541    let (y_plane, uv_plane) = dest.split_at_mut(width * height);
542    let (u_plane, v_plane) = uv_plane.split_at_mut(uv_plane.len() / 2);
543
544    // Convert YUYV to I420
545    for y in 0..height {
546        for x in (0..width).step_by(2) {
547            let base_index = (y * width + x) * 2;
548            let y0 = yuyv[base_index];
549            let u = yuyv[base_index + 1];
550            let y1 = yuyv[base_index + 2];
551            let v = yuyv[base_index + 3];
552
553            y_plane[y * width + x] = y0;
554            y_plane[y * width + x + 1] = y1;
555
556            if y % 2 == 0 {
557                u_plane[y / 2 * (width / 2) + x / 2] = u;
558                v_plane[y / 2 * (width / 2) + x / 2] = v;
559            }
560        }
561    }
562
563    Ok(())
564}