dev(narugo): save the damn code

imgutils/metadata/__init__.py

@@ -1,2 +1,3 @@
-from .geninfo import read_geninfo_gif, read_geninfo_parameters, read_geninfo_exif
+from .geninfo import read_geninfo_parameters, read_geninfo_exif, read_geninfo_gif, \
+    write_geninfo_parameters, write_geninfo_exif, write_geninfo_gif
 from .lsb import read_lsb_raw_bytes, read_lsb_metadata, write_lsb_raw_bytes, write_lsb_metadata, LSBReadError

imgutils/metadata/geninfo.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 import piexif
+from PIL.PngImagePlugin import PngInfo
 from piexif.helper import UserComment
 
 from ..data import ImageTyping, load_image
@@ -41,3 +42,26 @@ def read_geninfo_gif(image: ImageTyping) -> Optional[str]:
         return infos["comment"].decode("utf8", errors="ignore")
     else:
         return None
+
+
+def write_geninfo_parameters(image: ImageTyping, dst_filename: str, geninfo: str, **kwargs):
+    pnginfo = PngInfo()
+    pnginfo.add_text('parameters', geninfo)
+
+    image = load_image(image, force_background=None, mode=None)
+    image.save(dst_filename, pnginfo=pnginfo, *kwargs)
+
+
+def write_geninfo_exif(image: ImageTyping, dst_filename: str, geninfo: str, **kwargs):
+    exif_dict = {
+        "Exif": {piexif.ExifIFD.UserComment: UserComment.dump(geninfo, encoding="unicode")}}
+    exif_bytes = piexif.dump(exif_dict)
+
+    image = load_image(image, force_background=None, mode=None)
+    image.save(dst_filename, exif=exif_bytes, *kwargs)
+
+
+def write_geninfo_gif(image: ImageTyping, dst_filename: str, geninfo: str, **kwargs):
+    image = load_image(image, force_background=None, mode=None)
+    image.info['comment'] = geninfo.encode('utf-8')
+    image.save(dst_filename, *kwargs)

imgutils/sd/nai.py

@@ -13,6 +13,7 @@
 """
 
 import json
+import mimetypes
 import os
 import warnings
 from dataclasses import dataclass
@@ -23,7 +24,9 @@
 
 from ..data import load_image, ImageTyping
 from ..metadata import read_lsb_metadata, write_lsb_metadata, LSBReadError, read_geninfo_parameters, \
-    read_geninfo_exif, read_geninfo_gif
+    read_geninfo_exif, read_geninfo_gif, write_geninfo_exif, write_geninfo_gif
+
+mimetypes.add_type('image/webp', '.webp')
 
 
 @dataclass
@@ -54,6 +57,21 @@ class NAIMetadata:
     generation_time: Optional[float] = None
     description: Optional[str] = None
 
+    @property
+    def json(self) -> dict:
+        data = {
+            'Software': self.software,
+            'Source': self.source,
+            'Comment': json.dumps(self.parameters),
+        }
+        if self.title is not None:
+            data['Title'] = self.title
+        if self.generation_time is not None:
+            data['Generation time'] = json.dumps(self.generation_time)
+        if self.description is not None:
+            data['Description'] = self.description
+        return data
+
     @property
     def pnginfo(self) -> PngInfo:
         """
@@ -66,16 +84,8 @@ def pnginfo(self) -> PngInfo:
         :rtype: PngInfo
         """
         info = PngInfo()
-        info.add_text('Software', self.software)
-        info.add_text('Source', self.source)
-        if self.title is not None:
-            info.add_text('Title', self.title)
-        if self.generation_time is not None:
-            info.add_text('Generation time', json.dumps(self.generation_time)),
-        if self.description is not None:
-            info.add_text('Description', self.description)
-        if self.parameters is not None:
-            info.add_text('Comment', json.dumps(self.parameters))
+        for key, value in self.json.items():
+            info.add_text(key, value)
         return info
 
 
@@ -158,81 +168,60 @@ def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]:
         )
 
 
-def _get_pnginfo(metadata: Union[NAIMetadata, PngInfo]) -> PngInfo:
-    """
-    Convert metadata to PngInfo object.
+def add_naimeta_to_image(image: ImageTyping, metadata: NAIMetadata) -> Image.Image:
+    image = load_image(image, mode=None, force_background=None)
+    return write_lsb_metadata(image, data=metadata.pnginfo)
 
-    This function takes either a NAIMetadata object or a PngInfo object and returns a PngInfo object.
 
-    :param metadata: The metadata to convert.
-    :type metadata: Union[NAIMetadata, PngInfo]
+def _save_png_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetadata, **kwargs):
+    image.save(dst_file, pnginfo=metadata.pnginfo, **kwargs)
 
-    :return: A PngInfo object.
-    :rtype: PngInfo
 
-    :raises TypeError: If the metadata is neither NAIMetadata nor PngInfo.
-    """
-    if isinstance(metadata, NAIMetadata):
-        pnginfo = metadata.pnginfo
-    elif isinstance(metadata, PngInfo):
-        pnginfo = metadata
-    else:
-        raise TypeError(f'Unknown metadata type for NAI - {metadata!r}.')  # pragma: no cover
-    return pnginfo
+def _save_exif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetadata, **kwargs):
+    write_geninfo_exif(image, dst_file, json.dumps(metadata.json), **kwargs)
 
 
-def add_naimeta_to_image(image: ImageTyping, metadata: Union[NAIMetadata, PngInfo]) -> Image.Image:
-    """
-    Add NAI metadata to an image.
+def _save_gif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetadata, **kwargs):
+    write_geninfo_gif(image, dst_file, json.dumps(metadata.json), **kwargs)
 
-    This function injects the provided metadata into the image using LSB injection.
 
-    :param image: The input image.
-    :type image: ImageTyping
-    :param metadata: The metadata to add to the image.
-    :type metadata: Union[NAIMetadata, PngInfo]
+_FN_IMG_SAVE = {
+    'image/png': _save_png_with_naimeta,
+    'image/jpeg': _save_exif_with_naimeta,
+    'image/webp': _save_exif_with_naimeta,
+    'image/tiff': _save_exif_with_naimeta,
+    'image/gif': _save_gif_with_naimeta,
+}
+_LSB_ALLOWED_TYPES = {'image/png', 'image/tiff', 'image/gif', 'image/bmp'}
 
-    :return: The image with added metadata.
-    :rtype: Image.Image
-    """
-    pnginfo = _get_pnginfo(metadata)
-    image = load_image(image, mode=None, force_background=None)
-    return write_lsb_metadata(image, data=pnginfo)
-
-
-def save_image_with_naimeta(image: ImageTyping, dst_file: Union[str, os.PathLike],
-                            metadata: Union[NAIMetadata, PngInfo],
-                            add_lsb_meta: bool = True, save_pnginfo: bool = True, **kwargs) -> Image.Image:
-    """
-    Save an image with NAI metadata.
 
-    This function saves the given image to a file, optionally adding NAI metadata using LSB injection
-    and/or saving it as PNG metadata.
+def save_image_with_naimeta(
+        image: ImageTyping, dst_file: Union[str, os.PathLike], metadata: NAIMetadata,
+        add_lsb_meta: Union[str, bool] = 'auto', save_metainfo: bool = True, **kwargs) -> Image.Image:
+    mimetype, _ = mimetypes.guess_type(dst_file)
+    if add_lsb_meta == 'auto':
+        if mimetype in _LSB_ALLOWED_TYPES:
+            add_lsb_meta = True
+        else:
+            add_lsb_meta = False
+    else:
+        if add_lsb_meta and mimetype not in _LSB_ALLOWED_TYPES:
+            raise ValueError('LSB metadata cannot be saved to lossy image format, '
+                             'add_lsb_meta will be disabled. '
+                             f'Only {", ".join(sorted(_LSB_ALLOWED_TYPES))} images supported.')
+    if not add_lsb_meta and not save_metainfo:
+        warnings.warn(f'Both LSB meta and pnginfo is disabled, no metadata will be saved to {dst_file!r}.')
 
