dev(narugo): just code transfer, ci skip (e618b226) · Commits · git-mirror / Imgutils

imgutils/sd/nai.py

+13 −10

Original line number	Diff line number	Diff line
		@@ -189,15 +189,14 @@ _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'}
		_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: bool = True, **kwargs) -> Image.Image:
		add_lsb_meta: Union[str, bool] = 'auto', save_metainfo: Union[str, bool] = 'auto', **kwargs) -> Image.Image:
		mimetype, _ = mimetypes.guess_type(dst_file)
		if add_lsb_meta == 'auto':
		if mimetype in _LSB_ALLOWED_TYPES:
		@@ -206,9 +205,18 @@ def save_image_with_naimeta(
		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, '
		raise ValueError('LSB metadata cannot be saved to lossy image format or RGBA-incompatible format, '
		'add_lsb_meta will be disabled. '
		f'Only {", ".join(sorted(_LSB_ALLOWED_TYPES))} images supported.')
		if save_metainfo == 'auto':
		if mimetype in _FN_IMG_SAVE:
		save_metainfo = True
		else:
		save_metainfo = False
		else:
		if save_metainfo and 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}.')
		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}.')

		@@ -216,11 +224,6 @@ def save_image_with_naimeta(
		if add_lsb_meta:
		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)

requirements.txt

+2 −1

Original line number	Diff line number	Diff line
		@@ -16,3 +16,4 @@ deprecation>=2.0.0
		hfutils>=0.2.3
		filelock
		bchlib
		piexif
		No newline at end of file

test/sd/test_nai.py

+43 −49

Original line number	Diff line number	Diff line
		@@ -178,18 +178,6 @@ 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_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)
		assert get_naimeta_from_image('image.png') == pytest.approx(nai3_meta_without_title)

		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_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):
		with isolated_directory():
		with pytest.warns(Warning):
		@@ -270,48 +258,54 @@ class TestSDNai:
		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',),
		@pytest.mark.parametrize(['ext', 'okay'], [
		('.png', True),
		('.webp', False),
		('.jpg', False),
		('.jpeg', False),
		('.tiff', True),
		('.gif', False),
		])
		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):
		def test_save_image_with_naimeta_lsb_true(self, nai3_clear_file, nai3_meta_without_title,
		ext, okay):
		with isolated_directory():
		if okay:
		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)
		else:
		with 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'], [
		('.webp',),
		('.jpg',),
		('.jpeg',),
		('.tiff',),
		('.gif',),
		@pytest.mark.parametrize(['ext', 'warns', 'okay'], [
		('.png', False, True),
		('.webp', False, True),
		('.jpg', False, True),
		('.jpeg', False, True),
		('.tiff', True, False),
		('.gif', False, True),
		])
		def test_save_image_with_naimeta_metainfo_only_exifs(self, nai3_clear_file, nai3_meta_without_title, ext):
		with isolated_directory(), pytest.warns(None):
		def test_save_image_with_naimeta_metainfo_only(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, add_lsb_meta=False)
		assert get_naimeta_from_image(f'image{ext}') == pytest.approx(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',),
		('.tiff',),
		('.gif',),
		@pytest.mark.parametrize(['ext', 'warns', 'okay'], [
		('.png', False, True),
		('.webp', True, False),
		('.jpg', True, False),
		('.jpeg', True, False),
		('.tiff', False, True),
		('.gif', True, False),
		])
		def test_save_image_with_naimeta_lsbmeta_only_exifs(self, nai3_clear_file, nai3_meta_without_title, ext):
		with isolated_directory(), pytest.warns(Warning):
		def test_save_image_with_naimeta_lsbmeta_only(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, save_metainfo=False)
		assert get_naimeta_from_image(f'image{ext}') is None
		assert get_naimeta_from_image(f'image{ext}') == \
		(pytest.approx(nai3_meta_without_title) if okay else None)