Skip to content

Commit

Permalink
dev(narugo): just code transfer, ci skip
Browse files Browse the repository at this point in the history
  • Loading branch information
@narugo1992
narugo1992 committed Sep 10, 2024
1 parent f255ce4 commit e618b22
Show file tree
Hide file tree
Showing 3 changed files with 58 additions and 60 deletions.
23 changes: 13 additions & 10 deletions imgutils/sd/nai.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -189,15 +189,14 @@ def _save_gif_with_naimeta(image: Image.Image, dst_file: Union[str, os.PathLike]
'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:
Expand All @@ -206,22 +205,26 @@ 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}.')

image = load_image(image, mode=None, force_background=None)
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)
_FN_IMG_SAVE[mimetype](image, dst_file, metadata, **kwargs)
else:
image.save(dst_file, **kwargs)
return image
3 changes: 2 additions & 1 deletion requirements.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,4 +15,5 @@ pyclipper
deprecation>=2.0.0
hfutils>=0.2.3
filelock
bchlib
bchlib
piexif
92 changes: 43 additions & 49 deletions test/sd/test_nai.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -178,18 +178,6 @@ def test_add_naimeta_to_image_rgba(self, nai3_clear_rgba_image, nai3_meta_withou
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):
Expand Down Expand Up @@ -270,48 +258,54 @@ def test_save_image_with_naimeta_rgba(self, nai3_clear_rgba_file, nai3_meta_with
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():
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',),
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', '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)

@pytest.mark.parametrize(['ext'], [
('.webp',),
('.jpg',),
('.jpeg',),
('.tiff',),
('.gif',),
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', 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)

0 comments on commit e618b22

Please sign in to comment.