Loading docs/source/api_doc/sd/index.rst +1 −0 Original line number Diff line number Diff line Loading @@ -11,4 +11,5 @@ imgutils.sd metadata model nai docs/source/api_doc/sd/nai.rst 0 → 100644 +65 −0 Original line number Diff line number Diff line imgutils.sd.nai ==================================== .. currentmodule:: imgutils.sd.nai .. automodule:: imgutils.sd.nai NAIMetadata ------------------------------------------------ .. autoclass:: NAIMetadata :members: __init__, pnginfo get_naimeta_from_image ------------------------------------------------ .. autofunction:: get_naimeta_from_image add_naimeta_to_image ------------------------------------------------ .. autofunction:: add_naimeta_to_image save_image_with_naimeta ------------------------------------------------ .. autofunction:: save_image_with_naimeta LSBExtractor ------------------------------------------------ .. autoclass:: LSBExtractor :members: __init__, get_one_byte, get_next_n_bytes, read_32bit_integer ImageLsbDataExtractor ------------------------------------------------ .. autoclass:: ImageLsbDataExtractor :members: __init__, extract_data serialize_metadata ------------------------------------------------ .. autofunction:: serialize_metadata inject_data ------------------------------------------------ .. autofunction:: inject_data imgutils/sd/nai/__init__.py +14 −0 Original line number Diff line number Diff line """ This module provides functionality for handling LSB (Least Significant Bit) data extraction and injection, as well as managing Novel AI (NAI) metadata in images. The module includes the following main components: 1. LSB extraction from images 2. Data injection into images 3. NAI metadata handling (extraction, creation, addition, and saving) This module is particularly useful for working with steganography in images and managing metadata for AI-generated images. """ from .extract import LSBExtractor, ImageLsbDataExtractor from .inject import serialize_metadata, inject_data from .metadata import get_naimeta_from_image, NAIMetadata, add_naimeta_to_image, save_image_with_naimeta imgutils/sd/nai/extract.py +98 −2 Original line number Diff line number Diff line """ This module provides functionality for extracting hidden metadata from images using LSB (Least Significant Bit) steganography. It includes two main classes: 1. LSBExtractor: Extracts bits and bytes from image data. 2. ImageLsbDataExtractor: Uses LSBExtractor to extract and decode hidden JSON data from images. The module is based on the implementation from the NovelAI project (https://github.com/NovelAI/novelai-image-metadata). Usage: >>> from PIL import Image >>> >>> # Load an image >>> image = Image.open('path_to_image.png') >>> >>> # Create an extractor >>> extractor = ImageLsbDataExtractor() >>> >>> # Extract metadata >>> metadata = extractor.extract_data(image) >>> >>> # Process the extracted metadata >>> print(metadata) """ import gzip import json Loading @@ -5,9 +32,24 @@ import numpy as np from PIL import Image # MIT: https://github.com/NovelAI/novelai-image-metadata/blob/main/nai_meta.py class LSBExtractor(object): """ A class for extracting data hidden in the least significant bits of image pixels. This class provides methods to extract individual bits, bytes, and multi-byte values from image data using LSB steganography techniques. :param data: The image data as a numpy array. :type data: np.ndarray """ def __init__(self, data: np.ndarray): """ Initialize the LSBExtractor with image data. :param data: The image data as a numpy array. :type data: np.ndarray """ self.data = data self.rows, self.cols, self.dim = data.shape self.bits = 0 Loading @@ -16,6 +58,12 @@ class LSBExtractor(object): self.col = 0 def _extract_next_bit(self): """ Extract the next bit from the image data. This method updates the internal state of the extractor, moving to the next pixel as necessary. """ if self.row < self.rows and self.col < self.cols: bit = self.data[self.row, self.col, self.dim - 1] & 1 self.bits += 1 Loading @@ -27,6 +75,12 @@ class LSBExtractor(object): self.col += 1 def get_one_byte(self): """ Extract and return one byte of data. :return: A single byte of extracted data. :rtype: bytearray """ while self.bits < 8: self._extract_next_bit() byte = bytearray([self.byte]) Loading @@ -35,6 +89,14 @@ class LSBExtractor(object): return byte def get_next_n_bytes(self, n): """ Extract and return the next n bytes of data. :param n: The number of bytes to extract. :type n: int :return: The extracted bytes. :rtype: bytearray """ bytes_list = bytearray() for _ in range(n): byte = self.get_one_byte() Loading @@ -44,6 +106,12 @@ class LSBExtractor(object): return bytes_list def read_32bit_integer(self): """ Extract and return a 32-bit integer from the image data. :return: The extracted 32-bit integer, or None if not enough data is available. :rtype: int or None """ bytes_list = self.get_next_n_bytes(4) if len(bytes_list) == 4: integer_value = int.from_bytes(bytes_list, byteorder='big') Loading @@ -52,14 +120,42 @@ class LSBExtractor(object): return None # MIT: https://github.com/NovelAI/novelai-image-metadata/blob/main/nai_meta.py class ImageLsbDataExtractor(object): """ A class for extracting hidden JSON data from images using LSB steganography. This class uses the LSBExtractor to read hidden data from an image, expecting a specific magic number and format for the hidden data. :param magic: The magic string used to identify the start of the hidden data. :type magic: str """ def __init__(self, magic: str = "stealth_pngcomp"): """ Initialize the ImageLsbDataExtractor with a magic string. :param magic: The magic string used to identify the start of the hidden data. :type magic: str """ self._magic_bytes = magic.encode('utf-8') def extract_data(self, image: Image.Image) -> dict: """ Extract hidden JSON data from the given image. This method reads the LSB data from the image, verifies the magic number, and extracts, decompresses, and decodes the hidden JSON data. :param image: The input image. :type image: Image.Image :return: The extracted JSON data as a dictionary. :rtype: dict :raises ValueError: If the image is not in RGBA mode or if the magic number doesn't match. """ if image.mode != 'RGBA': raise ValueError(f'Image should be in RGBA mode, but {image.mode!r} found.') # noinspection PyTypeChecker image = np.array(image) reader = LSBExtractor(image) Loading imgutils/sd/nai/inject.py +110 −6 Original line number Diff line number Diff line # MIT: https://github.com/NovelAI/novelai-image-metadata/blob/main/nai_meta_writer.py """ This module provides functionality for injecting metadata and error correction codes into PNG images. It includes classes and functions for: - Bit shuffling and error correction encoding - LSB (Least Significant Bit) injection of data into image pixels - Serializing PNG metadata - Injecting encoded data and metadata into PNG images The module uses techniques like BCH error correction, bit manipulation, and LSB steganography to embed data robustly into image files. """ import gzip import json Loading @@ -14,6 +27,18 @@ code_block_len = 1920 def bit_shuffle(data_bytes, w, h): """ Shuffle the bits of input data into a specific pattern based on image dimensions. :param data_bytes: Input data bytes to be shuffled :type data_bytes: bytes :param w: Width of the image :type w: int :param h: Height of the image :type h: int :return: Tuple containing shuffled data, dimension, rest tile size, and rest dimension :rtype: tuple(bytearray, int, int, int) """ bits = np.frombuffer(data_bytes, dtype=np.uint8) bit_fac = 1 bits = bits.reshape((h, w, 3 * bit_fac)) Loading Loading @@ -43,6 +68,21 @@ def bit_shuffle(data_bytes, w, h): def split_byte_ranges(data_bytes, n, w, h): """ Split the input data bytes into chunks after shuffling. :param data_bytes: Input data bytes :type data_bytes: bytes :param n: Size of each chunk :type n: int :param w: Width of the image :type w: int :param h: Height of the image :type h: int :return: Tuple containing list of chunks, dimension, rest size, and rest dimension :rtype: tuple(list, int, int, int) """ # noinspection PyUnresolvedReferences data_bytes, dim, rest_size, rest_dim = bit_shuffle(data_bytes.copy(), w, h) chunks = [] for i in range(0, len(data_bytes), n): Loading @@ -51,34 +91,82 @@ def split_byte_ranges(data_bytes, n, w, h): def pad(data_bytes): """ Pad the input data bytes to a fixed length of 2019 bytes. :param data_bytes: Input data bytes :type data_bytes: bytes :return: Padded data bytes :rtype: bytearray """ return bytearray(data_bytes + b'\x00' * (2019 - len(data_bytes))) # Returns codes for the data in data_bytes def fec_encode(data_bytes, w, h): """ Perform Forward Error Correction (FEC) encoding on the input data. :param data_bytes: Input data bytes :type data_bytes: bytes :param w: Width of the image :type w: int :param h: Height of the image :type h: int :return: FEC encoded data :rtype: bytes """ # noinspection PyArgumentList encoder = bchlib.BCH(16, prim_poly=17475) # import galois # encoder = galois.BCH(16383, 16383-224, d=17, c=224) chunks = [bytearray(encoder.encode(pad(x))) for x in split_byte_ranges(data_bytes, 2019, w, h)[0]] return b''.join(chunks) class LSBInjector: """ A class for injecting data into the least significant bits of image pixels. """ def __init__(self, data): """ Initialize the LSBInjector with image data. :param data: Image data :type data: numpy.ndarray """ self.data = data self.buffer = bytearray() def put_32bit_integer(self, integer_value): """ Add a 32-bit integer to the buffer. :param integer_value: Integer to be added :type integer_value: int """ self.buffer.extend(integer_value.to_bytes(4, byteorder='big')) def put_bytes(self, bytes_list): """ Add bytes to the buffer. :param bytes_list: Bytes to be added :type bytes_list: bytes """ self.buffer.extend(bytes_list) def put_string(self, string): """ Add a string to the buffer (encoded as UTF-8). :param string: String to be added :type string: str """ self.put_bytes(string.encode('utf-8')) def finalize(self): """ Finalize the injection process by embedding the buffer data into the image's least significant bits. """ buffer = np.frombuffer(self.buffer, dtype=np.uint8) buffer = np.unpackbits(buffer) data = self.data[..., -1].T Loading @@ -93,7 +181,14 @@ class LSBInjector: def serialize_metadata(metadata: PngInfo) -> bytes: # Extract metadata from PNG chunks """ Serialize PNG metadata into a compressed byte string. :param metadata: PNG metadata :type metadata: PIL.PngImagePlugin.PngInfo :return: Serialized and compressed metadata :rtype: bytes """ data = { k: v for k, v in [ Loading @@ -104,12 +199,21 @@ def serialize_metadata(metadata: PngInfo) -> bytes: if data[0] == b"tEXt" or data[0] == b"iTXt" ] } # Encode and compress data using gzip data_encoded = json.dumps(data) return gzip.compress(bytes(data_encoded, "utf-8")) def inject_data(image: Image.Image, data: PngInfo) -> Image.Image: """ Inject metadata and error correction data into an image. :param image: Input image :type image: PIL.Image.Image :param data: PNG metadata to be injected :type data: PIL.PngImagePlugin.PngInfo :return: Image with injected data :rtype: PIL.Image.Image """ # noinspection PyTypeChecker rgb = np.array(image.convert('RGB')) image = image.convert('RGBA') Loading Loading
docs/source/api_doc/sd/index.rst +1 −0 Original line number Diff line number Diff line Loading @@ -11,4 +11,5 @@ imgutils.sd metadata model nai
docs/source/api_doc/sd/nai.rst 0 → 100644 +65 −0 Original line number Diff line number Diff line imgutils.sd.nai ==================================== .. currentmodule:: imgutils.sd.nai .. automodule:: imgutils.sd.nai NAIMetadata ------------------------------------------------ .. autoclass:: NAIMetadata :members: __init__, pnginfo get_naimeta_from_image ------------------------------------------------ .. autofunction:: get_naimeta_from_image add_naimeta_to_image ------------------------------------------------ .. autofunction:: add_naimeta_to_image save_image_with_naimeta ------------------------------------------------ .. autofunction:: save_image_with_naimeta LSBExtractor ------------------------------------------------ .. autoclass:: LSBExtractor :members: __init__, get_one_byte, get_next_n_bytes, read_32bit_integer ImageLsbDataExtractor ------------------------------------------------ .. autoclass:: ImageLsbDataExtractor :members: __init__, extract_data serialize_metadata ------------------------------------------------ .. autofunction:: serialize_metadata inject_data ------------------------------------------------ .. autofunction:: inject_data
imgutils/sd/nai/__init__.py +14 −0 Original line number Diff line number Diff line """ This module provides functionality for handling LSB (Least Significant Bit) data extraction and injection, as well as managing Novel AI (NAI) metadata in images. The module includes the following main components: 1. LSB extraction from images 2. Data injection into images 3. NAI metadata handling (extraction, creation, addition, and saving) This module is particularly useful for working with steganography in images and managing metadata for AI-generated images. """ from .extract import LSBExtractor, ImageLsbDataExtractor from .inject import serialize_metadata, inject_data from .metadata import get_naimeta_from_image, NAIMetadata, add_naimeta_to_image, save_image_with_naimeta
imgutils/sd/nai/extract.py +98 −2 Original line number Diff line number Diff line """ This module provides functionality for extracting hidden metadata from images using LSB (Least Significant Bit) steganography. It includes two main classes: 1. LSBExtractor: Extracts bits and bytes from image data. 2. ImageLsbDataExtractor: Uses LSBExtractor to extract and decode hidden JSON data from images. The module is based on the implementation from the NovelAI project (https://github.com/NovelAI/novelai-image-metadata). Usage: >>> from PIL import Image >>> >>> # Load an image >>> image = Image.open('path_to_image.png') >>> >>> # Create an extractor >>> extractor = ImageLsbDataExtractor() >>> >>> # Extract metadata >>> metadata = extractor.extract_data(image) >>> >>> # Process the extracted metadata >>> print(metadata) """ import gzip import json Loading @@ -5,9 +32,24 @@ import numpy as np from PIL import Image # MIT: https://github.com/NovelAI/novelai-image-metadata/blob/main/nai_meta.py class LSBExtractor(object): """ A class for extracting data hidden in the least significant bits of image pixels. This class provides methods to extract individual bits, bytes, and multi-byte values from image data using LSB steganography techniques. :param data: The image data as a numpy array. :type data: np.ndarray """ def __init__(self, data: np.ndarray): """ Initialize the LSBExtractor with image data. :param data: The image data as a numpy array. :type data: np.ndarray """ self.data = data self.rows, self.cols, self.dim = data.shape self.bits = 0 Loading @@ -16,6 +58,12 @@ class LSBExtractor(object): self.col = 0 def _extract_next_bit(self): """ Extract the next bit from the image data. This method updates the internal state of the extractor, moving to the next pixel as necessary. """ if self.row < self.rows and self.col < self.cols: bit = self.data[self.row, self.col, self.dim - 1] & 1 self.bits += 1 Loading @@ -27,6 +75,12 @@ class LSBExtractor(object): self.col += 1 def get_one_byte(self): """ Extract and return one byte of data. :return: A single byte of extracted data. :rtype: bytearray """ while self.bits < 8: self._extract_next_bit() byte = bytearray([self.byte]) Loading @@ -35,6 +89,14 @@ class LSBExtractor(object): return byte def get_next_n_bytes(self, n): """ Extract and return the next n bytes of data. :param n: The number of bytes to extract. :type n: int :return: The extracted bytes. :rtype: bytearray """ bytes_list = bytearray() for _ in range(n): byte = self.get_one_byte() Loading @@ -44,6 +106,12 @@ class LSBExtractor(object): return bytes_list def read_32bit_integer(self): """ Extract and return a 32-bit integer from the image data. :return: The extracted 32-bit integer, or None if not enough data is available. :rtype: int or None """ bytes_list = self.get_next_n_bytes(4) if len(bytes_list) == 4: integer_value = int.from_bytes(bytes_list, byteorder='big') Loading @@ -52,14 +120,42 @@ class LSBExtractor(object): return None # MIT: https://github.com/NovelAI/novelai-image-metadata/blob/main/nai_meta.py class ImageLsbDataExtractor(object): """ A class for extracting hidden JSON data from images using LSB steganography. This class uses the LSBExtractor to read hidden data from an image, expecting a specific magic number and format for the hidden data. :param magic: The magic string used to identify the start of the hidden data. :type magic: str """ def __init__(self, magic: str = "stealth_pngcomp"): """ Initialize the ImageLsbDataExtractor with a magic string. :param magic: The magic string used to identify the start of the hidden data. :type magic: str """ self._magic_bytes = magic.encode('utf-8') def extract_data(self, image: Image.Image) -> dict: """ Extract hidden JSON data from the given image. This method reads the LSB data from the image, verifies the magic number, and extracts, decompresses, and decodes the hidden JSON data. :param image: The input image. :type image: Image.Image :return: The extracted JSON data as a dictionary. :rtype: dict :raises ValueError: If the image is not in RGBA mode or if the magic number doesn't match. """ if image.mode != 'RGBA': raise ValueError(f'Image should be in RGBA mode, but {image.mode!r} found.') # noinspection PyTypeChecker image = np.array(image) reader = LSBExtractor(image) Loading
imgutils/sd/nai/inject.py +110 −6 Original line number Diff line number Diff line # MIT: https://github.com/NovelAI/novelai-image-metadata/blob/main/nai_meta_writer.py """ This module provides functionality for injecting metadata and error correction codes into PNG images. It includes classes and functions for: - Bit shuffling and error correction encoding - LSB (Least Significant Bit) injection of data into image pixels - Serializing PNG metadata - Injecting encoded data and metadata into PNG images The module uses techniques like BCH error correction, bit manipulation, and LSB steganography to embed data robustly into image files. """ import gzip import json Loading @@ -14,6 +27,18 @@ code_block_len = 1920 def bit_shuffle(data_bytes, w, h): """ Shuffle the bits of input data into a specific pattern based on image dimensions. :param data_bytes: Input data bytes to be shuffled :type data_bytes: bytes :param w: Width of the image :type w: int :param h: Height of the image :type h: int :return: Tuple containing shuffled data, dimension, rest tile size, and rest dimension :rtype: tuple(bytearray, int, int, int) """ bits = np.frombuffer(data_bytes, dtype=np.uint8) bit_fac = 1 bits = bits.reshape((h, w, 3 * bit_fac)) Loading Loading @@ -43,6 +68,21 @@ def bit_shuffle(data_bytes, w, h): def split_byte_ranges(data_bytes, n, w, h): """ Split the input data bytes into chunks after shuffling. :param data_bytes: Input data bytes :type data_bytes: bytes :param n: Size of each chunk :type n: int :param w: Width of the image :type w: int :param h: Height of the image :type h: int :return: Tuple containing list of chunks, dimension, rest size, and rest dimension :rtype: tuple(list, int, int, int) """ # noinspection PyUnresolvedReferences data_bytes, dim, rest_size, rest_dim = bit_shuffle(data_bytes.copy(), w, h) chunks = [] for i in range(0, len(data_bytes), n): Loading @@ -51,34 +91,82 @@ def split_byte_ranges(data_bytes, n, w, h): def pad(data_bytes): """ Pad the input data bytes to a fixed length of 2019 bytes. :param data_bytes: Input data bytes :type data_bytes: bytes :return: Padded data bytes :rtype: bytearray """ return bytearray(data_bytes + b'\x00' * (2019 - len(data_bytes))) # Returns codes for the data in data_bytes def fec_encode(data_bytes, w, h): """ Perform Forward Error Correction (FEC) encoding on the input data. :param data_bytes: Input data bytes :type data_bytes: bytes :param w: Width of the image :type w: int :param h: Height of the image :type h: int :return: FEC encoded data :rtype: bytes """ # noinspection PyArgumentList encoder = bchlib.BCH(16, prim_poly=17475) # import galois # encoder = galois.BCH(16383, 16383-224, d=17, c=224) chunks = [bytearray(encoder.encode(pad(x))) for x in split_byte_ranges(data_bytes, 2019, w, h)[0]] return b''.join(chunks) class LSBInjector: """ A class for injecting data into the least significant bits of image pixels. """ def __init__(self, data): """ Initialize the LSBInjector with image data. :param data: Image data :type data: numpy.ndarray """ self.data = data self.buffer = bytearray() def put_32bit_integer(self, integer_value): """ Add a 32-bit integer to the buffer. :param integer_value: Integer to be added :type integer_value: int """ self.buffer.extend(integer_value.to_bytes(4, byteorder='big')) def put_bytes(self, bytes_list): """ Add bytes to the buffer. :param bytes_list: Bytes to be added :type bytes_list: bytes """ self.buffer.extend(bytes_list) def put_string(self, string): """ Add a string to the buffer (encoded as UTF-8). :param string: String to be added :type string: str """ self.put_bytes(string.encode('utf-8')) def finalize(self): """ Finalize the injection process by embedding the buffer data into the image's least significant bits. """ buffer = np.frombuffer(self.buffer, dtype=np.uint8) buffer = np.unpackbits(buffer) data = self.data[..., -1].T Loading @@ -93,7 +181,14 @@ class LSBInjector: def serialize_metadata(metadata: PngInfo) -> bytes: # Extract metadata from PNG chunks """ Serialize PNG metadata into a compressed byte string. :param metadata: PNG metadata :type metadata: PIL.PngImagePlugin.PngInfo :return: Serialized and compressed metadata :rtype: bytes """ data = { k: v for k, v in [ Loading @@ -104,12 +199,21 @@ def serialize_metadata(metadata: PngInfo) -> bytes: if data[0] == b"tEXt" or data[0] == b"iTXt" ] } # Encode and compress data using gzip data_encoded = json.dumps(data) return gzip.compress(bytes(data_encoded, "utf-8")) def inject_data(image: Image.Image, data: PngInfo) -> Image.Image: """ Inject metadata and error correction data into an image. :param image: Input image :type image: PIL.Image.Image :param data: PNG metadata to be injected :type data: PIL.PngImagePlugin.PngInfo :return: Image with injected data :rtype: PIL.Image.Image """ # noinspection PyTypeChecker rgb = np.array(image.convert('RGB')) image = image.convert('RGBA') Loading
