styx_codec/decoder/raw/
nv12.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::{YuvBiPlanarImage, YuvConversionMode, 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        // When color is marked as `Srgb` in our metadata it usually means "full-range output".
16        // libcamera often reports `sYCC` for PiSP outputs; `sYCC` uses a Rec.601 matrix.
17        ColorSpace::Srgb => (YuvRange::Full, YuvStandardMatrix::Bt601),
18        // Unknown is far more likely to be limited-range BT.709 than full-range BT.601.
19        ColorSpace::Unknown => (YuvRange::Limited, YuvStandardMatrix::Bt709),
20    }
21}
22
23/// CPU NV12 (Y plane + interleaved UV) → RGB24 decoder.
24pub struct Nv12ToRgbDecoder {
25    descriptor: CodecDescriptor,
26    pool: BufferPool,
27}
28
29impl Nv12ToRgbDecoder {
30    pub fn new(max_width: u32, max_height: u32) -> Self {
31        let bytes = max_width as usize * max_height as usize * 3;
32        Self::with_pool(BufferPool::with_limits(2, bytes, 4))
33    }
34
35    pub fn with_pool(pool: BufferPool) -> Self {
36        Self {
37            descriptor: CodecDescriptor {
38                kind: CodecKind::Decoder,
39                input: FourCc::new(*b"NV12"),
40                output: FourCc::new(*b"RG24"),
41                name: "yuv2rgb",
42                impl_name: "nv12-cpu",
43            },
44            pool,
45        }
46    }
47
48    /// Decode into a caller-provided tightly-packed RGB24 buffer.
49    ///
50    /// `dst` must be at least `width * height * 3` bytes.
51    pub fn decode_into(&self, input: &FrameLease, dst: &mut [u8]) -> Result<FrameMeta, CodecError> {
52        let meta = input.meta();
53        if meta.format.code != self.descriptor.input {
54            return Err(CodecError::FormatMismatch {
55                expected: self.descriptor.input,
56                actual: meta.format.code,
57            });
58        }
59        let planes = input.planes();
60        let (y_plane_data, uv_plane_data, y_stride, uv_stride) = if planes.len() >= 2 {
61            let y_plane = &planes[0];
62            let uv_plane = &planes[1];
63            let width = meta.format.resolution.width.get() as usize;
64            let height = meta.format.resolution.height.get() as usize;
65            let chroma_width = width.div_ceil(2);
66            let y_stride = y_plane.stride().max(width);
67            let uv_stride = uv_plane.stride().max(chroma_width * 2);
68            let chroma_height = height.div_ceil(2);
69
70            let y_required = y_stride
71                .checked_mul(height)
72                .ok_or_else(|| CodecError::Codec("nv12 y stride overflow".into()))?;
73            let uv_required = uv_stride
74                .checked_mul(chroma_height)
75                .ok_or_else(|| CodecError::Codec("nv12 uv stride overflow".into()))?;
76            if y_plane.data().len() < y_required || uv_plane.data().len() < uv_required {
77                return Err(CodecError::Codec("nv12 plane buffer too short".into()));
78            }
79
80            (
81                &y_plane.data()[..y_required],
82                &uv_plane.data()[..uv_required],
83                y_stride,
84                uv_stride,
85            )
86        } else if planes.len() == 1 {
87            let plane = &planes[0];
88            let width = meta.format.resolution.width.get() as usize;
89            let height = meta.format.resolution.height.get() as usize;
90            let chroma_width = width.div_ceil(2);
91            let stride = plane.stride().max(width);
92            let chroma_height = height.div_ceil(2);
93            let y_required = stride
94                .checked_mul(height)
95                .ok_or_else(|| CodecError::Codec("nv12 y stride overflow".into()))?;
96            let uv_required = stride
97                .checked_mul(chroma_height)
98                .ok_or_else(|| CodecError::Codec("nv12 uv stride overflow".into()))?;
99            let total_required = y_required
100                .checked_add(uv_required)
101                .ok_or_else(|| CodecError::Codec("nv12 plane length overflow".into()))?;
102            if plane.data().len() < total_required {
103                return Err(CodecError::Codec(
104                    "nv12 packed plane buffer too short".into(),
105                ));
106            }
107            let uv_min_stride = (chroma_width * 2).max(1);
108            let uv_stride = stride.max(uv_min_stride);
109            (
110                &plane.data()[..y_required],
111                &plane.data()[y_required..y_required + uv_required],
112                stride,
113                uv_stride,
114            )
115        } else {
116            return Err(CodecError::Codec("nv12 frame missing planes".into()));
117        };
118
119        let width = meta.format.resolution.width.get() as usize;
120        let height = meta.format.resolution.height.get() as usize;
121        let chroma_width = width.div_ceil(2);
122        let row_bytes = width
123            .checked_mul(3)
124            .ok_or_else(|| CodecError::Codec("nv12 output overflow".into()))?;
125        let out_len = row_bytes
126            .checked_mul(height)
127            .ok_or_else(|| CodecError::Codec("nv12 output overflow".into()))?;
128        if dst.len() < out_len {
129            return Err(CodecError::Codec("nv12 dst buffer too short".into()));
130        }
131        let dst = &mut dst[..out_len];
132
133        let color = meta.format.color;
134        let (range, matrix) = map_colorspace(color);
135        let bi = YuvBiPlanarImage {
136            y_plane: y_plane_data,
137            y_stride: y_stride as u32,
138            uv_plane: uv_plane_data,
139            uv_stride: uv_stride as u32,
140            width: width as u32,
141            height: height as u32,
142        };
143        if yuvutils_rs::yuv_nv12_to_rgb(
144            &bi,
145            dst,
146            row_bytes as u32,
147            range,
148            matrix,
149            YuvConversionMode::Balanced,
150        )
151        .is_err()
152        {
153            dst.par_chunks_mut(row_bytes)
154                .enumerate()
155                .for_each(|(y, dst_line)| {
156                    let y_line = &y_plane_data[y * y_stride..][..width];
157                    let uv_line = &uv_plane_data[(y / 2) * uv_stride..][..chroma_width * 2];
158                    let pair_count = width / 2;
159                    for pair in 0..pair_count {
160                        let y_base = pair * 2;
161                        let uv_idx = pair * 2;
162                        let di = pair * 6;
163                        let y0 = unsafe { *y_line.get_unchecked(y_base) as i32 };
164                        let y1 = unsafe { *y_line.get_unchecked(y_base + 1) as i32 };
165                        let u = unsafe { *uv_line.get_unchecked(uv_idx) as i32 };
166                        let v = unsafe { *uv_line.get_unchecked(uv_idx + 1) as i32 };
167                        let (r0, g0, b0) = yuv_to_rgb(y0, u, v, color);
168                        let (r1, g1, b1) = yuv_to_rgb(y1, u, v, color);
169                        unsafe {
170                            *dst_line.get_unchecked_mut(di) = r0;
171                            *dst_line.get_unchecked_mut(di + 1) = g0;
172                            *dst_line.get_unchecked_mut(di + 2) = b0;
173                            *dst_line.get_unchecked_mut(di + 3) = r1;
174                            *dst_line.get_unchecked_mut(di + 4) = g1;
175                            *dst_line.get_unchecked_mut(di + 5) = b1;
176                        }
177                    }
178                    if width % 2 == 1 && width >= 1 {
179                        let last_idx = width - 1;
180                        let uv_idx = (last_idx / 2) * 2;
181                        let di = last_idx * 3;
182                        let y_val = unsafe { *y_line.get_unchecked(last_idx) as i32 };
183                        let u = unsafe { *uv_line.get_unchecked(uv_idx) as i32 };
184                        let v = unsafe { *uv_line.get_unchecked(uv_idx + 1) as i32 };
185                        let (r, g, b) = yuv_to_rgb(y_val, u, v, color);
186                        unsafe {
187                            *dst_line.get_unchecked_mut(di) = r;
188                            *dst_line.get_unchecked_mut(di + 1) = g;
189                            *dst_line.get_unchecked_mut(di + 2) = b;
190                        }
191                    }
192                });
193        }
194
195        Ok(FrameMeta::new(
196            MediaFormat::new(
197                self.descriptor.output,
198                meta.format.resolution,
199                meta.format.color,
200            ),
201            meta.timestamp,
202        ))
203    }
204}
205
206impl Codec for Nv12ToRgbDecoder {
207    fn descriptor(&self) -> &CodecDescriptor {
208        &self.descriptor
209    }
210
211    fn process(&self, input: FrameLease) -> Result<FrameLease, CodecError> {
212        let layout = plane_layout_from_dims(
213            input.meta().format.resolution.width,
214            input.meta().format.resolution.height,
215            3,
216        );
217        let mut buf = self.pool.lease();
218        unsafe { buf.resize_uninit(layout.len) };
219        let meta = self.decode_into(&input, buf.as_mut_slice())?;
220        Ok(unsafe { FrameLease::single_plane_uninit(meta, buf, layout.len, layout.stride) })
221    }
222}
223
224#[cfg(feature = "image")]
225impl ImageDecode for Nv12ToRgbDecoder {
226    fn decode_image(&self, frame: FrameLease) -> Result<image::DynamicImage, CodecError> {
227        process_to_dynamic(self, frame)
228    }
229}
230
231/// CPU NV12 (Y plane + interleaved UV) → BGR24 decoder.
232pub struct Nv12ToBgrDecoder {
233    descriptor: CodecDescriptor,
234    pool: BufferPool,
235}
236
237impl Nv12ToBgrDecoder {
238    pub fn new(max_width: u32, max_height: u32) -> Self {
239        let bytes = max_width as usize * max_height as usize * 3;
240        Self::with_pool(BufferPool::with_limits(2, bytes, 4))
241    }
242
243    pub fn with_pool(pool: BufferPool) -> Self {
244        Self {
245            descriptor: CodecDescriptor {
246                kind: CodecKind::Decoder,
247                input: FourCc::new(*b"NV12"),
248                output: FourCc::new(*b"BG24"),
249                name: "yuv2bgr",
250                impl_name: "nv12-cpu",
251            },
252            pool,
253        }
254    }
255
256    /// Decode into a caller-provided tightly-packed BGR24 buffer.
257    ///
258    /// `dst` must be at least `width * height * 3` bytes.
259    pub fn decode_into(&self, input: &FrameLease, dst: &mut [u8]) -> Result<FrameMeta, CodecError> {
260        let meta = input.meta();
261        if meta.format.code != self.descriptor.input {
262            return Err(CodecError::FormatMismatch {
263                expected: self.descriptor.input,
264                actual: meta.format.code,
265            });
266        }
267        let planes = input.planes();
268        let (y_plane_data, uv_plane_data, y_stride, uv_stride) = if planes.len() >= 2 {
269            let y_plane = &planes[0];
270            let uv_plane = &planes[1];
271            let width = meta.format.resolution.width.get() as usize;
272            let height = meta.format.resolution.height.get() as usize;
273            let chroma_width = width.div_ceil(2);
274            let y_stride = y_plane.stride().max(width);
275            let uv_stride = uv_plane.stride().max(chroma_width * 2);
276            let chroma_height = height.div_ceil(2);
277
278            let y_required = y_stride
279                .checked_mul(height)
280                .ok_or_else(|| CodecError::Codec("nv12 y stride overflow".into()))?;
281            let uv_required = uv_stride
282                .checked_mul(chroma_height)
283                .ok_or_else(|| CodecError::Codec("nv12 uv stride overflow".into()))?;
284            if y_plane.data().len() < y_required || uv_plane.data().len() < uv_required {
285                return Err(CodecError::Codec("nv12 plane buffer too short".into()));
286            }
287
288            (
289                &y_plane.data()[..y_required],
290                &uv_plane.data()[..uv_required],
291                y_stride,
292                uv_stride,
293            )
294        } else if planes.len() == 1 {
295            let plane = &planes[0];
296            let width = meta.format.resolution.width.get() as usize;
297            let height = meta.format.resolution.height.get() as usize;
298            let chroma_width = width.div_ceil(2);
299            let stride = plane.stride().max(width);
300            let chroma_height = height.div_ceil(2);
301            let y_required = stride
302                .checked_mul(height)
303                .ok_or_else(|| CodecError::Codec("nv12 y stride overflow".into()))?;
304            let uv_required = stride
305                .checked_mul(chroma_height)
306                .ok_or_else(|| CodecError::Codec("nv12 uv stride overflow".into()))?;
307            let total_required = y_required
308                .checked_add(uv_required)
309                .ok_or_else(|| CodecError::Codec("nv12 plane length overflow".into()))?;
310            if plane.data().len() < total_required {
311                return Err(CodecError::Codec(
312                    "nv12 packed plane buffer too short".into(),
313                ));
314            }
315            let uv_min_stride = (chroma_width * 2).max(1);
316            let uv_stride = stride.max(uv_min_stride);
317            (
318                &plane.data()[..y_required],
319                &plane.data()[y_required..y_required + uv_required],
320                stride,
321                uv_stride,
322            )
323        } else {
324            return Err(CodecError::Codec("nv12 frame missing planes".into()));
325        };
326
327        let width = meta.format.resolution.width.get() as usize;
328        let height = meta.format.resolution.height.get() as usize;
329        let chroma_width = width.div_ceil(2);
330        let row_bytes = width
331            .checked_mul(3)
332            .ok_or_else(|| CodecError::Codec("nv12 output overflow".into()))?;
333        let out_len = row_bytes
334            .checked_mul(height)
335            .ok_or_else(|| CodecError::Codec("nv12 output overflow".into()))?;
336        if dst.len() < out_len {
337            return Err(CodecError::Codec("nv12 dst buffer too short".into()));
338        }
339        let dst = &mut dst[..out_len];
340
341        let color = meta.format.color;
342        let (range, matrix) = map_colorspace(color);
343        let bi = YuvBiPlanarImage {
344            y_plane: y_plane_data,
345            y_stride: y_stride as u32,
346            uv_plane: uv_plane_data,
347            uv_stride: uv_stride as u32,
348            width: width as u32,
349            height: height as u32,
350        };
351        if yuvutils_rs::yuv_nv12_to_bgr(
352            &bi,
353            dst,
354            row_bytes as u32,
355            range,
356            matrix,
357            YuvConversionMode::Balanced,
358        )
359        .is_err()
360        {
361            dst.par_chunks_mut(row_bytes)
362                .enumerate()
363                .for_each(|(y, dst_line)| {
364                    let y_line = &y_plane_data[y * y_stride..][..width];
365                    let uv_line = &uv_plane_data[(y / 2) * uv_stride..][..chroma_width * 2];
366                    let pair_count = width / 2;
367                    for pair in 0..pair_count {
368                        let y_base = pair * 2;
369                        let uv_idx = pair * 2;
370                        let di = pair * 6;
371                        let y0 = unsafe { *y_line.get_unchecked(y_base) as i32 };
372                        let y1 = unsafe { *y_line.get_unchecked(y_base + 1) as i32 };
373                        let u = unsafe { *uv_line.get_unchecked(uv_idx) as i32 };
374                        let v = unsafe { *uv_line.get_unchecked(uv_idx + 1) as i32 };
375                        let (r0, g0, b0) = yuv_to_rgb(y0, u, v, color);
376                        let (r1, g1, b1) = yuv_to_rgb(y1, u, v, color);
377                        unsafe {
378                            *dst_line.get_unchecked_mut(di) = b0;
379                            *dst_line.get_unchecked_mut(di + 1) = g0;
380                            *dst_line.get_unchecked_mut(di + 2) = r0;
381                            *dst_line.get_unchecked_mut(di + 3) = b1;
382                            *dst_line.get_unchecked_mut(di + 4) = g1;
383                            *dst_line.get_unchecked_mut(di + 5) = r1;
384                        }
385                    }
386                    if width % 2 == 1 && width >= 1 {
387                        let last_idx = width - 1;
388                        let uv_idx = (last_idx / 2) * 2;
389                        let di = last_idx * 3;
390                        let y_val = unsafe { *y_line.get_unchecked(last_idx) as i32 };
391                        let u = unsafe { *uv_line.get_unchecked(uv_idx) as i32 };
392                        let v = unsafe { *uv_line.get_unchecked(uv_idx + 1) as i32 };
393                        let (r, g, b) = yuv_to_rgb(y_val, u, v, color);
394                        unsafe {
395                            *dst_line.get_unchecked_mut(di) = b;
396                            *dst_line.get_unchecked_mut(di + 1) = g;
397                            *dst_line.get_unchecked_mut(di + 2) = r;
398                        }
399                    }
400                });
401        }
402
403        Ok(FrameMeta::new(
404            MediaFormat::new(
405                self.descriptor.output,
406                meta.format.resolution,
407                meta.format.color,
408            ),
409            meta.timestamp,
410        ))
411    }
412}
413
414impl Codec for Nv12ToBgrDecoder {
415    fn descriptor(&self) -> &CodecDescriptor {
416        &self.descriptor
417    }
418
419    fn process(&self, input: FrameLease) -> Result<FrameLease, CodecError> {
420        let layout = plane_layout_from_dims(
421            input.meta().format.resolution.width,
422            input.meta().format.resolution.height,
423            3,
424        );
425        let mut buf = self.pool.lease();
426        unsafe { buf.resize_uninit(layout.len) };
427        let meta = self.decode_into(&input, buf.as_mut_slice())?;
428        Ok(unsafe { FrameLease::single_plane_uninit(meta, buf, layout.len, layout.stride) })
429    }
430}
styx_codec/decoder/raw/nv12.rs

styx_codec/decoder/raw/
nv12.rs
