1use crate::Size;
4
5#[derive(Clone, Copy, Debug, PartialEq, Eq)]
21#[non_exhaustive]
22pub enum Color {
23 Bt601Limited,
25 Bt601Full,
27 Bt709Limited,
29 Bt709Full,
31}
32
33impl Color {
34 pub fn infer(size: Size) -> Self {
41 match size.height <= 576 {
42 true => Color::Bt601Limited,
43 false => Color::Bt709Limited,
44 }
45 }
46
47 #[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 pub(crate) fn limited(self) -> bool {
68 matches!(self, Color::Bt601Limited | Color::Bt709Limited)
69 }
70
71 #[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 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 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#[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 #[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 #[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 #[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 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}