Distinguish between HEIF and HEIC

Signed-off-by: RSUU <[email protected]>
Roughsketch · May 21, 2024 · 2c0e0df · 2c0e0df
1 parent 6d94263
commit 2c0e0df
Show file tree

Hide file tree

Showing 7 changed files with 99 additions and 58 deletions.
diff --git a/src/formats/heif.rs → src/container/heif.rs b/src/formats/heif.rs → src/container/heif.rs
@@ -1,8 +1,15 @@
 use crate::util::*;
 use crate::{ImageError, ImageResult, ImageSize};
 
+use std::convert::TryInto;
 use std::io::{BufRead, Seek, SeekFrom};
 
+pub enum Heif {
+    Avif,
+    Heic,
+    Unknown,
+}
+
 pub fn size<R: BufRead + Seek>(reader: &mut R) -> ImageResult<ImageSize> {
     reader.seek(SeekFrom::Start(0))?;
     //  Read the ftyp header size
@@ -75,34 +82,39 @@ pub fn size<R: BufRead + Seek>(reader: &mut R) -> ImageResult<ImageSize> {
     })
 }
 
-pub fn matches(header: &[u8]) -> bool {
+pub fn matches(header: &[u8]) -> Option<Heif> {
     if header.len() < 12 || &header[4..8] != b"ftyp" {
-        return false;
+        return None;
     }
 
-    let header_brand = &header[8..12];
+    let header_brand: [u8; 4] = header[8..12].try_into().unwrap();
 
     // Since other non-heif files may contain ftype in the header
     // we try to use brands to distinguish image files specifically.
     // List of brands from here: https://mp4ra.org/#/brands
-    let valid_brands = [
-        // HEIF specific
+    #[rustfmt::skip]
+    // HEIC specific
+    let heic_brands = [
         b"avci", b"avcs", b"heic", b"heim",
         b"heis", b"heix", b"hevc", b"hevm",
         b"hevs", b"hevx", b"jpeg", b"jpgs",
         b"mif1", b"msf1", b"mif2", b"pred",
-        // AVIF specific
+    ];
+
+    #[rustfmt::skip]
+    // AVIF specific
+    let avif_brands = [
         b"avif", b"avio", b"avis", b"MA1A",
         b"MA1B",
     ];
 
-    for brand in valid_brands {
-        if brand == header_brand {
-            return true;
-        }
-    }
-
-    false
+    Some(if heic_brands.contains(&&header_brand) {
+        Heif::Heic
+    } else if avif_brands.contains(&&header_brand) {
+        Heif::Avif
+    } else {
+        Heif::Unknown
+    })
 }
 
 fn skip_to_tag<R: BufRead + Seek>(reader: &mut R, tag: &[u8]) -> ImageResult<u32> {

diff --git a/src/container/mod.rs b/src/container/mod.rs
@@ -0,0 +1 @@
+pub mod heif;
diff --git a/src/formats/avif.rs b/src/formats/avif.rs
diff --git a/src/formats/mod.rs b/src/formats/mod.rs
@@ -1,12 +1,10 @@
 pub mod aesprite;
-pub mod avif;
 pub mod bmp;
 pub mod dds;
 pub mod exr;
 pub mod farbfeld;
 pub mod gif;
 pub mod hdr;
-pub mod heif;
 pub mod ico;
 pub mod jpeg;
 pub mod jxl;
@@ -20,7 +18,7 @@ pub mod tiff;
 pub mod vtf;
 pub mod webp;
 
-use crate::{ImageError, ImageResult, ImageType};
+use crate::{container, ImageError, ImageResult, ImageType};
 use std::io::{BufRead, Seek};
 
 pub fn image_type<R: BufRead + Seek>(reader: &mut R) -> ImageResult<ImageType> {
@@ -56,12 +54,15 @@ pub fn image_type<R: BufRead + Seek>(reader: &mut R) -> ImageResult<ImageType> {
         return Ok(ImageType::Webp);
     }
 
-    if heif::matches(&header) {
-        return Ok(ImageType::Heif);
-    }
+    if let Some(ty) = container::heif::matches(&header) {
+        use container::heif::Heif;
 
-    if avif::matches(&header) {
-        return Ok(ImageType::Avif);
+        return Ok(match ty {
+            Heif::Avif => ImageType::Avif,
+            Heif::Heic => ImageType::Heic,
+            // Unknown format in HEIF.
+            _ => ImageType::Heif,
+        });
     }
 
     if jxl::matches(&header) {

diff --git a/src/lib.rs b/src/lib.rs
@@ -4,10 +4,11 @@ use std::fs::File;
 use std::io::{BufRead, BufReader, Cursor, Seek};
 use std::path::Path;
 
+mod container;
+mod formats;
 mod util;
 
-mod formats;
-use formats::*;
+use {container::*, formats::*};
 
 /// An Error type used in failure cases.
 #[derive(Debug)]
@@ -63,6 +64,8 @@ pub enum ImageType {
     /// Radiance HDR
     Hdr,
     /// High Efficiency Image File Format
+    Heic,
+    /// Image Container Format
     Heif,
     /// Icon file
     Ico,
@@ -248,14 +251,12 @@ pub fn reader_size<R: BufRead + Seek>(mut reader: R) -> ImageResult<ImageSize> {
 fn dispatch_header<R: BufRead + Seek>(reader: &mut R) -> ImageResult<ImageSize> {
     match formats::image_type(reader)? {
         ImageType::Aseprite => aesprite::size(reader),
-        ImageType::Avif => heif::size(reader), // AVIF uses HEIF size on purpose
         ImageType::Bmp => bmp::size(reader),
         ImageType::Dds => dds::size(reader),
         ImageType::Exr => exr::size(reader),
         ImageType::Farbfeld => farbfeld::size(reader),
         ImageType::Gif => gif::size(reader),
         ImageType::Hdr => hdr::size(reader),
-        ImageType::Heif => heif::size(reader),
         ImageType::Ico => ico::size(reader),
         ImageType::Jpeg => jpeg::size(reader),
         ImageType::Jxl => jxl::size(reader),
@@ -268,5 +269,8 @@ fn dispatch_header<R: BufRead + Seek>(reader: &mut R) -> ImageResult<ImageSize>
         ImageType::Tiff => tiff::size(reader),
         ImageType::Vtf => vtf::size(reader),
         ImageType::Webp => webp::size(reader),
+
+        // AVIF and HEIC uses HEIF size on purpose
+        ImageType::Heif | ImageType::Heic | ImageType::Avif => heif::size(reader),
     }
 }
diff --git a/tests/avif.rs b/tests/avif.rs
@@ -1,14 +1,38 @@
 #[cfg(test)]
-use imagesize::{size, ImageSize};
+use imagesize::{image_type, size, ImageSize, ImageType};
 
 #[test]
 fn avif_test() {
     let dim = size("tests/images/avif/test.avif").unwrap();
-    assert_eq!(dim, ImageSize { width: 1204, height: 800 });
+    assert_eq!(
+        dim,
+        ImageSize {
+            width: 1204,
+            height: 800
+        }
+    );
 }
 
 #[test]
 fn avif_multi_picks_largest() {
     let dim = size("tests/images/avif/test.avifs").unwrap();
-    assert_eq!(dim, ImageSize { width: 159, height: 159 });
+    assert_eq!(
+        dim,
+        ImageSize {
+            width: 159,
+            height: 159
+        }
+    );
+}
+
+#[test]
+fn avif_type() {
+    use std::{fs::File, io::Read};
+
+    let mut f = File::open("tests/images/avif/test.avif").unwrap();
+    let mut buf = vec![];
+    f.read_to_end(&mut buf).unwrap();
+
+    let ty = image_type(&buf).unwrap();
+    assert_eq!(ty, ImageType::Avif);
 }
diff --git a/tests/heic.rs b/tests/heic.rs
@@ -1,14 +1,38 @@
 #[cfg(test)]
-use imagesize::{size, ImageSize};
+use imagesize::{image_type, size, ImageSize, ImageType};
 
 #[test]
-fn heif_test() {
+fn heic_test() {
     let dim = size("tests/images/heic/test.heic").unwrap();
-    assert_eq!(dim, ImageSize { width: 1280, height: 720 });
+    assert_eq!(
+        dim,
+        ImageSize {
+            width: 1280,
+            height: 720
+        }
+    );
 }
 
 #[test]
-fn heif_multi_picks_largest() {
+fn heic_multi_picks_largest() {
     let dim = size("tests/images/heic/IMG_0007.heic").unwrap();
-    assert_eq!(dim, ImageSize { width: 2448, height: 3264 });
+    assert_eq!(
+        dim,
+        ImageSize {
+            width: 2448,
+            height: 3264
+        }
+    );
+}
+
+#[test]
+fn heic_type() {
+    use std::{fs::File, io::Read};
+
+    let mut f = File::open("tests/images/heic/test.heic").unwrap();
+    let mut buf = vec![];
+    f.read_to_end(&mut buf).unwrap();
+
+    let ty = image_type(&buf).unwrap();
+    assert_eq!(ty, ImageType::Heic);
 }