dev(narugo): add support for webp metadata read (5f39609a) · Commits · git-mirror / Imgutils

imgutils/metadata/init.py

+1 −0

Original line number	Diff line number	Diff line
		from .geninfo import read_geninfo_gif, read_geninfo_parameters, read_geninfo_exif
		from .lsb import read_lsb_raw_bytes, read_lsb_metadata, write_lsb_raw_bytes, write_lsb_metadata, LSBReadError

imgutils/metadata/geninfo.py

0 → 100644

+43 −0

Original line number	Diff line number	Diff line
		from typing import Optional

		import piexif
		from piexif.helper import UserComment

		from ..data import ImageTyping, load_image


		def read_geninfo_parameters(image: ImageTyping) -> Optional[str]:
		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]:
		image = load_image(image, mode=None, force_background=None)
		infos = image.info or {}
		if "exif" in infos:
		exif_data = infos["exif"]
		try:
		exif = piexif.load(exif_data)
		except OSError:
		# memory / exif was not valid so piexif tried to read from a file
		exif = None

		exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b"")
		try:
		exif_comment = UserComment.load(exif_comment)
		except ValueError:
		exif_comment = exif_comment.decode("utf8", errors="ignore")

		return exif_comment
		else:
		return None


		def read_geninfo_gif(image: ImageTyping) -> Optional[str]:
		image = load_image(image, mode=None, force_background=None)
		infos = image.info or {}
		if "comment" in infos: # for gif
		return infos["comment"].decode("utf8", errors="ignore")
		else:
		return None

imgutils/sd/metadata.py

+4 −1

Original line number	Diff line number	Diff line
		@@ -15,6 +15,7 @@ from typing import Dict, Any, Optional
		from PIL.PngImagePlugin import PngInfo

		from ..data import ImageTyping, load_image
		from ..metadata import read_geninfo_parameters, read_geninfo_exif

		_PARAM_PATTERN = re.compile(r'\s(?P<key>[\w ]+):\s(?P<value>"(?:\\.\|[^\\"])+"\|[^,]*)(?:,\|$)')
		_SIZE_PATTERN = re.compile(r"^(?P<size1>-?\d+)\sx\s(?P<size2>-?\d+)$")
		@@ -256,7 +257,9 @@ def get_sdmeta_from_image(image: ImageTyping) -> Optional[SDMetaData]:
		<class 'imgutils.sd.metadata.SDMetaData'>
		"""
		image = load_image(image, mode=None, force_background=None)
		pnginfo_text = image.info.get('parameters')
		pnginfo_text = (read_geninfo_parameters(image) or
		read_geninfo_exif(image) or
		read_geninfo_parameters(image))
		if pnginfo_text:
		return parse_sdmeta_from_text(pnginfo_text)
		else:

imgutils/sd/nai.py

+42 −9

Original line number	Diff line number	Diff line
		@@ -21,8 +21,9 @@ from typing import Optional, Union
		from PIL import Image
		from PIL.PngImagePlugin import PngInfo

		from imgutils.data import load_image, ImageTyping
		from imgutils.metadata import read_lsb_metadata, write_lsb_metadata, LSBReadError
		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


		@dataclass
		@@ -78,6 +79,17 @@ class NAIMetadata:
		return info


		class _InvalidNAIMetaError(Exception):
		pass


		def _naimeta_validate(data):
		if isinstance(data, dict) and data.get('Software') and data.get('Source') and data.get('Comment'):
		return data
		else:
		raise _InvalidNAIMetaError


		def _get_naimeta_raw(image: ImageTyping) -> dict:
		"""
		Extract raw NAI metadata from an image.
		@@ -93,9 +105,29 @@ def _get_naimeta_raw(image: ImageTyping) -> dict:
		"""
		image = load_image(image, force_background=None, mode=None)
		try:
		return read_lsb_metadata(image)
		except LSBReadError:
		return image.info or {}
		return _naimeta_validate(read_lsb_metadata(image))
		except (LSBReadError, _InvalidNAIMetaError):
		pass

		try:
		return _naimeta_validate(image.info or {})
		except (LSBReadError, _InvalidNAIMetaError):
		pass

		try:
		return _naimeta_validate(json.loads(read_geninfo_parameters(image)))
		except (TypeError, json.JSONDecodeError, _InvalidNAIMetaError):
		pass

		try:
		return _naimeta_validate(json.loads(read_geninfo_exif(image)))
		except (TypeError, json.JSONDecodeError, _InvalidNAIMetaError):
		pass

		try:
		return _naimeta_validate(json.loads(read_geninfo_gif(image)))
		except (TypeError, json.JSONDecodeError, _InvalidNAIMetaError):
		raise _InvalidNAIMetaError


		def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]:
		@@ -111,8 +143,11 @@ def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]:
		:return: A NAIMetadata object if successful, None otherwise.
		:rtype: Optional[NAIMetadata]
		"""
		try:
		data = _get_naimeta_raw(image)
		if data.get('Software') and data.get('Source') and data.get('Comment'):
		except _InvalidNAIMetaError:
		return None
		else:
		return NAIMetadata(
		software=data['Software'],
		source=data['Source'],
		@@ -121,8 +156,6 @@ def get_naimeta_from_image(image: ImageTyping) -> Optional[NAIMetadata]:
		generation_time=float(data['Generation time']) if data.get('Generation time') else None,
		description=data.get('Description'),
		)
		else:
		return None


		def _get_pnginfo(metadata: Union[NAIMetadata, PngInfo]) -> PngInfo:

