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