From 7e27df835c32bdc8b7c4a2c7570df5a02c7c4b9d Mon Sep 17 00:00:00 2001 From: Ruin0x11 Date: Fri, 7 May 2021 21:20:42 -0700 Subject: [PATCH] Add AVIF/HEIF/HEIC upload support --- client/html/post_merge_side.tpl | 1 + client/html/post_readonly_sidebar.tpl | 1 + client/js/views/post_upload_view.js | 2 + server/szurubooru/func/images.py | 16 ++++- server/szurubooru/func/mime.py | 14 +++- .../assets/{heic-mif1.heic => heic-heix.heic} | Bin 2806 -> 2806 bytes .../szurubooru/tests/assets/heif-similar.heif | Bin 3480 -> 3480 bytes server/szurubooru/tests/assets/heif.heif | Bin 3489 -> 3489 bytes server/szurubooru/tests/func/test_mime.py | 30 +++++++- server/szurubooru/tests/func/test_posts.py | 67 ++++++++++++++++++ 10 files changed, 127 insertions(+), 4 deletions(-) rename server/szurubooru/tests/assets/{heic-mif1.heic => heic-heix.heic} (96%) diff --git a/client/html/post_merge_side.tpl b/client/html/post_merge_side.tpl index ad50a2e..372b5ed 100644 --- a/client/html/post_merge_side.tpl +++ b/client/html/post_merge_side.tpl @@ -37,6 +37,7 @@ 'image/png': 'PNG', 'image/webp': 'WEBP', 'image/avif': 'AVIF', + 'image/heif': 'HEIF', 'image/heic': 'HEIC', 'video/webm': 'WEBM', 'video/mp4': 'MPEG-4', diff --git a/client/html/post_readonly_sidebar.tpl b/client/html/post_readonly_sidebar.tpl index 01a086d..cce33bf 100644 --- a/client/html/post_readonly_sidebar.tpl +++ b/client/html/post_readonly_sidebar.tpl @@ -10,6 +10,7 @@ 'image/png': 'PNG', 'image/webp': 'WEBP', 'image/avif': 'AVIF', + 'image/heif': 'HEIF', 'image/heic': 'HEIC', 'video/webm': 'WEBM', 'video/mp4': 'MPEG-4', diff --git a/client/js/views/post_upload_view.js b/client/js/views/post_upload_view.js index 8e40435..8c506a8 100644 --- a/client/js/views/post_upload_view.js +++ b/client/js/views/post_upload_view.js @@ -16,6 +16,7 @@ function _mimeTypeToPostType(mimeType) { "image/png": "image", "image/webp": "image", "image/avif": "image", + "image/heif": "image", "image/heic": "image", "video/mp4": "video", "video/webm": "video", @@ -112,6 +113,7 @@ class Url extends Uploadable { gif: "image/gif", webp: "image/webp", avif: "image/avif", + heif: "image/heif", heic: "image/heic", mp4: "video/mp4", webm: "video/webm", diff --git a/server/szurubooru/func/images.py b/server/szurubooru/func/images.py index 6413ac8..101bba8 100644 --- a/server/szurubooru/func/images.py +++ b/server/szurubooru/func/images.py @@ -4,7 +4,9 @@ import math import re import shlex import subprocess +from io import BytesIO from typing import List +from PIL import Image as PILImage from szurubooru import errors from szurubooru.func import mime, util @@ -12,6 +14,13 @@ from szurubooru.func import mime, util logger = logging.getLogger(__name__) +def convert_heif_to_png(content: bytes) -> bytes: + img = PILImage.open(BytesIO(content)) + img_byte_arr = BytesIO() + img.save(img_byte_arr, format='PNG') + return img_byte_arr.getvalue() + + class Image: def __init__(self, content: bytes) -> None: self.content = content @@ -252,7 +261,12 @@ class Image: ignore_error_if_data: bool = False, get_logs: bool = False, ) -> bytes: - extension = mime.get_extension(mime.get_mime_type(self.content)) + mime_type = mime.get_mime_type(self.content) + if mime.is_heif(mime_type): + # FFmpeg does not support HEIF. + # https://trac.ffmpeg.org/ticket/6521 + self.content = convert_heif_to_png(self.content) + extension = mime.get_extension(mime_type) assert extension with util.create_temp_file(suffix="." + extension) as handle: handle.write(self.content) diff --git a/server/szurubooru/func/mime.py b/server/szurubooru/func/mime.py index f1d9fbb..45eb110 100644 --- a/server/szurubooru/func/mime.py +++ b/server/szurubooru/func/mime.py @@ -24,7 +24,10 @@ def get_mime_type(content: bytes) -> str: if content[4:12] in (b"ftypavif", b"ftypavis"): return "image/avif" - if content[4:12] in (b"ftypheic", b"ftypmif1"): + if content[4:12] == b"ftypmif1": + return "image/heif" + + if content[4:12] in (b"ftypheic", b"ftypheix"): return "image/heic" if content[0:4] == b"\x1A\x45\xDF\xA3": @@ -44,6 +47,7 @@ def get_extension(mime_type: str) -> Optional[str]: "image/png": "png", "image/webp": "webp", "image/avif": "avif", + "image/heif": "heif", "image/heic": "heic", "video/mp4": "mp4", "video/webm": "webm", @@ -67,6 +71,7 @@ def is_image(mime_type: str) -> bool: "image/gif", "image/webp", "image/avif", + "image/heif", "image/heic", ) @@ -77,3 +82,10 @@ def is_animated_gif(content: bytes) -> bool: get_mime_type(content) == "image/gif" and len(re.findall(pattern, content)) > 1 ) + +def is_heif(mime_type: str) -> bool: + return mime_type.lower() in ( + "image/heif", + "image/heic", + "image/avif", + ) diff --git a/server/szurubooru/tests/assets/heic-mif1.heic b/server/szurubooru/tests/assets/heic-heix.heic similarity index 96% rename from server/szurubooru/tests/assets/heic-mif1.heic rename to server/szurubooru/tests/assets/heic-heix.heic index b370c6b9d298761d826dd8245b1dd6b62a503548..27d1268159f6c5c05f3ed1bda22e2c03acf661f7 100644 GIT binary patch delta 21 ccmew+`c0IFfq_9Ht)#LbBQ>*PBhM=?08Ok1T>t<8 delta 21 ccmew+`c0IFfq_9Ht)#LbH#5y}BhM=?08D!a8~^|S diff --git a/server/szurubooru/tests/assets/heif-similar.heif b/server/szurubooru/tests/assets/heif-similar.heif index 26acdc9f0c8406e03340766d9639105096dcf8b0..54385f22617aa5b7deb25dcff65bbbba76a19502 100644 GIT binary patch delta 21 ccmbOsJwuv@fq_9Ht)#LbH#5y}BTo-606q2w5dZ)H delta 21 ccmbOsJwuv@fq_9Ht)#LbBQ-O5BTo-606xA3Jpcdz diff --git a/server/szurubooru/tests/assets/heif.heif b/server/szurubooru/tests/assets/heif.heif index ad1fd77c9f40772d189b96535ebefc4a43a0449e..7fec72b96cc5d352e9f08b1358ce15a820d5ce66 100644 GIT binary patch delta 21 ccmZ1|y-=Emfq_9Ht)#LbH#5y}BhOS`06%aBEC2ui delta 21 ccmZ1|y-=Emfq_9Ht)#LbBQ-O5BhOS`06;hfSO5S3 diff --git a/server/szurubooru/tests/func/test_mime.py b/server/szurubooru/tests/func/test_mime.py index d5238a1..da0d183 100644 --- a/server/szurubooru/tests/func/test_mime.py +++ b/server/szurubooru/tests/func/test_mime.py @@ -15,9 +15,9 @@ from szurubooru.func import mime ("webp.webp", "image/webp"), ("avif.avif", "image/avif"), ("avif-avis.avif", "image/avif"), + ("heif.heif", "image/heif"), ("heic.heic", "image/heic"), - ("heic-mif1.heic", "image/heic"), - ("heif.heif", "image/heic"), + ("heic-heix.heic", "image/heic"), ("text.txt", "application/octet-stream"), ], ) @@ -40,6 +40,7 @@ def test_get_mime_type_for_empty_file(): ("image/gif", "gif"), ("image/webp", "webp"), ("image/avif", "avif"), + ("image/heif", "heif"), ("image/heic", "heic"), ("application/octet-stream", "dat"), ], @@ -84,11 +85,13 @@ def test_is_video(input_mime_type, expected_state): ("image/jpeg", True), ("image/avif", True), ("image/heic", True), + ("image/heif", True), ("IMAGE/GIF", True), ("IMAGE/PNG", True), ("IMAGE/JPEG", True), ("IMAGE/AVIF", True), ("IMAGE/HEIC", True), + ("IMAGE/HEIF", True), ("image/anything_else", False), ("not an image", False), ], @@ -106,3 +109,26 @@ def test_is_image(input_mime_type, expected_state): ) def test_is_animated_gif(read_asset, input_path, expected_state): assert mime.is_animated_gif(read_asset(input_path)) == expected_state + + +@pytest.mark.parametrize( + "input_mime_type,expected_state", + [ + ("image/gif", False), + ("image/png", False), + ("image/jpeg", False), + ("image/avif", True), + ("image/heic", True), + ("image/heif", True), + ("IMAGE/GIF", False), + ("IMAGE/PNG", False), + ("IMAGE/JPEG", False), + ("IMAGE/AVIF", True), + ("IMAGE/HEIC", True), + ("IMAGE/HEIF", True), + ("image/anything_else", False), + ("not an image", False), + ], +) +def test_is_heif(input_mime_type, expected_state): + assert mime.is_heif(input_mime_type) == expected_state diff --git a/server/szurubooru/tests/func/test_posts.py b/server/szurubooru/tests/func/test_posts.py index 6139555..609949f 100644 --- a/server/szurubooru/tests/func/test_posts.py +++ b/server/szurubooru/tests/func/test_posts.py @@ -392,6 +392,41 @@ def test_update_post_source_with_too_long_string(): model.Post.TYPE_IMAGE, "1_244c8840887984c4.gif", ), + ( + False, + "avif.avif", + "image/avif", + model.Post.TYPE_IMAGE, + "1_244c8840887984c4.avif", + ), + ( + False, + "avif-avis.avif", + "image/avif", + model.Post.TYPE_IMAGE, + "1_244c8840887984c4.avif", + ), + ( + False, + "heic.heic", + "image/heic", + model.Post.TYPE_IMAGE, + "1_244c8840887984c4.heic", + ), + ( + False, + "heic-heix.heic", + "image/heic", + model.Post.TYPE_IMAGE, + "1_244c8840887984c4.heic", + ), + ( + False, + "heif.heif", + "image/heif", + model.Post.TYPE_IMAGE, + "1_244c8840887984c4.heif", + ), ( False, "gif-animated.gif", @@ -699,6 +734,38 @@ def test_update_post_content_leaving_custom_thumbnail( assert os.path.exists(generated_path) +@pytest.mark.parametrize("filename", ("avif.avif", "heic.heic", "heif.heif")) +def test_update_post_content_convert_heif_to_png_when_processing( + tmpdir, config_injector, read_asset, post_factory, filename +): + config_injector( + { + "data_dir": str(tmpdir.mkdir("data")), + "thumbnails": { + "post_width": 300, + "post_height": 300, + }, + "secret": "test", + "allow_broken_uploads": False, + } + ) + post = post_factory(id=1) + db.session.add(post) + posts.update_post_content(post, read_asset(filename)) + posts.update_post_thumbnail(post, read_asset(filename)) + db.session.flush() + generated_path = ( + "{}/data/generated-thumbnails/".format(tmpdir) + + "1_244c8840887984c4.jpg" + ) + source_path = ( + "{}/data/posts/custom-thumbnails/".format(tmpdir) + + "1_244c8840887984c4.dat" + ) + assert os.path.exists(source_path) + assert os.path.exists(generated_path) + + def test_update_post_tags(tag_factory): post = model.Post() with patch("szurubooru.func.tags.get_or_create_tags_by_names"):