1use std::error::Error;
2use std::io::{Seek, SeekFrom};
3
4use image::error::{DecodingError, ImageFormatHint};
5use image::hooks::GenericReader;
6use image::{ColorType, ImageError, ImageResult};
7
8use crate::{ColorSpace, HeifContext, HeifError, ImageHandle, LibHeif, RgbChroma, StreamReader};
9
10macro_rules! magick {
11 ($v1:literal, $v2:literal, $v3:literal, $v4:literal) => {
12 [
13 &[
15 0, 0, 0, 0, b'f', b't', b'y', b'p', $v1, $v2, $v3, $v4, 0, 0, 0, 0,
16 ],
17 &[
19 0, 0, 0, 0, b'f', b't', b'y', b'p', 0, 0, 0, 0, $v1, $v2, $v3, $v4,
20 ],
21 ]
22 };
23}
24
25static MASKS: [&[u8]; 2] = [
26 &[
27 0, 0, 0, 0, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0,
28 ],
29 &[
30 0, 0, 0, 0, 0xff, 0xff, 0xff, 0xff, 0, 0, 0, 0, 0xff, 0xff, 0xff, 0xff,
31 ],
32];
33
34static HEIC_BRAND: [&[u8]; 2] = magick!(b'h', b'e', b'i', b'c');
38static HEIX_BRAND: [&[u8]; 2] = magick!(b'h', b'e', b'i', b'x');
42static AVIF_BRAND: [&[u8]; 2] = magick!(b'a', b'v', b'i', b'f');
44static JPGS_BRAND: [&[u8]; 2] = magick!(b'j', b'p', b'g', b's');
46static J2KI_BRAND: [&[u8]; 2] = magick!(b'j', b'2', b'k', b'i');
48static MIF1_BRAND: [&[u8]; 2] = magick!(b'm', b'i', b'f', b'1');
52static MIF2_BRAND: [&[u8]; 2] = magick!(b'm', b'i', b'f', b'2');
58
59pub fn register_heif_decoding_hook() -> bool {
61 let registered = image::hooks::register_decoding_hook(
62 "heif".into(),
63 Box::new(|r| Ok(Box::new(HeifDecoder::new(r)?))),
64 );
65 if registered {
66 for brands in [MIF1_BRAND, MIF2_BRAND, JPGS_BRAND, J2KI_BRAND] {
67 image::hooks::register_format_detection_hook("heif".into(), brands[0], Some(MASKS[0]));
68 image::hooks::register_format_detection_hook("heif".into(), brands[1], Some(MASKS[1]));
69 }
70 }
71 registered
72}
73
74pub fn register_heic_decoding_hook() -> bool {
76 let registered = image::hooks::register_decoding_hook(
77 "heic".into(),
78 Box::new(|r| Ok(Box::new(HeifDecoder::new(r)?))),
79 );
80 if registered {
81 for brands in [HEIC_BRAND, HEIX_BRAND] {
82 image::hooks::register_format_detection_hook("heic".into(), brands[0], Some(MASKS[0]));
83 image::hooks::register_format_detection_hook("heic".into(), brands[1], Some(MASKS[1]));
84 }
85 }
86 registered
87}
88
89pub fn register_avif_decoding_hook() -> bool {
91 let registered = image::hooks::register_decoding_hook(
92 "avif".into(),
93 Box::new(|r| Ok(Box::new(HeifDecoder::new(r)?))),
94 );
95 if registered {
96 image::hooks::register_format_detection_hook("avif".into(), AVIF_BRAND[0], Some(MASKS[0]));
97 image::hooks::register_format_detection_hook("avif".into(), AVIF_BRAND[1], Some(MASKS[1]));
98 }
99 registered
100}
101
102pub fn register_all_decoding_hooks() {
105 register_heif_decoding_hook();
106 register_heic_decoding_hook();
107 register_avif_decoding_hook();
108}
109
110fn image_error(err: impl Into<Box<dyn Error + Send + Sync>>) -> ImageError {
111 ImageError::Decoding(DecodingError::new(
112 ImageFormatHint::Name("heif".into()),
113 err,
114 ))
115}
116
117impl From<HeifError> for ImageError {
118 fn from(e: HeifError) -> Self {
119 image_error(e)
120 }
121}
122
123struct HeifDecoder<'a> {
124 _context: HeifContext<'a>,
125 image_handle: ImageHandle,
126 color_type: ColorType,
127}
128
129impl<'a> HeifDecoder<'a> {
130 fn new(mut reader: GenericReader<'a>) -> ImageResult<HeifDecoder<'a>> {
131 reader.seek(SeekFrom::End(0))?;
132 let total_size = reader.stream_position()?;
133 reader.seek(SeekFrom::Start(0))?;
134 let stream_reader = StreamReader::new(reader, total_size);
135 let context = HeifContext::read_from_reader(Box::new(stream_reader))?;
136 let image_handle = context.primary_image_handle()?;
137 let color_type = get_color_type(&image_handle)?;
138
139 Ok(Self {
140 _context: context,
141 image_handle,
142 color_type,
143 })
144 }
145}
146
147fn get_color_type(image_handle: &ImageHandle) -> ImageResult<ColorType> {
148 let has_alpha = image_handle.has_alpha_channel();
149 let color_space = image_handle.preferred_decoding_colorspace()?;
150 let (is_monochrome, is_hdr) = match color_space {
151 ColorSpace::YCbCr(_) => {
152 let bit_depth = image_handle.luma_bits_per_pixel();
153 (false, bit_depth > 8)
154 }
155 ColorSpace::Rgb(chroma) => match chroma {
156 RgbChroma::C444 | RgbChroma::Rgb | RgbChroma::Rgba => (false, false),
157 RgbChroma::HdrRgbBe
158 | RgbChroma::HdrRgbaBe
159 | RgbChroma::HdrRgbLe
160 | RgbChroma::HdrRgbaLe => (false, true),
161 },
162 ColorSpace::Monochrome => {
163 let bit_depth = image_handle.luma_bits_per_pixel();
164 (true, bit_depth > 8)
165 }
166 ColorSpace::Undefined => {
167 let bit_depth = image_handle.luma_bits_per_pixel();
168 (false, bit_depth > 8)
169 }
170 #[cfg(feature = "v1_19")]
171 ColorSpace::NonVisual => {
172 return Err(image_error("Container doesn't have image data."));
173 }
174 };
175 let color_type = match (is_monochrome, has_alpha, is_hdr) {
176 (false, false, false) => ColorType::Rgb8,
177 (false, false, true) => ColorType::Rgb16,
178 (false, true, false) => ColorType::Rgba8,
179 (false, true, true) => ColorType::Rgba16,
180 (true, false, false) => ColorType::L8,
181 (true, false, true) => ColorType::L16,
182 (true, true, false) => ColorType::La8,
183 (true, true, true) => ColorType::La16,
184 };
185 Ok(color_type)
186}
187
188fn get_color_space(color_type: ColorType) -> ColorSpace {
189 let is_target_little_endian = u16::from_ne_bytes([1, 0]) == 1;
190 match color_type {
191 ColorType::L8 | ColorType::La8 | ColorType::L16 | ColorType::La16 => ColorSpace::Monochrome,
192 ColorType::Rgb8 => ColorSpace::Rgb(RgbChroma::Rgb),
193 ColorType::Rgba8 => ColorSpace::Rgb(RgbChroma::Rgba),
194 ColorType::Rgb16 => {
195 if is_target_little_endian {
196 ColorSpace::Rgb(RgbChroma::HdrRgbLe)
197 } else {
198 ColorSpace::Rgb(RgbChroma::HdrRgbBe)
199 }
200 }
201 ColorType::Rgba16 => {
202 if is_target_little_endian {
203 ColorSpace::Rgb(RgbChroma::HdrRgbaLe)
204 } else {
205 ColorSpace::Rgb(RgbChroma::HdrRgbaBe)
206 }
207 }
208 _ => ColorSpace::Rgb(RgbChroma::Rgb),
209 }
210}
211
212impl<'a> image::ImageDecoder for HeifDecoder<'a> {
213 fn dimensions(&self) -> (u32, u32) {
214 (self.image_handle.width(), self.image_handle.height())
215 }
216
217 fn color_type(&self) -> ColorType {
218 self.color_type
219 }
220
221 fn read_image(self, buf: &mut [u8]) -> ImageResult<()>
222 where
223 Self: Sized,
224 {
225 let color_space = get_color_space(self.color_type);
226 let img = LibHeif::new().decode(&self.image_handle, color_space, None)?;
227 if !matches!(img.color_space(), Some(c) if c == color_space) {
228 return Err(image_error("Color space mismatch."));
229 }
230 let planes = img.planes();
231 let Some(plane) = planes.interleaved else {
232 return Err(image_error("Image is not interleaved."));
233 };
234
235 let row_size = plane.width as usize * (plane.storage_bits_per_pixel / 8) as usize;
236 if row_size > plane.stride {
237 return Err(image_error("Row size is greater than stride."));
238 }
239 let dst_rows = buf.chunks_exact_mut(row_size);
240 let src_rows = plane
241 .data
242 .chunks_exact(plane.stride)
243 .take(plane.height as usize)
244 .map(|row| &row[..row_size]);
245 for (dst_row, src_row) in dst_rows.zip(src_rows) {
246 dst_row.copy_from_slice(src_row);
247 }
248 Ok(())
249 }
250
251 fn read_image_boxed(self: Box<Self>, buf: &mut [u8]) -> ImageResult<()> {
252 (*self).read_image(buf)
253 }
254}