scripts/watermark.py

"""Файл, хранящий класс WaterMark и функции для шифрования"""

import numpy as np
import cv2

from cv2.typing import MatLike

from watermarkcore import WaterMarkCore


class WaterMark:
    """Класс-обёртка для работы с водяными знаками"""

    def __init__(self):
        self.wm_core = WaterMarkCore()

    def embed_and_save(
        self,
        src_img_filename: str,
        wm_content: str,
        password: str,
        img_path: str,
        compression_ratio: int = 100,
    ) -> int:
        """Внедрить водяной знак в изображение и сохранить его"""

        img, wm_size = self.embed(src_img_filename, wm_content, password)

        cv2.imwrite(
            img_path,
            img,
            params=[cv2.IMWRITE_JPEG_QUALITY, compression_ratio],
        )

        return wm_size

    def embed(
        self, src_img_filename: str, wm_content: str, password: str
    ) -> tuple[MatLike, int]:
        """Внедрить водяной знак в изображение"""

        self.prepare_img_filename(src_img_filename)
        wm_size = self.prepare_wm(wm_content, password)
        img = self.wm_core.embed(password)

        # print(
        #     "\nFill percentage without duplication",
        #     wm_size
        #     / (
        #         img.shape[0]
        #         * img.shape[1]
        #         / (
        #             self.wm_core.BLOCK_SHAPE[0]
        #             * self.wm_core.BLOCK_SHAPE[1]
        #             * np.dtype(np.float32).itemsize
        #         )
        #     ),
        # )

        return (img, wm_size)

    def prepare_img_filename(self, filename: str) -> MatLike:
        """Подготовить изображение из файла"""

        img = cv2.imread(filename, flags=cv2.IMREAD_UNCHANGED)
        assert img is not None, f"Image file \"'{filename}'\" is unreadable!"

        self.wm_core.prepare_img_arr(img)

    def prepare_wm(
        self, wm_content: str, password: int, encoding: str = "utf-8"
    ) -> int:
        """Зашифровать и подготовить водяной знак"""

        wm_bit = encrypt_wm_content(wm_content, password, encoding)
        self.wm_core.prepare_wm(wm_bit)

        return wm_bit.size

    def extract_img_filename(self, filename: str, password: int, wm_size: int) -> str:
        """Извлечь водяной знак из файла"""

        img = cv2.imread(filename, flags=cv2.IMREAD_COLOR)
        assert img is not None, f"Image file \"'{filename}'\" is unreadable!"

        return self.extract_img(img, password, wm_size)

    def extract_img(self, img: MatLike, password: int, wm_size: int) -> str:
        """Извлечь водяной знак из изображения"""

        wm_avg = self.wm_core.extract_with_kmeans(img, wm_size, password)

        return decrypt_wm_content(wm_avg, password, wm_size)


def encrypt_wm_content(
    wm_content: str, password: int, encoding: str = "utf-8"
) -> np.ndarray[bool]:
    """Зашифровать водяной знак"""

    bits = txt_to_bits(wm_content, encoding)
    shuffle_array(bits, password)

    return bits


def decrypt_wm_content(
    wm_bits: np.ndarray[bool], password: int, wm_size: int, encoding: str = "utf-8"
) -> str:
    """Дешифровать водяной знак"""

    deshuffle_array(wm_bits, password, wm_size)
    wm_content = bits_to_txt(wm_bits, encoding)

    return wm_content


def shuffle_array(array: np.ndarray, seed: int):
    """Перемешать случайные блоки массива битов"""

    random_state = np.random.RandomState(seed)
    random_state.shuffle(array)


def deshuffle_array(array: np.ndarray, seed: int, wm_size: int):
    """Восстановить блоки перемешанного массива"""

    indexes_for_deshuffle = np.arange(wm_size)
    shuffle_array(indexes_for_deshuffle, seed)
    array[indexes_for_deshuffle] = array.copy()


def txt_to_bits(txt: str, encoding: str = "utf-8") -> np.ndarray[bool]:
    """Преобразовать водяной знак ("abc") в массив битов [True, False].
    txt -> hex_str -> num -> bin_str -> bin_chars -> bits"""

    hex_str = hex_encode(txt, encoding)
    num = hex_str_to_num(hex_str)
    bin_str = num_to_bin_str(num)
    bin_chars = str_to_chars(bin_str)
    bits = bin_chars_to_bits(bin_chars)

    return bits


def bits_to_txt(bits: np.ndarray[bool], encoding: str = "utf-8"):
    """Преобразовать массив битов [True, False] в водяной знак ("abc").
    bits -> bin_chars -> bin_str -> num -> hex_str -> txt"""

    bin_chars = bits_to_bin_chars(bits)
    bin_str = chars_to_str(bin_chars)
    num = bin_str_to_num(bin_str)
    hex_str = num_to_hex_str(num)
    txt = hex_decode(hex_str, encoding)

    return txt


def hex_encode(txt: str, encoding: str = "utf-8") -> str:
    """Побайтово кодировать строку ("abc") в заданной кодировке (по умолчанию: "utf-8"),
    а затем преобразовать в строку с 16-ми представлениями этих байтов ("123F")"""

    encoded_txt = txt.encode(encoding)

    return encoded_txt.hex()


def hex_decode(hex_str: str, encoding: str = "utf-8") -> str:
    """Преобразовать строку с 16-ми представлениями символов в байты,
    а затем побайтово декодировать строку из символов ("abc") в заданной кодировке (по умолчанию: "utf-8").
    Заменяет нечитаемые символы символом (�)"""

    hex_bytes = bytes.fromhex(hex_str)

    return hex_bytes.decode(encoding, errors="replace")


def hex_str_to_num(hex_str: str) -> int:
    """Преобразовать строку с 16-м числом ("67F") в 10-е число (123)"""

    return int(hex_str, base=16)


def num_to_hex_str(num: int) -> str:
    """Преобразовать 10-е число (123) в строку с 16-м числом ("67F")"""

    return hex(num)[2:]


def num_to_bin_str(num: int) -> str:
    """Преобразовать 10-е число (123)
    в 2-е число в виде строки без указания системы счисления в начале строки ("10101")
    """

    return bin(num)[2:]


def bin_str_to_num(bin_str: str) -> int:
    """Преобразовать 2-е число в виде строки без указания системы счисления в начале строки ("10101")
    в 10-е число (123)"""

    return int(bin_str, base=2)


def str_to_chars(txt: str) -> np.ndarray[str]:
    """Преобразовать строку ('abc') в массив символов ['a', 'b', 'c']"""

    return np.array(list(txt))


def chars_to_str(bin_chars: list[str]) -> str:
    """Преобразовать массив символов ['a', 'b', 'c'] в строку ('abc')"""

    return "".join(bin_chars)


def bin_chars_to_bits(bin_chars: np.ndarray[str]) -> np.ndarray[bool]:
    """Преобразовать бинарные символы ['1', '0'] в массив битов [True, False]"""

    return bin_chars == "1"


def bits_to_bin_chars(bits: np.ndarray[bool]) -> list[str]:
    """Преобразовать массив битов [True, False] в бинарные символы ['1', '0']"""
    return [str(int(bit)) for bit in bits]