Add AVIF/HEIC detection

ffmpeg doesn't support HEIC decoding yet...

client/html/post_merge_side.tpl

@@ -36,6 +36,8 @@
                     'image/jpeg': 'JPEG',
                     'image/png': 'PNG',
                     'image/webp': 'WEBP',
+                    'image/avif': 'AVIF',
+                    'image/heic': 'HEIC',
                     'video/webm': 'WEBM',
                     'video/mp4': 'MPEG-4',
                     'application/x-shockwave-flash': 'SWF',

client/html/post_readonly_sidebar.tpl

@@ -9,6 +9,8 @@
                     'image/jpeg': 'JPEG',
                     'image/png': 'PNG',
                     'image/webp': 'WEBP',
+                    'image/avif': 'AVIF',
+                    'image/heic': 'HEIC',
                     'video/webm': 'WEBM',
                     'video/mp4': 'MPEG-4',
                     'application/x-shockwave-flash': 'SWF',

client/js/views/post_upload_view.js

@@ -15,6 +15,8 @@ function _mimeTypeToPostType(mimeType) {
             "image/jpeg": "image",
             "image/png": "image",
             "image/webp": "image",
+            "image/avif": "image",
+            "image/heic": "image",
             "video/mp4": "video",
             "video/webm": "video",
         }[mimeType] || "unknown"
@@ -109,6 +111,8 @@ class Url extends Uploadable {
             png: "image/png",
             gif: "image/gif",
             webp: "image/webp",
+            avif: "image/avif",
+            heic: "image/heic",
             mp4: "video/mp4",
             webm: "video/webm",
         };

server/requirements.txt

@@ -9,4 +9,6 @@ pillow>=4.3.0
 pynacl>=1.2.1
 pytz>=2018.3
 pyRFC3339>=1.0
+pillow-avif-plugin>=1.1.0
+pyheif-pillow-opener>=0.1.0
 youtube_dl

server/szurubooru/func/image_hash.py

@@ -6,6 +6,10 @@ from typing import Any, Callable, List, Optional, Set, Tuple
 
 import numpy as np
 from PIL import Image
+import pillow_avif
+import pyheif
+from pyheif_pillow_opener import register_heif_opener
+register_heif_opener()
 
 from szurubooru import config, errors
 

server/szurubooru/func/mime.py

@@ -21,6 +21,12 @@ def get_mime_type(content: bytes) -> str:
     if content[8:12] == b"WEBP":
         return "image/webp"
 
+    if content[4:12] in (b"ftypavif", b"ftypavis"):
+        return "image/avif"
+
+    if content[4:12] in (b"ftypheic", b"ftypmif1"):
+        return "image/heic"
+
     if content[0:4] == b"\x1A\x45\xDF\xA3":
         return "video/webm"
 
@@ -37,6 +43,8 @@ def get_extension(mime_type: str) -> Optional[str]:
         "image/jpeg": "jpg",
         "image/png": "png",
         "image/webp": "webp",
+        "image/avif": "avif",
+        "image/heic": "heic",
         "video/mp4": "mp4",
         "video/webm": "webm",
         "application/octet-stream": "dat",
@@ -58,6 +66,8 @@ def is_image(mime_type: str) -> bool:
         "image/png",
         "image/gif",
         "image/webp",
+        "image/avif",
+        "image/heic",
     )
 
 

server/szurubooru/func/posts.py

@@ -650,7 +650,8 @@ def update_post_content(post: model.Post, content: Optional[bytes]) -> None:
         image = images.Image(content)
         post.canvas_width = image.width
         post.canvas_height = image.height
-    except errors.ProcessingError:
+    except errors.ProcessingError as ex:
+        logger.exception(ex)
         if not config.config["allow_broken_uploads"]:
             raise InvalidPostContentError("Unable to process image metadata")
         else:

server/szurubooru/tests/func/test_image_hash.py

@@ -27,3 +27,53 @@ def test_signature_functions(read_asset, config_injector):
     words2 = image_hash.generate_words(sig2)
     words_match = sum(word1 == word2 for word1, word2 in zip(words1, words2))
     assert words_match == 18
+
+
+def test_signature_heif(read_asset, config_injector):
+    sig1 = image_hash.generate_signature(read_asset("heif.heif"))
+    sig2 = image_hash.generate_signature(read_asset("heif-similar.heif"))
+
+    sig1_repacked = image_hash.unpack_signature(
+        image_hash.pack_signature(sig1)
+    )
+    sig2_repacked = image_hash.unpack_signature(
+        image_hash.pack_signature(sig2)
+    )
+    assert array_equal(sig1, sig1_repacked)
+    assert array_equal(sig2, sig2_repacked)
+
+    dist1 = image_hash.normalized_distance([sig1], sig2)
+    assert abs(dist1[0] - 0.136777724290135) < 1e-8
+
+    dist2 = image_hash.normalized_distance([sig2], sig2)
+    assert abs(dist2[0]) < 1e-8
+
+    words1 = image_hash.generate_words(sig1)
+    words2 = image_hash.generate_words(sig2)
+    words_match = sum(word1 == word2 for word1, word2 in zip(words1, words2))
+    assert words_match == 43
+
+
+def test_signature_avif(read_asset, config_injector):
+    sig1 = image_hash.generate_signature(read_asset("avif.avif"))
+    sig2 = image_hash.generate_signature(read_asset("avif-similar.avif"))
+
+    sig1_repacked = image_hash.unpack_signature(
+        image_hash.pack_signature(sig1)
+    )
+    sig2_repacked = image_hash.unpack_signature(
+        image_hash.pack_signature(sig2)
+    )
+    assert array_equal(sig1, sig1_repacked)
+    assert array_equal(sig2, sig2_repacked)
+
+    dist1 = image_hash.normalized_distance([sig1], sig2)
+    assert abs(dist1[0] - 0.22628712858355998) < 1e-8
+
+    dist2 = image_hash.normalized_distance([sig2], sig2)
+    assert abs(dist2[0]) < 1e-8
+
+    words1 = image_hash.generate_words(sig1)
+    words2 = image_hash.generate_words(sig2)
+    words_match = sum(word1 == word2 for word1, word2 in zip(words1, words2))
+    assert words_match == 12

server/szurubooru/tests/func/test_mime.py

@@ -13,6 +13,11 @@ from szurubooru.func import mime
         ("jpeg.jpg", "image/jpeg"),
         ("gif.gif", "image/gif"),
         ("webp.webp", "image/webp"),
+        ("avif.avif", "image/avif"),
+        ("avif-avis.avif", "image/avif"),
+        ("heic.heic", "image/heic"),
+        ("heic-mif1.heic", "image/heic"),
+        ("heif.heif", "image/heic"),
         ("text.txt", "application/octet-stream"),
     ],
 )
@@ -34,6 +39,8 @@ def test_get_mime_type_for_empty_file():
         ("image/jpeg", "jpg"),
         ("image/gif", "gif"),
         ("image/webp", "webp"),
+        ("image/avif", "avif"),
+        ("image/heic", "heic"),
         ("application/octet-stream", "dat"),
     ],
 )
@@ -75,9 +82,13 @@ def test_is_video(input_mime_type, expected_state):
         ("image/gif", True),
         ("image/png", True),
         ("image/jpeg", True),
+        ("image/avif", True),
+        ("image/heic", True),
         ("IMAGE/GIF", True),
         ("IMAGE/PNG", True),
         ("IMAGE/JPEG", True),
+        ("IMAGE/AVIF", True),
+        ("IMAGE/HEIC", True),
         ("image/anything_else", False),
         ("not an image", False),
     ],