Skip to main content

edgefirst_tensor/
format.rs

1// SPDX-FileCopyrightText: Copyright 2025 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4use serde::{Deserialize, Serialize};
5use std::fmt;
6
7/// Pixel format identifier.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
9#[repr(u8)]
10#[non_exhaustive]
11pub enum PixelFormat {
12    /// Packed RGB [H, W, 3]
13    Rgb = 1,
14    /// Packed RGBA [H, W, 4]
15    Rgba,
16    /// Packed BGRA [H, W, 4]
17    Bgra,
18    /// Grayscale [H, W, 1]
19    Grey,
20    /// Packed YUV 4:2:2, YUYV byte order [H, W, 2]
21    Yuyv,
22    /// Packed YUV 4:2:2, VYUY byte order [H, W, 2]
23    Vyuy,
24    /// Semi-planar YUV 4:2:0 [H*3/2, W] or multiplane [H, W] + [H/2, W]
25    Nv12,
26    /// Semi-planar YUV 4:2:2 [H*2, W] or multiplane [H, W] + [H, W]
27    Nv16,
28    /// Planar RGB, channels-first [3, H, W]
29    PlanarRgb,
30    /// Planar RGBA, channels-first [4, H, W]
31    PlanarRgba,
32}
33
34/// Memory layout category.
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
36#[non_exhaustive]
37pub enum PixelLayout {
38    /// Interleaved channels: [H, W, C]
39    Packed,
40    /// Channels-first: [C, H, W]
41    Planar,
42    /// Luma plane + interleaved chroma plane
43    SemiPlanar,
44}
45
46/// FourCC code constants (V4L2/DRM compatible).
47const FOURCC_RGB: u32 = u32::from_le_bytes(*b"RGB ");
48const FOURCC_RGBA: u32 = u32::from_le_bytes(*b"RGBA");
49const FOURCC_BGRA: u32 = u32::from_le_bytes(*b"BGRA");
50const FOURCC_GREY: u32 = u32::from_le_bytes(*b"Y800");
51const FOURCC_YUYV: u32 = u32::from_le_bytes(*b"YUYV");
52const FOURCC_VYUY: u32 = u32::from_le_bytes(*b"VYUY");
53const FOURCC_NV12: u32 = u32::from_le_bytes(*b"NV12");
54const FOURCC_NV16: u32 = u32::from_le_bytes(*b"NV16");
55
56impl PixelFormat {
57    /// Returns the number of channels for this pixel format.
58    ///
59    /// For semi-planar formats (NV12, NV16), this returns 1 (the luma channel
60    /// count for the primary plane). For packed formats, this is the total
61    /// number of interleaved components per pixel.
62    pub const fn channels(&self) -> usize {
63        match self {
64            Self::Rgb | Self::PlanarRgb => 3,
65            Self::Rgba | Self::Bgra | Self::PlanarRgba => 4,
66            Self::Grey | Self::Nv12 | Self::Nv16 => 1,
67            Self::Yuyv | Self::Vyuy => 2,
68        }
69    }
70
71    /// Returns the memory layout category for this pixel format.
72    pub const fn layout(&self) -> PixelLayout {
73        match self {
74            Self::Rgb | Self::Rgba | Self::Bgra | Self::Grey | Self::Yuyv | Self::Vyuy => {
75                PixelLayout::Packed
76            }
77            Self::PlanarRgb | Self::PlanarRgba => PixelLayout::Planar,
78            Self::Nv12 | Self::Nv16 => PixelLayout::SemiPlanar,
79        }
80    }
81
82    /// Returns `true` if this format encodes YUV (luma/chroma) data.
83    pub const fn is_yuv(&self) -> bool {
84        matches!(self, Self::Yuyv | Self::Vyuy | Self::Nv12 | Self::Nv16)
85    }
86
87    /// Returns `true` if this format includes an alpha channel.
88    pub const fn has_alpha(&self) -> bool {
89        matches!(self, Self::Rgba | Self::Bgra | Self::PlanarRgba)
90    }
91
92    /// Returns the V4L2/DRM FourCC code for this format, or `0` for formats
93    /// that have no standard FourCC representation (e.g., `PlanarRgb`).
94    pub const fn to_fourcc(&self) -> u32 {
95        match self {
96            Self::Rgb => FOURCC_RGB,
97            Self::Rgba => FOURCC_RGBA,
98            Self::Bgra => FOURCC_BGRA,
99            Self::Grey => FOURCC_GREY,
100            Self::Yuyv => FOURCC_YUYV,
101            Self::Vyuy => FOURCC_VYUY,
102            Self::Nv12 => FOURCC_NV12,
103            Self::Nv16 => FOURCC_NV16,
104            Self::PlanarRgb | Self::PlanarRgba => 0,
105        }
106    }
107
108    /// Converts a V4L2/DRM FourCC code to a `PixelFormat`, returning `None`
109    /// for unrecognized or zero codes.
110    pub const fn from_fourcc(fourcc: u32) -> Option<Self> {
111        match fourcc {
112            FOURCC_RGB => Some(Self::Rgb),
113            FOURCC_RGBA => Some(Self::Rgba),
114            FOURCC_BGRA => Some(Self::Bgra),
115            FOURCC_GREY => Some(Self::Grey),
116            FOURCC_YUYV => Some(Self::Yuyv),
117            FOURCC_VYUY => Some(Self::Vyuy),
118            FOURCC_NV12 => Some(Self::Nv12),
119            FOURCC_NV16 => Some(Self::Nv16),
120            _ => None,
121        }
122    }
123}
124
125impl fmt::Display for PixelFormat {
126    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
127        let fcc = self.to_fourcc();
128        if fcc != 0 {
129            let bytes = fcc.to_le_bytes();
130            for &b in &bytes {
131                if b == b' ' {
132                    break;
133                }
134                write!(f, "{}", b as char)?;
135            }
136            Ok(())
137        } else {
138            write!(f, "{self:?}")
139        }
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn channels() {
149        assert_eq!(PixelFormat::Rgb.channels(), 3);
150        assert_eq!(PixelFormat::Rgba.channels(), 4);
151        assert_eq!(PixelFormat::Bgra.channels(), 4);
152        assert_eq!(PixelFormat::Grey.channels(), 1);
153        assert_eq!(PixelFormat::Yuyv.channels(), 2);
154        assert_eq!(PixelFormat::Vyuy.channels(), 2);
155        assert_eq!(PixelFormat::Nv12.channels(), 1);
156        assert_eq!(PixelFormat::Nv16.channels(), 1);
157        assert_eq!(PixelFormat::PlanarRgb.channels(), 3);
158        assert_eq!(PixelFormat::PlanarRgba.channels(), 4);
159    }
160
161    #[test]
162    fn layout() {
163        assert_eq!(PixelFormat::Rgb.layout(), PixelLayout::Packed);
164        assert_eq!(PixelFormat::Rgba.layout(), PixelLayout::Packed);
165        assert_eq!(PixelFormat::Bgra.layout(), PixelLayout::Packed);
166        assert_eq!(PixelFormat::Grey.layout(), PixelLayout::Packed);
167        assert_eq!(PixelFormat::Yuyv.layout(), PixelLayout::Packed);
168        assert_eq!(PixelFormat::Vyuy.layout(), PixelLayout::Packed);
169        assert_eq!(PixelFormat::Nv12.layout(), PixelLayout::SemiPlanar);
170        assert_eq!(PixelFormat::Nv16.layout(), PixelLayout::SemiPlanar);
171        assert_eq!(PixelFormat::PlanarRgb.layout(), PixelLayout::Planar);
172        assert_eq!(PixelFormat::PlanarRgba.layout(), PixelLayout::Planar);
173    }
174
175    #[test]
176    fn is_yuv() {
177        assert!(!PixelFormat::Rgb.is_yuv());
178        assert!(!PixelFormat::Rgba.is_yuv());
179        assert!(PixelFormat::Yuyv.is_yuv());
180        assert!(PixelFormat::Vyuy.is_yuv());
181        assert!(PixelFormat::Nv12.is_yuv());
182        assert!(PixelFormat::Nv16.is_yuv());
183        assert!(!PixelFormat::PlanarRgb.is_yuv());
184    }
185
186    #[test]
187    fn has_alpha() {
188        assert!(!PixelFormat::Rgb.has_alpha());
189        assert!(PixelFormat::Rgba.has_alpha());
190        assert!(PixelFormat::Bgra.has_alpha());
191        assert!(!PixelFormat::Grey.has_alpha());
192        assert!(!PixelFormat::Yuyv.has_alpha());
193        assert!(!PixelFormat::PlanarRgb.has_alpha());
194        assert!(PixelFormat::PlanarRgba.has_alpha());
195    }
196
197    #[test]
198    fn fourcc_roundtrip() {
199        for fmt in [
200            PixelFormat::Rgb,
201            PixelFormat::Rgba,
202            PixelFormat::Bgra,
203            PixelFormat::Grey,
204            PixelFormat::Yuyv,
205            PixelFormat::Vyuy,
206            PixelFormat::Nv12,
207            PixelFormat::Nv16,
208        ] {
209            let fcc = fmt.to_fourcc();
210            assert_ne!(fcc, 0, "{fmt:?} should have a fourcc code");
211            assert_eq!(
212                PixelFormat::from_fourcc(fcc),
213                Some(fmt),
214                "roundtrip failed for {fmt:?}"
215            );
216        }
217    }
218
219    #[test]
220    fn fourcc_planar_returns_zero() {
221        assert_eq!(PixelFormat::PlanarRgb.to_fourcc(), 0);
222        assert_eq!(PixelFormat::PlanarRgba.to_fourcc(), 0);
223    }
224
225    #[test]
226    fn from_fourcc_unknown() {
227        assert_eq!(PixelFormat::from_fourcc(0), None);
228        assert_eq!(PixelFormat::from_fourcc(0xDEADBEEF), None);
229    }
230
231    #[test]
232    fn display_fourcc_formats() {
233        assert_eq!(format!("{}", PixelFormat::Rgba), "RGBA");
234        assert_eq!(format!("{}", PixelFormat::Nv12), "NV12");
235        assert_eq!(format!("{}", PixelFormat::Yuyv), "YUYV");
236        // Grey uses V4L2 FourCC "Y800", not "GREY"
237        assert_eq!(format!("{}", PixelFormat::Grey), "Y800");
238    }
239
240    #[test]
241    fn display_planar_formats() {
242        assert_eq!(format!("{}", PixelFormat::PlanarRgb), "PlanarRgb");
243        assert_eq!(format!("{}", PixelFormat::PlanarRgba), "PlanarRgba");
244    }
245
246    #[test]
247    fn repr_starts_at_one() {
248        assert_eq!(PixelFormat::Rgb as u8, 1);
249    }
250
251    #[test]
252    fn serde_roundtrip() {
253        let fmt = PixelFormat::Nv12;
254        let json = serde_json::to_string(&fmt).unwrap();
255        let back: PixelFormat = serde_json::from_str(&json).unwrap();
256        assert_eq!(fmt, back);
257    }
258}