styx_codec/decoder/raw/
yuv.rs

1use styx_core::prelude::*;
2
3use crate::decoder::raw::yuv_to_rgb;
4#[cfg(feature = "image")]
5use crate::decoder::{ImageDecode, process_to_dynamic};
6use crate::{Codec, CodecDescriptor, CodecError, CodecKind};
7use rayon::prelude::*;
8use yuvutils_rs::{
9    YuvBiPlanarImage, YuvConversionMode, YuvPackedImage, YuvPlanarImage, YuvRange,
10    YuvStandardMatrix,
11};
12
13#[inline(always)]
14fn map_colorspace(color: ColorSpace) -> (YuvRange, YuvStandardMatrix) {
15    match color {
16        ColorSpace::Bt709 => (YuvRange::Limited, YuvStandardMatrix::Bt709),
17        ColorSpace::Bt2020 => (YuvRange::Limited, YuvStandardMatrix::Bt2020),
18        ColorSpace::Srgb => (YuvRange::Full, YuvStandardMatrix::Bt601),
19        ColorSpace::Unknown => (YuvRange::Limited, YuvStandardMatrix::Bt709),
20    }
21}
22
23#[derive(Clone, Copy)]
24enum UvOrder {
25    Uv,
26    Vu,
27}
28
29/// NVXX (Y plane + interleaved chroma) → RGB24 decoder.
30pub struct NvToRgbDecoder {
31    descriptor: CodecDescriptor,
32    pool: BufferPool,
33    uv_order: UvOrder,
34    subsample_h: usize,
35    subsample_v: usize,
36}
37
38impl NvToRgbDecoder {
39    pub fn new(
40        input: FourCc,
41        impl_name: &'static str,
42        subsample_h: usize,
43        subsample_v: usize,
44        uv_is_uv: bool,
45        max_width: u32,
46        max_height: u32,
47    ) -> Self {
48        let bytes = max_width as usize * max_height as usize * 3;
49        Self::with_pool(
50            input,
51            impl_name,
52            subsample_h,
53            subsample_v,
54            uv_is_uv,
55            BufferPool::with_limits(2, bytes, 4),
56        )
57    }
58
59    pub fn with_pool(
60        input: FourCc,
61        impl_name: &'static str,
62        subsample_h: usize,
63        subsample_v: usize,
64        uv_is_uv: bool,
65        pool: BufferPool,
66    ) -> Self {
67        Self {
68            descriptor: CodecDescriptor {
69                kind: CodecKind::Decoder,
70                input,
71                output: FourCc::new(*b"RG24"),
72                name: "yuv2rgb",
73                impl_name,
74            },
75            pool,
76            uv_order: if uv_is_uv { UvOrder::Uv } else { UvOrder::Vu },
77            subsample_h: subsample_h.max(1),
78            subsample_v: subsample_v.max(1),
79        }
80    }
81
82    /// Decode into a caller-provided tightly-packed RGB24 buffer.
83    ///
84    /// `dst` must be at least `width * height * 3` bytes.
85    pub fn decode_into(&self, input: &FrameLease, dst: &mut [u8]) -> Result<FrameMeta, CodecError> {
86        let meta = input.meta();
87        if meta.format.code != self.descriptor.input {
88            return Err(CodecError::FormatMismatch {
89                expected: self.descriptor.input,
90                actual: meta.format.code,
91            });
92        }
93        let planes = input.planes();
94        if planes.len() < 2 {
95            return Err(CodecError::Codec("nv frame missing planes".into()));
96        }
97        let y_plane = &planes[0];
98        let uv_plane = &planes[1];
99
100        let width = meta.format.resolution.width.get() as usize;
101        let height = meta.format.resolution.height.get() as usize;
102        let chroma_width = width.div_ceil(self.subsample_h);
103        let chroma_height = height.div_ceil(self.subsample_v);
104        let y_stride = y_plane.stride().max(width);
105        let uv_stride = uv_plane.stride().max(chroma_width * 2);
106
107        let y_required = y_stride
108            .checked_mul(height)
109            .ok_or_else(|| CodecError::Codec("nv y stride overflow".into()))?;
110        let uv_required = uv_stride
111            .checked_mul(chroma_height)
112            .ok_or_else(|| CodecError::Codec("nv uv stride overflow".into()))?;
113        if y_plane.data().len() < y_required || uv_plane.data().len() < uv_required {
114            return Err(CodecError::Codec("nv plane buffer too short".into()));
115        }
116
117        let row_bytes = width
118            .checked_mul(3)
119            .ok_or_else(|| CodecError::Codec("nv output overflow".into()))?;
120        let out_len = row_bytes
121            .checked_mul(height)
122            .ok_or_else(|| CodecError::Codec("nv output overflow".into()))?;
123        if dst.len() < out_len {
124            return Err(CodecError::Codec("nv dst buffer too short".into()));
125        }
126        let dst = &mut dst[..out_len];
127
128        let (range, matrix) = map_colorspace(meta.format.color);
129        let y_plane_data = &y_plane.data()[..y_required];
130        let uv_plane_data = &uv_plane.data()[..uv_required];
131        let bi = YuvBiPlanarImage {
132            y_plane: y_plane_data,
133            y_stride: y_stride as u32,
134            uv_plane: uv_plane_data,
135            uv_stride: uv_stride as u32,
136            width: width as u32,
137            height: height as u32,
138        };
139        let mode = YuvConversionMode::Balanced;
140        let input_code = self.descriptor.input.to_u32().to_le_bytes();
141        let ok = match &input_code {
142            b"NV21" => Some(yuvutils_rs::yuv_nv21_to_rgb(
143                &bi,
144                dst,
145                row_bytes as u32,
146                range,
147                matrix,
148                mode,
149            )),
150            b"NV16" => Some(yuvutils_rs::yuv_nv16_to_rgb(
151                &bi,
152                dst,
153                row_bytes as u32,
154                range,
155                matrix,
156                mode,
157            )),
158            b"NV61" => Some(yuvutils_rs::yuv_nv61_to_rgb(
159                &bi,
160                dst,
161                row_bytes as u32,
162                range,
163                matrix,
164                mode,
165            )),
166            b"NV24" => Some(yuvutils_rs::yuv_nv24_to_rgb(
167                &bi,
168                dst,
169                row_bytes as u32,
170                range,
171                matrix,
172                mode,
173            )),
174            b"NV42" => Some(yuvutils_rs::yuv_nv42_to_rgb(
175                &bi,
176                dst,
177                row_bytes as u32,
178                range,
179                matrix,
180                mode,
181            )),
182            _ => None,
183        };
184        if ok.is_some_and(|r| r.is_ok()) {
185            return Ok(FrameMeta::new(
186                MediaFormat::new(
187                    self.descriptor.output,
188                    meta.format.resolution,
189                    meta.format.color,
190                ),
191                meta.timestamp,
192            ));
193        }
194
195        let color = meta.format.color;
196        let uv_order = self.uv_order;
197        let subsample_h = self.subsample_h;
198        let subsample_v = self.subsample_v;
199        dst.par_chunks_mut(row_bytes)
200            .enumerate()
201            .for_each(|(y, dst_line)| {
202                let y_line = &y_plane_data[y * y_stride..][..width];
203                let uv_line =
204                    &uv_plane_data[(y / subsample_v) * uv_stride..][..chroma_width * 2];
205                for x in 0..width {
206                    let uv_base = (x / subsample_h) * 2;
207                    let (u, v) = unsafe {
208                        match uv_order {
209                            UvOrder::Uv => (
210                                *uv_line.get_unchecked(uv_base) as i32,
211                                *uv_line.get_unchecked(uv_base + 1) as i32,
212                            ),
213                            UvOrder::Vu => (
214                                *uv_line.get_unchecked(uv_base + 1) as i32,
215                                *uv_line.get_unchecked(uv_base) as i32,
216                            ),
217                        }
218                    };
219                    let y_val = unsafe { *y_line.get_unchecked(x) as i32 };
220                    let (r, g, b) = yuv_to_rgb(y_val, u, v, color);
221                    let di = x * 3;
222                    unsafe {
223                        *dst_line.get_unchecked_mut(di) = r;
224                        *dst_line.get_unchecked_mut(di + 1) = g;
225                        *dst_line.get_unchecked_mut(di + 2) = b;
226                    }
227                }
228            });
229
230        Ok(FrameMeta::new(
231            MediaFormat::new(
232                self.descriptor.output,
233                meta.format.resolution,
234                meta.format.color,
235            ),
236            meta.timestamp,
237        ))
238    }
239}
240
241impl Codec for NvToRgbDecoder {
242    fn descriptor(&self) -> &CodecDescriptor {
243        &self.descriptor
244    }
245
246    fn process(&self, input: FrameLease) -> Result<FrameLease, CodecError> {
247        let layout = plane_layout_from_dims(
248            input.meta().format.resolution.width,
249            input.meta().format.resolution.height,
250            3,
251        );
252        let mut buf = self.pool.lease();
253        unsafe { buf.resize_uninit(layout.len) };
254        let meta = self.decode_into(&input, buf.as_mut_slice())?;
255        Ok(unsafe { FrameLease::single_plane_uninit(meta, buf, layout.len, layout.stride) })
256    }
257}
258
259#[cfg(feature = "image")]
260impl ImageDecode for NvToRgbDecoder {
261    fn decode_image(&self, frame: FrameLease) -> Result<image::DynamicImage, CodecError> {
262        process_to_dynamic(self, frame)
263    }
264}
265
266/// Planar YUV → RGB24 decoder (Y + separate U/V planes).
267pub struct PlanarYuvToRgbDecoder {
268    descriptor: CodecDescriptor,
269    pool: BufferPool,
270    u_first: bool,
271    subsample_h: usize,
272    subsample_v: usize,
273}
274
275impl PlanarYuvToRgbDecoder {
276    pub fn new(
277        input: FourCc,
278        impl_name: &'static str,
279        subsample_h: usize,
280        subsample_v: usize,
281        u_first: bool,
282        max_width: u32,
283        max_height: u32,
284    ) -> Self {
285        let bytes = max_width as usize * max_height as usize * 3;
286        Self::with_pool(
287            input,
288            impl_name,
289            subsample_h,
290            subsample_v,
291            u_first,
292            BufferPool::with_limits(2, bytes, 4),
293        )
294    }
295
296    pub fn with_pool(
297        input: FourCc,
298        impl_name: &'static str,
299        subsample_h: usize,
300        subsample_v: usize,
301        u_first: bool,
302        pool: BufferPool,
303    ) -> Self {
304        Self {
305            descriptor: CodecDescriptor {
306                kind: CodecKind::Decoder,
307                input,
308                output: FourCc::new(*b"RG24"),
309                name: "yuv2rgb",
310                impl_name,
311            },
312            pool,
313            u_first,
314            subsample_h: subsample_h.max(1),
315            subsample_v: subsample_v.max(1),
316        }
317    }
318
319    /// Decode into a caller-provided tightly-packed RGB24 buffer.
320    ///
321    /// `dst` must be at least `width * height * 3` bytes.
322    pub fn decode_into(&self, input: &FrameLease, dst: &mut [u8]) -> Result<FrameMeta, CodecError> {
323        let meta = input.meta();
324        if meta.format.code != self.descriptor.input {
325            return Err(CodecError::FormatMismatch {
326                expected: self.descriptor.input,
327                actual: meta.format.code,
328            });
329        }
330        let planes = input.planes();
331        if planes.len() < 3 {
332            return Err(CodecError::Codec("planar yuv frame missing planes".into()));
333        }
334
335        let y_plane = &planes[0];
336        let (u_plane, v_plane) = if self.u_first {
337            (&planes[1], &planes[2])
338        } else {
339            (&planes[2], &planes[1])
340        };
341
342        let width = meta.format.resolution.width.get() as usize;
343        let height = meta.format.resolution.height.get() as usize;
344        let chroma_width = width.div_ceil(self.subsample_h);
345        let chroma_height = height.div_ceil(self.subsample_v);
346        let y_stride = y_plane.stride().max(width);
347        let u_stride = u_plane.stride().max(chroma_width);
348        let v_stride = v_plane.stride().max(chroma_width);
349
350        let y_required = y_stride
351            .checked_mul(height)
352            .ok_or_else(|| CodecError::Codec("planar yuv y stride overflow".into()))?;
353        let u_required = u_stride
354            .checked_mul(chroma_height)
355            .ok_or_else(|| CodecError::Codec("planar yuv u stride overflow".into()))?;
356        let v_required = v_stride
357            .checked_mul(chroma_height)
358            .ok_or_else(|| CodecError::Codec("planar yuv v stride overflow".into()))?;
359        if y_plane.data().len() < y_required
360            || u_plane.data().len() < u_required
361            || v_plane.data().len() < v_required
362        {
363            return Err(CodecError::Codec(
364                "planar yuv plane buffer too short".into(),
365            ));
366        }
367
368        let row_bytes = width
369            .checked_mul(3)
370            .ok_or_else(|| CodecError::Codec("planar yuv output overflow".into()))?;
371        let out_len = row_bytes
372            .checked_mul(height)
373            .ok_or_else(|| CodecError::Codec("planar yuv output overflow".into()))?;
374        if dst.len() < out_len {
375            return Err(CodecError::Codec("planar yuv dst buffer too short".into()));
376        }
377        let dst = &mut dst[..out_len];
378
379        let (range, matrix) = map_colorspace(meta.format.color);
380        let y_plane_data = &y_plane.data()[..y_required];
381        let u_plane_data = &u_plane.data()[..u_required];
382        let v_plane_data = &v_plane.data()[..v_required];
383        let planar = YuvPlanarImage {
384            y_plane: y_plane_data,
385            y_stride: y_stride as u32,
386            u_plane: u_plane_data,
387            u_stride: u_stride as u32,
388            v_plane: v_plane_data,
389            v_stride: v_stride as u32,
390            width: width as u32,
391            height: height as u32,
392        };
393        let ok = match (self.subsample_h, self.subsample_v) {
394            (2, 2) => Some(yuvutils_rs::yuv420_to_rgb(
395                &planar,
396                dst,
397                row_bytes as u32,
398                range,
399                matrix,
400            )),
401            (2, 1) => Some(yuvutils_rs::yuv422_to_rgb(
402                &planar,
403                dst,
404                row_bytes as u32,
405                range,
406                matrix,
407            )),
408            (1, 1) => Some(yuvutils_rs::yuv444_to_rgb(
409                &planar,
410                dst,
411                row_bytes as u32,
412                range,
413                matrix,
414            )),
415            _ => None,
416        };
417        if ok.is_some_and(|r| r.is_ok()) {
418            return Ok(FrameMeta::new(
419                MediaFormat::new(
420                    self.descriptor.output,
421                    meta.format.resolution,
422                    meta.format.color,
423                ),
424                meta.timestamp,
425            ));
426        }
427
428        let color = meta.format.color;
429        let subsample_h = self.subsample_h;
430        let subsample_v = self.subsample_v;
431        dst.par_chunks_mut(row_bytes)
432            .enumerate()
433            .for_each(|(y, dst_line)| {
434                let y_line = &y_plane_data[y * y_stride..][..width];
435                let u_line =
436                    &u_plane_data[(y / subsample_v) * u_stride..][..chroma_width];
437                let v_line =
438                    &v_plane_data[(y / subsample_v) * v_stride..][..chroma_width];
439                for x in 0..width {
440                    let chroma_idx = x / subsample_h;
441                    let u = unsafe { *u_line.get_unchecked(chroma_idx) as i32 };
442                    let v = unsafe { *v_line.get_unchecked(chroma_idx) as i32 };
443                    let y_val = unsafe { *y_line.get_unchecked(x) as i32 };
444                    let (r, g, b) = yuv_to_rgb(y_val, u, v, color);
445                    let di = x * 3;
446                    unsafe {
447                        *dst_line.get_unchecked_mut(di) = r;
448                        *dst_line.get_unchecked_mut(di + 1) = g;
449                        *dst_line.get_unchecked_mut(di + 2) = b;
450                    }
451                }
452            });
453
454        Ok(FrameMeta::new(
455            MediaFormat::new(
456                self.descriptor.output,
457                meta.format.resolution,
458                meta.format.color,
459            ),
460            meta.timestamp,
461        ))
462    }
463}
464
465impl Codec for PlanarYuvToRgbDecoder {
466    fn descriptor(&self) -> &CodecDescriptor {
467        &self.descriptor
468    }
469
470    fn process(&self, input: FrameLease) -> Result<FrameLease, CodecError> {
471        let layout = plane_layout_from_dims(
472            input.meta().format.resolution.width,
473            input.meta().format.resolution.height,
474            3,
475        );
476        let mut buf = self.pool.lease();
477        unsafe { buf.resize_uninit(layout.len) };
478        let meta = self.decode_into(&input, buf.as_mut_slice())?;
479        Ok(unsafe { FrameLease::single_plane_uninit(meta, buf, layout.len, layout.stride) })
480    }
481}
482
483#[cfg(feature = "image")]
484impl ImageDecode for PlanarYuvToRgbDecoder {
485    fn decode_image(&self, frame: FrameLease) -> Result<image::DynamicImage, CodecError> {
486        process_to_dynamic(self, frame)
487    }
488}
489
490/// Packed YUV422 → RGB24 decoder with configurable byte order.
491pub struct Packed422ToRgbDecoder {
492    descriptor: CodecDescriptor,
493    pool: BufferPool,
494    byte_order: [usize; 4],
495}
496
497impl Packed422ToRgbDecoder {
498    pub fn new(
499        input: FourCc,
500        impl_name: &'static str,
501        byte_order: [usize; 4],
502        max_width: u32,
503        max_height: u32,
504    ) -> Self {
505        let bytes = max_width as usize * max_height as usize * 3;
506        Self::with_pool(
507            input,
508            impl_name,
509            byte_order,
510            BufferPool::with_limits(2, bytes, 4),
511        )
512    }
513
514    pub fn with_pool(
515        input: FourCc,
516        impl_name: &'static str,
517        byte_order: [usize; 4],
518        pool: BufferPool,
519    ) -> Self {
520        Self {
521            descriptor: CodecDescriptor {
522                kind: CodecKind::Decoder,
523                input,
524                output: FourCc::new(*b"RG24"),
525                name: "yuv2rgb",
526                impl_name,
527            },
528            pool,
529            byte_order,
530        }
531    }
532
533    /// Decode into a caller-provided tightly-packed RGB24 buffer.
534    ///
535    /// `dst` must be at least `width * height * 3` bytes.
536    pub fn decode_into(&self, input: &FrameLease, dst: &mut [u8]) -> Result<FrameMeta, CodecError> {
537        let meta = input.meta();
538        if meta.format.code != self.descriptor.input {
539            return Err(CodecError::FormatMismatch {
540                expected: self.descriptor.input,
541                actual: meta.format.code,
542            });
543        }
544        let plane = input
545            .planes()
546            .into_iter()
547            .next()
548            .ok_or_else(|| CodecError::Codec("yuv422 frame missing plane".into()))?;
549
550        let width = meta.format.resolution.width.get() as usize;
551        let height = meta.format.resolution.height.get() as usize;
552        let stride = plane.stride().max(width * 2);
553        let required = stride
554            .checked_mul(height)
555            .ok_or_else(|| CodecError::Codec("yuv422 stride overflow".into()))?;
556        if plane.data().len() < required {
557            return Err(CodecError::Codec("yuv422 plane buffer too short".into()));
558        }
559
560        let row_bytes = width
561            .checked_mul(3)
562            .ok_or_else(|| CodecError::Codec("yuv422 output overflow".into()))?;
563        let out_len = row_bytes
564            .checked_mul(height)
565            .ok_or_else(|| CodecError::Codec("yuv422 output overflow".into()))?;
566        if dst.len() < out_len {
567            return Err(CodecError::Codec("yuv422 dst buffer too short".into()));
568        }
569        let dst = &mut dst[..out_len];
570
571        let src = plane.data();
572        let src_required = &src[..required];
573        let (range, matrix) = map_colorspace(meta.format.color);
574        let packed = YuvPackedImage {
575            yuy: src_required,
576            yuy_stride: stride as u32,
577            width: width as u32,
578            height: height as u32,
579        };
580        let input_code = self.descriptor.input.to_u32().to_le_bytes();
581        let ok = match &input_code {
582            b"UYVY" => Some(yuvutils_rs::uyvy422_to_rgb(
583                &packed,
584                dst,
585                row_bytes as u32,
586                range,
587                matrix,
588            )),
589            b"YUYV" => Some(yuvutils_rs::yuyv422_to_rgb(
590                &packed,
591                dst,
592                row_bytes as u32,
593                range,
594                matrix,
595            )),
596            _ => None,
597        };
598        if ok.is_some_and(|r| r.is_ok()) {
599            return Ok(FrameMeta::new(
600                MediaFormat::new(
601                    self.descriptor.output,
602                    meta.format.resolution,
603                    meta.format.color,
604                ),
605                meta.timestamp,
606            ));
607        }
608
609        let color = meta.format.color;
610        let byte_order = self.byte_order;
611        dst.par_chunks_mut(row_bytes)
612            .enumerate()
613            .for_each(|(y, dst_line)| {
614                let src_line = &src_required[y * stride..][..width * 2];
615                let pair_count = width / 2;
616                for pair in 0..pair_count {
617                    let si = pair * 4;
618                    let di = pair * 6;
619                    let y0 = unsafe { *src_line.get_unchecked(si + byte_order[0]) as i32 };
620                    let u = unsafe { *src_line.get_unchecked(si + byte_order[1]) as i32 };
621                    let y1 = unsafe { *src_line.get_unchecked(si + byte_order[2]) as i32 };
622                    let v = unsafe { *src_line.get_unchecked(si + byte_order[3]) as i32 };
623                    let (r0, g0, b0) = yuv_to_rgb(y0, u, v, color);
624                    let (r1, g1, b1) = yuv_to_rgb(y1, u, v, color);
625                    unsafe {
626                        *dst_line.get_unchecked_mut(di) = r0;
627                        *dst_line.get_unchecked_mut(di + 1) = g0;
628                        *dst_line.get_unchecked_mut(di + 2) = b0;
629                        *dst_line.get_unchecked_mut(di + 3) = r1;
630                        *dst_line.get_unchecked_mut(di + 4) = g1;
631                        *dst_line.get_unchecked_mut(di + 5) = b1;
632                    }
633                }
634                if width % 2 == 1 && width >= 1 {
635                    let last_x = width - 1;
636                    let si = (last_x / 2) * 4;
637                    let di = last_x * 3;
638                    let y_val = unsafe { *src_line.get_unchecked(si + byte_order[0]) as i32 };
639                    let u = unsafe { *src_line.get_unchecked(si + byte_order[1]) as i32 };
640                    let v = unsafe { *src_line.get_unchecked(si + byte_order[3]) as i32 };
641                    let (r, g, b) = yuv_to_rgb(y_val, u, v, color);
642                    unsafe {
643                        *dst_line.get_unchecked_mut(di) = r;
644                        *dst_line.get_unchecked_mut(di + 1) = g;
645                        *dst_line.get_unchecked_mut(di + 2) = b;
646                    }
647                }
648            });
649
650        Ok(FrameMeta::new(
651            MediaFormat::new(
652                self.descriptor.output,
653                meta.format.resolution,
654                meta.format.color,
655            ),
656            meta.timestamp,
657        ))
658    }
659}
660
661impl Codec for Packed422ToRgbDecoder {
662    fn descriptor(&self) -> &CodecDescriptor {
663        &self.descriptor
664    }
665
666    fn process(&self, input: FrameLease) -> Result<FrameLease, CodecError> {
667        let layout = plane_layout_from_dims(
668            input.meta().format.resolution.width,
669            input.meta().format.resolution.height,
670            3,
671        );
672        let mut buf = self.pool.lease();
673        unsafe { buf.resize_uninit(layout.len) };
674        let meta = self.decode_into(&input, buf.as_mut_slice())?;
675        Ok(unsafe { FrameLease::single_plane_uninit(meta, buf, layout.len, layout.stride) })
676    }
677}
678
679#[cfg(feature = "image")]
680impl ImageDecode for Packed422ToRgbDecoder {
681    fn decode_image(&self, frame: FrameLease) -> Result<image::DynamicImage, CodecError> {
682        process_to_dynamic(self, frame)
683    }
684}
styx_codec/decoder/raw/yuv.rs

styx_codec/decoder/raw/
yuv.rs
