diff --git a/playground/image_metadata/read_nai_tag.py b/playground/image_metadata/read_nai_tag.py index 291fa1e..0967ab8 100644 --- a/playground/image_metadata/read_nai_tag.py +++ b/playground/image_metadata/read_nai_tag.py @@ -17,3 +17,13 @@ print(meta.Title) print(meta.Description) print(meta.Comment) + +image = Path(__file__).parent.joinpath("sample-0317.png") +try: + meta = ImageMetadata.load_image(image) +except ValueError: + raise LookupError("Cant find a MetaData") + +print(meta.Title) +print(meta.Description) +print(meta.Comment) diff --git a/playground/image_metadata/sample-0317.png b/playground/image_metadata/sample-0317.png new file mode 100644 index 0000000..39d86dc Binary files /dev/null and b/playground/image_metadata/sample-0317.png differ diff --git a/src/novelai_python/tool/image_metadata/__init__.py b/src/novelai_python/tool/image_metadata/__init__.py index 68619e4..1d32483 100644 --- a/src/novelai_python/tool/image_metadata/__init__.py +++ b/src/novelai_python/tool/image_metadata/__init__.py @@ -85,6 +85,8 @@ def reset_alpha(input_img: BytesIO) -> BytesIO: :param input_img: :return: """ + if isinstance(input_img, BytesIO): + input_img.seek(0) image = Image.open(input_img).convert('RGBA') data = np.array(image) data[..., 3] = 254 @@ -130,6 +132,8 @@ def load_image(cls, :return: ImageMetadata :raises ValidationError: Data extraction failed """ + if isinstance(image_io, BytesIO): + image_io.seek(0) try: image_data = ImageLsbDataExtractor().extract_data(image_io) model = cls(**image_data)