yuvxyb/
yuv.rs

1use std::fmt;
2use std::mem::size_of;
3
4use av_data::pixel::{ColorPrimaries, MatrixCoefficients, TransferCharacteristic};
5use v_frame::{frame::Frame, plane::Plane, prelude::Pixel};
6
7use crate::{yuv_rgb::rgb_to_yuv, ConversionError, LinearRgb, Rgb, Xyb};
8
9/// Contains a YCbCr image in a color space defined by [`YuvConfig`].
10///
11/// YCbCr (often called YUV) representations of images are transformed from a "true" RGB image.
12/// These transformations are dependent on the original RGB color space and use one luminance (Y)
13/// as well as two chrominance (U, V) components to represent a color.
14///
15/// The data in this structure supports chroma subsampling (e.g. 4:2:0), multiple bit depths,
16/// and limited range support.
17#[derive(Debug, Clone)]
18pub struct Yuv<T: Pixel> {
19    data: Frame<T>,
20    config: YuvConfig,
21}
22
23/// Contains the configuration data for a YCbCr image.
24///
25/// This includes color space information, bit depth, chroma subsampling and whether the data
26/// is full or limited range.
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub struct YuvConfig {
29    pub bit_depth: u8,
30    pub subsampling_x: u8,
31    pub subsampling_y: u8,
32    pub full_range: bool,
33    pub matrix_coefficients: MatrixCoefficients,
34    pub transfer_characteristics: TransferCharacteristic,
35    pub color_primaries: ColorPrimaries,
36}
37
38/// Error type for when creating a [`Yuv`] struct goes wrong.
39#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum YuvError {
41    /// The configured subsampling does not match
42    /// the actual subsampling in the frame data.
43    SubsamplingMismatch,
44
45    /// The width of the luminance plane is not
46    /// compatible with the configured subsampling.
47    ///
48    /// For example, for 4:2:0 chroma subsampling, the width
49    /// of the luminance plane must be divisible by 2.
50    InvalidLumaWidth,
51
52    /// The height of the luminance plane is not
53    /// compatible with the configured subsampling.
54    ///
55    /// For example, for 4:2:0 chroma subsampling, the height
56    /// of the luminance plane must be divisible by 2.
57    InvalidLumaHeight,
58
59    /// The supplied data contains values which are outside
60    /// the valid range for the configured bit depth.
61    ///
62    /// This error can not occur for bit depths of 8 and 16,
63    /// as the valid range of values matches the available
64    /// range in the underlying data type exactly.
65    InvalidData,
66}
67
68impl std::error::Error for YuvError {}
69
70impl fmt::Display for YuvError {
71    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
72        match *self {
73            Self::SubsamplingMismatch => write!(f, "Configured subsampling does not match subsampling of frame data."),
74            Self::InvalidLumaWidth => write!(f, "The frame width does not support the configured subsampling."),
75            Self::InvalidLumaHeight => write!(f, "The frame height does not support the configured subsampling."),
76            Self::InvalidData => write!(f, "Data contains values which are not valid for the configured bit depth."),
77        }
78    }
79}
80
81impl<T: Pixel> Yuv<T> {
82    /// Create a new [`Yuv`] with the given data and configuration.
83    ///
84    /// # Errors
85    /// - If luma plane length does not match `width * height`
86    /// - If chroma plane lengths do not match `(width * height) >>
87    ///   (subsampling_x + subsampling_y)`
88    /// - If chroma subsampling is enabled and dimensions are not a multiple of
89    ///   2
90    /// - If chroma sampling set in `config` does not match subsampling in the
91    ///   frame data
92    /// - If `data` contains values which are not valid for the specified bit
93    ///   depth (note: out-of-range values for limited range are allowed)
94    // Clippy complains about T::to_u16 maybe panicking, but it can be assumed
95    // to never panic because the Pixel trait is only implemented by u8 and
96    // u16, both of which will successfully return a u16 from to_u16.
97    #[allow(clippy::missing_panics_doc)]
98    pub fn new(data: Frame<T>, config: YuvConfig) -> Result<Self, YuvError> {
99        if config.subsampling_x != data.planes[1].cfg.xdec as u8
100            || config.subsampling_x != data.planes[2].cfg.xdec as u8
101            || config.subsampling_y != data.planes[1].cfg.ydec as u8
102            || config.subsampling_y != data.planes[2].cfg.ydec as u8
103        {
104            return Err(YuvError::SubsamplingMismatch);
105        }
106
107        let width = data.planes[0].cfg.width;
108        let height = data.planes[0].cfg.height;
109        if width % (1 << config.subsampling_x) != 0 {
110            return Err(YuvError::InvalidLumaWidth);
111        }
112        if height % (1 << config.subsampling_y) != 0 {
113            return Err(YuvError::InvalidLumaHeight);
114        }
115        if size_of::<T>() == 2 && config.bit_depth < 16 {
116            let max_value = u16::MAX >> (16 - config.bit_depth);
117            if data.planes.iter().any(|plane| {
118                plane
119                    .iter()
120                    .any(|pix| pix.to_u16().expect("This is a u16") > max_value)
121            }) {
122                return Err(YuvError::InvalidData);
123            }
124        }
125
126        Ok(Self {
127            data,
128            config: config.fix_unspecified_data(width, height),
129        })
130    }
131
132    #[must_use]
133    #[inline]
134    pub const fn data(&self) -> &[Plane<T>] {
135        &self.data.planes
136    }
137
138    #[must_use]
139    #[inline]
140    pub const fn width(&self) -> usize {
141        self.data.planes[0].cfg.width
142    }
143
144    #[must_use]
145    #[inline]
146    pub const fn height(&self) -> usize {
147        self.data.planes[0].cfg.height
148    }
149
150    #[must_use]
151    #[inline]
152    pub const fn config(&self) -> YuvConfig {
153        self.config
154    }
155}
156
157impl YuvConfig {
158    pub(crate) fn fix_unspecified_data(mut self, width: usize, height: usize) -> Self {
159        if self.matrix_coefficients == MatrixCoefficients::Unspecified {
160            self.matrix_coefficients = guess_matrix_coefficients(width, height);
161            log::warn!(
162                "Matrix coefficients not specified. Guessing {}",
163                self.matrix_coefficients
164            );
165        }
166
167        if self.color_primaries == ColorPrimaries::Unspecified {
168            self.color_primaries = guess_color_primaries(self.matrix_coefficients, width, height);
169            log::warn!(
170                "Color primaries not specified. Guessing {}",
171                self.color_primaries
172            );
173        }
174
175        if self.transfer_characteristics == TransferCharacteristic::Unspecified {
176            self.transfer_characteristics = TransferCharacteristic::BT1886;
177            log::warn!(
178                "Transfer characteristics not specified. Guessing {}",
179                self.transfer_characteristics
180            );
181        }
182
183        self
184    }
185}
186
187impl<T: Pixel> TryFrom<(Xyb, YuvConfig)> for Yuv<T> {
188    type Error = ConversionError;
189
190    /// # Errors
191    /// - If the `YuvConfig` would produce an invalid image
192    fn try_from(other: (Xyb, YuvConfig)) -> Result<Self, Self::Error> {
193        let lrgb = LinearRgb::from(other.0);
194        Self::try_from((lrgb, other.1))
195    }
196}
197
198impl<T: Pixel> TryFrom<(Rgb, YuvConfig)> for Yuv<T> {
199    type Error = ConversionError;
200
201    fn try_from(other: (Rgb, YuvConfig)) -> Result<Self, Self::Error> {
202        Self::try_from((&other.0, other.1))
203    }
204}
205
206impl<T: Pixel> TryFrom<(LinearRgb, YuvConfig)> for Yuv<T> {
207    type Error = ConversionError;
208
209    fn try_from(other: (LinearRgb, YuvConfig)) -> Result<Self, Self::Error> {
210        let config = other.1;
211        let rgb = Rgb::try_from((
212            other.0,
213            config.transfer_characteristics,
214            config.color_primaries,
215        ))?;
216        Self::try_from((&rgb, config))
217    }
218}
219
220impl<T: Pixel> TryFrom<(&Rgb, YuvConfig)> for Yuv<T> {
221    type Error = ConversionError;
222
223    fn try_from(other: (&Rgb, YuvConfig)) -> Result<Self, Self::Error> {
224        let rgb = other.0;
225        let config = other.1;
226        rgb_to_yuv(rgb.data(), rgb.width(), rgb.height(), config)
227    }
228}
229
230// Heuristic taken from mpv
231const fn guess_matrix_coefficients(width: usize, height: usize) -> MatrixCoefficients {
232    if width >= 1280 || height > 576 {
233        MatrixCoefficients::BT709
234    } else if height == 576 {
235        MatrixCoefficients::BT470BG
236    } else {
237        MatrixCoefficients::ST170M
238    }
239}
240
241// Heuristic taken from mpv
242fn guess_color_primaries(
243    matrix: MatrixCoefficients,
244    width: usize,
245    height: usize,
246) -> ColorPrimaries {
247    if matrix == MatrixCoefficients::BT2020NonConstantLuminance
248        || matrix == MatrixCoefficients::BT2020ConstantLuminance
249    {
250        ColorPrimaries::BT2020
251    } else if matrix == MatrixCoefficients::BT709 || width >= 1280 || height > 576 {
252        ColorPrimaries::BT709
253    } else if height == 576 {
254        ColorPrimaries::BT470BG
255    } else if height == 480 || height == 488 {
256        ColorPrimaries::ST170M
257    } else {
258        ColorPrimaries::BT709
259    }
260}