libheif_rs/integration/
image.rs

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