styx_codec/decoder/raw/
yuv420p.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::{YuvPlanarImage, YuvRange, YuvStandardMatrix};
9
10#[inline(always)]
11fn map_colorspace(color: ColorSpace) -> (YuvRange, YuvStandardMatrix) {
12    match color {
13        ColorSpace::Bt709 => (YuvRange::Limited, YuvStandardMatrix::Bt709),
14        ColorSpace::Bt2020 => (YuvRange::Limited, YuvStandardMatrix::Bt2020),
15        // See `nv12.rs`: `Srgb` is used as a "full-range output" hint for camera YUV.
16        ColorSpace::Srgb => (YuvRange::Full, YuvStandardMatrix::Bt601),
17        ColorSpace::Unknown => (YuvRange::Limited, YuvStandardMatrix::Bt709),
18    }
19}
20
21#[derive(Clone, Copy)]
22enum UvOrder {
23    Uv,
24    Vu,
25}
26
27/// CPU planar 4:2:0 (Y + U + V planes) → RGB24 decoder for formats such as `YU12`/`YV12`.
28pub struct Yuv420pToRgbDecoder {
29    descriptor: CodecDescriptor,
30    pool: BufferPool,
31    uv_order: UvOrder,
32}
33
34impl Yuv420pToRgbDecoder {
35    pub fn new(
36        input: FourCc,
37        impl_name: &'static str,
38        uv_is_uv: bool,
39        max_width: u32,
40        max_height: u32,
41    ) -> Self {
42        let bytes = max_width as usize * max_height as usize * 3;
43        Self::with_pool(
44            input,
45            impl_name,
46            uv_is_uv,
47            BufferPool::with_limits(2, bytes, 4),
48        )
49    }
50
51    pub fn with_pool(
52        input: FourCc,
53        impl_name: &'static str,
54        uv_is_uv: bool,
55        pool: BufferPool,
56    ) -> Self {
57        Self {
58            descriptor: CodecDescriptor {
59                kind: CodecKind::Decoder,
60                input,
61                output: FourCc::new(*b"RG24"),
62                name: "yuv2rgb",
63                impl_name,
64            },
65            pool,
66            uv_order: if uv_is_uv { UvOrder::Uv } else { UvOrder::Vu },
67        }
68    }
69
70    /// Decode into a caller-provided tightly-packed RGB24 buffer.
71    ///
72    /// `dst` must be at least `width * height * 3` bytes.
73    pub fn decode_into(&self, input: &FrameLease, dst: &mut [u8]) -> Result<FrameMeta, CodecError> {
74        let meta = input.meta();
75        if meta.format.code != self.descriptor.input {
76            return Err(CodecError::FormatMismatch {
77                expected: self.descriptor.input,
78                actual: meta.format.code,
79            });
80        }
81
82        let planes = input.planes();
83        if planes.len() < 3 {
84            return Err(CodecError::Codec("yuv420p frame missing planes".into()));
85        }
86        let y_plane = &planes[0];
87        let (u_plane, v_plane) = match self.uv_order {
88            UvOrder::Uv => (&planes[1], &planes[2]),
89            UvOrder::Vu => (&planes[2], &planes[1]),
90        };
91
92        let width = meta.format.resolution.width.get() as usize;
93        let height = meta.format.resolution.height.get() as usize;
94        let color = meta.format.color;
95        let y_stride = y_plane.stride().max(width);
96        let chroma_width = width.div_ceil(2);
97        let chroma_height = height.div_ceil(2);
98        let u_stride = u_plane.stride().max(chroma_width);
99        let v_stride = v_plane.stride().max(chroma_width);
100
101        let y_required = y_stride
102            .checked_mul(height)
103            .ok_or_else(|| CodecError::Codec("yuv420p y stride overflow".into()))?;
104        let u_required = u_stride
105            .checked_mul(chroma_height)
106            .ok_or_else(|| CodecError::Codec("yuv420p u stride overflow".into()))?;
107        let v_required = v_stride
108            .checked_mul(chroma_height)
109            .ok_or_else(|| CodecError::Codec("yuv420p v stride overflow".into()))?;
110        if y_plane.data().len() < y_required
111            || u_plane.data().len() < u_required
112            || v_plane.data().len() < v_required
113        {
114            return Err(CodecError::Codec("yuv420p plane buffer too short".into()));
115        }
116
117        let row_bytes = width
118            .checked_mul(3)
119            .ok_or_else(|| CodecError::Codec("yuv420p output overflow".into()))?;
120        let out_len = row_bytes
121            .checked_mul(height)
122            .ok_or_else(|| CodecError::Codec("yuv420p output overflow".into()))?;
123        if dst.len() < out_len {
124            return Err(CodecError::Codec("yuv420p dst buffer too short".into()));
125        }
126        let dst = &mut dst[..out_len];
127
128        let y_data = &y_plane.data()[..y_required];
129        let u_data = &u_plane.data()[..u_required];
130        let v_data = &v_plane.data()[..v_required];
131        let planar = YuvPlanarImage {
132            y_plane: y_data,
133            y_stride: y_stride as u32,
134            u_plane: u_data,
135            u_stride: u_stride as u32,
136            v_plane: v_data,
137            v_stride: v_stride as u32,
138            width: width as u32,
139            height: height as u32,
140        };
141        let (range, matrix) = map_colorspace(color);
142        if yuvutils_rs::yuv420_to_rgb(&planar, dst, row_bytes as u32, range, matrix).is_err() {
143            dst.par_chunks_mut(row_bytes)
144                .enumerate()
145                .for_each(|(y, dst_line)| {
146                    let y_line = &y_plane.data()[y * y_stride..][..width];
147                    let u_line = &u_plane.data()[(y / 2) * u_stride..][..chroma_width];
148                    let v_line = &v_plane.data()[(y / 2) * v_stride..][..chroma_width];
149
150                    let pair_count = width / 2;
151                    for pair in 0..pair_count {
152                        let y_base = pair * 2;
153                        let di = pair * 6;
154                        let y0 = unsafe { *y_line.get_unchecked(y_base) as i32 };
155                        let y1 = unsafe { *y_line.get_unchecked(y_base + 1) as i32 };
156                        let u_val = unsafe { *u_line.get_unchecked(pair) as i32 };
157                        let v_val = unsafe { *v_line.get_unchecked(pair) as i32 };
158                        let (r0, g0, b0) = yuv_to_rgb(y0, u_val, v_val, color);
159                        let (r1, g1, b1) = yuv_to_rgb(y1, u_val, v_val, color);
160                        unsafe {
161                            *dst_line.get_unchecked_mut(di) = r0;
162                            *dst_line.get_unchecked_mut(di + 1) = g0;
163                            *dst_line.get_unchecked_mut(di + 2) = b0;
164                            *dst_line.get_unchecked_mut(di + 3) = r1;
165                            *dst_line.get_unchecked_mut(di + 4) = g1;
166                            *dst_line.get_unchecked_mut(di + 5) = b1;
167                        }
168                    }
169                    if width % 2 == 1 && width >= 1 {
170                        let last_idx = width - 1;
171                        let chroma_idx = last_idx / 2;
172                        let di = last_idx * 3;
173                        let y_val = unsafe { *y_line.get_unchecked(last_idx) as i32 };
174                        let u_val = unsafe { *u_line.get_unchecked(chroma_idx) as i32 };
175                        let v_val = unsafe { *v_line.get_unchecked(chroma_idx) as i32 };
176                        let (r, g, b) = yuv_to_rgb(y_val, u_val, v_val, color);
177                        unsafe {
178                            *dst_line.get_unchecked_mut(di) = r;
179                            *dst_line.get_unchecked_mut(di + 1) = g;
180                            *dst_line.get_unchecked_mut(di + 2) = b;
181                        }
182                    }
183                });
184        }
185
186        Ok(FrameMeta::new(
187            MediaFormat::new(
188                self.descriptor.output,
189                meta.format.resolution,
190                meta.format.color,
191            ),
192            meta.timestamp,
193        ))
194    }
195}
196
197impl Codec for Yuv420pToRgbDecoder {
198    fn descriptor(&self) -> &CodecDescriptor {
199        &self.descriptor
200    }
201
202    fn process(&self, input: FrameLease) -> Result<FrameLease, CodecError> {
203        let layout = plane_layout_from_dims(
204            input.meta().format.resolution.width,
205            input.meta().format.resolution.height,
206            3,
207        );
208        let mut buf = self.pool.lease();
209        unsafe { buf.resize_uninit(layout.len) };
210        let meta = self.decode_into(&input, buf.as_mut_slice())?;
211        Ok(unsafe { FrameLease::single_plane_uninit(meta, buf, layout.len, layout.stride) })
212    }
213}
214
215#[cfg(feature = "image")]
216impl ImageDecode for Yuv420pToRgbDecoder {
217    fn decode_image(&self, frame: FrameLease) -> Result<image::DynamicImage, CodecError> {
218        process_to_dynamic(self, frame)
219    }
220}
styx_codec/decoder/raw/yuv420p.rs

styx_codec/decoder/raw/
yuv420p.rs
