dev(narugo): save the damn code (f255ce4b) · Commits · git-mirror / Imgutils

imgutils/metadata/init.py

+2 −1

Original line number	Diff line number	Diff line
		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

+24 −0

Original line number	Diff line number	Diff line
		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

+64 −75

Original line number	Diff line number	Diff line
		@@ -13,6 +13,7 @@ This module is particularly useful for working with AI-generated images and thei
		"""

		import json
		import mimetypes
		import os
		import warnings
		from dataclasses import dataclass
		@@ -23,7 +24,9 @@ from PIL.PngImagePlugin import PngInfo

		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 @@ class NAIMetadata:
		: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.

		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 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)

		: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_png_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: NAIMetadata, **kwargs):
		image.save(dst_file, pnginfo=metadata.pnginfo, **kwargs)


		def add_naimeta_to_image(image: ImageTyping, metadata: Union[NAIMetadata, PngInfo]) -> Image.Image:
		"""
		Add NAI metadata to an image.
		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)

		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]
		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)

		: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)

		_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'}

		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 = 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

+80 −16

Original line number	Diff line number	Diff line
		@@ -178,17 +178,7 @@ class TestSDNai:
		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 @@ class TestSDNai:
		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 @@ class TestSDNai:
		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 @@ class TestSDNai:
		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 @@ class TestSDNai:
		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 @@ class TestSDNai:
		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 @@ class TestSDNai:

		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