Commit 71d451fd authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): add unittest for sd meta writing

parent 3226c2b1

docs/source/api_doc/sd/metadata.rst

+7 −0
Original line number Diff line number Diff line
@@ -28,3 +28,10 @@ get_sdmeta_from_image 






save_image_with_sdmeta

------------------------------------------



.. autofunction:: save_image_with_sdmeta

docs/source/api_doc/sd/nai.rst

+11 −2
Original line number Diff line number Diff line
@@ -6,14 +6,23 @@ imgutils.sd.nai 
.. automodule:: imgutils.sd.nai





NAIMetadata

NAIMetaData

------------------------------------------------



.. autoclass:: NAIMetadata

.. autoclass:: NAIMetaData

    :members: __init__, pnginfo







NAIMetadata

------------------------------------------------



.. class:: NAIMetadata



    Alias of :class:`NAIMetaData`







get_naimeta_from_image

------------------------------------------------

imgutils/sd/__init__.py

+2 −2
Original line number Diff line number Diff line
@@ -17,6 +17,6 @@ Submodules: 
For detailed usage of each function and class, please refer to their individual docstrings.

"""



from .metadata import parse_sdmeta_from_text, get_sdmeta_from_image, SDMetaData

from .metadata import parse_sdmeta_from_text, get_sdmeta_from_image, SDMetaData, save_image_with_sdmeta

from .model import read_metadata, save_with_metadata

from .nai import get_naimeta_from_image, NAIMetadata, add_naimeta_to_image, save_image_with_naimeta
 
from .nai import get_naimeta_from_image, NAIMetaData, add_naimeta_to_image, save_image_with_naimeta, NAIMetadata

imgutils/sd/metadata.py

+41 −3
Original line number Diff line number Diff line
@@ -7,15 +7,19 @@ Overview: 
"""

import io

import json

import mimetypes

import os

import re

import textwrap

from dataclasses import dataclass

from typing import Dict, Any, Optional

from typing import Dict, Any, Optional, Union



from PIL import Image

from PIL.PngImagePlugin import PngInfo



from ..data import ImageTyping, load_image

from ..metadata import read_geninfo_parameters, read_geninfo_exif

from ..metadata import read_geninfo_parameters, read_geninfo_exif, write_geninfo_exif, write_geninfo_gif, \

    read_geninfo_gif



_PARAM_PATTERN = re.compile(r'\s*(?P<key>[\w ]+):\s*(?P<value>"(?:\\.|[^\\"])+"|[^,]*)(?:,|$)')

_SIZE_PATTERN = re.compile(r"^(?P<size1>-?\d+)\s*x\s*(?P<size2>-?\d+)$")
@@ -105,6 +109,10 @@ class SDMetaData: 


            return sio.getvalue().strip()



    @property

    def text(self) -> str:

        return self._sdmeta_text()



    @property

    def pnginfo(self) -> PngInfo:

        """
@@ -259,8 +267,38 @@ def get_sdmeta_from_image(image: ImageTyping) -> Optional[SDMetaData]: 
    image = load_image(image, mode=None, force_background=None)

    pnginfo_text = (read_geninfo_parameters(image) or

                    read_geninfo_exif(image) or

                    read_geninfo_parameters(image))

                    read_geninfo_gif(image))

    if pnginfo_text:

        return parse_sdmeta_from_text(pnginfo_text)

    else:

        return None





def _save_png_with_sdmeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: SDMetaData, **kwargs):

    image.save(dst_file, pnginfo=metadata.pnginfo, **kwargs)





def _save_exif_with_sdmeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: SDMetaData, **kwargs):

    write_geninfo_exif(image, dst_file, metadata.text, **kwargs)





def _save_gif_with_sdmeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: SDMetaData, **kwargs):

    write_geninfo_gif(image, dst_file, metadata.text, **kwargs)





_FN_IMG_SAVE = {

    'image/png': _save_png_with_sdmeta,

    'image/jpeg': _save_exif_with_sdmeta,

    'image/webp': _save_exif_with_sdmeta,

    'image/gif': _save_gif_with_sdmeta,

}





def save_image_with_sdmeta(image: ImageTyping, dst_file: Union[str, os.PathLike], metadata: SDMetaData, **kwargs):

    mimetype, _ = mimetypes.guess_type(str(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}.')



    image = load_image(image, mode=None, force_background=None)

    _FN_IMG_SAVE[mimetype](image, dst_file, metadata, **kwargs)

imgutils/sd/nai.py

+18 −15
Original line number Diff line number Diff line
@@ -30,7 +30,7 @@ mimetypes.add_type('image/webp', '.webp') 




@dataclass

class NAIMetadata:

class NAIMetaData:

    """

    A dataclass representing Novel AI (NAI) metadata.
@@ -95,6 +95,9 @@ class NAIMetadata: 
        return info





NAIMetadata = NAIMetaData





