1 //! Encoding of AVIF images. 2 /// 3 /// The [AVIF] specification defines an image derivative of the AV1 bitstream, an open video codec. 4 /// 5 /// [AVIF]: https://aomediacodec.github.io/av1-avif/ 6 use std::borrow::Cow; 7 use std::io::Write; 8 use std::cmp::min; 9 10 use crate::{ColorType, ImageBuffer, ImageFormat, Pixel}; 11 use crate::{ImageError, ImageResult}; 12 use crate::buffer::ConvertBuffer; 13 use crate::color::{FromColor, Luma, LumaA, Bgr, Bgra, Rgb, Rgba}; 14 use crate::error::{EncodingError, ParameterError, ParameterErrorKind, UnsupportedError, UnsupportedErrorKind}; 15 16 use bytemuck::{Pod, PodCastError, try_cast_slice, try_cast_slice_mut}; 17 use num_traits::Zero; 18 use ravif::{Img, Config, RGBA8, encode_rgba}; 19 use rgb::AsPixels; 20 21 /// AVIF Encoder. 22 /// 23 /// Writes one image into the chosen output. 24 pub struct AvifEncoder<W> { 25 inner: W, 26 fallback: Vec<u8>, 27 config: Config 28 } 29 30 /// An enumeration over supported AVIF color spaces 31 #[derive(Debug, Copy, Clone, PartialEq, Eq)] 32 pub enum ColorSpace { 33 /// sRGB colorspace 34 Srgb, 35 /// BT.709 colorspace 36 Bt709, 37 38 #[doc(hidden)] 39 __NonExhaustive(crate::utils::NonExhaustiveMarker), 40 } 41 42 impl ColorSpace { to_ravif(self) -> ravif::ColorSpace43 fn to_ravif(self) -> ravif::ColorSpace { 44 match self { 45 Self::Srgb => ravif::ColorSpace::RGB, 46 Self::Bt709 => ravif::ColorSpace::YCbCr, 47 Self::__NonExhaustive(marker) => match marker._private {}, 48 } 49 } 50 } 51 52 impl<W: Write> AvifEncoder<W> { 53 /// Create a new encoder that writes its output to `w`. new(w: W) -> Self54 pub fn new(w: W) -> Self { 55 AvifEncoder::new_with_speed_quality(w, 1, 100) 56 } 57 58 /// Create a new encoder with specified speed and quality, that writes its output to `w`. 59 /// `speed` accepts a value in the range 0-10, where 0 is the slowest and 10 is the fastest. 60 /// `quality` accepts a value in the range 0-100, where 0 is the worst and 100 is the best. new_with_speed_quality(w: W, speed: u8, quality: u8) -> Self61 pub fn new_with_speed_quality(w: W, speed: u8, quality: u8) -> Self { 62 // Clamp quality and speed to range 63 let quality = min(quality, 100); 64 let speed = min(speed, 10); 65 66 AvifEncoder { 67 inner: w, 68 fallback: vec![], 69 config: Config { 70 quality, 71 alpha_quality: quality, 72 speed, 73 premultiplied_alpha: false, 74 color_space: ravif::ColorSpace::RGB, 75 } 76 } 77 } 78 79 /// Encode with the specified `color_space`. with_colorspace(mut self, color_space: ColorSpace) -> Self80 pub fn with_colorspace(mut self, color_space: ColorSpace) -> Self { 81 self.config.color_space = color_space.to_ravif(); 82 self 83 } 84 85 /// Encode image data with the indicated color type. 86 /// 87 /// The encoder currently requires all data to be RGBA8, it will be converted internally if 88 /// necessary. When data is suitably aligned, i.e. u16 channels to two bytes, then the 89 /// conversion may be more efficient. write_image(mut self, data: &[u8], width: u32, height: u32, color: ColorType) -> ImageResult<()>90 pub fn write_image(mut self, data: &[u8], width: u32, height: u32, color: ColorType) -> ImageResult<()> { 91 self.set_color(color); 92 let config = self.config; 93 // `ravif` needs strongly typed data so let's convert. We can either use a temporarily 94 // owned version in our own buffer or zero-copy if possible by using the input buffer. 95 // This requires going through `rgb`. 96 let buffer = self.encode_as_img(data, width, height, color)?; 97 let (data, _color_size, _alpha_size) = encode_rgba(buffer, &config) 98 .map_err(|err| ImageError::Encoding( 99 EncodingError::new(ImageFormat::Avif.into(), err) 100 ))?; 101 self.inner.write_all(&data)?; 102 Ok(()) 103 } 104 105 // Does not currently do anything. Mirrors behaviour of old config function. set_color(&mut self, _color: ColorType)106 fn set_color(&mut self, _color: ColorType) { 107 // self.config.color_space = ColorSpace::RGB; 108 } 109 encode_as_img<'buf>(&'buf mut self, data: &'buf [u8], width: u32, height: u32, color: ColorType) -> ImageResult<Img<&'buf [RGBA8]>>110 fn encode_as_img<'buf>(&'buf mut self, data: &'buf [u8], width: u32, height: u32, color: ColorType) 111 -> ImageResult<Img<&'buf [RGBA8]>> 112 { 113 // Error wrapping utility for color dependent buffer dimensions. 114 fn try_from_raw<P: Pixel + 'static>(data: &[P::Subpixel], width: u32, height: u32) 115 -> ImageResult<ImageBuffer<P, &[P::Subpixel]>> 116 { 117 ImageBuffer::from_raw(width, height, data).ok_or_else(|| { 118 ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch)) 119 }) 120 }; 121 122 // Convert to target color type using few buffer allocations. 123 fn convert_into<'buf, P>(buf: &'buf mut Vec<u8>, image: ImageBuffer<P, &[P::Subpixel]>) 124 -> Img<&'buf [RGBA8]> 125 where 126 P: Pixel + 'static, 127 Rgba<u8>: FromColor<P>, 128 { 129 let (width, height) = image.dimensions(); 130 // TODO: conversion re-using the target buffer? 131 let image: ImageBuffer<Rgba<u8>, _> = image.convert(); 132 *buf = image.into_raw(); 133 Img::new(buf.as_pixels(), width as usize, height as usize) 134 } 135 136 // Cast the input slice using few buffer allocations if possible. 137 // In particular try not to allocate if the caller did the infallible reverse. 138 fn cast_buffer<Channel>(buf: &[u8]) -> ImageResult<Cow<[Channel]>> 139 where 140 Channel: Pod + Zero, 141 { 142 match try_cast_slice(buf) { 143 Ok(slice) => Ok(Cow::Borrowed(slice)), 144 Err(PodCastError::OutputSliceWouldHaveSlop) => { 145 Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch))) 146 } 147 Err(PodCastError::TargetAlignmentGreaterAndInputNotAligned) => { 148 // Sad, but let's allocate. 149 // bytemuck checks alignment _before_ slop but size mismatch before this.. 150 if buf.len() % std::mem::size_of::<Channel>() != 0 { 151 Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch))) 152 } else { 153 let len = buf.len() / std::mem::size_of::<Channel>(); 154 let mut data = vec![Channel::zero(); len]; 155 let view = try_cast_slice_mut::<_, u8>(data.as_mut_slice()).unwrap(); 156 view.copy_from_slice(buf); 157 Ok(Cow::Owned(data)) 158 } 159 } 160 Err(err) => { // Are you trying to encode a ZST?? 161 Err(ImageError::Parameter(ParameterError::from_kind( 162 ParameterErrorKind::Generic(format!("{:?}", err)) 163 ))) 164 } 165 } 166 } 167 168 match color { 169 ColorType::Rgba8 => { 170 // ravif doesn't do any checks but has some asserts, so we do the checks. 171 let img = try_from_raw::<Rgba<u8>>(data, width, height)?; 172 // Now, internally ravif uses u32 but it takes usize. We could do some checked 173 // conversion but instead we use that a non-empty image must be addressable. 174 if img.pixels().len() == 0 { 175 return Err(ImageError::Parameter(ParameterError::from_kind(ParameterErrorKind::DimensionMismatch))); 176 } 177 178 Ok(Img::new(rgb::AsPixels::as_pixels(data), width as usize, height as usize)) 179 }, 180 // we need a separate buffer.. 181 ColorType::L8 => { 182 let image = try_from_raw::<Luma<u8>>(data, width, height)?; 183 Ok(convert_into(&mut self.fallback, image)) 184 } 185 ColorType::La8 => { 186 let image = try_from_raw::<LumaA<u8>>(data, width, height)?; 187 Ok(convert_into(&mut self.fallback, image)) 188 } 189 ColorType::Rgb8 => { 190 let image = try_from_raw::<Rgb<u8>>(data, width, height)?; 191 Ok(convert_into(&mut self.fallback, image)) 192 } 193 ColorType::Bgr8 => { 194 let image = try_from_raw::<Bgr<u8>>(data, width, height)?; 195 Ok(convert_into(&mut self.fallback, image)) 196 } 197 ColorType::Bgra8 => { 198 let image = try_from_raw::<Bgra<u8>>(data, width, height)?; 199 Ok(convert_into(&mut self.fallback, image)) 200 } 201 // we need to really convert data.. 202 ColorType::L16 => { 203 let buffer = cast_buffer(data)?; 204 let image = try_from_raw::<Luma<u16>>(&buffer, width, height)?; 205 Ok(convert_into(&mut self.fallback, image)) 206 } 207 ColorType::La16 => { 208 let buffer = cast_buffer(data)?; 209 let image = try_from_raw::<LumaA<u16>>(&buffer, width, height)?; 210 Ok(convert_into(&mut self.fallback, image)) 211 } 212 ColorType::Rgb16 => { 213 let buffer = cast_buffer(data)?; 214 let image = try_from_raw::<Rgb<u16>>(&buffer, width, height)?; 215 Ok(convert_into(&mut self.fallback, image)) 216 } 217 ColorType::Rgba16 => { 218 let buffer = cast_buffer(data)?; 219 let image = try_from_raw::<Rgba<u16>>(&buffer, width, height)?; 220 Ok(convert_into(&mut self.fallback, image)) 221 } 222 // for cases we do not support at all? 223 _ => Err(ImageError::Unsupported(UnsupportedError::from_format_and_kind( 224 ImageFormat::Avif.into(), 225 UnsupportedErrorKind::Color(color.into()), 226 ))) 227 } 228 } 229 } 230