nai-extract-metadata.py

from PIL import Image
import numpy as np
import gzip
import json

from io import BytesIO

def byteize(alpha):
    alpha = alpha.T.reshape((-1,))
    alpha = alpha[:(alpha.shape[0] // 8) * 8]
    alpha = np.bitwise_and(alpha, 1)
    alpha = alpha.reshape((-1, 8))
    alpha = np.packbits(alpha, axis=1)
    return alpha



class LSBExtractor:
    def __init__(self, data):
        self.data = byteize(data[..., -1])
        self.pos = 0

    def get_next_n_bytes(self, n):
        n_bytes = self.data[self.pos:self.pos + n]
        self.pos += n
        return bytearray(n_bytes)

    def read_32bit_integer(self):
        bytes_list = self.get_next_n_bytes(4)
        return int.from_bytes(bytes_list, 'big') if len(bytes_list) == 4 else None


def extract_image_metadata(image: Image.Image) -> dict:
    # Convert PIL Image to numpy array
    image_array = np.array(image.convert("RGBA"))
    
    # Verify image format
    if image_array.shape[-1] != 4 or len(image_array.shape) != 3:
        raise ValueError("Image must be in RGBA format")

    # Extract data
    reader = LSBExtractor(image_array)
    magic = "stealth_pngcomp"
    read_magic = reader.get_next_n_bytes(len(magic)).decode("utf-8")
    
    if magic != read_magic:
        raise ValueError("Invalid magic number")

    # Get and decompress JSON data
    json_len = reader.read_32bit_integer() // 8
    compressed_json = reader.get_next_n_bytes(json_len)
    json_data = json.loads(gzip.decompress(compressed_json).decode("utf-8"))

    # Handle nested Comment field if present
    if "Comment" in json_data and isinstance(json_data["Comment"], str):
        json_data["Comment"] = json.loads(json_data["Comment"])

    return json_data