class _InvalidNAIMetaError(Exception):

    """

    Custom exception raised when NAI metadata is invalid.
@@ -158,7 +161,7 @@ def _get_naimeta_raw(image: ImageTyping) -> dict: 
        raise _InvalidNAIMetaError





def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]:

def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetaData]:

    """

    Extract and create a NAIMetadata object from an image.
@@ -168,14 +171,14 @@ def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]: 
    :param image: The input image.

    :type image: ImageTyping

    :return: A NAIMetadata object if successful, None otherwise.

    :rtype: Optional[NAIMetadata]

    :rtype: Optional[NAIMetaData]

    """

    try:

        data = _get_naimeta_raw(image)

    except _InvalidNAIMetaError:

        return None

    else:

        return NAIMetadata(

        return NAIMetaData(

            software=data['Software'],

            source=data['Source'],

            parameters=json.loads(data['Comment']),
@@ -185,14 +188,14 @@ def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]: 
        )





def add_naimeta_to_image(image: ImageTyping, metadata: NAIMetadata) -> Image.Image:

def add_naimeta_to_image(image: ImageTyping, metadata: NAIMetaData) -> Image.Image:

    """

    Add NAI metadata to an image using LSB (Least Significant Bit) encoding.



    :param image: The input image.

    :type image: ImageTyping

    :param metadata: The NAIMetadata object to add to the image.

    :type metadata: NAIMetadata

    :type metadata: NAIMetaData

    :return: The image with added metadata.

    :rtype: Image.Image

    """
@@ -200,7 +203,7 @@ def add_naimeta_to_image(image: ImageTyping, metadata: NAIMetadata) -> Image.Ima 
    return write_lsb_metadata(image, data=metadata.pnginfo)





def _save_png_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetadata, **kwargs):

def _save_png_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetaData, **kwargs):

    """

    Save a PNG image with NAI metadata.
@@ -209,13 +212,13 @@ def _save_png_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike] 
    :param dst_file: The destination file path.

    :type dst_file: Union[str, os.PathLike]

    :param metadata: The NAIMetadata object to include in the image.

    :type metadata: NAIMetadata

    :type metadata: NAIMetaData

    :param kwargs: Additional keyword arguments for image saving.

    """

    image.save(dst_file, pnginfo=metadata.pnginfo, **kwargs)





def _save_exif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetadata, **kwargs):

def _save_exif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetaData, **kwargs):

    """

    Save an image with NAI metadata in EXIF format.
@@ -224,13 +227,13 @@ def _save_exif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike 
    :param dst_file: The destination file path.

    :type dst_file: Union[str, os.PathLike]

    :param metadata: The NAIMetadata object to include in the image.

    :type metadata: NAIMetadata

    :type metadata: NAIMetaData

    :param kwargs: Additional keyword arguments for image saving.

    """

    write_geninfo_exif(image, dst_file, json.dumps(metadata.json), **kwargs)





def _save_gif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetadata, **kwargs):

def _save_gif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetaData, **kwargs):

    """

    Save a GIF image with NAI metadata.
@@ -239,7 +242,7 @@ def _save_gif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike] 
    :param dst_file: The destination file path.

    :type dst_file: Union[str, os.PathLike]

    :param metadata: The NAIMetadata object to include in the image.

    :type metadata: NAIMetadata

    :type metadata: NAIMetaData

    :param kwargs: Additional keyword arguments for image saving.

    """

    write_geninfo_gif(image, dst_file, json.dumps(metadata.json), **kwargs)
@@ -255,7 +258,7 @@ _LSB_ALLOWED_TYPES = {'image/png', 'image/tiff'} 




def save_image_with_naimeta(

        image: ImageTyping, dst_file: Union[str, os.PathLike], metadata: NAIMetadata,

        image: ImageTyping, dst_file: Union[str, os.PathLike], metadata: NAIMetaData,

        add_lsb_meta: Union[str, bool] = 'auto', save_metainfo: Union[str, bool] = 'auto', **kwargs) -> Image.Image:

    """

    Save an image with NAI metadata.
@@ -268,7 +271,7 @@ def save_image_with_naimeta( 
    :param dst_file: The destination file path.

    :type dst_file: Union[str, os.PathLike]

    :param metadata: The NAIMetadata object to include in the image.

    :type metadata: NAIMetadata

    :type metadata: NAIMetaData

    :param add_lsb_meta: Whether to add LSB metadata. Can be 'auto', True, or False.

    :type add_lsb_meta: Union[str, bool]

    :param save_metainfo: Whether to save metainfo. Can be 'auto', True, or False.
@@ -279,7 +282,7 @@ def save_image_with_naimeta( 
    :raises ValueError: If LSB metadata cannot be saved to the specified image format.

    :raises SystemError: If the image format is not supported for saving metainfo.

    """

    mimetype, _ = mimetypes.guess_type(dst_file)

    mimetype, _ = mimetypes.guess_type(str(dst_file))

    if add_lsb_meta == 'auto':

        if mimetype in _LSB_ALLOWED_TYPES:

            add_lsb_meta = True