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

styx_codec/decoder/raw/
i420.rs
