moving add_thumbnails method to thumbnails file

Author: bigcat88

pillow_heif/__init__.py

@@ -43,4 +43,4 @@
     read_heif,
 )
 from .misc import get_file_mimetype, getxmp, set_orientation
-from .thumbnails import thumbnail
+from .thumbnails import add_thumbnails, thumbnail

pillow_heif/heif.py

@@ -348,38 +348,6 @@ def scale(self, width: int, height: int):
         self._img_to_img_data_dict(scaled_heif_img)
         return self
 
-    def add_thumbnails(self, boxes: Union[List[int], int]) -> None:
-        """Add thumbnail(s) to an image.
-
-        .. note:: Method creates thumbnails without image data, they will be encoded during `save` operation.
-
-        :param boxes: int or list of ints determining size of thumbnail(s) to generate for image.
-
-        :returns: None"""
-
-        if isinstance(boxes, list):
-            boxes_list = boxes
-        else:
-            boxes_list = [boxes]
-        self.load()
-        for box in boxes_list:
-            if box <= 3:
-                continue
-            if self.size[0] <= box and self.size[1] <= box:
-                continue
-            if self.size[0] > self.size[1]:
-                thumb_height = int(self.size[1] * box / self.size[0])
-                thumb_width = box
-            else:
-                thumb_width = int(self.size[0] * box / self.size[1])
-                thumb_height = box
-            thumb_height = thumb_height - 1 if (thumb_height & 1) else thumb_height
-            thumb_width = thumb_width - 1 if (thumb_width & 1) else thumb_width
-            if max((thumb_height, thumb_width)) in [max(i.size) for i in self.thumbnails]:
-                continue
-            __heif_ctx = HeifCtxAsDict(self.mode, (thumb_width, thumb_height), None, stride=0)
-            self.thumbnails.append(HeifThumbnail(__heif_ctx, self))
-
     def copy_thumbnails(self, thumbnails: List[HeifThumbnail]):
         """Private. For use only in ``add_from_pillow`` and ``add_from_heif``."""
 
@@ -634,15 +602,6 @@ def add_from_heif(self, heif_image, load_one=False, ignore_primary=True, **kwarg
                 break
         return self
 
-    def add_thumbnails(self, boxes: Union[List[int], int]) -> None:
-        """Add thumbnail(s) to the primary image.
-
-        :param boxes: int or list of ints determining size of thumbnail(s) to generate.
-
-        :returns: None"""
-
-        self.images[self.primary_index()].add_thumbnails(boxes)
-
     def save(self, fp, **kwargs) -> None:
         """Saves image under the given fp.
 

pillow_heif/thumbnails.py

@@ -1,7 +1,12 @@
 """Functions for work with thumbnails."""
 
+from typing import List, Union
+
+from PIL import Image
+
 from .as_opener import HeifImageFile
-from .heif import HeifFile, HeifImage
+from .heif import HeifFile, HeifImage, HeifThumbnail
+from .private import HeifCtxAsDict
 
 
 def thumbnail(im, min_box: int = 0):
@@ -28,6 +33,55 @@ def thumbnail(im, min_box: int = 0):
     return thumb if thumb else im
 
 
+def add_thumbnails(im: Union[HeifFile, HeifImage, Image.Image], boxes: Union[List[int], int]) -> None:
+    """Add thumbnail(s) to an image(s).
+
+    .. note:: Method creates thumbnails without image data, they will be encoded during `save` operation.
+
+    :param im: For ``HeifFile`` will add thumbnail(s) to all images in it.
+        For ``HeifImage`` or a Pillow image it will add thumbnails only to specified image.
+    :param boxes: int or list of ints determining size of thumbnail(s) to generate for image.
+
+    :returns: None"""
+
+    if isinstance(im, HeifFile):
+        for i in im:
+            _add_thumbnails(i, boxes)
+    else:
+        _add_thumbnails(im, boxes)
+
+
+def _add_thumbnails(im: Union[HeifImage, Image.Image], boxes: Union[List[int], int]) -> None:
+    if isinstance(boxes, list):
+        boxes_list = boxes
+    else:
+        boxes_list = [boxes]
+    for box in boxes_list:
+        if box <= 3:
+            continue
+        if im.size[0] <= box and im.size[1] <= box:
+            continue
+        if im.size[0] > im.size[1]:
+            thumb_height = int(im.size[1] * box / im.size[0])
+            thumb_width = box
+        else:
+            thumb_width = int(im.size[0] * box / im.size[1])
+            thumb_height = box
+        thumb_height = thumb_height - 1 if (thumb_height & 1) else thumb_height
+        thumb_width = thumb_width - 1 if (thumb_width & 1) else thumb_width
+        thumbnails = im.thumbnails if isinstance(im, HeifImage) else im.info.get("thumbnails", [])
+        if max((thumb_height, thumb_width)) in [max(i.size) for i in thumbnails]:
+            continue
+        __heif_ctx = HeifCtxAsDict(im.mode, (thumb_width, thumb_height), None, stride=0)
+        if isinstance(im, HeifImage):
+            im.thumbnails.append(HeifThumbnail(__heif_ctx, im))
+        else:
+            if im.info.get("thumbnails", None) is None:
+                im.info["thumbnails"] = [HeifThumbnail(__heif_ctx, __heif_ctx)]
+            else:
+                im.info["thumbnails"].append(HeifThumbnail(__heif_ctx, __heif_ctx))
+
+
 def _choose_thumbnail(thumbnails: list, min_box: int):
     for thumb in thumbnails:
         if thumb.size[0] >= min_box or thumb.size[1] >= min_box:

tests/helpers.py

@@ -4,7 +4,7 @@
 
 from PIL import Image, ImageMath, ImageOps
 
-from pillow_heif import HeifFile, HeifImage, HeifThumbnail
+from pillow_heif import HeifFile, HeifImage, HeifThumbnail, add_thumbnails
 
 try:
     import numpy as np
@@ -128,7 +128,7 @@ def create_heif(size: tuple = None, thumb_boxes: list = None, n_images=1, **kwar
     for i in range(n_images):
         im_heif.add_from_pillow(Image.effect_mandelbrot(size, (-3, -2.5, 2, 2.5), 100))
         size = (int(size[0] / 2), int(size[1] / 2))
-        im_heif[i].add_thumbnails(thumb_boxes)
+        add_thumbnails(im_heif[i], thumb_boxes)
     heif_buf = BytesIO()
     im_heif.save(heif_buf, **kwargs)
     return heif_buf

tests/numpy_test.py

@@ -15,7 +15,7 @@
 # Creating HEIF file in memory with 1 image and 3 thumbnails.
 im_pillow = Image.effect_mandelbrot((512, 512), (-3, -2.5, 2, 2.5), 100)
 im_heif = pillow_heif.from_pillow(im_pillow)
-im_heif.add_thumbnails(boxes=[192, 128, 64])
+pillow_heif.add_thumbnails(im_heif, boxes=[192, 128, 64])
 heif_buf = BytesIO()
 im_heif.save(heif_buf)
 

tests/opencv_test.py

@@ -8,7 +8,8 @@
 from helpers import compare_hashes, gradient_rgb_bytes, gradient_rgba_bytes
 from PIL import Image
 
-from pillow_heif import from_bytes, open_heif, options, register_heif_opener
+from pillow_heif import HeifImagePlugin  # noqa
+from pillow_heif import from_bytes, open_heif, options
 
 if not options().hevc_enc:
     pytest.skip("No HEVC encoder.", allow_module_level=True)
@@ -17,7 +18,6 @@
 cv2 = pytest.importorskip("cv2", reason="OpenCV not installed")
 
 os.chdir(os.path.dirname(os.path.abspath(__file__)))
-register_heif_opener()
 
 
 @pytest.mark.parametrize("enc_bits", (10, 12))

tests/thumbnails_test.py

@@ -19,8 +19,8 @@ def create_thumbnail_heif(size):
     _ = Image.effect_mandelbrot(size, (-3, -2.5, 2, 2.5), 100)
     im_heif = pillow_heif.from_pillow(_)
     im_heif.add_from_pillow(_.crop((0, 0, 256, 256)))
-    im_heif[0].add_thumbnails(boxes=[128, 64])
-    im_heif[1].add_thumbnails(boxes=[128, 64])
+    pillow_heif.add_thumbnails(im_heif[0], boxes=[128, 64])
+    pillow_heif.add_thumbnails(im_heif[1], boxes=[128, 64])
     im_heif.add_from_pillow(_.crop((0, 0, 192, 192)))
     _heif_buf = BytesIO()
     exif = Image.Exif()
@@ -216,12 +216,16 @@ def test_heif_thumbnail_references():
     ),
 )
 @pytest.mark.parametrize("heif_file_buf", (heif_buf, create_thumbnail_heif((317, 311))))
-def test_heif_image_add_thumbs(thumbs, expected_after, heif_file_buf):
+@pytest.mark.parametrize("method", ("HeifFile", "HeifImage"))
+def test_heif_add_thumbs(thumbs, expected_after, heif_file_buf, method):
     output = BytesIO()
     heif_file = pillow_heif.open_heif(heif_file_buf)
-    heif_file[0].add_thumbnails(thumbs)
-    heif_file[1].add_thumbnails(thumbs)
-    heif_file[2].add_thumbnails(thumbs)
+    if method == "HeifFile":
+        pillow_heif.add_thumbnails(heif_file, thumbs)
+    else:
+        pillow_heif.add_thumbnails(heif_file[0], thumbs)
+        pillow_heif.add_thumbnails(heif_file[1], thumbs)
+        pillow_heif.add_thumbnails(heif_file[2], thumbs)
     heif_file.save(output, quality=-1)
     out_heif = pillow_heif.open_heif(output)
     for i in range(3):
@@ -232,25 +236,48 @@ def test_heif_image_add_thumbs(thumbs, expected_after, heif_file_buf):
 @pytest.mark.parametrize(
     "thumbs,expected_after",
     (
-        (-1, 2),
-        (64, 2),
-        (96, 3),
-        ([96], 3),
-        ([0, 84], 3),
-        ([96, 84], 4),
+        (-1, [2, 2, 0]),
+        ([0], [2, 2, 0]),
+        (64, [2, 2, 1]),
+        ([2048], [2, 2, 0]),
+        (96, [3, 3, 1]),
+        ([0, 84], [3, 3, 1]),
+        ([96, 84], [4, 4, 2]),
     ),
 )
-@pytest.mark.parametrize("heif_file_buf", (heif_buf, create_thumbnail_heif((537, 511))))
-def test_heif_primary_add_thumbs(thumbs, expected_after, heif_file_buf):
-    heif_file = pillow_heif.open_heif(heif_file_buf)
-    heif_file.add_thumbnails(thumbs)
+def test_pillow_add_thumbs(thumbs, expected_after):
     output = BytesIO()
-    heif_file.save(output, quality=-1)
+    im = Image.open(heif_buf)
+    for frame in ImageSequence.Iterator(im):
+        pillow_heif.add_thumbnails(frame, thumbs)
+    im.save(output, format="HEIF", quality=-1, save_all=True)
+    out_heif = pillow_heif.open_heif(output)
+    for i in range(3):
+        assert len(out_heif[i].thumbnails) == expected_after[i]
+    compare_hashes([out_heif[0].to_pillow(), out_heif[0].thumbnails[0].to_pillow()], hash_size=8, max_difference=4)
+
+
+@pytest.mark.parametrize(
+    "thumbs,expected_after",
+    (
+        (-1, 0),
+        ([0], 0),
+        (64, 1),
+        ([2048], 0),
+        ([0, 128], 1),
+        ([64, 128], 2),
+    ),
+)
+def test_pillow_add_thumbs_empty_info(thumbs, expected_after):
+    output = BytesIO()
+    im = Image.open(heif_buf)
+    im.info.pop("thumbnails")
+    pillow_heif.add_thumbnails(im, thumbs)
+    im.save(output, format="HEIF", quality=-1)
     out_heif = pillow_heif.open_heif(output)
-    assert len(out_heif[0].thumbnails) == len(heif_file[0].thumbnails)
-    assert len(out_heif[1].thumbnails) == expected_after
     assert len(out_heif.thumbnails) == expected_after
-    assert len(out_heif[2].thumbnails) == len(heif_file[2].thumbnails)
+    if expected_after >= 1:
+        compare_hashes([out_heif[0].to_pillow(), out_heif.thumbnails[0].to_pillow()], hash_size=8)
 
 
 def test_heif_remove_thumbs():