styx_codec/decoder/
mod.rs

1//! Decoder namespace with per-format modules.
2
3#[cfg(feature = "image")]
4use crate::decoder::raw::yuv_to_rgb;
5#[cfg(feature = "image")]
6use crate::{Codec, CodecError};
7#[cfg(feature = "image")]
8use image::{DynamicImage, GenericImageView};
9#[cfg(feature = "image")]
10use rayon::prelude::*;
11#[cfg(feature = "image")]
12use std::cell::RefCell;
13#[cfg(feature = "image")]
14use styx_core::prelude::{
15    BufferPool, ColorSpace, FourCc, FrameLease, FrameMeta, MediaFormat, Resolution,
16};
17#[cfg(feature = "image")]
18use yuvutils_rs::{
19    YuvBiPlanarImage, YuvConversionMode, YuvPackedImage, YuvPlanarImage, YuvRange,
20    YuvStandardMatrix,
21};
22
23#[cfg(feature = "codec-ffmpeg")]
24pub mod ffmpeg;
25pub mod mjpeg;
26pub mod raw;
27
28/// Trait to retrieve a `DynamicImage` from any decoder.
29#[cfg(feature = "image")]
30pub trait ImageDecode {
31    fn decode_image(&self, frame: FrameLease) -> Result<DynamicImage, CodecError>;
32}
33
34#[cfg(feature = "image")]
35pub(crate) fn process_to_dynamic<D: Codec>(
36    decoder: &D,
37    frame: FrameLease,
38) -> Result<DynamicImage, CodecError> {
39    // Prefer converting the original input frame directly when possible to avoid unnecessary
40    // format conversions (e.g. routing BGRA through an RGB24 decoder) when the goal is a
41    // `DynamicImage`.
42    match frame_lease_to_dynamic_image(frame) {
43        Ok(img) => Ok(img),
44        Err(frame) => {
45            if let Some(img) = frame_to_dynamic_image(&frame) {
46                return Ok(img);
47            }
48            let decoded = decoder.process(frame)?;
49            match frame_lease_to_dynamic_image(decoded) {
50                Ok(img) => Ok(img),
51                Err(decoded) => frame_to_dynamic_image(&decoded)
52                    .ok_or_else(|| CodecError::Codec("unable to convert to DynamicImage".into())),
53            }
54        }
55    }
56}
57
58#[cfg(feature = "image")]
59thread_local! {
60    static PACKED_FRAME_POOLS: RefCell<Vec<(usize, BufferPool)>> = const { RefCell::new(Vec::new()) };
61}
62
63#[cfg(feature = "image")]
64const PACKED_FRAME_POOL_SLOTS: usize = 4;
65
66#[cfg(feature = "image")]
67fn packed_frame_pool(len: usize) -> BufferPool {
68    PACKED_FRAME_POOLS.with(|pools| {
69        let mut pools = pools.borrow_mut();
70        if let Some(pos) = pools.iter().position(|(k, _)| *k == len) {
71            let (k, pool) = pools.remove(pos);
72            pools.insert(0, (k, pool.clone()));
73            return pool;
74        }
75
76        let pool = BufferPool::with_limits(2, len, 2);
77        pools.insert(0, (len, pool.clone()));
78        if pools.len() > PACKED_FRAME_POOL_SLOTS {
79            pools.truncate(PACKED_FRAME_POOL_SLOTS);
80        }
81        pool
82    })
83}
84
85/// Drop per-thread packed frame pools (helps release peak allocations between streams).
86#[cfg(feature = "image")]
87pub fn clear_packed_frame_pools() {
88    PACKED_FRAME_POOLS.with(|pools| pools.borrow_mut().clear());
89}
90
91/// Clear packed frame pools for the current thread and all Rayon worker threads.
92///
93/// The decoder fast-path uses `rayon` for some pixel conversions; since Rayon keeps a global
94/// thread pool alive for the lifetime of the process, per-thread buffers can otherwise retain
95/// peak allocations even after streams are stopped or codecs are switched.
96#[cfg(feature = "image")]
97pub fn clear_packed_frame_pools_all_threads() {
98    clear_packed_frame_pools();
99    rayon::broadcast(|_| clear_packed_frame_pools());
100}
101
102/// Convert an owned `FrameLease` into a `DynamicImage`.
103///
104/// This is preferred over [`frame_to_dynamic_image`] when the caller can give up the frame, as it
105/// can avoid copies for tightly-packed CPU frames (e.g. `RG24`, `RGBA`, `R8  `).
106///
107/// This function intentionally only handles packed CPU formats that can be wrapped cheaply.
108/// For planar formats (e.g. `NV12`, `YUYV`) and raw sensor formats (e.g. Bayer), prefer routing
109/// through the codec registry so the fastest available decoder can be selected.
110///
111/// If the format is not supported by this fast-path, the original frame is returned as
112/// `Err(frame)` so the caller can route it through a decoder.
113#[cfg(feature = "image")]
114#[allow(clippy::result_large_err)]
115pub fn frame_lease_to_dynamic_image(frame: FrameLease) -> Result<DynamicImage, FrameLease> {
116    let code = frame.meta().format.code;
117
118    // Only return Ok() for formats we can represent correctly as a `DynamicImage` without
119    // performing pixel-format conversions. Everything else should return Err(frame) so callers
120    // can route through `frame_to_dynamic_image` (conversion/copy) or a codec.
121    let is_packed = code == FourCc::new(*b"R8  ")
122        || code == FourCc::new(*b"GREY")
123        || code == FourCc::new(*b"NV12")
124        || code == FourCc::new(*b"NV21")
125        || code == FourCc::new(*b"RG24")
126        || code == FourCc::new(*b"RGBA");
127
128    if !is_packed {
129        return Err(frame);
130    }
131
132    // External-backed frames do not own their underlying CPU buffer; we must copy.
133    if frame.is_external() {
134        return frame_to_dynamic_image(&frame).ok_or(frame);
135    }
136
137    let meta = frame.meta();
138    let width = meta.format.resolution.width.get();
139    let height = meta.format.resolution.height.get();
140
141    // For NV12/NV21, treat plane0 as luma and ignore UV. We intentionally avoid trying to "steal"
142    // the underlying buffer because these formats are typically multi-plane.
143    if code == FourCc::new(*b"NV12") || code == FourCc::new(*b"NV21") {
144        let planes = frame.planes();
145        if planes.is_empty() {
146            drop(planes);
147            return Err(frame);
148        }
149        let plane = &planes[0];
150        let stride = plane.stride().max(width as usize);
151        let expected = width as usize;
152        let required = stride.saturating_mul(height as usize);
153        if plane.data().len() < required {
154            drop(planes);
155            return Err(frame);
156        }
157        let out = if stride == expected {
158            let required = expected.saturating_mul(height as usize);
159            plane.data()[..required].to_vec()
160        } else {
161            let required = expected.saturating_mul(height as usize);
162            let mut out = vec![0u8; required];
163            let dst: *mut u8 = out.as_mut_ptr();
164            let src: *const u8 = plane.data().as_ptr();
165            for y in 0..height as usize {
166                let src_off = y.saturating_mul(stride);
167                let dst_off = y.saturating_mul(expected);
168                unsafe {
169                    std::ptr::copy_nonoverlapping(src.add(src_off), dst.add(dst_off), expected);
170                }
171            }
172            out
173        };
174        drop(planes);
175        let Some(img) = image::GrayImage::from_raw(width, height, out) else {
176            return Err(frame);
177        };
178        return Ok(DynamicImage::ImageLuma8(img));
179    }
180
181    let (bytes_per_pixel, wrap) = if code == FourCc::new(*b"R8  ") || code == FourCc::new(*b"GREY")
182    {
183        (1usize, 0u8)
184    } else if code == FourCc::new(*b"RG24") {
185        (3usize, 1u8)
186    } else if code == FourCc::new(*b"RGBA") {
187        (4usize, 2u8)
188    } else {
189        return Err(frame);
190    };
191    let expected_stride = (width as usize).saturating_mul(bytes_per_pixel);
192
193    let planes = frame.planes();
194    if planes.is_empty() {
195        drop(planes);
196        return Err(frame);
197    }
198    let plane_stride = planes[0].stride();
199    let plane_len = planes[0].data().len();
200    drop(planes);
201    let stride = plane_stride.max(expected_stride);
202    let required = stride.saturating_mul(height as usize);
203    if plane_len < required {
204        return Err(frame);
205    }
206
207    // Some capture backends allocate slightly larger buffers than strictly required (e.g. padding
208    // or alignment), even when the visible image is tightly packed. If stride matches exactly,
209    // we can still take/move the buffer by truncating down to the required length.
210    let can_take_zero_copy = plane_stride == expected_stride && plane_len >= required;
211    if can_take_zero_copy {
212        let (_meta, layouts, mut buffers) = frame.into_parts();
213        let layout = *layouts
214            .first()
215            .expect("frame has at least one plane layout");
216        let buf = buffers
217            .pop()
218            .expect("non-external packed frame has an owned buffer");
219
220        let stride = layout.stride.max(expected_stride);
221        let required = stride.saturating_mul(height as usize);
222        if layout.offset == 0
223            && stride == expected_stride
224            && layout.len >= required
225            && buf.len() >= required
226        {
227            let mut buf = buf;
228            buf.truncate(required);
229            match wrap {
230                0 => {
231                    let img =
232                        image::GrayImage::from_raw(width, height, buf).expect("length validated");
233                    return Ok(DynamicImage::ImageLuma8(img));
234                }
235                1 => {
236                    let img =
237                        image::RgbImage::from_raw(width, height, buf).expect("length validated");
238                    return Ok(DynamicImage::ImageRgb8(img));
239                }
240                _ => {
241                    let img =
242                        image::RgbaImage::from_raw(width, height, buf).expect("length validated");
243                    return Ok(DynamicImage::ImageRgba8(img));
244                }
245            }
246        }
247
248        let data = buf
249            .get(layout.offset..layout.offset.saturating_add(layout.len))
250            .unwrap_or(&[]);
251        return Ok(copy_packed_to_image(
252            code,
253            width,
254            height,
255            expected_stride,
256            stride,
257            data,
258        ));
259    }
260
261    let planes = frame.planes();
262    let plane = &planes[0];
263    Ok(copy_packed_to_image(
264        code,
265        width,
266        height,
267        expected_stride,
268        stride,
269        plane.data(),
270    ))
271}
272
273#[cfg(feature = "image")]
274fn copy_packed_to_image(
275    code: FourCc,
276    width: u32,
277    height: u32,
278    expected_stride: usize,
279    stride: usize,
280    data: &[u8],
281) -> DynamicImage {
282    fn copy_strided(
283        out: &mut Vec<u8>,
284        expected_stride: usize,
285        stride: usize,
286        height: u32,
287        data: &[u8],
288    ) {
289        let height = height as usize;
290        let required = expected_stride.saturating_mul(height);
291        out.clear();
292        out.resize(required, 0);
293        let dst = out.as_mut_ptr();
294        let src = data.as_ptr();
295        for y in 0..height {
296            let src_off = y.saturating_mul(stride);
297            let dst_off = y.saturating_mul(expected_stride);
298            unsafe {
299                std::ptr::copy_nonoverlapping(src.add(src_off), dst.add(dst_off), expected_stride);
300            }
301        }
302    }
303
304    match code {
305        c if c == FourCc::new(*b"R8  ") || c == FourCc::new(*b"GREY") => {
306            let mut out = Vec::new();
307            copy_strided(&mut out, expected_stride, stride, height, data);
308            DynamicImage::ImageLuma8(
309                image::GrayImage::from_raw(width, height, out).expect("length validated"),
310            )
311        }
312        c if c == FourCc::new(*b"RG24") => {
313            let mut out = Vec::new();
314            copy_strided(&mut out, expected_stride, stride, height, data);
315            DynamicImage::ImageRgb8(
316                image::RgbImage::from_raw(width, height, out).expect("length validated"),
317            )
318        }
319        c if c == FourCc::new(*b"RGBA") => {
320            let mut out = Vec::new();
321            copy_strided(&mut out, expected_stride, stride, height, data);
322            DynamicImage::ImageRgba8(
323                image::RgbaImage::from_raw(width, height, out).expect("length validated"),
324            )
325        }
326        _ => unreachable!("copy_packed_to_image only called for supported packed formats"),
327    }
328}
329
330/// Convert a `DynamicImage` to a packed `RG24` frame for codec input.
331#[cfg(feature = "image")]
332pub fn dynamic_image_to_rg24_frame(img: DynamicImage, timestamp: u64) -> Option<FrameLease> {
333    match img {
334        DynamicImage::ImageRgb8(rgb) => {
335            let (width, height) = rgb.dimensions();
336            let res = Resolution::new(width, height)?;
337            let stride = (width as usize) * 3;
338            let len = stride.checked_mul(height as usize)?;
339            let raw = rgb.into_raw();
340            if raw.len() != len {
341                return None;
342            }
343            let pool = packed_frame_pool(len);
344            let mut buf = pool.lease();
345            buf.replace_owned(raw);
346            let format = MediaFormat::new(FourCc::new(*b"RG24"), res, ColorSpace::Srgb);
347            Some(FrameLease::single_plane(
348                FrameMeta::new(format, timestamp),
349                buf,
350                len,
351                stride,
352            ))
353        }
354        other => {
355            let rgb = other.into_rgb8();
356            let (width, height) = rgb.dimensions();
357            let res = Resolution::new(width, height)?;
358            let stride = (width as usize) * 3;
359            let len = stride.checked_mul(height as usize)?;
360            let pool = packed_frame_pool(len);
361            let mut buf = pool.lease();
362            buf.resize(len);
363            buf.as_mut_slice().copy_from_slice(&rgb);
364            let format = MediaFormat::new(FourCc::new(*b"RG24"), res, ColorSpace::Srgb);
365            Some(FrameLease::single_plane(
366                FrameMeta::new(format, timestamp),
367                buf,
368                len,
369                stride,
370            ))
371        }
372    }
373}
374
375/// Convert a `DynamicImage` reference to a packed `RG24` frame for codec input.
376///
377/// This avoids cloning the full image buffer when the caller already holds the image in an `Arc`.
378#[cfg(feature = "image")]
379pub fn dynamic_image_ref_to_rg24_frame(img: &DynamicImage, timestamp: u64) -> Option<FrameLease> {
380    let (width, height) = img.dimensions();
381    let res = Resolution::new(width, height)?;
382    let stride = (width as usize) * 3;
383    let len = stride.checked_mul(height as usize)?;
384    let pool = packed_frame_pool(len);
385    let mut buf = pool.lease();
386    buf.resize(len);
387    if let Some(rgb) = img.as_rgb8() {
388        let raw = rgb.as_raw();
389        if raw.len() < len {
390            return None;
391        }
392        buf.as_mut_slice().copy_from_slice(&raw[..len]);
393    } else {
394        let rgb = img.to_rgb8();
395        let raw = rgb.as_raw();
396        if raw.len() < len {
397            return None;
398        }
399        buf.as_mut_slice().copy_from_slice(&raw[..len]);
400    }
401
402    let format = MediaFormat::new(FourCc::new(*b"RG24"), res, ColorSpace::Srgb);
403    Some(FrameLease::single_plane(
404        FrameMeta::new(format, timestamp),
405        buf,
406        len,
407        stride,
408    ))
409}
410
411#[cfg(feature = "image")]
412pub fn frame_to_dynamic_image(frame: &FrameLease) -> Option<DynamicImage> {
413    let meta = frame.meta();
414    let res = meta.format.resolution;
415    let width = res.width.get();
416    let height = res.height.get();
417    let color = meta.format.color;
418    let code = meta.format.code;
419    let planes = frame.planes();
420
421    #[inline(always)]
422    fn map_colorspace(color: ColorSpace) -> (YuvRange, YuvStandardMatrix) {
423        match color {
424            ColorSpace::Bt709 => (YuvRange::Limited, YuvStandardMatrix::Bt709),
425            ColorSpace::Bt2020 => (YuvRange::Limited, YuvStandardMatrix::Bt2020),
426            // In our metadata `Srgb` is used as a "full-range output" hint for camera YUV.
427            ColorSpace::Srgb => (YuvRange::Full, YuvStandardMatrix::Bt601),
428            ColorSpace::Unknown => (YuvRange::Limited, YuvStandardMatrix::Bt709),
429        }
430    }
431
432    fn copy_tightly_packed(src: &[u8], len: usize) -> Vec<u8> {
433        let mut out = vec![0u8; len];
434        out.copy_from_slice(&src[..len]);
435        out
436    }
437
438    fn copy_strided_packed_external(
439        plane_data: &[u8],
440        src_stride: usize,
441        dst_stride: usize,
442        height: usize,
443    ) -> Vec<u8> {
444        let required_src = src_stride.saturating_mul(height);
445
446        // For external-backed (DMA) buffers, many small strided reads can be dramatically slower on
447        // some SoCs. Prefer staging the full plane (including padding) into a contiguous Vec and
448        // then repacking from cacheable memory when the plane is large enough.
449        //
450        // NOTE: This is intentionally not tied strictly to "heavy padding"; even modest padding can
451        // still trigger slow strided DMA reads at full-frame resolutions.
452        const STAGE_THRESHOLD_BYTES: usize = 256 * 1024;
453        let use_contiguous_stage =
454            plane_data.len() >= required_src && required_src >= STAGE_THRESHOLD_BYTES;
455
456        if use_contiguous_stage {
457            let mut staged = vec![0u8; required_src];
458            staged.copy_from_slice(&plane_data[..required_src]);
459            return copy_strided_packed(&staged, src_stride, dst_stride, height);
460        }
461
462        copy_strided_packed(plane_data, src_stride, dst_stride, height)
463    }
464
465    fn copy_strided_packed(plane_data: &[u8], src_stride: usize, dst_stride: usize, height: usize) -> Vec<u8> {
466        let required_dst = dst_stride.saturating_mul(height);
467        let mut out: Vec<u8> = vec![0u8; required_dst];
468        let dst: *mut u8 = out.as_mut_ptr();
469        let src: *const u8 = plane_data.as_ptr();
470        for y in 0..height {
471            let src_off = y.saturating_mul(src_stride);
472            let dst_off = y.saturating_mul(dst_stride);
473            unsafe {
474                std::ptr::copy_nonoverlapping(src.add(src_off), dst.add(dst_off), dst_stride);
475            }
476        }
477        out
478    }
479
480    #[cfg(target_arch = "aarch64")]
481    #[inline(always)]
482    unsafe fn bgr_row_to_rgb24_neon(src: &[u8], dst: &mut [u8], width: usize) {
483        use std::arch::aarch64::{uint8x16x3_t, vld3q_u8, vst3q_u8};
484        debug_assert!(src.len() >= width * 3);
485        debug_assert!(dst.len() >= width * 3);
486
487        let src_ptr = src.as_ptr();
488        let dst_ptr = dst.as_mut_ptr();
489        let mut x = 0usize;
490        while x + 16 <= width {
491            unsafe {
492                let bgr = vld3q_u8(src_ptr.add(x * 3));
493                let rgb = uint8x16x3_t(bgr.2, bgr.1, bgr.0);
494                vst3q_u8(dst_ptr.add(x * 3), rgb);
495            }
496            x += 16;
497        }
498        for x in x..width {
499            unsafe {
500                let si = x * 3;
501                let di = x * 3;
502                let b = *src_ptr.add(si);
503                let g = *src_ptr.add(si + 1);
504                let r = *src_ptr.add(si + 2);
505                *dst_ptr.add(di) = r;
506                *dst_ptr.add(di + 1) = g;
507                *dst_ptr.add(di + 2) = b;
508            }
509        }
510    }
511
512    #[cfg(target_arch = "aarch64")]
513    #[inline(always)]
514    unsafe fn bgra_row_to_rgba_neon(src: &[u8], dst: &mut [u8], width: usize) {
515        use std::arch::aarch64::{uint8x16x4_t, vld4q_u8, vst4q_u8};
516        debug_assert!(src.len() >= width * 4);
517        debug_assert!(dst.len() >= width * 4);
518
519        let src_ptr = src.as_ptr();
520        let dst_ptr = dst.as_mut_ptr();
521        let mut x = 0usize;
522        while x + 16 <= width {
523            unsafe {
524                let bgra = vld4q_u8(src_ptr.add(x * 4));
525                let rgba = uint8x16x4_t(bgra.2, bgra.1, bgra.0, bgra.3);
526                vst4q_u8(dst_ptr.add(x * 4), rgba);
527            }
528            x += 16;
529        }
530        for x in x..width {
531            unsafe {
532                let si = x * 4;
533                let di = x * 4;
534                dst_ptr.add(di).write(*src_ptr.add(si + 2));
535                dst_ptr.add(di + 1).write(*src_ptr.add(si + 1));
536                dst_ptr.add(di + 2).write(*src_ptr.add(si));
537                dst_ptr.add(di + 3).write(*src_ptr.add(si + 3));
538            }
539        }
540    }
541
542    #[inline(always)]
543    fn convert_strided_bgr_to_rgb(width: usize, height: usize, src: &[u8], src_stride: usize) -> Vec<u8> {
544        let dst_stride = width * 3;
545        let required = dst_stride.saturating_mul(height);
546        let mut out = vec![0u8; required];
547
548        let out_ptr: *mut u8 = out.as_mut_ptr();
549        for y in 0..height {
550            let src_line = &src[y * src_stride..][..width * 3];
551            let dst_line = unsafe {
552                std::slice::from_raw_parts_mut(out_ptr.add(y * dst_stride), dst_stride)
553            };
554
555            #[cfg(target_arch = "aarch64")]
556            unsafe {
557                bgr_row_to_rgb24_neon(src_line, dst_line, width);
558                continue;
559            }
560
561            #[cfg(not(target_arch = "aarch64"))]
562            {
563                for (dst_px, src_px) in dst_line.chunks_exact_mut(3).zip(src_line.chunks_exact(3)) {
564                    dst_px[0] = src_px[2];
565                    dst_px[1] = src_px[1];
566                    dst_px[2] = src_px[0];
567                }
568            }
569        }
570        out
571    }
572
573    #[inline(always)]
574    fn convert_strided_bgra_to_rgba(width: usize, height: usize, src: &[u8], src_stride: usize) -> Vec<u8> {
575        let dst_stride = width * 4;
576        let required = dst_stride.saturating_mul(height);
577        let mut out = vec![0u8; required];
578
579        let out_ptr: *mut u8 = out.as_mut_ptr();
580        for y in 0..height {
581            let src_line = &src[y * src_stride..][..width * 4];
582            let dst_line = unsafe {
583                std::slice::from_raw_parts_mut(out_ptr.add(y * dst_stride), dst_stride)
584            };
585
586            #[cfg(target_arch = "aarch64")]
587            unsafe {
588                bgra_row_to_rgba_neon(src_line, dst_line, width);
589                continue;
590            }
591
592            #[cfg(not(target_arch = "aarch64"))]
593            {
594                for (dst_px, src_px) in dst_line.chunks_exact_mut(4).zip(src_line.chunks_exact(4)) {
595                    dst_px[0] = src_px[2];
596                    dst_px[1] = src_px[1];
597                    dst_px[2] = src_px[0];
598                    dst_px[3] = src_px[3];
599                }
600            }
601        }
602        out
603    }
604
605    match code {
606	        c if c == FourCc::new(*b"R8  ") || c == FourCc::new(*b"GREY") => {
607	            let plane = planes.into_iter().next()?;
608	            let stride = plane.stride().max(width as usize);
609	            let required = stride.checked_mul(height as usize)?;
610	            if plane.data().len() < required {
611	                return None;
612	            }
613	            let expected = (width as usize).saturating_mul(1);
614	            if stride == expected {
615	                let required = expected.checked_mul(height as usize)?;
616	                let out = copy_tightly_packed(plane.data(), required);
617	                return image::GrayImage::from_raw(width, height, out).map(DynamicImage::ImageLuma8);
618	            }
619	            let out = copy_strided_packed_external(plane.data(), stride, expected, height as usize);
620	            image::GrayImage::from_raw(width, height, out).map(DynamicImage::ImageLuma8)
621	        }
622        c if c == FourCc::new(*b"RG24") => {
623            let plane = planes.into_iter().next()?;
624            let stride = plane.stride().max(width as usize * 3);
625            let required = stride.checked_mul(height as usize)?;
626            if plane.data().len() < required {
627                return None;
628            }
629            let expected = (width as usize).saturating_mul(3);
630            if stride == expected {
631                let required = expected.checked_mul(height as usize)?;
632                let out = copy_tightly_packed(plane.data(), required);
633                return image::RgbImage::from_raw(width, height, out).map(DynamicImage::ImageRgb8);
634            }
635            let out = copy_strided_packed_external(plane.data(), stride, expected, height as usize);
636            image::RgbImage::from_raw(width, height, out).map(DynamicImage::ImageRgb8)
637        }
638        c if c == FourCc::new(*b"RGBA") => {
639            let plane = planes.into_iter().next()?;
640            let stride = plane.stride().max(width as usize * 4);
641            let required = stride.checked_mul(height as usize)?;
642            if plane.data().len() < required {
643                return None;
644            }
645            let expected = (width as usize).saturating_mul(4);
646            if stride == expected {
647                let required = expected.checked_mul(height as usize)?;
648                let out = copy_tightly_packed(plane.data(), required);
649                return image::RgbaImage::from_raw(width, height, out).map(DynamicImage::ImageRgba8);
650            }
651	            let out = copy_strided_packed_external(plane.data(), stride, expected, height as usize);
652	            image::RgbaImage::from_raw(width, height, out).map(DynamicImage::ImageRgba8)
653	        }
654        c if c == FourCc::new(*b"BGR3") => {
655            let plane = planes.into_iter().next()?;
656            let stride = plane.stride().max(width as usize * 3);
657            let required = stride.checked_mul(height as usize)?;
658            if plane.data().len() < required {
659                return None;
660            }
661            let out = convert_strided_bgr_to_rgb(
662                width as usize,
663                height as usize,
664                &plane.data()[..required],
665                stride,
666            );
667            image::RgbImage::from_raw(width, height, out).map(DynamicImage::ImageRgb8)
668        }
669        c if c == FourCc::new(*b"BG24") => {
670            let plane = planes.into_iter().next()?;
671            let stride = plane.stride().max(width as usize * 3);
672            let required = stride.checked_mul(height as usize)?;
673            if plane.data().len() < required {
674                return None;
675            }
676            let out = convert_strided_bgr_to_rgb(
677                width as usize,
678                height as usize,
679                &plane.data()[..required],
680                stride,
681            );
682            image::RgbImage::from_raw(width, height, out).map(DynamicImage::ImageRgb8)
683        }
684        c if c == FourCc::new(*b"BGRA") => {
685            let plane = planes.into_iter().next()?;
686            let stride = plane.stride().max(width as usize * 4);
687            let required = stride.checked_mul(height as usize)?;
688            if plane.data().len() < required {
689                return None;
690            }
691            let out = convert_strided_bgra_to_rgba(
692                width as usize,
693                height as usize,
694                &plane.data()[..required],
695                stride,
696            );
697            image::RgbaImage::from_raw(width, height, out).map(DynamicImage::ImageRgba8)
698        }
699        c if c == FourCc::new(*b"XB24") => {
700            let plane = planes.into_iter().next()?;
701            let stride = plane.stride().max(width as usize * 4);
702            let required = stride.checked_mul(height as usize)?;
703            if plane.data().len() < required {
704                return None;
705            }
706            let dst_stride = width as usize * 3;
707            let len = dst_stride.checked_mul(height as usize)?;
708            let mut out = vec![0u8; len];
709            let src = &plane.data()[..required];
710            out.par_chunks_mut(dst_stride)
711                .enumerate()
712                .for_each(|(y, dst_line)| {
713                    let start = y * stride;
714                    let src_line = &src[start..start + (width as usize * 4)];
715                    for (dst_px, src_px) in dst_line.chunks_exact_mut(3).zip(src_line.chunks_exact(4))
716                    {
717                        dst_px[0] = src_px[2];
718                        dst_px[1] = src_px[1];
719                        dst_px[2] = src_px[0];
720                    }
721                });
722            image::RgbImage::from_raw(width, height, out).map(DynamicImage::ImageRgb8)
723        }
724        c if c == FourCc::new(*b"XR24") => {
725            let plane = planes.into_iter().next()?;
726            let stride = plane.stride().max(width as usize * 4);
727            let required = stride.checked_mul(height as usize)?;
728            if plane.data().len() < required {
729                return None;
730            }
731            let dst_stride = width as usize * 3;
732            let len = dst_stride.checked_mul(height as usize)?;
733            let mut out = vec![0u8; len];
734            let src = &plane.data()[..required];
735            out.par_chunks_mut(dst_stride)
736                .enumerate()
737                .for_each(|(y, dst_line)| {
738                    let start = y * stride;
739                    let src_line = &src[start..start + (width as usize * 4)];
740                    for (dst_px, src_px) in dst_line.chunks_exact_mut(3).zip(src_line.chunks_exact(4))
741                    {
742                        dst_px[0] = src_px[0];
743                        dst_px[1] = src_px[1];
744                        dst_px[2] = src_px[2];
745                    }
746                });
747            image::RgbImage::from_raw(width, height, out).map(DynamicImage::ImageRgb8)
748        }
749        c if c == FourCc::new(*b"YUYV") => {
750            let plane = planes.into_iter().next()?;
751            let stride = plane.stride().max((width as usize) * 2);
752            let required = stride.checked_mul(height as usize)?;
753            if plane.data().len() < required {
754                return None;
755            }
756            let dst_stride = (width as usize) * 3;
757            let rgb_len = dst_stride.checked_mul(height as usize)?;
758            let mut rgb = vec![0u8; rgb_len];
759
760            let packed = YuvPackedImage {
761                yuy: &plane.data()[..required],
762                yuy_stride: stride as u32,
763                width,
764                height,
765            };
766            let (range, matrix) = map_colorspace(color);
767            if yuvutils_rs::yuyv422_to_rgb(&packed, &mut rgb, dst_stride as u32, range, matrix)
768                .is_err()
769            {
770                // Scalar fallback (threaded) writing into preallocated output.
771                let src = &plane.data()[..required];
772                rgb.par_chunks_mut(dst_stride)
773                    .enumerate()
774                    .for_each(|(y, dst_line)| {
775                        let line = &src[y * stride..][..(width as usize) * 2];
776                        let pair_count = (width as usize) / 2;
777                        for pair in 0..pair_count {
778                            let si = pair * 4;
779                            let di = pair * 6;
780                            let y0 = line[si] as i32;
781                            let u = line[si + 1] as i32;
782                            let y1 = line[si + 2] as i32;
783                            let v = line[si + 3] as i32;
784                            let (r0, g0, b0) = yuv_to_rgb(y0, u, v, color);
785                            let (r1, g1, b1) = yuv_to_rgb(y1, u, v, color);
786                            dst_line[di] = r0;
787                            dst_line[di + 1] = g0;
788                            dst_line[di + 2] = b0;
789                            dst_line[di + 3] = r1;
790                            dst_line[di + 4] = g1;
791                            dst_line[di + 5] = b1;
792                        }
793                        if (width as usize) % 2 == 1 && (width as usize) >= 1 {
794                            let last_x = (width as usize) - 1;
795                            let si = (last_x / 2) * 4;
796                            let di = last_x * 3;
797                            let yv = line[si] as i32;
798                            let u = line[si + 1] as i32;
799                            let v = line[si + 3] as i32;
800                            let (r, g, b) = yuv_to_rgb(yv, u, v, color);
801                            dst_line[di] = r;
802                            dst_line[di + 1] = g;
803                            dst_line[di + 2] = b;
804                        }
805                    });
806            }
807            image::RgbImage::from_raw(width, height, rgb).map(DynamicImage::ImageRgb8)
808        }
809        c if c == FourCc::new(*b"NV12") || c == FourCc::new(*b"NV21") => {
810            if planes.len() < 2 {
811                return None;
812            }
813            let y_plane = &planes[0];
814            let uv_plane = &planes[1];
815            let y_stride = y_plane.stride().max(width as usize);
816            let chroma_width = (width as usize).div_ceil(2);
817            let uv_stride = uv_plane.stride().max(chroma_width * 2);
818            let chroma_height = (height as usize).div_ceil(2);
819            let y_required = y_stride.checked_mul(height as usize)?;
820            let uv_required = uv_stride.checked_mul(chroma_height)?;
821            if y_plane.data().len() < y_required || uv_plane.data().len() < uv_required {
822                return None;
823            }
824
825            let dst_stride = (width as usize) * 3;
826            let rgb_len = dst_stride.checked_mul(height as usize)?;
827            let mut rgb = vec![0u8; rgb_len];
828
829            let bi = YuvBiPlanarImage {
830                y_plane: &y_plane.data()[..y_required],
831                y_stride: y_stride as u32,
832                uv_plane: &uv_plane.data()[..uv_required],
833                uv_stride: uv_stride as u32,
834                width,
835                height,
836            };
837            let (range, matrix) = map_colorspace(color);
838            let mode = YuvConversionMode::Balanced;
839            let is_nv12 = code == FourCc::new(*b"NV12");
840            let ok = if is_nv12 {
841                yuvutils_rs::yuv_nv12_to_rgb(&bi, &mut rgb, dst_stride as u32, range, matrix, mode)
842            } else {
843                yuvutils_rs::yuv_nv21_to_rgb(&bi, &mut rgb, dst_stride as u32, range, matrix, mode)
844            };
845            if ok.is_err() {
846                // Scalar fallback (threaded) writing into preallocated output.
847                let y_data = &y_plane.data()[..y_required];
848                let uv_data = &uv_plane.data()[..uv_required];
849                let chroma_width = (width as usize).div_ceil(2);
850                rgb.par_chunks_mut(dst_stride)
851                    .enumerate()
852                    .for_each(|(y, dst_line)| {
853                        let y_line = &y_data[y * y_stride..][..width as usize];
854                        let uv_line = &uv_data[(y / 2) * uv_stride..][..chroma_width * 2];
855                        for (x, yv) in y_line.iter().enumerate() {
856                            let uv_idx = (x / 2) * 2;
857                            let (u, v) = if is_nv12 {
858                                (uv_line[uv_idx] as i32, uv_line[uv_idx + 1] as i32)
859                            } else {
860                                (uv_line[uv_idx + 1] as i32, uv_line[uv_idx] as i32)
861                            };
862                            let (r, g, b) = yuv_to_rgb(*yv as i32, u, v, color);
863                            let di = x * 3;
864                            dst_line[di] = r;
865                            dst_line[di + 1] = g;
866                            dst_line[di + 2] = b;
867                        }
868                    });
869            }
870            image::RgbImage::from_raw(width, height, rgb).map(DynamicImage::ImageRgb8)
871        }
872        c if c == FourCc::new(*b"I420") => {
873            if planes.len() < 3 {
874                return None;
875            }
876            let y_plane = &planes[0];
877            let u_plane = &planes[1];
878            let v_plane = &planes[2];
879            let y_stride = y_plane.stride().max(width as usize);
880            let chroma_width = (width as usize).div_ceil(2);
881            let chroma_height = (height as usize).div_ceil(2);
882            let u_stride = u_plane.stride().max(chroma_width);
883            let v_stride = v_plane.stride().max(chroma_width);
884            let y_required = y_stride.checked_mul(height as usize)?;
885            let u_required = u_stride.checked_mul(chroma_height)?;
886            let v_required = v_stride.checked_mul(chroma_height)?;
887            if y_plane.data().len() < y_required
888                || u_plane.data().len() < u_required
889                || v_plane.data().len() < v_required
890            {
891                return None;
892            }
893            let dst_stride = (width as usize) * 3;
894            let rgb_len = dst_stride.checked_mul(height as usize)?;
895            let mut rgb = vec![0u8; rgb_len];
896
897            let planar = YuvPlanarImage {
898                y_plane: &y_plane.data()[..y_required],
899                y_stride: y_stride as u32,
900                u_plane: &u_plane.data()[..u_required],
901                u_stride: u_stride as u32,
902                v_plane: &v_plane.data()[..v_required],
903                v_stride: v_stride as u32,
904                width,
905                height,
906            };
907            let (range, matrix) = map_colorspace(color);
908            if yuvutils_rs::yuv420_to_rgb(&planar, &mut rgb, dst_stride as u32, range, matrix)
909                .is_err()
910            {
911                // Scalar fallback (threaded) writing into preallocated output.
912                let y_data = &y_plane.data()[..y_required];
913                let u_data = &u_plane.data()[..u_required];
914                let v_data = &v_plane.data()[..v_required];
915                let chroma_width = (width as usize).div_ceil(2);
916                rgb.par_chunks_mut(dst_stride)
917                    .enumerate()
918                    .for_each(|(y, dst_line)| {
919                        let y_line = &y_data[y * y_stride..][..width as usize];
920                        let u_line = &u_data[(y / 2) * u_stride..][..chroma_width];
921                        let v_line = &v_data[(y / 2) * v_stride..][..chroma_width];
922                        for (x, yv) in y_line.iter().enumerate() {
923                            let u = u_line[x / 2] as i32;
924                            let v = v_line[x / 2] as i32;
925                            let (r, g, b) = yuv_to_rgb(*yv as i32, u, v, color);
926                            let di = x * 3;
927                            dst_line[di] = r;
928                            dst_line[di + 1] = g;
929                            dst_line[di + 2] = b;
930                        }
931                    });
932            }
933            image::RgbImage::from_raw(width, height, rgb).map(DynamicImage::ImageRgb8)
934        }
935        _ => None,
936    }
937}
938
939#[cfg(all(test, feature = "image"))]
940mod tests {
941    use super::*;
942
943    #[test]
944    fn rejects_short_rgb_buffer() {
945        let res = Resolution::new(2, 2).unwrap();
946        let format = MediaFormat::new(FourCc::new(*b"RG24"), res, ColorSpace::Srgb);
947        let stride = res.width.get() as usize * 3;
948        let len = stride * res.height.get() as usize - 1; // intentionally short
949        let mut buf = BufferPool::with_limits(1, len, 1).lease();
950        buf.resize(len);
951        let frame = FrameLease::single_plane(FrameMeta::new(format, 0), buf, len, stride);
952        assert!(frame_to_dynamic_image(&frame).is_none());
953    }
954
955    struct NoopCodec {
956        desc: crate::CodecDescriptor,
957    }
958
959    impl NoopCodec {
960        fn new(input: FourCc) -> Self {
961            Self {
962                desc: crate::CodecDescriptor {
963                    kind: crate::CodecKind::Decoder,
964                    input,
965                    output: input,
966                    name: "noop",
967                    impl_name: "noop",
968                },
969            }
970        }
971    }
972
973    impl crate::Codec for NoopCodec {
974        fn descriptor(&self) -> &crate::CodecDescriptor {
975            &self.desc
976        }
977
978        fn process(&self, input: FrameLease) -> Result<FrameLease, crate::CodecError> {
979            Ok(input)
980        }
981    }
982
983    #[test]
984    fn process_to_dynamic_prefers_input_conversion() {
985        // BGRA can't be zero-copy wrapped as RGBA; we should still be able to convert to an image
986        // without forcing an intermediate RG24 decode.
987        let res = Resolution::new(2, 1).unwrap();
988        let format = MediaFormat::new(FourCc::new(*b"BGRA"), res, ColorSpace::Srgb);
989        let stride = res.width.get() as usize * 4;
990        let len = stride * res.height.get() as usize;
991        let mut buf = BufferPool::with_limits(1, len, 1).lease();
992        buf.resize(len);
993        // Two pixels: blue then red.
994        buf.as_mut_slice().copy_from_slice(&[255, 0, 0, 255, 0, 0, 255, 255]);
995        let frame = FrameLease::single_plane(FrameMeta::new(format, 0), buf, len, stride);
996
997        let codec = NoopCodec::new(FourCc::new(*b"BGRA"));
998        let out = process_to_dynamic(&codec, frame).unwrap();
999        let rgba = out.into_rgba8();
1000        assert_eq!(rgba.as_raw(), &[0, 0, 255, 255, 255, 0, 0, 255]);
1001    }
1002}
styx_codec/decoder/mod.rs

styx_codec/decoder/
mod.rs