test/sd/test_nai.py

+39 −0

Original line number	Diff line number	Diff line
		@@ -6,6 +6,11 @@ from imgutils.sd import get_naimeta_from_image, NAIMetadata, add_naimeta_to_imag
		from ..testings import get_testfile


		@pytest.fixture()
		def nai3_webp_file():
		return get_testfile('nai3_webp.webp')


		@pytest.fixture()
		def nai3_file():
		return get_testfile('nai3.png')
		@@ -40,6 +45,37 @@ def nai3_clear_rgba_image():
		return image


		@pytest.fixture()
		def nai3_webp_meta():
		return NAIMetadata(
		software='NovelAI',
		source='Stable Diffusion XL C1E1DE52',
		parameters={
		'prompt': '2girls,side-by-side,nekomata okayu,shiina mahiru,symmetrical pose,general,masterpiece,, '
		'best quality, amazing quality, very aesthetic, absurdres',
		'steps': 28, 'height': 832, 'width': 1216, 'scale': 5.0, 'uncond_scale': 0.0, 'cfg_rescale': 0.0,
		'seed': 210306140, 'n_samples': 1, 'hide_debug_overlay': False, 'noise_schedule': 'native',
		'legacy_v3_extend': False, 'reference_information_extracted_multiple': [],
		'reference_strength_multiple': [], 'sampler': 'k_euler_ancestral', 'controlnet_strength': 1.0,
		'controlnet_model': None, 'dynamic_thresholding': False, 'dynamic_thresholding_percentile': 0.999,
		'dynamic_thresholding_mimic_scale': 10.0, 'sm': False, 'sm_dyn': False, 'skip_cfg_above_sigma': None,
		'skip_cfg_below_sigma': 0.0, 'lora_unet_weights': None, 'lora_clip_weights': None,
		'deliberate_euler_ancestral_bug': True, 'prefer_brownian': False,
		'cfg_sched_eligibility': 'enable_for_post_summer_samplers', 'explike_fine_detail': False,
		'minimize_sigma_inf': False, 'uncond_per_vibe': True, 'wonky_vibe_correlation': True, 'version': 1,
		'uc': 'lowres, {bad}, error, fewer, extra, missing, worst quality, jpeg artifacts, bad quality, '
		'watermark, unfinished, displeasing, chromatic aberration, signature, extra digits, '
		'artistic error, username, scan, [abstract], ',
		'request_type': 'PromptGenerateRequest',
		'signed_hash': 'nM6vZLFGJWW7SH2xc4lpRY9sJGbPQKXaUzhUVX/u2NvCAyLg9abn90XBCiNmwqh1hK5hk+o7wYHkPJvhkfAnBg=='
		},
		title=None,
		generation_time=6.494704299024306,
		description='2girls,side-by-side,nekomata okayu,shiina mahiru,symmetrical pose,general,masterpiece,, '
		'best quality, amazing quality, very aesthetic, absurdres'
		)


		@pytest.fixture()
		def nai3_meta_without_title():
		return NAIMetadata(
		@@ -212,3 +248,6 @@ class TestSDNai:
		])
		def test_image_error_with_wrong_format(self, file):
		assert get_naimeta_from_image(get_testfile(file)) is None

		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)