-    :param image: The input image.
-    :type image: ImageTyping
-    :param dst_file: The destination file path.
-    :type dst_file: Union[str, os.PathLike]
-    :param metadata: The metadata to add to the image.
-    :type metadata: Union[NAIMetadata, PngInfo]
-    :param add_lsb_meta: Whether to add metadata using LSB injection. Defaults to True.
-    :type add_lsb_meta: bool
-    :param save_pnginfo: Whether to save metadata as PNG metadata. Defaults to True.
-    :type save_pnginfo: bool
-    :param kwargs: Additional keyword arguments to pass to the image save function.
-
-    :return: The saved image.
-    :rtype: Image.Image
-
-    :raises Warning: If both LSB meta and pnginfo are disabled.
-    """
-    pnginfo = _get_pnginfo(metadata)
     image = load_image(image, mode=None, force_background=None)
-    if not add_lsb_meta and not save_pnginfo:
-        warnings.warn(f'Both LSB meta and pnginfo is disabled, no metadata will be saved to {dst_file!r}.')
     if add_lsb_meta:
-        image = add_naimeta_to_image(image, metadata=pnginfo)
-    if save_pnginfo:
-        kwargs['pnginfo'] = pnginfo
-    image.save(dst_file, **kwargs)
+        image = add_naimeta_to_image(image, metadata=metadata)
+    if save_metainfo:
+        mimetype, _ = mimetypes.guess_type(dst_file)
+        if mimetype not in _FN_IMG_SAVE:
+            raise SystemError(f'Not supported to save as a {mimetype!r} type, '
+                              f'supported mimetypes are {sorted(_FN_IMG_SAVE.keys())!r}.')
+        else:
+            _FN_IMG_SAVE[mimetype](image, dst_file, metadata, **kwargs)
+    else:
+        image.save(dst_file, **kwargs)
     return image

test/sd/test_nai.py

@@ -178,17 +178,7 @@ def test_add_naimeta_to_image_rgba(self, nai3_clear_rgba_image, nai3_meta_withou
         image = add_naimeta_to_image(nai3_clear_rgba_image, metadata=nai3_meta_without_title)
         assert get_naimeta_from_image(image) == pytest.approx(nai3_meta_without_title)
 
-    def test_save_image_with_naimeta(self, nai3_clear_file, nai3_meta_without_title):
-        with isolated_directory():
-            save_image_with_naimeta(nai3_clear_file, 'image.png', metadata=nai3_meta_without_title)
-            assert get_naimeta_from_image('image.png') == pytest.approx(nai3_meta_without_title)
-
-    def test_save_image_with_naimeta_rgba(self, nai3_clear_rgba_file, nai3_meta_without_title):
-        with isolated_directory():
-            save_image_with_naimeta(nai3_clear_rgba_file, 'image.png', metadata=nai3_meta_without_title)
-            assert get_naimeta_from_image('image.png') == pytest.approx(nai3_meta_without_title)
-
-    def test_save_image_with_naimeta_pnginfo_only(self, nai3_clear_file, nai3_meta_without_title):
+    def test_save_image_with_naimeta_metainfo_only(self, nai3_clear_file, nai3_meta_without_title):
         with isolated_directory():
             save_image_with_naimeta(nai3_clear_file, 'image.png',
                                     metadata=nai3_meta_without_title, add_lsb_meta=False)
@@ -197,7 +187,7 @@ def test_save_image_with_naimeta_pnginfo_only(self, nai3_clear_file, nai3_meta_w
     def test_save_image_with_naimeta_lsbmeta_only(self, nai3_clear_file, nai3_meta_without_title):
         with isolated_directory():
             save_image_with_naimeta(nai3_clear_file, 'image.png',
-                                    metadata=nai3_meta_without_title, save_pnginfo=False)
+                                    metadata=nai3_meta_without_title, save_metainfo=False)
             assert get_naimeta_from_image('image.png') == pytest.approx(nai3_meta_without_title)
 
     def test_save_image_with_naimeta_both_no(self, nai3_clear_file, nai3_meta_without_title):
@@ -206,7 +196,7 @@ def test_save_image_with_naimeta_both_no(self, nai3_clear_file, nai3_meta_withou
                 save_image_with_naimeta(
                     nai3_clear_file, 'image.png',
                     metadata=nai3_meta_without_title,
-                    save_pnginfo=False, add_lsb_meta=False,
+                    save_metainfo=False, add_lsb_meta=False,
                 )
             assert get_naimeta_from_image('image.png') is None
 
@@ -220,7 +210,7 @@ def test_save_image_with_naimeta_rgba_with_title(self, nai3_clear_rgba_file, nai
             save_image_with_naimeta(nai3_clear_rgba_file, 'image.png', metadata=nai3_meta)
             assert get_naimeta_from_image('image.png') == pytest.approx(nai3_meta)
 
-    def test_save_image_with_naimeta_pnginfo_only_with_title(self, nai3_clear_file, nai3_meta):
+    def test_save_image_with_naimeta_metainfo_only_with_title(self, nai3_clear_file, nai3_meta):
         with isolated_directory():
             save_image_with_naimeta(nai3_clear_file, 'image.png',
                                     metadata=nai3_meta, add_lsb_meta=False)
@@ -229,7 +219,7 @@ def test_save_image_with_naimeta_pnginfo_only_with_title(self, nai3_clear_file,
     def test_save_image_with_naimeta_lsbmeta_only_with_title(self, nai3_clear_file, nai3_meta):
         with isolated_directory():
             save_image_with_naimeta(nai3_clear_file, 'image.png',
-                                    metadata=nai3_meta, save_pnginfo=False)
+                                    metadata=nai3_meta, save_metainfo=False)
             assert get_naimeta_from_image('image.png') == pytest.approx(nai3_meta)
 
     def test_save_image_with_naimeta_both_no_with_title(self, nai3_clear_file, nai3_meta):
@@ -238,7 +228,7 @@ def test_save_image_with_naimeta_both_no_with_title(self, nai3_clear_file, nai3_
                 save_image_with_naimeta(
                     nai3_clear_file, 'image.png',
                     metadata=nai3_meta,
-                    save_pnginfo=False, add_lsb_meta=False,
+                    save_metainfo=False, add_lsb_meta=False,
                 )
             assert get_naimeta_from_image('image.png') is None
 
@@ -251,3 +241,77 @@ def test_image_error_with_wrong_format(self, file):
 
     def test_get_naimeta_from_image_webp(self, nai3_webp_file, nai3_webp_meta):
         assert get_naimeta_from_image(nai3_webp_file) == pytest.approx(nai3_webp_meta)
+
+    @pytest.mark.parametrize(['ext', 'warns', 'okay'], [
+        ('.png', False, True),
+        ('.webp', False, True),
+        ('.jpg', False, True),
+        ('.jpeg', False, True),
+        ('.tiff', False, True),
+        ('.gif', False, True),
+    ])
+    def test_save_image_with_naimeta(self, nai3_clear_file, nai3_meta_without_title,
+                                     ext, warns, okay):
+        with isolated_directory(), pytest.warns(Warning if warns else None):
+            save_image_with_naimeta(nai3_clear_file, f'image{ext}', metadata=nai3_meta_without_title)
+            assert get_naimeta_from_image(f'image{ext}') == \
+                   (pytest.approx(nai3_meta_without_title) if okay else None)
+
+    @pytest.mark.parametrize(['ext', 'warns', 'okay'], [
+        ('.png', False, True),
+        ('.webp', False, True),
+        ('.tiff', False, True),
+        ('.gif', False, True),
+    ])
+    def test_save_image_with_naimeta_rgba(self, nai3_clear_rgba_file, nai3_meta_without_title,
+                                          ext, warns, okay):
+        with isolated_directory(), pytest.warns(Warning if warns else None):
+            save_image_with_naimeta(nai3_clear_rgba_file, f'image{ext}', metadata=nai3_meta_without_title)
+            assert get_naimeta_from_image(f'image{ext}') == \
+                   (pytest.approx(nai3_meta_without_title) if okay else None)
+
+    @pytest.mark.parametrize(['ext'], [
+        ('.webp',),
+        ('.jpg',),
+        ('.jpeg',),
+    ])
+    def test_save_image_with_naimeta_exifs_lsb_true_lossy(self, nai3_clear_file, nai3_meta_without_title, ext):
+        with isolated_directory(), pytest.raises(ValueError):
+            save_image_with_naimeta(nai3_clear_file, f'image{ext}',
+                                    add_lsb_meta=True, metadata=nai3_meta_without_title)
+
+    @pytest.mark.parametrize(['ext'], [
+        ('.tiff',),
+        ('.gif',),
+    ])
+    def test_save_image_with_naimeta_exifs_lsb_true_non_lossy(self, nai3_clear_file, nai3_meta_without_title, ext):
+        with isolated_directory():
+            save_image_with_naimeta(nai3_clear_file, f'image{ext}',
+                                    add_lsb_meta=True, metadata=nai3_meta_without_title)
+            assert get_naimeta_from_image(f'image{ext}') == pytest.approx(nai3_meta_without_title)
+
+    @pytest.mark.parametrize(['ext'], [
+        ('.webp',),
+        ('.jpg',),
+        ('.jpeg',),
+        ('.tiff',),
+        ('.gif',),
+    ])
+    def test_save_image_with_naimeta_metainfo_only_exifs(self, nai3_clear_file, nai3_meta_without_title, ext):
+        with isolated_directory(), pytest.warns(None):
+            save_image_with_naimeta(nai3_clear_file, f'image{ext}',
+                                    metadata=nai3_meta_without_title, add_lsb_meta=False)
+            assert get_naimeta_from_image(f'image{ext}') == pytest.approx(nai3_meta_without_title)
+
+    @pytest.mark.parametrize(['ext'], [
+        ('.webp',),
+        ('.jpg',),
+        ('.jpeg',),
+        ('.tiff',),
+        ('.gif',),
+    ])
+    def test_save_image_with_naimeta_lsbmeta_only_exifs(self, nai3_clear_file, nai3_meta_without_title, ext):
+        with isolated_directory(), pytest.warns(Warning):
+            save_image_with_naimeta(nai3_clear_file, f'image{ext}',
+                                    metadata=nai3_meta_without_title, save_metainfo=False)
+            assert get_naimeta_from_image(f'image{ext}') is None