Commit 01c95b36 authored by narugo1992's avatar narugo1992
Browse files

dev(narugo): add more unittest

parent ae999643

imgutils/sd/__init__.py

+1 −1
Original line number Diff line number Diff line
@@ -4,4 +4,4 @@ Overview: 
"""

from .metadata import parse_sdmeta_from_text, get_sdmeta_from_image, SDMetaData

from .model import read_metadata, save_with_metadata

from .nai import get_naimeta_from_image, NAIMetadata
 
from .nai import get_naimeta_from_image, NAIMetadata, add_naimeta_to_image, save_image_with_naimeta

imgutils/sd/nai/inject.py

+6 −16
Original line number Diff line number Diff line
@@ -13,15 +13,10 @@ block_length = 2019 
code_block_len = 1920





def bit_shuffle(data_bytes, w, h, use_bytes=False):

def bit_shuffle(data_bytes, w, h):

    bits = np.frombuffer(data_bytes, dtype=np.uint8)

    bit_fac = 8

    if use_bytes:

    bit_fac = 1

    else:

        bits = np.unpackbits(bits)

    bits = bits.reshape((h, w, 3 * bit_fac))

    code_block_len = 1920

    flat_tile_len = (w * h * 3) // code_block_len

    tile_w = 32

    if flat_tile_len // tile_w > 100:
@@ -44,13 +39,11 @@ def bit_shuffle(data_bytes, w, h, use_bytes=False): 
    bits = np.concatenate((easy_tiles, rest_tiles), axis=0)

    dim = bits.shape[-1]

    bits = bits.reshape((-1,))

    if not use_bytes:

        bits = np.packbits(bits)

    return bytearray(bits.tobytes()), dim, rest_tiles.shape[0], rest_dim





def split_byte_ranges(data_bytes, n, w, h):

    data_bytes, dim, rest_size, rest_dim = bit_shuffle(data_bytes.copy(), w, h, use_bytes=True)

    data_bytes, dim, rest_size, rest_dim = bit_shuffle(data_bytes.copy(), w, h)

    chunks = []

    for i in range(0, len(data_bytes), n):

        chunks.append(data_bytes[i:i + n])
@@ -63,6 +56,7 @@ def pad(data_bytes): 


# Returns codes for the data in data_bytes

def fec_encode(data_bytes, w, h):

    # noinspection PyArgumentList

    encoder = bchlib.BCH(16, prim_poly=17475)

    # import galois

    # encoder = galois.BCH(16383, 16383-224, d=17, c=224)
@@ -75,9 +69,6 @@ class LSBInjector: 
        self.data = data

        self.buffer = bytearray()



    def put_byte(self, byte):

        self.buffer.append(byte)



    def put_32bit_integer(self, integer_value):

        self.buffer.extend(integer_value.to_bytes(4, byteorder='big'))
@@ -113,18 +104,17 @@ def serialize_metadata(metadata: PngInfo) -> bytes: 
            if data[0] == b"tEXt" or data[0] == b"iTXt"

        ]

    }

    # Save space by getting rid of reduntant metadata (Title is static)

    if "Title" in data:

        del data["Title"]

    # 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:

    # noinspection PyTypeChecker

    rgb = np.array(image.convert('RGB'))

    image = image.convert('RGBA')

    w, h = image.size

    # noinspection PyTypeChecker

    pixels = np.array(image)

    injector = LSBInjector(pixels)

    injector.put_string("stealth_pngcomp")

imgutils/sd/nai/metadata.py

+4 −1
Original line number Diff line number Diff line 
import json

import os

import warnings

from dataclasses import dataclass

from typing import Optional, Union
@@ -65,7 +66,7 @@ def _get_pnginfo(metadata: Union[NAIMetadata, PngInfo]) -> PngInfo: 
    elif isinstance(metadata, PngInfo):

        pnginfo = metadata

    else:

        raise TypeError(f'Unknown metadata type for NAI - {metadata!r}.')

        raise TypeError(f'Unknown metadata type for NAI - {metadata!r}.')  # pragma: no cover

    return pnginfo
@@ -80,6 +81,8 @@ def save_image_with_naimeta(image: ImageTyping, dst_file: Union[str, os.PathLike 
                            add_lsb_meta: bool = True, save_pnginfo: bool = True, **kwargs) -> Image.Image:

    pnginfo = _get_pnginfo(metadata)

    image = load_image(image, mode=None, force_background=None)

    if not add_lsb_meta and not save_pnginfo:

        warnings.warn(f'Both LSB meta and pnginfo is disabled, no metadata will be saved to {dst_file!r}.')

    if add_lsb_meta:

        image = add_naimeta_to_image(image, metadata=pnginfo)

    if save_pnginfo:

test/sd/test_nai.py

+59 −1
Original line number Diff line number Diff line 
import pytest

from hbutils.testing import isolated_directory



from imgutils.sd import get_naimeta_from_image, NAIMetadata

from imgutils.data import load_image

from imgutils.sd import get_naimeta_from_image, NAIMetadata, add_naimeta_to_image, save_image_with_naimeta

from ..testings import get_testfile
@@ -24,6 +26,20 @@ def nai3_clear_rgba_file(): 
    return get_testfile('nai3_clear_rgba.png')





@pytest.fixture()

def nai3_clear_rgb_image():

    image = load_image(get_testfile('nai3_clear.png'))

    image.load()

    return image





@pytest.fixture()

def nai3_clear_rgba_image():

    image = load_image(get_testfile('nai3_clear_rgba.png'))

    image.load()

    return image





@pytest.fixture()

def nai3_meta_without_title():

    return NAIMetadata(
@@ -76,3 +92,45 @@ class TestSDNai: 


    def test_get_naimeta_from_image_cleared_rgba(self, nai3_clear_rgba_file, nai3_meta_without_title):

        assert get_naimeta_from_image(nai3_clear_rgba_file) is None



    def test_add_naimeta_to_image(self, nai3_clear_rgb_image, nai3_meta_without_title):

        assert get_naimeta_from_image(nai3_clear_rgb_image) is None

        image = add_naimeta_to_image(nai3_clear_rgb_image, metadata=nai3_meta_without_title)

        assert get_naimeta_from_image(image) == pytest.approx(nai3_meta_without_title)



    def test_add_naimeta_to_image_rgba(self, nai3_clear_rgba_image, nai3_meta_without_title):

        assert get_naimeta_from_image(nai3_clear_rgba_image) is None

        image = add_naimeta_to_image(nai3_clear_rgba_image, metadata=nai3_meta_without_title)

        assert get_naimeta_from_image(image) == pytest.approx(nai3_meta_without_title)



    def test_save_image_with_naimeta(self, nai3_clear_file, nai3_meta_without_title):

        with isolated_directory():

            save_image_with_naimeta(nai3_clear_file, 'image.png', metadata=nai3_meta_without_title)

            assert get_naimeta_from_image('image.png') == pytest.approx(nai3_meta_without_title)



    def test_save_image_with_naimeta_rgba(self, nai3_clear_rgba_file, nai3_meta_without_title):

        with isolated_directory():

            save_image_with_naimeta(nai3_clear_rgba_file, 'image.png', metadata=nai3_meta_without_title)

            assert get_naimeta_from_image('image.png') == pytest.approx(nai3_meta_without_title)



    def test_save_image_with_naimeta_pnginfo_only(self, nai3_clear_file, nai3_meta_without_title):

        with isolated_directory():

            save_image_with_naimeta(nai3_clear_file, 'image.png',

                                    metadata=nai3_meta_without_title, add_lsb_meta=False)

            assert get_naimeta_from_image('image.png') == pytest.approx(nai3_meta_without_title)



    def test_save_image_with_naimeta_lsbmeta_only(self, nai3_clear_file, nai3_meta_without_title):

        with isolated_directory():

            save_image_with_naimeta(nai3_clear_file, 'image.png',

                                    metadata=nai3_meta_without_title, save_pnginfo=False)

            assert get_naimeta_from_image('image.png') == pytest.approx(nai3_meta_without_title)



    def test_save_image_with_naimeta_both_no(self, nai3_clear_file, nai3_meta_without_title):

        with isolated_directory():

            with pytest.warns(Warning):

                save_image_with_naimeta(

                    nai3_clear_file, 'image.png',

                    metadata=nai3_meta_without_title,

                    save_pnginfo=False, add_lsb_meta=False,

                )

            assert get_naimeta_from_image('image.png') is None