Loading imgutils/metadata/geninfo.py +101 −3 Original line number Diff line number Diff line """ This module provides functions for reading and writing generation information (geninfo) to image files. It supports different image formats including PNG, EXIF, and GIF. The module includes functions for: 1. Reading geninfo from image parameters, EXIF data, and GIF comments 2. Writing geninfo to image parameters, EXIF data, and GIF comments These functions are useful for storing and retrieving metadata about image generation, particularly in the context of AI-generated images. """ from typing import Optional import piexif Loading @@ -8,12 +21,38 @@ from ..data import ImageTyping, load_image def read_geninfo_parameters(image: ImageTyping) -> Optional[str]: """ Read generation information from image parameters. :param image: The input image. :type image: ImageTyping :return: The generation information if found, None otherwise. :rtype: Optional[str] This function loads the image and attempts to retrieve the 'parameters' information from the image metadata. It's commonly used for PNG images where generation information is stored in the image parameters. """ image = load_image(image, mode=None, force_background=None) infos = image.info or {} return infos.get('parameters') def read_geninfo_exif(image: ImageTyping) -> Optional[str]: """ Read generation information from EXIF data. :param image: The input image. :type image: ImageTyping :return: The generation information if found in EXIF data, None otherwise. :rtype: Optional[str] This function attempts to read generation information from the EXIF metadata of the image. It specifically looks for the UserComment field in the EXIF data. If the EXIF data is invalid or not present, it returns None. """ image = load_image(image, mode=None, force_background=None) infos = image.info or {} if "exif" in infos: Loading @@ -36,6 +75,19 @@ def read_geninfo_exif(image: ImageTyping) -> Optional[str]: def read_geninfo_gif(image: ImageTyping) -> Optional[str]: """ Read generation information from GIF comment. :param image: The input image. :type image: ImageTyping :return: The generation information if found in GIF comment, None otherwise. :rtype: Optional[str] This function is specifically designed to read generation information from GIF images. It looks for the 'comment' field in the image metadata, which is commonly used in GIF files to store additional information. """ image = load_image(image, mode=None, force_background=None) infos = image.info or {} if "comment" in infos: # for gif Loading @@ -45,23 +97,69 @@ def read_geninfo_gif(image: ImageTyping) -> Optional[str]: def write_geninfo_parameters(image: ImageTyping, dst_filename: str, geninfo: str, **kwargs): """ Write generation information to image parameters. :param image: The input image. :type image: ImageTyping :param dst_filename: The destination filename to save the image with geninfo. :type dst_filename: str :param geninfo: The generation information to write. :type geninfo: str :param kwargs: Additional keyword arguments to pass to the image save function. This function writes the provided generation information to the image parameters. It's commonly used for PNG images where generation information can be stored in the image metadata. The function creates a PngInfo object, adds the geninfo as 'parameters', and saves the image with this metadata. """ pnginfo = PngInfo() pnginfo.add_text('parameters', geninfo) image = load_image(image, force_background=None, mode=None) image.save(dst_filename, pnginfo=pnginfo, *kwargs) image.save(dst_filename, pnginfo=pnginfo, **kwargs) def write_geninfo_exif(image: ImageTyping, dst_filename: str, geninfo: str, **kwargs): """ Write generation information to EXIF data. :param image: The input image. :type image: ImageTyping :param dst_filename: The destination filename to save the image with geninfo. :type dst_filename: str :param geninfo: The generation information to write. :type geninfo: str :param kwargs: Additional keyword arguments to pass to the image save function. This function writes the provided generation information to the EXIF metadata of the image. It creates an EXIF dictionary with the geninfo stored in the UserComment field, converts it to bytes, and saves the image with this EXIF data. """ 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) image.save(dst_filename, exif=exif_bytes, **kwargs) def write_geninfo_gif(image: ImageTyping, dst_filename: str, geninfo: str, **kwargs): """ Write generation information to GIF comment. :param image: The input image. :type image: ImageTyping :param dst_filename: The destination filename to save the image with geninfo. :type dst_filename: str :param geninfo: The generation information to write. :type geninfo: str :param kwargs: Additional keyword arguments to pass to the image save function. This function is specifically designed to write generation information to GIF images. It adds the geninfo to the image's 'comment' field, which is a standard way of including metadata in GIF files. """ image = load_image(image, force_background=None, mode=None) image.info['comment'] = geninfo.encode('utf-8') image.save(dst_filename, *kwargs) image.save(dst_filename, **kwargs) imgutils/sd/nai.py +86 −4 Original line number Diff line number Diff line Loading @@ -59,6 +59,12 @@ class NAIMetadata: @property def json(self) -> dict: """ Convert the NAIMetadata to a JSON-compatible dictionary. :return: A dictionary representation of the metadata. :rtype: dict """ data = { 'Software': self.software, 'Source': self.source, Loading Loading @@ -90,10 +96,22 @@ class NAIMetadata: class _InvalidNAIMetaError(Exception): """ Custom exception raised when NAI metadata is invalid. """ pass def _naimeta_validate(data): """ Validate the NAI metadata. :param data: The metadata to validate. :type data: dict :return: The validated metadata. :rtype: dict :raises _InvalidNAIMetaError: If the metadata is invalid. """ if isinstance(data, dict) and data.get('Software') and data.get('Source') and data.get('Comment'): return data else: Loading @@ -104,14 +122,14 @@ def _get_naimeta_raw(image: ImageTyping) -> dict: """ Extract raw NAI metadata from an image. This function attempts to extract metadata from the image using LSB (Least Significant Bit) extraction. If that fails, it falls back to using the image's info dictionary. This function attempts to extract metadata from the image using various methods, including LSB extraction, image info dictionary, and other specific metadata formats. :param image: The input image. :type image: ImageTyping :return: A dictionary containing the raw metadata. :rtype: dict :raises _InvalidNAIMetaError: If no valid metadata is found. """ image = load_image(image, force_background=None, mode=None) try: Loading Loading @@ -149,7 +167,6 @@ 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] """ Loading @@ -169,19 +186,62 @@ def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]: 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 :return: The image with added metadata. :rtype: Image.Image """ image = load_image(image, mode=None, force_background=None) 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): """ Save a PNG image with NAI metadata. :param image: The image to save. :type image: Image.Image :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 :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): """ Save an image with NAI metadata in EXIF format. :param image: The image to save. :type image: Image.Image :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 :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): """ Save a GIF image with NAI metadata. :param image: The image to save. :type image: Image.Image :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 :param kwargs: Additional keyword arguments for image saving. """ write_geninfo_gif(image, dst_file, json.dumps(metadata.json), **kwargs) Loading @@ -197,6 +257,28 @@ _LSB_ALLOWED_TYPES = {'image/png', 'image/tiff'} def save_image_with_naimeta( 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. This function saves the given image with the provided NAI metadata. It can add LSB metadata and save metainfo based on the image format and user preferences. :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 NAIMetadata object to include in the image. :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. :type save_metainfo: Union[str, bool] :param kwargs: Additional keyword arguments for image saving. :return: The saved image. :rtype: Image.Image :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) if add_lsb_meta == 'auto': if mimetype in _LSB_ALLOWED_TYPES: Loading Loading
imgutils/metadata/geninfo.py +101 −3 Original line number Diff line number Diff line """ This module provides functions for reading and writing generation information (geninfo) to image files. It supports different image formats including PNG, EXIF, and GIF. The module includes functions for: 1. Reading geninfo from image parameters, EXIF data, and GIF comments 2. Writing geninfo to image parameters, EXIF data, and GIF comments These functions are useful for storing and retrieving metadata about image generation, particularly in the context of AI-generated images. """ from typing import Optional import piexif Loading @@ -8,12 +21,38 @@ from ..data import ImageTyping, load_image def read_geninfo_parameters(image: ImageTyping) -> Optional[str]: """ Read generation information from image parameters. :param image: The input image. :type image: ImageTyping :return: The generation information if found, None otherwise. :rtype: Optional[str] This function loads the image and attempts to retrieve the 'parameters' information from the image metadata. It's commonly used for PNG images where generation information is stored in the image parameters. """ image = load_image(image, mode=None, force_background=None) infos = image.info or {} return infos.get('parameters') def read_geninfo_exif(image: ImageTyping) -> Optional[str]: """ Read generation information from EXIF data. :param image: The input image. :type image: ImageTyping :return: The generation information if found in EXIF data, None otherwise. :rtype: Optional[str] This function attempts to read generation information from the EXIF metadata of the image. It specifically looks for the UserComment field in the EXIF data. If the EXIF data is invalid or not present, it returns None. """ image = load_image(image, mode=None, force_background=None) infos = image.info or {} if "exif" in infos: Loading @@ -36,6 +75,19 @@ def read_geninfo_exif(image: ImageTyping) -> Optional[str]: def read_geninfo_gif(image: ImageTyping) -> Optional[str]: """ Read generation information from GIF comment. :param image: The input image. :type image: ImageTyping :return: The generation information if found in GIF comment, None otherwise. :rtype: Optional[str] This function is specifically designed to read generation information from GIF images. It looks for the 'comment' field in the image metadata, which is commonly used in GIF files to store additional information. """ image = load_image(image, mode=None, force_background=None) infos = image.info or {} if "comment" in infos: # for gif Loading @@ -45,23 +97,69 @@ def read_geninfo_gif(image: ImageTyping) -> Optional[str]: def write_geninfo_parameters(image: ImageTyping, dst_filename: str, geninfo: str, **kwargs): """ Write generation information to image parameters. :param image: The input image. :type image: ImageTyping :param dst_filename: The destination filename to save the image with geninfo. :type dst_filename: str :param geninfo: The generation information to write. :type geninfo: str :param kwargs: Additional keyword arguments to pass to the image save function. This function writes the provided generation information to the image parameters. It's commonly used for PNG images where generation information can be stored in the image metadata. The function creates a PngInfo object, adds the geninfo as 'parameters', and saves the image with this metadata. """ pnginfo = PngInfo() pnginfo.add_text('parameters', geninfo) image = load_image(image, force_background=None, mode=None) image.save(dst_filename, pnginfo=pnginfo, *kwargs) image.save(dst_filename, pnginfo=pnginfo, **kwargs) def write_geninfo_exif(image: ImageTyping, dst_filename: str, geninfo: str, **kwargs): """ Write generation information to EXIF data. :param image: The input image. :type image: ImageTyping :param dst_filename: The destination filename to save the image with geninfo. :type dst_filename: str :param geninfo: The generation information to write. :type geninfo: str :param kwargs: Additional keyword arguments to pass to the image save function. This function writes the provided generation information to the EXIF metadata of the image. It creates an EXIF dictionary with the geninfo stored in the UserComment field, converts it to bytes, and saves the image with this EXIF data. """ 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) image.save(dst_filename, exif=exif_bytes, **kwargs) def write_geninfo_gif(image: ImageTyping, dst_filename: str, geninfo: str, **kwargs): """ Write generation information to GIF comment. :param image: The input image. :type image: ImageTyping :param dst_filename: The destination filename to save the image with geninfo. :type dst_filename: str :param geninfo: The generation information to write. :type geninfo: str :param kwargs: Additional keyword arguments to pass to the image save function. This function is specifically designed to write generation information to GIF images. It adds the geninfo to the image's 'comment' field, which is a standard way of including metadata in GIF files. """ image = load_image(image, force_background=None, mode=None) image.info['comment'] = geninfo.encode('utf-8') image.save(dst_filename, *kwargs) image.save(dst_filename, **kwargs)
imgutils/sd/nai.py +86 −4 Original line number Diff line number Diff line Loading @@ -59,6 +59,12 @@ class NAIMetadata: @property def json(self) -> dict: """ Convert the NAIMetadata to a JSON-compatible dictionary. :return: A dictionary representation of the metadata. :rtype: dict """ data = { 'Software': self.software, 'Source': self.source, Loading Loading @@ -90,10 +96,22 @@ class NAIMetadata: class _InvalidNAIMetaError(Exception): """ Custom exception raised when NAI metadata is invalid. """ pass def _naimeta_validate(data): """ Validate the NAI metadata. :param data: The metadata to validate. :type data: dict :return: The validated metadata. :rtype: dict :raises _InvalidNAIMetaError: If the metadata is invalid. """ if isinstance(data, dict) and data.get('Software') and data.get('Source') and data.get('Comment'): return data else: Loading @@ -104,14 +122,14 @@ def _get_naimeta_raw(image: ImageTyping) -> dict: """ Extract raw NAI metadata from an image. This function attempts to extract metadata from the image using LSB (Least Significant Bit) extraction. If that fails, it falls back to using the image's info dictionary. This function attempts to extract metadata from the image using various methods, including LSB extraction, image info dictionary, and other specific metadata formats. :param image: The input image. :type image: ImageTyping :return: A dictionary containing the raw metadata. :rtype: dict :raises _InvalidNAIMetaError: If no valid metadata is found. """ image = load_image(image, force_background=None, mode=None) try: Loading Loading @@ -149,7 +167,6 @@ 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] """ Loading @@ -169,19 +186,62 @@ def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]: 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 :return: The image with added metadata. :rtype: Image.Image """ image = load_image(image, mode=None, force_background=None) 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): """ Save a PNG image with NAI metadata. :param image: The image to save. :type image: Image.Image :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 :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): """ Save an image with NAI metadata in EXIF format. :param image: The image to save. :type image: Image.Image :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 :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): """ Save a GIF image with NAI metadata. :param image: The image to save. :type image: Image.Image :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 :param kwargs: Additional keyword arguments for image saving. """ write_geninfo_gif(image, dst_file, json.dumps(metadata.json), **kwargs) Loading @@ -197,6 +257,28 @@ _LSB_ALLOWED_TYPES = {'image/png', 'image/tiff'} def save_image_with_naimeta( 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. This function saves the given image with the provided NAI metadata. It can add LSB metadata and save metainfo based on the image format and user preferences. :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 NAIMetadata object to include in the image. :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. :type save_metainfo: Union[str, bool] :param kwargs: Additional keyword arguments for image saving. :return: The saved image. :rtype: Image.Image :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) if add_lsb_meta == 'auto': if mimetype in _LSB_ALLOWED_TYPES: Loading
