1use crate::{GpuError, Result};
13use rayon::prelude::*;
14
15#[derive(Debug, Clone, Copy)]
21pub struct YcbcrCoefficients {
22 pub kr: f64,
24 pub kg: f64,
26 pub kb: f64,
28}
29
30impl YcbcrCoefficients {
31 pub const BT601: Self = Self {
33 kr: 0.299,
34 kg: 0.587,
35 kb: 0.114,
36 };
37
38 pub const BT709: Self = Self {
40 kr: 0.2126,
41 kg: 0.7152,
42 kb: 0.0722,
43 };
44
45 pub const BT2020: Self = Self {
47 kr: 0.2627,
48 kg: 0.6780,
49 kb: 0.0593,
50 };
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55pub enum ChromaSubsampling {
56 Yuv420,
58 Yuv422,
60}
61
62impl ChromaSubsampling {
63 fn output_size(self, width: u32, height: u32) -> usize {
65 let w = width as usize;
66 let h = height as usize;
67 let y_size = w * h;
68 match self {
69 Self::Yuv420 => {
70 let uv_w = (w + 1) / 2;
71 let uv_h = (h + 1) / 2;
72 y_size + 2 * uv_w * uv_h
73 }
74 Self::Yuv422 => {
75 let uv_w = (w + 1) / 2;
76 y_size + 2 * uv_w * h
77 }
78 }
79 }
80}
81
82pub struct ChromaOps;
84
85impl ChromaOps {
86 pub fn rgba_to_ycbcr(
95 rgba: &[u8],
96 width: u32,
97 height: u32,
98 format: ChromaSubsampling,
99 coefficients: YcbcrCoefficients,
100 ) -> Result<Vec<u8>> {
101 let w = width as usize;
102 let h = height as usize;
103 let expected_input = w * h * 4;
104
105 if width == 0 || height == 0 {
106 return Err(GpuError::InvalidDimensions { width, height });
107 }
108 if rgba.len() < expected_input {
109 return Err(GpuError::InvalidBufferSize {
110 expected: expected_input,
111 actual: rgba.len(),
112 });
113 }
114
115 match format {
116 ChromaSubsampling::Yuv420 => Self::rgba_to_yuv420(rgba, w, h, coefficients),
117 ChromaSubsampling::Yuv422 => Self::rgba_to_yuv422(rgba, w, h, coefficients),
118 }
119 }
120
121 pub fn ycbcr_to_rgba(
130 ycbcr: &[u8],
131 width: u32,
132 height: u32,
133 format: ChromaSubsampling,
134 coefficients: YcbcrCoefficients,
135 ) -> Result<Vec<u8>> {
136 let w = width as usize;
137 let h = height as usize;
138
139 if width == 0 || height == 0 {
140 return Err(GpuError::InvalidDimensions { width, height });
141 }
142
143 let expected = format.output_size(width, height);
144 if ycbcr.len() < expected {
145 return Err(GpuError::InvalidBufferSize {
146 expected,
147 actual: ycbcr.len(),
148 });
149 }
150
151 match format {
152 ChromaSubsampling::Yuv420 => Self::yuv420_to_rgba(ycbcr, w, h, coefficients),
153 ChromaSubsampling::Yuv422 => Self::yuv422_to_rgba(ycbcr, w, h, coefficients),
154 }
155 }
156
157 fn rgba_to_yuv420(
162 rgba: &[u8],
163 w: usize,
164 h: usize,
165 coeff: YcbcrCoefficients,
166 ) -> Result<Vec<u8>> {
167 let y_size = w * h;
168 let uv_w = (w + 1) / 2;
169 let uv_h = (h + 1) / 2;
170 let uv_size = uv_w * uv_h;
171 let mut output = vec![0u8; y_size + 2 * uv_size];
172
173 let y_plane = &mut output[..y_size];
175 y_plane
176 .par_chunks_exact_mut(w)
177 .enumerate()
178 .for_each(|(y, row)| {
179 for x in 0..w {
180 let base = (y * w + x) * 4;
181 let r = rgba[base] as f64;
182 let g = rgba[base + 1] as f64;
183 let b = rgba[base + 2] as f64;
184 let luma = coeff.kr * r + coeff.kg * g + coeff.kb * b;
185 row[x] = luma.round().clamp(0.0, 255.0) as u8;
186 }
187 });
188
189 let cb_start = y_size;
191 let cr_start = y_size + uv_size;
192
193 for by in 0..uv_h {
194 for bx in 0..uv_w {
195 let mut sum_cb = 0.0_f64;
196 let mut sum_cr = 0.0_f64;
197 let mut count = 0u32;
198
199 for dy in 0..2_usize {
200 let sy = by * 2 + dy;
201 if sy >= h {
202 continue;
203 }
204 for dx in 0..2_usize {
205 let sx = bx * 2 + dx;
206 if sx >= w {
207 continue;
208 }
209 let base = (sy * w + sx) * 4;
210 let r = rgba[base] as f64;
211 let g = rgba[base + 1] as f64;
212 let b = rgba[base + 2] as f64;
213 let y_val = coeff.kr * r + coeff.kg * g + coeff.kb * b;
214 let denom_cb = 2.0 * (1.0 - coeff.kb);
216 let cb = if denom_cb.abs() > 1e-10 {
217 (b - y_val) / denom_cb + 128.0
218 } else {
219 128.0
220 };
221 let denom_cr = 2.0 * (1.0 - coeff.kr);
223 let cr = if denom_cr.abs() > 1e-10 {
224 (r - y_val) / denom_cr + 128.0
225 } else {
226 128.0
227 };
228 sum_cb += cb;
229 sum_cr += cr;
230 count += 1;
231 }
232 }
233
234 let uv_idx = by * uv_w + bx;
235 if count > 0 {
236 output[cb_start + uv_idx] =
237 (sum_cb / count as f64).round().clamp(0.0, 255.0) as u8;
238 output[cr_start + uv_idx] =
239 (sum_cr / count as f64).round().clamp(0.0, 255.0) as u8;
240 } else {
241 output[cb_start + uv_idx] = 128;
242 output[cr_start + uv_idx] = 128;
243 }
244 }
245 }
246
247 Ok(output)
248 }
249
250 fn yuv420_to_rgba(
255 ycbcr: &[u8],
256 w: usize,
257 h: usize,
258 coeff: YcbcrCoefficients,
259 ) -> Result<Vec<u8>> {
260 let y_size = w * h;
261 let uv_w = (w + 1) / 2;
262 let uv_h = (h + 1) / 2;
263 let uv_size = uv_w * uv_h;
264
265 let y_plane = &ycbcr[..y_size];
266 let cb_plane = &ycbcr[y_size..y_size + uv_size];
267 let cr_plane = &ycbcr[y_size + uv_size..y_size + 2 * uv_size];
268
269 let mut rgba = vec![0u8; w * h * 4];
270
271 rgba.par_chunks_exact_mut(w * 4)
272 .enumerate()
273 .for_each(|(py, row)| {
274 let uv_y = (py / 2).min(uv_h.saturating_sub(1));
275 for px in 0..w {
276 let uv_x = (px / 2).min(uv_w.saturating_sub(1));
277 let uv_idx = uv_y * uv_w + uv_x;
278
279 let y_val = y_plane[py * w + px] as f64;
280 let cb = cb_plane[uv_idx] as f64 - 128.0;
281 let cr = cr_plane[uv_idx] as f64 - 128.0;
282
283 let r = y_val + 2.0 * (1.0 - coeff.kr) * cr;
284 let b = y_val + 2.0 * (1.0 - coeff.kb) * cb;
285 let g = if coeff.kg.abs() > 1e-10 {
286 (y_val - coeff.kr * r - coeff.kb * b) / coeff.kg
287 } else {
288 y_val
289 };
290
291 let base = px * 4;
292 row[base] = r.round().clamp(0.0, 255.0) as u8;
293 row[base + 1] = g.round().clamp(0.0, 255.0) as u8;
294 row[base + 2] = b.round().clamp(0.0, 255.0) as u8;
295 row[base + 3] = 255;
296 }
297 });
298
299 Ok(rgba)
300 }
301
302 fn rgba_to_yuv422(
307 rgba: &[u8],
308 w: usize,
309 h: usize,
310 coeff: YcbcrCoefficients,
311 ) -> Result<Vec<u8>> {
312 let y_size = w * h;
313 let uv_w = (w + 1) / 2;
314 let uv_size = uv_w * h;
315 let mut output = vec![0u8; y_size + 2 * uv_size];
316
317 let y_plane = &mut output[..y_size];
319 y_plane
320 .par_chunks_exact_mut(w)
321 .enumerate()
322 .for_each(|(y, row)| {
323 for x in 0..w {
324 let base = (y * w + x) * 4;
325 let r = rgba[base] as f64;
326 let g = rgba[base + 1] as f64;
327 let b = rgba[base + 2] as f64;
328 let luma = coeff.kr * r + coeff.kg * g + coeff.kb * b;
329 row[x] = luma.round().clamp(0.0, 255.0) as u8;
330 }
331 });
332
333 let cb_start = y_size;
335 let cr_start = y_size + uv_size;
336
337 for y in 0..h {
338 for bx in 0..uv_w {
339 let mut sum_cb = 0.0_f64;
340 let mut sum_cr = 0.0_f64;
341 let mut count = 0u32;
342
343 for dx in 0..2_usize {
344 let sx = bx * 2 + dx;
345 if sx >= w {
346 continue;
347 }
348 let base = (y * w + sx) * 4;
349 let r = rgba[base] as f64;
350 let g = rgba[base + 1] as f64;
351 let b = rgba[base + 2] as f64;
352 let y_val = coeff.kr * r + coeff.kg * g + coeff.kb * b;
353
354 let denom_cb = 2.0 * (1.0 - coeff.kb);
355 let cb = if denom_cb.abs() > 1e-10 {
356 (b - y_val) / denom_cb + 128.0
357 } else {
358 128.0
359 };
360 let denom_cr = 2.0 * (1.0 - coeff.kr);
361 let cr = if denom_cr.abs() > 1e-10 {
362 (r - y_val) / denom_cr + 128.0
363 } else {
364 128.0
365 };
366
367 sum_cb += cb;
368 sum_cr += cr;
369 count += 1;
370 }
371
372 let uv_idx = y * uv_w + bx;
373 if count > 0 {
374 output[cb_start + uv_idx] =
375 (sum_cb / count as f64).round().clamp(0.0, 255.0) as u8;
376 output[cr_start + uv_idx] =
377 (sum_cr / count as f64).round().clamp(0.0, 255.0) as u8;
378 } else {
379 output[cb_start + uv_idx] = 128;
380 output[cr_start + uv_idx] = 128;
381 }
382 }
383 }
384
385 Ok(output)
386 }
387
388 fn yuv422_to_rgba(
393 ycbcr: &[u8],
394 w: usize,
395 h: usize,
396 coeff: YcbcrCoefficients,
397 ) -> Result<Vec<u8>> {
398 let y_size = w * h;
399 let uv_w = (w + 1) / 2;
400 let uv_size = uv_w * h;
401
402 let y_plane = &ycbcr[..y_size];
403 let cb_plane = &ycbcr[y_size..y_size + uv_size];
404 let cr_plane = &ycbcr[y_size + uv_size..y_size + 2 * uv_size];
405
406 let mut rgba = vec![0u8; w * h * 4];
407
408 rgba.par_chunks_exact_mut(w * 4)
409 .enumerate()
410 .for_each(|(py, row)| {
411 for px in 0..w {
412 let uv_x = (px / 2).min(uv_w.saturating_sub(1));
413 let uv_idx = py * uv_w + uv_x;
414
415 let y_val = y_plane[py * w + px] as f64;
416 let cb = cb_plane[uv_idx] as f64 - 128.0;
417 let cr = cr_plane[uv_idx] as f64 - 128.0;
418
419 let r = y_val + 2.0 * (1.0 - coeff.kr) * cr;
420 let b = y_val + 2.0 * (1.0 - coeff.kb) * cb;
421 let g = if coeff.kg.abs() > 1e-10 {
422 (y_val - coeff.kr * r - coeff.kb * b) / coeff.kg
423 } else {
424 y_val
425 };
426
427 let base = px * 4;
428 row[base] = r.round().clamp(0.0, 255.0) as u8;
429 row[base + 1] = g.round().clamp(0.0, 255.0) as u8;
430 row[base + 2] = b.round().clamp(0.0, 255.0) as u8;
431 row[base + 3] = 255;
432 }
433 });
434
435 Ok(rgba)
436 }
437}
438
439#[cfg(test)]
444mod tests {
445 use super::*;
446
447 fn solid_rgba(w: u32, h: u32, r: u8, g: u8, b: u8) -> Vec<u8> {
448 let n = (w as usize) * (h as usize);
449 let mut buf = Vec::with_capacity(n * 4);
450 for _ in 0..n {
451 buf.extend_from_slice(&[r, g, b, 255]);
452 }
453 buf
454 }
455
456 fn gradient_rgba(w: u32, h: u32) -> Vec<u8> {
457 let ww = w as usize;
458 let hh = h as usize;
459 let mut buf = Vec::with_capacity(ww * hh * 4);
460 for y in 0..hh {
461 for x in 0..ww {
462 let v = ((x + y) % 256) as u8;
463 buf.extend_from_slice(&[v, v / 2, 255 - v, 255]);
464 }
465 }
466 buf
467 }
468
469 #[test]
472 fn test_420_output_size() {
473 let rgba = solid_rgba(16, 16, 128, 128, 128);
474 let yuv = ChromaOps::rgba_to_ycbcr(
475 &rgba,
476 16,
477 16,
478 ChromaSubsampling::Yuv420,
479 YcbcrCoefficients::BT601,
480 )
481 .expect("conversion should succeed");
482 assert_eq!(yuv.len(), 384);
484 }
485
486 #[test]
487 fn test_420_roundtrip_grey() {
488 let rgba = solid_rgba(16, 16, 128, 128, 128);
489 let yuv = ChromaOps::rgba_to_ycbcr(
490 &rgba,
491 16,
492 16,
493 ChromaSubsampling::Yuv420,
494 YcbcrCoefficients::BT601,
495 )
496 .expect("forward");
497 let back = ChromaOps::ycbcr_to_rgba(
498 &yuv,
499 16,
500 16,
501 ChromaSubsampling::Yuv420,
502 YcbcrCoefficients::BT601,
503 )
504 .expect("inverse");
505 for i in 0..(16 * 16) {
507 let base = i * 4;
508 for c in 0..3 {
509 let diff = (rgba[base + c] as i32 - back[base + c] as i32).unsigned_abs();
510 assert!(diff <= 2, "pixel {i} channel {c}: diff={diff}");
511 }
512 }
513 }
514
515 #[test]
516 fn test_420_roundtrip_gradient() {
517 let rgba = gradient_rgba(32, 32);
518 let yuv = ChromaOps::rgba_to_ycbcr(
519 &rgba,
520 32,
521 32,
522 ChromaSubsampling::Yuv420,
523 YcbcrCoefficients::BT601,
524 )
525 .expect("forward");
526 let back = ChromaOps::ycbcr_to_rgba(
527 &yuv,
528 32,
529 32,
530 ChromaSubsampling::Yuv420,
531 YcbcrCoefficients::BT601,
532 )
533 .expect("inverse");
534 let max_diff: u32 = (0..(32 * 32))
536 .map(|i| {
537 let base = i * 4;
538 (0..3)
539 .map(|c| (rgba[base + c] as i32 - back[base + c] as i32).unsigned_abs())
540 .max()
541 .unwrap_or(0)
542 })
543 .max()
544 .unwrap_or(0);
545 assert!(
546 max_diff <= 10,
547 "max roundtrip diff={max_diff}, expected <= 10"
548 );
549 }
550
551 #[test]
552 fn test_420_white() {
553 let rgba = solid_rgba(4, 4, 255, 255, 255);
554 let yuv = ChromaOps::rgba_to_ycbcr(
555 &rgba,
556 4,
557 4,
558 ChromaSubsampling::Yuv420,
559 YcbcrCoefficients::BT601,
560 )
561 .expect("forward");
562 assert!(yuv[0] > 250, "Y for white should be ~255, got {}", yuv[0]);
564 let y_size = 4 * 4;
566 let cb = yuv[y_size];
567 let cr = yuv[y_size + 4]; assert!(
569 (cb as i32 - 128).unsigned_abs() <= 2,
570 "Cb for white should be ~128, got {cb}"
571 );
572 assert!(
573 (cr as i32 - 128).unsigned_abs() <= 2,
574 "Cr for white should be ~128, got {cr}"
575 );
576 }
577
578 #[test]
579 fn test_420_black() {
580 let rgba = solid_rgba(4, 4, 0, 0, 0);
581 let yuv = ChromaOps::rgba_to_ycbcr(
582 &rgba,
583 4,
584 4,
585 ChromaSubsampling::Yuv420,
586 YcbcrCoefficients::BT601,
587 )
588 .expect("forward");
589 assert_eq!(yuv[0], 0, "Y for black should be 0");
590 }
591
592 #[test]
593 fn test_420_odd_dimensions() {
594 let rgba = solid_rgba(15, 13, 100, 150, 200);
596 let yuv = ChromaOps::rgba_to_ycbcr(
597 &rgba,
598 15,
599 13,
600 ChromaSubsampling::Yuv420,
601 YcbcrCoefficients::BT601,
602 )
603 .expect("forward");
604 let expected = ChromaSubsampling::Yuv420.output_size(15, 13);
605 assert_eq!(yuv.len(), expected);
606 let back = ChromaOps::ycbcr_to_rgba(
608 &yuv,
609 15,
610 13,
611 ChromaSubsampling::Yuv420,
612 YcbcrCoefficients::BT601,
613 )
614 .expect("inverse");
615 assert_eq!(back.len(), 15 * 13 * 4);
616 }
617
618 #[test]
621 fn test_422_output_size() {
622 let rgba = solid_rgba(16, 16, 128, 128, 128);
623 let yuv = ChromaOps::rgba_to_ycbcr(
624 &rgba,
625 16,
626 16,
627 ChromaSubsampling::Yuv422,
628 YcbcrCoefficients::BT601,
629 )
630 .expect("conversion should succeed");
631 assert_eq!(yuv.len(), 512);
633 }
634
635 #[test]
636 fn test_422_roundtrip_grey() {
637 let rgba = solid_rgba(16, 16, 128, 128, 128);
638 let yuv = ChromaOps::rgba_to_ycbcr(
639 &rgba,
640 16,
641 16,
642 ChromaSubsampling::Yuv422,
643 YcbcrCoefficients::BT601,
644 )
645 .expect("forward");
646 let back = ChromaOps::ycbcr_to_rgba(
647 &yuv,
648 16,
649 16,
650 ChromaSubsampling::Yuv422,
651 YcbcrCoefficients::BT601,
652 )
653 .expect("inverse");
654 for i in 0..(16 * 16) {
655 let base = i * 4;
656 for c in 0..3 {
657 let diff = (rgba[base + c] as i32 - back[base + c] as i32).unsigned_abs();
658 assert!(diff <= 2, "pixel {i} channel {c}: diff={diff}");
659 }
660 }
661 }
662
663 #[test]
664 fn test_422_roundtrip_gradient() {
665 let rgba = gradient_rgba(32, 32);
666 let yuv = ChromaOps::rgba_to_ycbcr(
667 &rgba,
668 32,
669 32,
670 ChromaSubsampling::Yuv422,
671 YcbcrCoefficients::BT709,
672 )
673 .expect("forward");
674 let back = ChromaOps::ycbcr_to_rgba(
675 &yuv,
676 32,
677 32,
678 ChromaSubsampling::Yuv422,
679 YcbcrCoefficients::BT709,
680 )
681 .expect("inverse");
682 let max_diff: u32 = (0..(32 * 32))
683 .map(|i| {
684 let base = i * 4;
685 (0..3)
686 .map(|c| (rgba[base + c] as i32 - back[base + c] as i32).unsigned_abs())
687 .max()
688 .unwrap_or(0)
689 })
690 .max()
691 .unwrap_or(0);
692 assert!(
693 max_diff <= 8,
694 "max roundtrip diff={max_diff}, expected <= 8"
695 );
696 }
697
698 #[test]
699 fn test_422_odd_width() {
700 let rgba = solid_rgba(15, 8, 100, 200, 50);
701 let yuv = ChromaOps::rgba_to_ycbcr(
702 &rgba,
703 15,
704 8,
705 ChromaSubsampling::Yuv422,
706 YcbcrCoefficients::BT601,
707 )
708 .expect("forward");
709 let expected = ChromaSubsampling::Yuv422.output_size(15, 8);
710 assert_eq!(yuv.len(), expected);
711 }
712
713 #[test]
716 fn test_bt2020_roundtrip() {
717 let rgba = gradient_rgba(16, 16);
718 let yuv = ChromaOps::rgba_to_ycbcr(
719 &rgba,
720 16,
721 16,
722 ChromaSubsampling::Yuv420,
723 YcbcrCoefficients::BT2020,
724 )
725 .expect("forward");
726 let back = ChromaOps::ycbcr_to_rgba(
727 &yuv,
728 16,
729 16,
730 ChromaSubsampling::Yuv420,
731 YcbcrCoefficients::BT2020,
732 )
733 .expect("inverse");
734 let max_diff: u32 = (0..(16 * 16))
735 .map(|i| {
736 let base = i * 4;
737 (0..3)
738 .map(|c| (rgba[base + c] as i32 - back[base + c] as i32).unsigned_abs())
739 .max()
740 .unwrap_or(0)
741 })
742 .max()
743 .unwrap_or(0);
744 assert!(max_diff <= 10, "BT.2020 max roundtrip diff={max_diff}");
745 }
746
747 #[test]
750 fn test_zero_dimensions() {
751 let rgba = vec![];
752 assert!(ChromaOps::rgba_to_ycbcr(
753 &rgba,
754 0,
755 0,
756 ChromaSubsampling::Yuv420,
757 YcbcrCoefficients::BT601,
758 )
759 .is_err());
760 }
761
762 #[test]
763 fn test_buffer_too_small() {
764 let rgba = vec![0u8; 10];
765 assert!(ChromaOps::rgba_to_ycbcr(
766 &rgba,
767 16,
768 16,
769 ChromaSubsampling::Yuv420,
770 YcbcrCoefficients::BT601,
771 )
772 .is_err());
773 }
774
775 #[test]
776 fn test_inverse_buffer_too_small() {
777 let yuv = vec![0u8; 10];
778 assert!(ChromaOps::ycbcr_to_rgba(
779 &yuv,
780 16,
781 16,
782 ChromaSubsampling::Yuv420,
783 YcbcrCoefficients::BT601,
784 )
785 .is_err());
786 }
787
788 #[test]
791 fn test_output_size_420() {
792 assert_eq!(ChromaSubsampling::Yuv420.output_size(16, 16), 384);
793 assert_eq!(
794 ChromaSubsampling::Yuv420.output_size(15, 13),
795 15 * 13 + 2 * 8 * 7
796 );
797 }
798
799 #[test]
800 fn test_output_size_422() {
801 assert_eq!(ChromaSubsampling::Yuv422.output_size(16, 16), 512);
802 assert_eq!(
803 ChromaSubsampling::Yuv422.output_size(15, 8),
804 15 * 8 + 2 * 8 * 8
805 );
806 }
807}