ezk_image/color/
space.rs

1#![allow(clippy::too_many_arguments)]
2
3use crate::ColorTransfer;
4use crate::color::mat_idxs::*;
5use crate::vector::Vector;
6
7/// Color space used for RGB to YUV conversion
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum ColorSpace {
10    /// YUV Rec. ITU-R BT.601-7 625
11    BT601,
12
13    /// YUV Rec. ITU-R BT.709-6
14    BT709,
15
16    /// YUV Rec. ITU-R BT.2020-2
17    BT2020,
18
19    /// ICtCp Rec. ITU-R BT.2100-2 ICtCp (PQ transfer)
20    ICtCpPQ,
21
22    /// ICtCp Rec. ITU-R BT.2100-2 ICtCp (HLG transfer)
23    ICtCpHLG,
24}
25
26impl ColorSpace {
27    #[inline(always)]
28    pub(crate) unsafe fn yuv_to_rgb<V: Vector>(
29        &self,
30        transfer: ColorTransfer,
31        xyz_to_rgb: &'static [[f32; 3]; 3],
32        y: V,
33        u: V,
34        v: V,
35    ) -> (V, V, V) {
36        match self {
37            ColorSpace::BT601 => convert_yuv_to_rgb_matrix(&BT601_YUV_TO_RGB, y, u, v),
38            ColorSpace::BT709 => convert_yuv_to_rgb_matrix(&BT709_YUV_TO_RGB, y, u, v),
39            ColorSpace::BT2020 => convert_yuv_to_rgb_matrix(&BT2020_YUV_TO_RGB, y, u, v),
40            ColorSpace::ICtCpPQ => bt2100::pq_yuv_to_rgb(transfer, xyz_to_rgb, y, u, v),
41            ColorSpace::ICtCpHLG => bt2100::hlg_yuv_to_rgb(transfer, xyz_to_rgb, y, u, v),
42        }
43    }
44    #[inline(always)]
45    pub(crate) unsafe fn yx4_uv_to_rgb<V: Vector>(
46        &self,
47        transfer: ColorTransfer,
48        xyz_to_rgb: &'static [[f32; 3]; 3],
49        y00: V,
50        y01: V,
51        y10: V,
52        y11: V,
53        u: V,
54        v: V,
55    ) -> [[V; 3]; 4] {
56        match self {
57            ColorSpace::BT601 => {
58                convert_yx4_uv_to_rgb_matrix(&BT601_YUV_TO_RGB, y00, y01, y10, y11, u, v)
59            }
60            ColorSpace::BT709 => {
61                convert_yx4_uv_to_rgb_matrix(&BT709_YUV_TO_RGB, y00, y01, y10, y11, u, v)
62            }
63            ColorSpace::BT2020 => {
64                convert_yx4_uv_to_rgb_matrix(&BT2020_YUV_TO_RGB, y00, y01, y10, y11, u, v)
65            }
66            ColorSpace::ICtCpPQ => bt2100_yx4_uv_to_rgb(
67                bt2100::pq_yuv_to_rgb,
68                transfer,
69                xyz_to_rgb,
70                y00,
71                y01,
72                y10,
73                y11,
74                u,
75                v,
76            ),
77            ColorSpace::ICtCpHLG => bt2100_yx4_uv_to_rgb(
78                bt2100::hlg_yuv_to_rgb,
79                transfer,
80                xyz_to_rgb,
81                y00,
82                y01,
83                y10,
84                y11,
85                u,
86                v,
87            ),
88        }
89    }
90
91    #[inline(always)]
92    pub(crate) unsafe fn rgb_to_yuv<V: Vector>(
93        &self,
94        transfer: ColorTransfer,
95        rgb_to_xyz: &'static [[f32; 3]; 3],
96        r: V,
97        g: V,
98        b: V,
99    ) -> (V, V, V) {
100        match self {
101            ColorSpace::BT601 => convert_rgb_to_yuv_matrix(&BT601_RGB_TO_YUV, r, g, b),
102            ColorSpace::BT709 => convert_rgb_to_yuv_matrix(&BT709_RGB_TO_YUV, r, g, b),
103            ColorSpace::BT2020 => convert_rgb_to_yuv_matrix(&BT2020_RGB_TO_YUV, r, g, b),
104            ColorSpace::ICtCpPQ => bt2100::pq_rgb_to_yuv(transfer, rgb_to_xyz, r, g, b),
105            ColorSpace::ICtCpHLG => bt2100::hlg_rgb_to_yuv(transfer, rgb_to_xyz, r, g, b),
106        }
107    }
108    #[inline(always)]
109    pub(crate) unsafe fn rgbx4_to_yx4_uv<V: Vector>(
110        &self,
111        transfer: ColorTransfer,
112        rgb_to_xyz: &'static [[f32; 3]; 3],
113        r: [V; 4],
114        g: [V; 4],
115        b: [V; 4],
116    ) -> ([V; 4], V, V) {
117        match self {
118            ColorSpace::BT601 => rgbx4_to_yx4_uv_matrix(&BT601_RGB_TO_YUV, r, g, b),
119            ColorSpace::BT709 => rgbx4_to_yx4_uv_matrix(&BT709_RGB_TO_YUV, r, g, b),
120            ColorSpace::BT2020 => rgbx4_to_yx4_uv_matrix(&BT2020_RGB_TO_YUV, r, g, b),
121            ColorSpace::ICtCpPQ => {
122                bt2100_rgbx4_to_yx4_uv(bt2100::pq_rgb_to_yuv, transfer, rgb_to_xyz, r, g, b)
123            }
124            ColorSpace::ICtCpHLG => {
125                bt2100_rgbx4_to_yx4_uv(bt2100::hlg_rgb_to_yuv, transfer, rgb_to_xyz, r, g, b)
126            }
127        }
128    }
129}
130
131#[inline(always)]
132unsafe fn convert_yuv_to_rgb_matrix<V: Vector>(mat: &[[f32; 3]; 3], y: V, u: V, v: V) -> (V, V, V) {
133    let r = y.vadd(v.vmulf(mat[V][R]));
134    let g = y.vadd(v.vmulf(mat[V][G]).vadd(u.vmulf(mat[U][G])));
135    let b = y.vadd(u.vmulf(mat[U][B]));
136
137    (r, g, b)
138}
139
140#[inline(always)]
141unsafe fn convert_yx4_uv_to_rgb_matrix<V: Vector>(
142    mat: &[[f32; 3]; 3],
143    y00: V,
144    y01: V,
145    y10: V,
146    y11: V,
147    u: V,
148    v: V,
149) -> [[V; 3]; 4] {
150    #[inline(always)]
151    unsafe fn prepare_rgb<V: Vector>(mat: &[[f32; 3]; 3], u: V, v: V) -> (V, V, V) {
152        let r = v.vmulf(mat[V][R]);
153        let g = v.vmulf(mat[V][G]).vadd(u.vmulf(mat[U][G]));
154        let b = u.vmulf(mat[U][B]);
155
156        (r, g, b)
157    }
158
159    let (left_u, right_u) = u.zip(u);
160    let (left_v, right_v) = v.zip(v);
161
162    let (r_left, g_left, b_left) = prepare_rgb(mat, left_u, left_v);
163    let (r_right, g_right, b_right) = prepare_rgb(mat, right_u, right_v);
164
165    let r00 = y00.vadd(r_left);
166    let g00 = y00.vadd(g_left);
167    let b00 = y00.vadd(b_left);
168
169    let r01 = y01.vadd(r_right);
170    let g01 = y01.vadd(g_right);
171    let b01 = y01.vadd(b_right);
172
173    let r10 = y10.vadd(r_left);
174    let g10 = y10.vadd(g_left);
175    let b10 = y10.vadd(b_left);
176
177    let r11 = y11.vadd(r_right);
178    let g11 = y11.vadd(g_right);
179    let b11 = y11.vadd(b_right);
180
181    [
182        [r00, g00, b00],
183        [r01, g01, b01],
184        [r10, g10, b10],
185        [r11, g11, b11],
186    ]
187}
188
189#[inline(always)]
190unsafe fn convert_rgb_to_yuv_matrix<V: Vector>(mat: &[[f32; 3]; 3], r: V, g: V, b: V) -> (V, V, V) {
191    let y = r
192        .vmulf(mat[0][0])
193        .vadd(g.vmulf(mat[0][1]))
194        .vadd(b.vmulf(mat[0][2]));
195    let u = r
196        .vmulf(mat[1][0])
197        .vadd(g.vmulf(mat[1][1]))
198        .vadd(b.vmulf(mat[1][2]));
199    let v = r
200        .vmulf(mat[2][0])
201        .vadd(g.vmulf(mat[2][1]))
202        .vadd(b.vmulf(mat[2][2]));
203
204    (y, u, v)
205}
206
207#[inline(always)]
208unsafe fn rgbx4_to_yx4_uv_matrix<V: Vector>(
209    mat: &[[f32; 3]; 3],
210    r: [V; 4],
211    g: [V; 4],
212    b: [V; 4],
213) -> ([V; 4], V, V) {
214    #[inline(always)]
215    unsafe fn calc_y<V: Vector>(mat: &[[f32; 3]; 3], r: V, g: V, b: V) -> V {
216        r.vmulf(mat[Y][R])
217            .vadd(g.vmulf(mat[Y][G]))
218            .vadd(b.vmulf(mat[Y][B]))
219    }
220
221    #[inline(always)]
222    unsafe fn calc_u<V: Vector>(mat: &[[f32; 3]; 3], r: V, g: V, b: V) -> V {
223        r.vmulf(mat[U][R])
224            .vadd(g.vmulf(mat[U][G]))
225            .vadd(b.vmulf(mat[U][B]))
226    }
227
228    #[inline(always)]
229    unsafe fn calc_v<V: Vector>(mat: &[[f32; 3]; 3], r: V, g: V, b: V) -> V {
230        r.vmulf(mat[V][R])
231            .vadd(g.vmulf(mat[V][G]))
232            .vadd(b.vmulf(mat[V][B]))
233    }
234
235    // Calculate Y
236    let y00 = calc_y(mat, r[0], g[0], b[0]);
237    let y01 = calc_y(mat, r[1], g[1], b[1]);
238    let y10 = calc_y(mat, r[2], g[2], b[2]);
239    let y11 = calc_y(mat, r[3], g[3], b[3]);
240
241    // Calculate U/V pixel from 2x2 rgb pixel blocks
242    // Add the upper and lower pixels
243    let rgb0_r = r[0].vadd(r[2]);
244    let rgb0_g = g[0].vadd(g[2]);
245    let rgb0_b = b[0].vadd(b[2]);
246
247    let rgb1_r = r[1].vadd(r[3]);
248    let rgb1_g = g[1].vadd(g[3]);
249    let rgb1_b = b[1].vadd(b[3]);
250
251    let (rgb0_r, rgb1_r) = rgb0_r.unzip(rgb1_r);
252    let (rgb0_g, rgb1_g) = rgb0_g.unzip(rgb1_g);
253    let (rgb0_b, rgb1_b) = rgb0_b.unzip(rgb1_b);
254
255    let r = rgb0_r.vadd(rgb1_r);
256    let g = rgb0_g.vadd(rgb1_g);
257    let b = rgb0_b.vadd(rgb1_b);
258
259    let r = r.vmulf(0.25);
260    let g = g.vmulf(0.25);
261    let b = b.vmulf(0.25);
262
263    let u = calc_u(mat, r, g, b);
264    let v = calc_v(mat, r, g, b);
265
266    ([y00, y01, y10, y11], u, v)
267}
268
269#[rustfmt::skip]
270macro_rules! make_matrices {
271    ($($yuv_to_rgb:ident, $rgb_to_yuv:ident: $kr:expr, $kg:expr, $kb:expr;)*) => {
272        $(
273        pub(crate) const $yuv_to_rgb: [[f32; 3]; 3] = [
274            // R                 G                                  B
275            [1.0,                1.0,                               1.0            ], // Y
276            [0.0,               (-($kb / $kg)) * (2.0 - 2.0 * $kb), 2.0 - 2.0 * $kb], // U
277            [(2.0 - 2.0 * $kr), (-($kr / $kg)) * (2.0 - 2.0 * $kr), 0.0            ], // V
278        ];
279
280        pub(crate) const $rgb_to_yuv: [[f32; 3]; 3] = [
281            // Y                         U                           V
282            [$kr,                        $kg,                        $kb                       ], // R
283            [-0.5 * ($kr / (1.0 - $kb)), -0.5 * ($kg / (1.0 - $kb)), 0.5                       ], // G
284            [0.5,                        -0.5 * ($kg / (1.0 - $kr)), -0.5 * ($kb / (1.0 - $kr))], // B
285        ];
286        )*
287    };
288}
289
290make_matrices! {
291    BT601_YUV_TO_RGB, BT601_RGB_TO_YUV: 0.299,  0.587,  0.114;
292    BT709_YUV_TO_RGB, BT709_RGB_TO_YUV: 0.2126, 0.7152, 0.0722;
293    BT2020_YUV_TO_RGB, BT2020_RGB_TO_YUV: 0.2627, 0.322, 0.0593;
294}
295
296type BT2100ConvertFn<V> = unsafe fn(ColorTransfer, &[[f32; 3]; 3], V, V, V) -> (V, V, V);
297
298#[inline(always)]
299unsafe fn bt2100_yx4_uv_to_rgb<V: Vector>(
300    f: BT2100ConvertFn<V>,
301    transfer: ColorTransfer,
302    xyz_to_rgb: &[[f32; 3]; 3],
303    y00: V,
304    y01: V,
305    y10: V,
306    y11: V,
307    u: V,
308    v: V,
309) -> [[V; 3]; 4] {
310    let (left_u, right_u) = u.zip(u);
311    let (left_v, right_v) = v.zip(v);
312
313    let rgb00 = f(transfer, xyz_to_rgb, y00, left_u, left_v);
314    let rgb01 = f(transfer, xyz_to_rgb, y01, right_u, right_v);
315    let rgb10 = f(transfer, xyz_to_rgb, y10, left_u, left_v);
316    let rgb11 = f(transfer, xyz_to_rgb, y11, right_u, right_v);
317
318    [
319        [rgb00.0, rgb00.1, rgb00.2],
320        [rgb01.0, rgb01.1, rgb01.2],
321        [rgb10.0, rgb10.1, rgb10.2],
322        [rgb11.0, rgb11.1, rgb11.2],
323    ]
324}
325
326#[inline(always)]
327unsafe fn bt2100_rgbx4_to_yx4_uv<V: Vector>(
328    f: BT2100ConvertFn<V>,
329    transfer: ColorTransfer,
330    rgb_to_xyz: &[[f32; 3]; 3],
331    r: [V; 4],
332    g: [V; 4],
333    b: [V; 4],
334) -> ([V; 4], V, V) {
335    let yuv00 = f(transfer, rgb_to_xyz, r[0], g[0], b[0]);
336    let yuv01 = f(transfer, rgb_to_xyz, r[1], g[1], b[1]);
337    let yuv10 = f(transfer, rgb_to_xyz, r[2], g[2], b[2]);
338    let yuv11 = f(transfer, rgb_to_xyz, r[3], g[3], b[3]);
339
340    let u = yuv00.1.vadd(yuv01.1).vadd(yuv10.1).vadd(yuv11.1);
341    let v = yuv00.2.vadd(yuv01.2).vadd(yuv10.2).vadd(yuv11.2);
342
343    let u = u.vdivf(4.0);
344    let v = v.vdivf(4.0);
345
346    ([yuv00.0, yuv01.0, yuv10.0, yuv11.0], u, v)
347}
348
349mod bt2100 {
350    use crate::ColorTransfer;
351    use crate::color::primaries::{rgb_to_xyz, xyz_to_rgb};
352    use crate::color::transfer::{bt2100_hlg, bt2100_pq};
353    use crate::vector::Vector;
354
355    #[inline(always)]
356    unsafe fn xyz_to_lms<V: Vector>(x: V, y: V, z: V) -> (V, V, V) {
357        let l = x.vmulf(0.3593).vadd(y.vmulf(0.6976)).vadd(z.vmulf(-0.0359));
358        let m = x.vmulf(-0.1921).vadd(y.vmulf(1.1005)).vadd(z.vmulf(0.0754));
359        let s = x.vmulf(0.0071).vadd(y.vmulf(0.0748)).vadd(z.vmulf(0.8433));
360
361        (l, m, s)
362    }
363
364    #[inline(always)]
365    unsafe fn lms_to_xyz<V: Vector>(l: V, m: V, s: V) -> (V, V, V) {
366        let x = l
367            .vmulf(2.070_034_5)
368            .vadd(m.vmulf(-1.326_231_2))
369            .vadd(s.vmulf(0.206_702_32));
370        let y = l
371            .vmulf(0.364_749_8)
372            .vadd(m.vmulf(0.680_545_7))
373            .vadd(s.vmulf(-0.045_320_32));
374        let z = l
375            .vmulf(-0.049_781_25)
376            .vadd(m.vmulf(-0.049_197_882))
377            .vadd(s.vmulf(1.188_097_2));
378        (x, y, z)
379    }
380
381    #[inline(always)]
382    pub(super) unsafe fn pq_rgb_to_yuv<V: Vector>(
383        transfer: ColorTransfer,
384        rgb_to_xyz_mat: &[[f32; 3]; 3],
385        mut r: V,
386        mut g: V,
387        mut b: V,
388    ) -> (V, V, V) {
389        transfer.scaled_to_linear_v(&mut [&mut r, &mut g, &mut b]);
390
391        let [x, y, z] = rgb_to_xyz(rgb_to_xyz_mat, r, g, b);
392
393        let (l, m, s) = xyz_to_lms(x, y, z);
394
395        let l = bt2100_pq::linear_to_scaled(l);
396        let m = bt2100_pq::linear_to_scaled(m);
397        let s = bt2100_pq::linear_to_scaled(s);
398
399        // I = 0.5L’ + 0.5M’
400        let y = l.vmulf(0.5).vadd(m.vmulf(0.5));
401
402        // Ct = (6610L’ - 13613M’ + 7003S’ ) / 4096
403        let u = l
404            .vmulf(6610.0)
405            .vadd(m.vmulf(-13613.0))
406            .vadd(s.vmulf(7003.0))
407            .vdivf(4096.0);
408
409        // Cp = (17933L’ - 17390M’ - 543S’ ) / 4096
410        let v = l
411            .vmulf(17933.0)
412            .vadd(m.vmulf(-17390.0))
413            .vadd(s.vmulf(-543.0))
414            .vdivf(4096.0);
415
416        (y, u, v)
417    }
418
419    #[inline(always)]
420    pub(super) unsafe fn pq_yuv_to_rgb<V: Vector>(
421        transfer: ColorTransfer,
422        xyz_to_rgb_mat: &[[f32; 3]; 3],
423        y: V,
424        u: V,
425        v: V,
426    ) -> (V, V, V) {
427        // l = y + u * 0.008609037 + v * 0.111029625
428        let l = y.vadd(u.vmulf(0.008609037)).vadd(v.vmulf(0.111029625));
429
430        // m = y + u * -0.008609037 + v * -0.111029625
431        let m = y.vadd(u.vmulf(-0.008609037)).vadd(v.vmulf(-0.111029625));
432
433        // s = y + u * 0.5600313 + v * -0.32062715
434        let s = y.vadd(u.vmulf(0.5600313)).vadd(v.vmulf(-0.32062715));
435
436        let l = bt2100_pq::scaled_to_linear(l);
437        let m = bt2100_pq::scaled_to_linear(m);
438        let s = bt2100_pq::scaled_to_linear(s);
439
440        let (x, y, z) = lms_to_xyz(l, m, s);
441
442        let [mut r, mut g, mut b] = xyz_to_rgb(xyz_to_rgb_mat, x, y, z);
443
444        transfer.linear_to_scaled_v(&mut [&mut r, &mut g, &mut b]);
445
446        (r, g, b)
447    }
448
449    #[inline(always)]
450    pub(super) unsafe fn hlg_rgb_to_yuv<V: Vector>(
451        transfer: ColorTransfer,
452        rgb_to_xyz_mat: &[[f32; 3]; 3],
453        mut r: V,
454        mut g: V,
455        mut b: V,
456    ) -> (V, V, V) {
457        transfer.scaled_to_linear_v(&mut [&mut r, &mut g, &mut b]);
458
459        let [x, y, z] = rgb_to_xyz(rgb_to_xyz_mat, r, g, b);
460
461        let (l, m, s) = xyz_to_lms(x, y, z);
462
463        let l = bt2100_hlg::linear_to_scaled(l);
464        let m = bt2100_hlg::linear_to_scaled(m);
465        let s = bt2100_hlg::linear_to_scaled(s);
466
467        // y = 0.5 * (l + m)
468        let y = l.vadd(m).vmulf(0.5);
469
470        // u = 0.88500977 * l - 1.8225098 * m + 0.9375 * s
471        let u = l
472            .vmulf(0.88500977)
473            .vsub(m.vmulf(1.8225098))
474            .vadd(s.vmulf(0.9375));
475
476        // v = 2.319336 * l - 2.2490234 * m - 0.0703125 * s
477        let v = l
478            .vmulf(2.319336)
479            .vsub(m.vmulf(2.2490234))
480            .vsub(s.vmulf(0.0703125));
481
482        (y, u, v)
483    }
484
485    #[inline(always)]
486    pub(super) unsafe fn hlg_yuv_to_rgb<V: Vector>(
487        transfer: ColorTransfer,
488        xyz_to_rgb_mat: &[[f32; 3]; 3],
489        y: V,
490        u: V,
491        v: V,
492    ) -> (V, V, V) {
493        // l = y + u * 0.01571858 + v * 0.20958106
494        let l = y
495            .vadd(u.vmul(V::splat(0.01571858)))
496            .vadd(v.vmul(V::splat(0.20958106)));
497
498        // m = y + u * -0.01571858 + v * -0.20958106
499        let m = y
500            .vadd(u.vmul(V::splat(-0.01571858)))
501            .vadd(v.vmul(V::splat(-0.20958106)));
502
503        // s = y + u * 1.0212711 + v * -0.6052745
504        let s = y
505            .vadd(u.vmul(V::splat(1.0212711)))
506            .vadd(v.vmul(V::splat(-0.6052745)));
507
508        let l = bt2100_hlg::scaled_to_linear(l);
509        let m = bt2100_hlg::scaled_to_linear(m);
510        let s = bt2100_hlg::scaled_to_linear(s);
511
512        let (x, y, z) = lms_to_xyz(l, m, s);
513
514        let [mut r, mut g, mut b] = xyz_to_rgb(xyz_to_rgb_mat, x, y, z);
515
516        transfer.linear_to_scaled_v(&mut [&mut r, &mut g, &mut b]);
517
518        (r, g, b)
519    }
520}
ezk_image/color/space.rs

ezk_image/color/
space.rs
