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#[derive(Debug, Clone)]
18pub struct Yuv<T: Pixel> {
19 data: Frame<T>,
20 config: YuvConfig,
21}
22
23#[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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
40pub enum YuvError {
41 SubsamplingMismatch,
44
45 InvalidLumaWidth,
51
52 InvalidLumaHeight,
58
59 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 #[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 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
230const 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
241fn 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}