Skip to main content

aom_decode/
aom.rs

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