✨ fix: Fix reading image metadata from file

Fix reading image metadata from a file by handling the case where image data is passed as BytesIO object.
LlmKira · Mar 16, 2024 · c7d2294 · c7d2294
1 parent d989e24
commit c7d2294
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 0 deletions.
diff --git 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
diff --git 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)