aom_decode/
aom.rs

1use crate::Error;
2use crate::Result;
3use libaom_sys::*;
4use std::ffi::c_int;
5use std::ffi::CStr;
6use std::fmt;
7use std::marker::PhantomData;
8use std::mem::MaybeUninit;
9use std::num::NonZeroU32;
10use std::ptr;
11use std::ptr::NonNull;
12use yuv::color::{
13    ChromaSamplePosition, ChromaSampling, ColorPrimaries, Depth, MatrixCoefficients, Range,
14    TransferCharacteristics,
15};
16
17/// AOM decoder context
18pub struct Decoder {
19    ctx: aom_codec_ctx,
20}
21
22/// Result of [`Decoder::frame_meta()`]
23#[derive(Debug, PartialEq)]
24#[non_exhaustive]
25pub struct FrameMeta {
26    /// The width the current frame is decoded at (note this may be different
27    /// to the display width of the current frame)
28    pub frame_width: u32,
29    /// The height the current frame is decoded at (note this may be different
30    /// to the display height of the current frame)
31    pub frame_height: u32,
32
33    /// The intended display width of the current frame (note this may be
34    /// different to the width the frame is decoded at)
35    pub display_width: u32,
36    /// The intended display height of the current frame (note this may be
37    /// different to the height the frame is decoded at).
38    pub display_height: u32,
39}
40
41/// Configuration for the decoer. For now it's just number of threads.
42///
43/// You can use `std::thread::available_parallelism().map(|v| v.get()).unwrap_or(4)`.
44#[derive(Debug, Clone)]
45pub struct Config {
46    pub threads: usize,
47}
48
49impl Decoder {
50    /// Create a new decoder
51    pub fn new(cfg: &Config) -> Result<Self> {
52        let cfg = aom_codec_dec_cfg {
53            w: 0,
54            h: 0,
55            threads: cfg.threads as _,
56            allow_lowbitdepth: 1,
57        };
58        unsafe {
59            let mut ctx = MaybeUninit::uninit();
60            let res = aom_codec_dec_init_ver(
61                ctx.as_mut_ptr(),
62                aom_codec_av1_dx(),
63                &cfg,
64                0,
65                AOM_DECODER_ABI_VERSION as i32,
66            );
67            if let Some(code) = NonZeroU32::new(res) {
68                Err(Error::AOM(code, None))
69            } else {
70                Ok(Self {
71                    ctx: ctx.assume_init(),
72                })
73            }
74        }
75    }
76
77    /// This parses the AV1 data independently of `decode_frame`
78    #[inline]
79    pub fn frame_meta<'a>(&'a mut self, av1_data: &[u8]) -> Result<FrameMeta> {
80        let res = unsafe {
81            aom_codec_decode(
82                &mut self.ctx,
83                av1_data.as_ptr(),
84                av1_data.len(),
85                ptr::null_mut(),
86            )
87        };
88        self.is_err(res)?;
89
90        #[cold]
91        fn overflow_err<E>(_: E) -> Error {
92            Error::Unsupported("overflow")
93        }
94
95        let mut frame_w_h = [0 as c_int; 2];
96        let res = unsafe {
97            aom_codec_control(
98                &mut self.ctx,
99                AV1D_GET_FRAME_SIZE.try_into().map_err(overflow_err)?,
100                frame_w_h.as_mut_ptr()
101            )
102        };
103        self.is_err(res)?;
104
105        let mut display_w_h = [0 as c_int; 2];
106        let res = unsafe {
107            aom_codec_control(
108                &mut self.ctx,
109                AV1D_GET_DISPLAY_SIZE.try_into().map_err(overflow_err)?,
110                display_w_h.as_mut_ptr()
111            )
112        };
113        self.is_err(res)?;
114
115        Ok(FrameMeta {
116            frame_width: frame_w_h[0].try_into().map_err(overflow_err)?,
117            frame_height: frame_w_h[1].try_into().map_err(overflow_err)?,
118
119            display_width: display_w_h[0].try_into().map_err(overflow_err)?,
120            display_height: display_w_h[1].try_into().map_err(overflow_err)?,
121        })
122    }
123
124    /// Take AV1-compressed data and decode a *single* frame into raw frame data (YUV pixels). This is for AVIF, and can't handle video.
125    ///
126    /// The returned frame is temporary. You must copy data out of it and drop it before decoding other files.
127    ///
128    /// See [yuv](https://lib.rs/yuv) crate for conversion to RGB.
129    #[inline]
130    pub fn decode_frame<'a>(&'a mut self, av1_data: &[u8]) -> Result<FrameTempRef<'a>> {
131        Ok(FrameTempRef(unsafe {
132            let res = aom_codec_decode(
133                &mut self.ctx,
134                av1_data.as_ptr(),
135                av1_data.len(),
136                ptr::null_mut(),
137            );
138            self.is_err(res)?;
139
140            let mut iter = ptr::null();
141            let res = aom_codec_get_frame(&mut self.ctx, &mut iter);
142            self.err_if_null(res)?
143        }, PhantomData))
144    }
145
146    #[inline]
147    fn is_err(&self, res: u32) -> Result<()> {
148        if let Some(code) = NonZeroU32::new(res) {
149            Err(Error::AOM(code, self.last_error_msg()))
150        } else {
151            Ok(())
152        }
153    }
154
155    #[inline]
156    fn err_if_null<T>(&self, ptr: *const T) -> Result<NonNull<T>, Error> {
157        self.err_if_null_mut(ptr.cast_mut())
158    }
159
160    fn err_if_null_mut<T>(&self, ptr: *mut T) -> Result<NonNull<T>, Error> {
161        if let Some(ptr) = NonNull::new(ptr) {
162            Ok(ptr)
163        } else {
164            Err(Error::AOM(NonZeroU32::new(libaom_sys::AOM_CODEC_ERROR).unwrap(), self.last_error_msg()))
165        }
166    }
167
168    fn last_error_msg(&self) -> Option<String> {
169        let s = unsafe {
170            let err = aom_codec_error(std::ptr::from_ref(&self.ctx).cast_mut());
171            if err.is_null() {
172                return None;
173            }
174            CStr::from_ptr(err).to_string_lossy()
175        };
176        Some(s.into_owned())
177    }
178}
179
180impl Drop for Decoder {
181    #[inline]
182    fn drop(&mut self) {
183        unsafe {
184            aom_codec_destroy(&mut self.ctx);
185        }
186    }
187}
188
189/// Iterates over rows of pixels in a Y/U/V plane
190pub struct RowsIter<'a, Pixel> {
191    plane_start: *const u8,
192    stride_bytes: isize,
193    w: usize, h: usize,
194    y: isize,
195    _data: PhantomData<&'a [Pixel]>,
196}
197
198impl<T> RowsIter<'_, T> {
199    #[inline(always)]
200    #[must_use]
201    pub const fn width(&self) -> usize {
202        self.w
203    }
204    #[inline(always)]
205    #[must_use]
206    pub const fn height(&self) -> usize {
207        self.h
208    }
209}
210
211impl<'a, Pixel> Iterator for RowsIter<'a, Pixel> {
212    type Item = &'a [Pixel];
213    #[inline]
214    fn next(&mut self) -> Option<Self::Item> {
215        let y = self.y;
216        if (self.y as usize) < self.h {
217            self.y += 1;
218            // stride may be negative for flipped images?
219            Some(unsafe {
220                debug_assert_eq!(0, self.stride_bytes.unsigned_abs() % std::mem::align_of::<Pixel>());
221                let ptr = self.plane_start.offset(y * self.stride_bytes);
222                std::slice::from_raw_parts(ptr.cast::<Pixel>(), self.w)
223            })
224        } else {
225            None
226        }
227    }
228}
229
230/// Iterators of frame's pixels (YUV planes)
231///
232/// It's an enum, because frames may have different pixel formats, and you get format-specific iterator in the enum
233pub enum RowsIters<'a> {
234    /// 8-bit YUV (YCbCr, YCgCo, etc.). This is a collection of 3 iterators. Each may have a different size depending on `chroma_sampling`.
235    ///
236    /// Consult `matrix_coefficients()` for meaning of these channels
237    YuvPlanes8 {
238        y: RowsIter<'a, u8>,
239        u: RowsIter<'a, u8>,
240        v: RowsIter<'a, u8>,
241        chroma_sampling: ChromaSampling,
242    },
243    /// 8-bit grayscale
244    Mono8(RowsIter<'a, u8>),
245    /// 10/12/16-bit color. Not subsampled. See `depth` field for actual range of pixels used.
246    ///
247    /// It's not `u16`, because it's not guaranteed to be aligned. Use `u16::from_ne_bytes()` or `ptr::read_unaligned`.
248    YuvPlanes16 {
249        y: RowsIter<'a, [u8; 2]>,
250        u: RowsIter<'a, [u8; 2]>,
251        v: RowsIter<'a, [u8; 2]>,
252        chroma_sampling: ChromaSampling,
253        depth: Depth,
254    },
255    /// 10/12/16-bit grayscale. It's not `u16`, because it's not guaranteed to be aligned. Use `u16::from_ne_bytes()` or `ptr::read_unaligned`.
256    Mono16(RowsIter<'a, [u8; 2]>, Depth),
257}
258
259/// Frame held in decoder's internal state. Must be dropped before the next call.
260pub struct FrameTempRef<'a>(NonNull<aom_image_t>, PhantomData<&'a mut Decoder>);
261
262impl FrameTempRef<'_> {
263    #[inline(always)]
264    const fn as_ref(&self) -> &aom_image_t {
265        unsafe { self.0.as_ref() }
266    }
267
268    #[inline]
269    unsafe fn single_plane_iter<T: Copy>(&self, plane_n: u8) -> Result<RowsIter<'_, T>> {
270        assert!(plane_n < 3);
271
272        let img = self.as_ref();
273        if img.bit_depth <= 8 {
274            assert_eq!(1, std::mem::size_of::<T>());
275        } else {
276            assert_eq!(2, std::mem::size_of::<T>());
277        }
278
279        let w = aom_img_plane_width(img, i32::from(plane_n)) as usize;
280        let h = aom_img_plane_height(img, i32::from(plane_n)) as usize;
281
282        let stride_bytes = img.stride[plane_n as usize] as isize;
283        if stride_bytes.unsigned_abs() * std::mem::size_of::<T>() < w {
284            return Err(Error::Unsupported("stride"));
285        }
286
287        Ok(RowsIter {
288            plane_start: img.planes[plane_n as usize],
289            stride_bytes,
290            w, h,
291            y: 0,
292            _data: PhantomData,
293        })
294    }
295
296    /// Access pixel data
297    ///
298    /// Iterator over rows of image data.
299    ///
300    /// The data can be grayscale (mono) or YUV (YCbCr), so the result is wrapped in an `enum`
301    pub fn rows_iter(&self) -> Result<RowsIters> {
302        let chroma_sampling = self.chroma_sampling()?;
303        let depth = self.depth()?;
304        let flipped_uv = match self.as_ref().fmt {
305            AOM_IMG_FMT_YV12 | AOM_IMG_FMT_AOMYV12 | AOM_IMG_FMT_YV1216 => true,
306            _ => false,
307        };
308        Ok(unsafe {match (chroma_sampling, depth) {
309            (ChromaSampling::Monochrome, Depth::Depth8) => RowsIters::Mono8(self.single_plane_iter(0)?),
310            (ChromaSampling::Monochrome, depth) => RowsIters::Mono16(self.single_plane_iter(0)?, depth),
311            (chroma_sampling, Depth::Depth8) => {
312                let y = self.single_plane_iter::<u8>(0)?;
313                let u = self.single_plane_iter(1)?;
314                let v = self.single_plane_iter(2)?;
315                let (u,v) = if flipped_uv {(v,u)} else {(u,v)};
316                RowsIters::YuvPlanes8 {y,u,v, chroma_sampling}
317            },
318            (chroma_sampling, depth) => {
319                let y = self.single_plane_iter::<[u8;2]>(0)?;
320                let u = self.single_plane_iter(1)?;
321                let v = self.single_plane_iter(2)?;
322                let (u,v) = if flipped_uv {(v,u)} else {(u,v)};
323                RowsIters::YuvPlanes16 {y,u,v, depth, chroma_sampling}
324            },
325        }})
326    }
327
328    /// Whether image uses chroma subsampling or not
329    #[inline]
330    pub const fn chroma_sampling(&self) -> Result<ChromaSampling> {
331        if self.as_ref().monochrome != 0 {
332            return Ok(ChromaSampling::Monochrome);
333        }
334        Ok(match self.as_ref().fmt {
335            AOM_IMG_FMT_YV12 |
336            AOM_IMG_FMT_I420 |
337            AOM_IMG_FMT_AOMYV12 |
338            AOM_IMG_FMT_AOMI420 |
339            AOM_IMG_FMT_I42016 |
340            AOM_IMG_FMT_YV1216 => ChromaSampling::Cs420,
341            AOM_IMG_FMT_I422 |
342            AOM_IMG_FMT_I42216 => ChromaSampling::Cs422,
343            AOM_IMG_FMT_I444 |
344            AOM_IMG_FMT_I44416 => ChromaSampling::Cs444,
345            _ => return Err(Error::Unsupported("Unknown image format")),
346        })
347    }
348
349    /// How many bits per pixel that is
350    #[inline]
351    pub const fn depth(&self) -> Result<Depth> {
352        Ok(match self.as_ref().bit_depth {
353            8 => Depth::Depth8,
354            10 => Depth::Depth10,
355            12 => Depth::Depth12,
356            16 => Depth::Depth16,
357            _ => return Err(Error::Unsupported("Bad depth")),
358        })
359    }
360
361    /// What flavor of RGB color this should be converted to
362    #[inline]
363    #[must_use]
364    pub const fn color_primaries(&self) -> Option<ColorPrimaries> {
365        Some(match self.as_ref().cp {
366            AOM_CICP_CP_BT_709 => ColorPrimaries::BT709,
367            AOM_CICP_CP_BT_470_M => ColorPrimaries::BT470M,
368            AOM_CICP_CP_BT_470_B_G => ColorPrimaries::BT470BG,
369            AOM_CICP_CP_BT_601 |
370            AOM_CICP_CP_SMPTE_240 => ColorPrimaries::BT601,
371            AOM_CICP_CP_GENERIC_FILM => ColorPrimaries::GenericFilm,
372            AOM_CICP_CP_BT_2020 => ColorPrimaries::BT2020,
373            AOM_CICP_CP_XYZ => ColorPrimaries::XYZ,
374            AOM_CICP_CP_SMPTE_431 => ColorPrimaries::SMPTE431,
375            AOM_CICP_CP_SMPTE_432 => ColorPrimaries::SMPTE432,
376            AOM_CICP_CP_EBU_3213 => ColorPrimaries::EBU3213,
377            _ => return None,
378        })
379    }
380
381    /// That's basically gamma correction
382    #[inline]
383    #[must_use]
384    pub const fn transfer_characteristics(&self) -> Option<TransferCharacteristics> {
385        Some(match self.as_ref().tc {
386            AOM_CICP_TC_BT_709 => TransferCharacteristics::BT709,
387            AOM_CICP_TC_BT_470_M => TransferCharacteristics::BT470M,
388            AOM_CICP_TC_BT_470_B_G => TransferCharacteristics::BT470BG,
389            AOM_CICP_TC_BT_601 => TransferCharacteristics::BT601,
390            AOM_CICP_TC_SMPTE_240 => TransferCharacteristics::SMPTE240,
391            AOM_CICP_TC_LINEAR => TransferCharacteristics::Linear,
392            AOM_CICP_TC_LOG_100 => TransferCharacteristics::Log100,
393            AOM_CICP_TC_LOG_100_SQRT10 => TransferCharacteristics::Log100Sqrt10,
394            AOM_CICP_TC_IEC_61966 => TransferCharacteristics::IEC61966,
395            AOM_CICP_TC_BT_1361 => TransferCharacteristics::BT1361,
396            AOM_CICP_TC_SRGB => TransferCharacteristics::SRGB,
397            AOM_CICP_TC_BT_2020_10_BIT => TransferCharacteristics::BT2020_10Bit,
398            AOM_CICP_TC_BT_2020_12_BIT => TransferCharacteristics::BT2020_12Bit,
399            AOM_CICP_TC_SMPTE_2084 => TransferCharacteristics::SMPTE2084,
400            AOM_CICP_TC_SMPTE_428 => TransferCharacteristics::SMPTE428,
401            AOM_CICP_TC_HLG => TransferCharacteristics::HLG,
402            _ => return None,
403        })
404    }
405
406    /// Flavor of YUV used for the pixels
407    ///
408    /// See [yuv](https://lib.rs/yuv) crate for conversion to RGB.
409    #[inline]
410    #[must_use]
411    pub const fn matrix_coefficients(&self) -> Option<MatrixCoefficients> {
412        Some(match self.as_ref().mc {
413            AOM_CICP_MC_IDENTITY => MatrixCoefficients::Identity,
414            AOM_CICP_MC_BT_709 => MatrixCoefficients::BT709,
415            AOM_CICP_MC_FCC => MatrixCoefficients::FCC,
416            AOM_CICP_MC_BT_470_B_G => MatrixCoefficients::BT470BG,
417            AOM_CICP_MC_BT_601 => MatrixCoefficients::BT601,
418            AOM_CICP_MC_SMPTE_240 => MatrixCoefficients::SMPTE240,
419            AOM_CICP_MC_SMPTE_YCGCO => MatrixCoefficients::YCgCo,
420            AOM_CICP_MC_BT_2020_NCL |
421            AOM_CICP_MC_BT_2020_CL => MatrixCoefficients::BT2020NCL,
422            AOM_CICP_MC_SMPTE_2085 => MatrixCoefficients::SMPTE2085,
423            AOM_CICP_MC_CHROMAT_NCL => MatrixCoefficients::ChromatNCL,
424            AOM_CICP_MC_CHROMAT_CL => MatrixCoefficients::ChromatCL,
425            AOM_CICP_MC_ICTCP => MatrixCoefficients::ICtCp,
426            _ => return None,
427        })
428    }
429
430    /// Whether pixels are in 0-255 or 16-235/240 range.
431    #[inline(always)]
432    #[must_use]
433    pub const fn range(&self) -> Range {
434        match self.as_ref().range {
435            AOM_CR_STUDIO_RANGE => Range::Limited,
436            _ => Range::Full,
437        }
438    }
439
440    /// Alignment of the chroma channels
441    ///
442    /// Routines in this library don't support this detail.
443    /// Also, chroma subsampling is useless in AV1, so please don't use it.
444    #[inline(always)]
445    #[must_use]
446    pub const fn chroma_sample_position(&self) -> Option<ChromaSamplePosition> {
447        match self.as_ref().csp {
448            AOM_CSP_VERTICAL => Some(ChromaSamplePosition::Vertical),
449            AOM_CSP_COLOCATED => Some(ChromaSamplePosition::Colocated),
450            _ => None,
451        }
452    }
453}
454
455impl fmt::Debug for FrameTempRef<'_> {
456    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
457        let img = self.as_ref();
458        f.debug_struct("FrameTempRef")
459            .field("chroma_sampling", &self.chroma_sampling())
460            .field("color_primaries", &self.color_primaries())
461            .field("transfer_characteristics", &self.transfer_characteristics())
462            .field("matrix_coefficients", &self.matrix_coefficients())
463            .field("monochrome", &img.monochrome)
464            .field("csp", &self.chroma_sample_position())
465            .field("range", &self.range())
466            .field("depth", &self.depth())
467            .field("width", &img.d_w)
468            .field("height", &img.d_h)
469            .finish()
470    }
471}