Skip to main content

moq_video/
color.rs

1//! [`Color`]: which YUV color space a frame's samples are in.
2
3use crate::Size;
4
5/// Which YUV color space a frame's samples are in.
6///
7/// Video carries luma and chroma, not RGB, and the matrix that converts between
8/// them differs by generation (BT.601 for standard definition, BT.709 for high
9/// definition) as does the numeric range (limited/studio swing pins luma to
10/// 16..235, full/full swing uses 0..255). Pairing samples with the wrong matrix
11/// is the classic tinted-video bug: it leaves grays untouched and skews
12/// saturated colors, so it survives a casual look at the picture.
13///
14/// [`Surface::color`](crate::Surface::color) reports it where the crate knows:
15/// when the crate did the conversion itself, or when the surface carries the
16/// answer (a macOS pixel buffer names its matrix, which VideoToolbox copies out
17/// of the stream's VUI). It is `None` for pixels that merely passed through with
18/// nothing naming their space, a camera's raw YUYV among them.
19/// [`Color::infer`] is the fallback then.
20#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21#[non_exhaustive]
22pub enum Color {
23	/// BT.601 (standard definition), limited range.
24	Bt601Limited,
25	/// BT.601 (standard definition), full range.
26	Bt601Full,
27	/// BT.709 (high definition), limited range.
28	Bt709Limited,
29	/// BT.709 (high definition), full range.
30	Bt709Full,
31}
32
33impl Color {
34	/// The conventional guess for a frame of this size: BT.601 up to standard
35	/// definition (576 lines), BT.709 above it, both limited range.
36	///
37	/// What a player does when the bitstream carries no VUI color description,
38	/// which is most of the time. A guess, so prefer a known [`Color`] whenever
39	/// one is available.
40	pub fn infer(size: Size) -> Self {
41		match size.height <= 576 {
42			true => Color::Bt601Limited,
43			false => Color::Bt709Limited,
44		}
45	}
46
47	/// The same matrix as `self` but in the given range, for a caller that knows
48	/// the range and not the matrix.
49	///
50	/// Only a surface whose pixel format spells out its range reaches this, which
51	/// is why it is macOS-only: CoreVideo's video-range and full-range NV12 name
52	/// theirs.
53	#[cfg(target_os = "macos")]
54	pub(crate) fn with_range(self, limited: bool) -> Self {
55		match (self, limited) {
56			(Color::Bt601Limited | Color::Bt601Full, true) => Color::Bt601Limited,
57			(Color::Bt601Limited | Color::Bt601Full, false) => Color::Bt601Full,
58			(_, true) => Color::Bt709Limited,
59			(_, false) => Color::Bt709Full,
60		}
61	}
62
63	/// Whether luma is 16..235 rather than 0..255.
64	///
65	/// The encoders need it for the VUI's `video_full_range_flag`, so unlike the
66	/// render module's `weights` it is not render-only.
67	pub(crate) fn limited(self) -> bool {
68		matches!(self, Color::Bt601Limited | Color::Bt709Limited)
69	}
70
71	/// The 8-bit RGB to Y'CbCr coefficients of this space, for a conversion the
72	/// crate runs itself (the GPU kernels).
73	///
74	/// Compiled for every test build so the cross-check against the `yuv` crate
75	/// runs without a GPU.
76	///
77	/// Display-referred RGB in: the samples are taken as already gamma-encoded,
78	/// which is what an 8-bit render target holds, so no transfer function is
79	/// applied on the way through. The offsets put chroma at 128 and, for limited
80	/// range, luma at 16; the scales fit the range (219/224 of 255 for limited,
81	/// all of it for full).
82	#[cfg(any(test, all(target_os = "linux", feature = "nvidia")))]
83	pub(crate) fn coefficients(self) -> Coefficients {
84		let (kr, kb) = match self {
85			Color::Bt601Limited | Color::Bt601Full => (0.299, 0.114),
86			Color::Bt709Limited | Color::Bt709Full => (0.2126, 0.0722),
87		};
88		let kg = 1.0 - kr - kb;
89		let (luma_scale, chroma_scale, luma_offset) = match self.limited() {
90			true => (219.0 / 255.0, 224.0 / 255.0, 16.0),
91			false => (1.0, 1.0, 0.0),
92		};
93		// Cb = (B - Y') / (2 (1 - Kb)) and Cr = (R - Y') / (2 (1 - Kr)), with Y'
94		// substituted so each channel is one weighted sum.
95		let cb = chroma_scale / (2.0 * (1.0 - kb));
96		let cr = chroma_scale / (2.0 * (1.0 - kr));
97		Coefficients {
98			y: [luma_scale * kr, luma_scale * kg, luma_scale * kb, luma_offset],
99			u: [-cb * kr, -cb * kg, cb * (1.0 - kb), 128.0],
100			v: [cr * (1.0 - kr), -cr * kg, -cr * kb, 128.0],
101		}
102	}
103
104	/// How the `yuv` crate names this color space, for the RGB conversions.
105	pub(crate) fn yuv(self) -> (yuv::YuvRange, yuv::YuvStandardMatrix) {
106		let range = match self.limited() {
107			true => yuv::YuvRange::Limited,
108			false => yuv::YuvRange::Full,
109		};
110		let matrix = match self {
111			Color::Bt601Limited | Color::Bt601Full => yuv::YuvStandardMatrix::Bt601,
112			Color::Bt709Limited | Color::Bt709Full => yuv::YuvStandardMatrix::Bt709,
113		};
114		(range, matrix)
115	}
116}
117
118/// The weights of one RGB to Y'CbCr conversion: each output sample is
119/// `[r, g, b, offset]` dotted with `(R, G, B, 1)`, all on the 0..255 scale.
120#[cfg(any(test, all(target_os = "linux", feature = "nvidia")))]
121#[derive(Clone, Copy, Debug, PartialEq)]
122pub(crate) struct Coefficients {
123	pub y: [f32; 4],
124	pub u: [f32; 4],
125	pub v: [f32; 4],
126}
127
128#[cfg(any(test, all(target_os = "linux", feature = "nvidia")))]
129impl Coefficients {
130	/// One pixel through the matrix, rounded and clamped the way the kernels do
131	/// it. The CPU reference for a GPU conversion, and what the tests compare
132	/// against the `yuv` crate.
133	#[cfg(test)]
134	pub(crate) fn apply(&self, rgb: [u8; 3]) -> [u8; 3] {
135		let dot = |w: [f32; 4]| {
136			let v = w[0] * rgb[0] as f32 + w[1] * rgb[1] as f32 + w[2] * rgb[2] as f32 + w[3];
137			v.round().clamp(0.0, 255.0) as u8
138		};
139		[dot(self.y), dot(self.u), dot(self.v)]
140	}
141}
142
143#[cfg(test)]
144mod tests {
145	use super::*;
146
147	#[test]
148	fn inference_splits_at_standard_definition() {
149		assert_eq!(Color::infer(Size::new(720, 480)), Color::Bt601Limited);
150		assert_eq!(Color::infer(Size::new(720, 576)), Color::Bt601Limited);
151		assert_eq!(Color::infer(Size::new(1280, 720)), Color::Bt709Limited);
152	}
153
154	#[cfg(target_os = "macos")]
155	#[test]
156	fn with_range_keeps_the_matrix() {
157		assert_eq!(Color::Bt709Limited.with_range(false), Color::Bt709Full);
158		assert_eq!(Color::Bt709Full.with_range(true), Color::Bt709Limited);
159		assert_eq!(Color::Bt601Limited.with_range(false), Color::Bt601Full);
160	}
161
162	/// The textbook 8-bit values for pure red: BT.709 limited is (63, 102, 240),
163	/// full range (54, 99, 255); BT.601 limited is (81, 90, 240).
164	#[test]
165	fn coefficients_match_the_textbook_values() {
166		let red = [255, 0, 0];
167		assert_eq!(Color::Bt709Limited.coefficients().apply(red), [63, 102, 240]);
168		assert_eq!(Color::Bt709Full.coefficients().apply(red), [54, 99, 255]);
169		assert_eq!(Color::Bt601Limited.coefficients().apply(red), [81, 90, 240]);
170		assert_eq!(Color::Bt601Full.coefficients().apply([255; 3]), [255, 128, 128]);
171		assert_eq!(Color::Bt709Limited.coefficients().apply([0; 3]), [16, 128, 128]);
172	}
173
174	/// The coefficients agree with the `yuv` crate's conversion, which every CPU
175	/// path uses, so a GPU frame converted with them decodes to the same picture
176	/// as the same pixels fed through `Surface::rgba`.
177	#[test]
178	fn coefficients_agree_with_the_yuv_crate() {
179		use yuv::{YuvChromaSubsampling, YuvConversionMode, YuvPlanarImageMut, rgba_to_yuv420};
180
181		let colors = [
182			Color::Bt601Limited,
183			Color::Bt601Full,
184			Color::Bt709Limited,
185			Color::Bt709Full,
186		];
187		let pixels: [[u8; 3]; 6] = [
188			[255, 0, 0],
189			[0, 255, 0],
190			[0, 0, 255],
191			[255, 255, 255],
192			[17, 200, 90],
193			[128, 128, 128],
194		];
195		for color in colors {
196			let (range, matrix) = color.yuv();
197			let coefficients = color.coefficients();
198			for rgb in pixels {
199				// A solid 2x2 block, so the crate's chroma subsampling changes nothing.
200				let rgba: Vec<u8> = std::iter::repeat_n([rgb[0], rgb[1], rgb[2], 255], 4)
201					.flatten()
202					.collect();
203				let mut planar = YuvPlanarImageMut::alloc(2, 2, YuvChromaSubsampling::Yuv420);
204				rgba_to_yuv420(&mut planar, &rgba, 8, range, matrix, YuvConversionMode::Balanced).unwrap();
205				let expected = [
206					planar.y_plane.borrow()[0],
207					planar.u_plane.borrow()[0],
208					planar.v_plane.borrow()[0],
209				];
210				let actual = coefficients.apply(rgb);
211				for (channel, (a, e)) in actual.iter().zip(expected).enumerate() {
212					assert!(
213						a.abs_diff(e) <= 1,
214						"{color:?} {rgb:?} channel {channel}: coefficients {actual:?}, yuv crate {expected:?}"
215					);
216				}
217			}
218		}
219	}
220}