Merge pull request #89 from LlmKira/dev

feat: implement image metadata extraction and verification

README.md

@@ -7,7 +7,8 @@
 
 ✨ NovelAI api python sdk with Pydantic, modern and user-friendly.
 
-The goal of this repository is to use Pydantic to build legitimate requests to access the NovelAI API service.
+The goal of this repository is to use Pydantic to build legitimate requests to access
+the [NovelAI API service](https://api.novelai.net/docs).
 
 > Python >= 3.9 is required.
 
@@ -22,6 +23,8 @@ The goal of this repository is to use Pydantic to build legitimate requests to a
 pip -U install novelai-python
 ```
 
+All API users must adhere to the NovelAI Terms of Service: https://novelai.net/terms.
+
 **More examples can be found in the [playground](https://github.com/LlmKira/novelai-python/tree/main/playground)
 directory, read code as documentation.**
 
@@ -38,6 +41,7 @@ from novelai_python.sdk.ai.generate_image.schema import PositionMap
 
 load_dotenv()
 session = ApiCredential(api_token=SecretStr(os.getenv("NOVELAI_JWT")))  # pst-***
+# For security reasons, storing user credentials in plaintext is strongly discouraged.
 
 prompt = "1girl, year 2023,dynamic angle,  best quality, amazing quality, very aesthetic, absurdres"
 
@@ -137,8 +141,12 @@ loop.run_until_complete(chat("Hello"))
 ```python
 from novelai_python.tool.random_prompt import RandomPromptGenerator
 
-prompt = RandomPromptGenerator(nsfw_enabled=False).random_prompt()
-print(prompt)
+generator = RandomPromptGenerator()
+for i in range(10):
+    print(generator.generate_common_tags(nsfw=False))
+    print(generator.generate_scene_tags())
+    print(generator.generate_scene_composition())
+    print(generator.get_holiday_themed_tags())
 ```
 
 #### 📦 Run A Server
@@ -203,8 +211,6 @@ https://dghs-imgutils.deepghs.org/main/api_doc/operate/censor.html
 
 ## 🙏 Acknowledgements
 
-[BackEnd](https://api.novelai.net/docs)
-
 [novelai-api](https://github.com/Aedial/novelai-api)
 
 [NovelAI-API](https://github.com/HanaokaYuzu/NovelAI-API)

playground/image_metadata/apply_lsb_to_image.py

@@ -0,0 +1,29 @@
+from pathlib import Path
+
+from PIL import Image
+
+from novelai_python.tool.image_metadata import ImageMetadata, ImageLsbDataExtractor, ImageVerifier
+
+'''
+    You can apply metadata to an image using the LSB method.
+'''
+
+image = Path(__file__).parent.joinpath("sample-0316.png")
+write_out = Path(__file__).parent.joinpath("sample-0316-out.png")
+try:
+    with Image.open(image) as img:
+        meta = ImageMetadata.load_image(img)
+    with Image.open(image) as img:
+        new_io = meta.apply_to_image(img, inject_lsb=True)
+    with open(write_out, 'wb') as f:
+        f.write(new_io.getvalue())
+    with Image.open(write_out) as img:
+        new_meta = ImageMetadata.load_image(img)
+        data = ImageLsbDataExtractor().extract_data(img)
+except ValueError:
+    raise LookupError("Cant find a MetaData")
+print(data)
+print(new_meta)
+print(new_meta.used_model)
+with Image.open(write_out) as img:
+    print(ImageVerifier().verify(img))

playground/image_metadata/read_lsb.py → playground/image_metadata/image_lsb_data_extractor.py

@@ -1,10 +1,13 @@
 from pathlib import Path
 
+from PIL import Image
+
 from novelai_python.tool.image_metadata.lsb_extractor import ImageLsbDataExtractor
 
 image = Path(__file__).parent.joinpath("sample-0316.png")
 try:
-    data = ImageLsbDataExtractor().extract_data(image)
+    with Image.open(image) as img:
+        data = ImageLsbDataExtractor().extract_data(img)
 except ValueError:
     raise LookupError("Cant find a MetaData")
 

playground/image_metadata/read_image_info.py

@@ -0,0 +1,33 @@
+import json
+from pathlib import Path
+
+from PIL import Image
+from loguru import logger
+
+"""
+   Read the metadata from the image using raw PIL
+   To conveniently read the metadata from the image, check read_image_nai_tag.py
+"""
+
+image_io = Path(__file__).parent.joinpath("sample-0316.png")
+with Image.open(image_io) as img:
+    title = img.info.get("Title", None)
+    prompt = img.info.get("Description", None)
+    comment = img.info.get("Comment", None)
+    print(img.info)
+
+    assert isinstance(comment, str), ValueError("Comment Empty")
+    try:
+        comment = json.loads(comment)
+    except Exception as e:
+        logger.debug(e)
+        comment = {}
+print(title)
+print(prompt)
+print(comment)
+
+"""
+AI generated image
+1girl, year 2023, dynamic angle, best quality, amazing quality, very aesthetic, absurdres, rating:general, amazing quality, very aesthetic, absurdres
+{'prompt': '1girl, year 2023, dynamic angle, best quality, amazing quality, very aesthetic, absurdres, rating:general, amazing quality, very aesthetic, absurdres', 'steps': 23, 'height': 1216, 'width': 832, 'scale': 6.0, 'uncond_scale': 0.0, 'cfg_rescale': 0.0, 'seed': 3685348292, 'n_samples': 1, 'noise_schedule': 'karras', 'legacy_v3_extend': False, 'reference_information_extracted_multiple': [], 'reference_strength_multiple': [], 'extra_passthrough_testing': {'prompt': None, 'uc': None, 'hide_debug_overlay': False, 'r': 0.0, 'eta': 1.0, 'negative_momentum': 0.0}, 'v4_prompt': {'caption': {'base_caption': '1girl, year 2023, dynamic angle, best quality, amazing quality, very aesthetic, absurdres, rating:general, amazing quality, very aesthetic, absurdres', 'char_captions': [{'char_caption': '1girl', 'centers': [{'x': 0.0, 'y': 0.0}]}, {'char_caption': '1boy', 'centers': [{'x': 0.9, 'y': 0.9}]}]}, 'use_coords': True, 'use_order': True}, 'v4_negative_prompt': {'caption': {'base_caption': 'lowres', 'char_captions': [{'char_caption': 'red hair', 'centers': [{'x': 0.0, 'y': 0.0}]}, {'char_caption': '', 'centers': [{'x': 0.9, 'y': 0.9}]}]}, 'use_coords': False, 'use_order': False}, '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': False, 'prefer_brownian': True, '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', 'request_type': 'PromptGenerateRequest', 'signed_hash': 'usuVOlHJg8QFGj4nXAkC7iWKlemqttUcuvjtvGRtPzBZWiHwa/XrcM3p928gsz0F97JMb70YoVYvBG+Cbtu/Bw=='}
+"""

playground/image_metadata/read_image_nai_tag.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+# @Time    : 2024/2/8 下午4:57
+# @Author  : sudoskys
+# @File    : read_nai_tag.py
+from pathlib import Path
+
+from PIL import Image
+
+from novelai_python.tool.image_metadata import ImageMetadata, ImageVerifier
+
+image = Path(__file__).parent.joinpath("sample-0316.png")
+with Image.open(image) as img:
+    image_clear = ImageMetadata.reset_alpha(
+        image=img,
+    )
+
+try:
+    with Image.open(image) as img:
+        meta_auto = ImageMetadata.load_image(img)
+        meta1 = ImageMetadata.load_from_watermark(img)
+        meta2 = ImageMetadata.load_from_pnginfo(img)
+except ValueError:
+    raise LookupError("Cant find a MetaData")
+
+print(meta1.Generation_time)  # Meatadata from watermark have no Generation_time...
+print(meta2.Generation_time)
+print(f"Description: {meta_auto.Description}")
+print(f"Comment: {meta_auto.Comment}")
+print(f"Request Method: {meta_auto.Comment.request_type}")
+print(f"Used image model: {meta_auto.used_model}")
+# Verify if the image is from NovelAI
+with Image.open(image) as img:
+    is_novelai, have_latent = ImageVerifier().verify(image=img)
+print(f"Is NovelAI: {is_novelai}")
+print(f"Have Latent: {have_latent}")

playground/image_metadata/verify.py → playground/image_metadata/verify_is_nai_generated.py

@@ -1,13 +1,13 @@
-import os
 from pathlib import Path
 
 from PIL import Image
 
-from novelai_python.tool.image_metadata import ImageMetadata
+from novelai_python.tool.image_metadata import ImageVerifier
 
 image = Path(__file__).parent.joinpath("sample-0316.png")
 try:
-    verify = ImageMetadata.verify_image_is_novelai(Image.open(image))
+    with Image.open(image) as img:
+        verify, have_latent = ImageVerifier().verify(image=img)
 except ValueError:
     raise LookupError("Cant find a MetaData")