diff --git a/imgutils/metadata/geninfo.py b/imgutils/metadata/geninfo.py index 2c418fb7f4..e06b9c57be 100644 --- a/imgutils/metadata/geninfo.py +++ b/imgutils/metadata/geninfo.py @@ -1,3 +1,16 @@ +""" +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 @@ -8,12 +21,38 @@ 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: @@ -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 @@ -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) diff --git a/imgutils/sd/nai.py b/imgutils/sd/nai.py index 03d17bd1f6..ede8bb36a8 100644 --- a/imgutils/sd/nai.py +++ b/imgutils/sd/nai.py @@ -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, @@ -90,10 +96,22 @@ def pnginfo(self) -> PngInfo: 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: @@ -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: @@ -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] """ @@ -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) @@ -197,6 +257,28 @@ def _save_gif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike] 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: