scripts/watermarkcore.py

"""Файл, хранящий класс WaterMarkCore и функцию-кластеризатор (метод k-средних)"""

from random import Random

import copy
import numpy as np
import cv2

from cv2 import dct, idct
from cv2.typing import MatLike
from numpy.linalg import svd
from pywt import dwt2, idwt2


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

    # Размер блока
    # Чем больше размер, тем выше стойкость, но тем больше искажается выходное изображение
    BLOCK_SHAPE: np.ndarray[np.int32] = np.array([3, 3])

    MIN_D = 2  # Нижний порог интервала квантования
    MAX_D = 9  # Верхний порог интервала квантования
    DUPLICATE_NUM = 10

    def __init__(self):
        self.ca_block = [np.array([])] * 3  # Результаты для: dct на канал
        self.block_num = 0  # Количество фрагментов информации, которые могут быть вставлены в исходное изображение
        self.wm_size = 0  # Высота и ширина водяного знака
        self.img_shape = None  # Размер изображения
        self.part_shape = None  # Это округленный двумерный размер ca, который используется для игнорирования смещения элементов справа и снизу при встраивании
        self.block_index = None  # Декартово произведение всех индексов блоков
        self.ca_block_shape = None  # Кортеж (кол-во блоков в высоту, кол-во блоков в ширину, высота блока, ширина блока)
        self.wm_bits = None  # Водяной знак (массив битов)

        # Аппроксимация от двумерного дискретного вейвлет-преобразования
        self.ca = [np.array([])] * 3

        # Коэффициенты от двумерного дискретного вейвлет-преобразования
        # (horizontal detail, vertical detail and diagonal detail coefficients respectively)
        self.hvd = [np.array([])] * 3

    def embed(self, password: int) -> MatLike:
        """Внедрить подготовленный водяной знак в подготовленное изображение"""

        self.init_block_index()
        embed_ca = copy.deepcopy(self.ca)
        embed_yuv = [np.array([])] * 3
        seed1 = str(password)[::2]
        seed2 = str(password)[1::2]

        for channel in range(3):
            random1 = Random(seed1)
            random2 = Random(seed2)
            blocks_args = [
                (
                    self.ca_block[channel][self.block_index[i]],
                    i,
                    random1.randint(WaterMarkCore.MIN_D, WaterMarkCore.MAX_D),
                    random2.randint(WaterMarkCore.MIN_D, WaterMarkCore.MAX_D),
                )
                for i in range(self.block_num)
            ]
            tmp = list(map(self.block_add_wm, blocks_args))

            for i in range(self.block_num):
                self.ca_block[channel][self.block_index[i]] = tmp[i]

            # Четырехмерное преобразование в двухмерное
            # Каждый канал хранит четырехмерный массив,
            # представляющий результат четырехмерной разбивки.
            # После четырехмерной разбивки иногда часть отсутствует из-за нецелочисленного деления.
            # ca_part - это часть ca, которая отсутствует в данном канале
            ca_part = np.concatenate(np.concatenate(self.ca_block[channel], 1), 1)

            # При 4-мерном разбиении правая и нижняя часть длинной полосы,
            # которая не является делимой, сохраняется,
            # а остальное - это основная часть данных,
            # которая после встраивания преобразуется в частотную область.
            embed_ca[channel][: self.part_shape[0], : self.part_shape[1]] = ca_part

            # инверсия
            embed_yuv[channel] = idwt2((embed_ca[channel], self.hvd[channel]), "haar")

        # Объединить 3 канала
        embed_img_yuv = np.stack(embed_yuv, axis=2)

        # Ранее, если оно не было целым числом 2, оно добавляло белую рамку, которая здесь удалена
        embed_img_yuv = embed_img_yuv[: self.img_shape[0], : self.img_shape[1]]
        embed_img = cv2.cvtColor(embed_img_yuv, cv2.COLOR_YUV2BGR)
        embed_img = np.clip(embed_img, a_min=0, a_max=255)

        return embed_img

    def extract_with_kmeans(
        self, img: MatLike, wm_size: int, password: int
    ) -> np.ndarray[bool]:
        """Извлечь кластеризированный водяной знак из изображения"""

        wm_avg = self.extract(img, wm_size, password)

        return one_dim_kmeans(wm_avg)

    def extract(
        self, img: MatLike, wm_size: int, password: int
    ) -> np.ndarray[np.float64]:
        """Извлечь водяной знак из изображения"""

        wm_raw_bits = self.extract_raw(img, password)

        return extract_avg(wm_raw_bits, wm_size)

    def extract_raw(self, img: MatLike, password: int) -> np.ndarray[np.float64]:
        """Извлечь необработанные биты, из каждого блока"""

        self.prepare_img_arr(img)
        self.init_block_index()
        wm_raw_bits = np.zeros(shape=(3, self.block_num))
        seed1 = str(password)[::2]
        seed2 = str(password)[1::2]

        for channel in range(3):
            random1 = Random(seed1)
            random2 = Random(seed2)
            blocks_args = [
                (
                    self.ca_block[channel][self.block_index[i]],
                    random1.randint(WaterMarkCore.MIN_D, WaterMarkCore.MAX_D),
                    random2.randint(WaterMarkCore.MIN_D, WaterMarkCore.MAX_D),
                )
                for i in range(self.block_num)
            ]
            wm_raw_bits[channel, :] = list(map(self.block_get_wm, blocks_args))

        return wm_raw_bits

    def prepare_img_arr(self, img: MatLike):
        """Подготовить изображение.
        Считывание изображения ->
        YUVise ->
        добавление белой границы, чтобы сделать пиксели равномерными ->
        4D чанкинг"""

        img = img.astype(np.float32)

        # shape (высота, ширина, кол-во цветов), img_shape (высота, ширина)
        self.img_shape = img.shape[:2]

        # Y (яркость) UV (цвет)
        img_yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)

        # Дополнить изображение чёрной рамкой в 1 пиксель в нижней и/или левой части,
        # чтобы ширина и высота были чётными
        img_yuv_with_border = cv2.copyMakeBorder(
            img_yuv,
            0,
            img.shape[0] % 2,
            0,
            img.shape[1] % 2,
            cv2.BORDER_CONSTANT,
            value=(0, 0, 0),
        )

        ca_shape = [(i + 1) // 2 for i in self.img_shape]  # центр
        self.ca_block_shape = (
            ca_shape[0] // WaterMarkCore.BLOCK_SHAPE[0],
            ca_shape[1] // WaterMarkCore.BLOCK_SHAPE[1],
            WaterMarkCore.BLOCK_SHAPE[0],
            WaterMarkCore.BLOCK_SHAPE[1],
        )

        # Размер/шаг в байтах для построения свёртки
        strides = np.dtype(np.float32).itemsize * np.array(
            [
                ca_shape[1] * WaterMarkCore.BLOCK_SHAPE[0],
                WaterMarkCore.BLOCK_SHAPE[1],
                ca_shape[1],
                1,  # Сдвиг по каналам
            ]
        )

        for channel in range(3):
            # haar, db1, bior1/1, rbio1/1
            self.ca[channel], self.hvd[channel] = dwt2(
                img_yuv_with_border[:, :, channel], "haar"
            )  # Двумерное дискретное вейвлет-преобразование (Тип вейвлет: "haar")
            # Переход к 4 измерениям
            self.ca_block[channel] = np.lib.stride_tricks.as_strided(
                self.ca[channel].astype(np.float32), self.ca_block_shape, strides
            )  # Построение свёртки исходного изображения (разделение на блоки)

    def prepare_wm(self, wm_bits: np.ndarray[bool]):
        """Подготовить водяной знака"""

        self.wm_bits = wm_bits
        self.wm_size = wm_bits.size

    def init_block_index(self):
        """Подготовить информацию о блоках"""

        self.block_num = self.ca_block_shape[0] * self.ca_block_shape[1]

        assert self.wm_size < self.block_num, IndexError(
            f"До {self.block_num * 3 / 128} кб встроенной информации, \
            более {self.wm_size * 3 / 128} кб информации с водяными знаками, переполнение"
        )

        self.part_shape = self.ca_block_shape[:2] * WaterMarkCore.BLOCK_SHAPE

        self.block_index = [
            (i, j)
            for i in range(self.ca_block_shape[0])
            for j in range(self.ca_block_shape[1])
        ]

    def block_add_wm(self, args) -> MatLike:
        """Внедрить информацию о водяном знаке в блок.
        d1, d2 ∊ N; d1, d2 > 1"""

        block: np.ndarray[np.ndarray[np.float64]]
        i: int
        d1: int
        d2: int
        block, i, d1, d2 = args

        if (
            WaterMarkCore.DUPLICATE_NUM is not None
            and i >= self.wm_size * WaterMarkCore.DUPLICATE_NUM
        ):
            return block

        bit: bool = self.wm_bits[i % self.wm_size]

        block_dct = dct(block)  # Дискретное косинус-преобразование

        u: np.ndarray[np.ndarray[np.float32]]
        s: np.ndarray[np.float32]
        v: np.ndarray[np.ndarray[np.float32]]
        u, s, v = svd(block_dct)  # Сингулярное разложение

        s[0] = quantization(s[0], d1, bit)
        s[1] = quantization(s[1], d2, bit)

        inverse_block_dct = isvd(u, s, v)
        inverse_block = idct(
            inverse_block_dct
        )  # Обратное дискретное косинус-преобразование

        return inverse_block

    def block_get_wm(self, args):
        """Извлечь информацию о водяном знаке из блока.
        d1, d2 ∊ N; d1, d2 > 1"""

        block: np.ndarray[np.ndarray[np.float64]]
        d1: int
        d2: int
        block, d1, d2 = args

        dct_block = dct(block)

        s: np.ndarray[np.float32]
        _, s, _ = svd(dct_block)

        bit1 = reverse_quantization(s[0], d1)
        bit2 = reverse_quantization(s[1], d2)

        # Первый столбец более устойчив к помехам
        # и имеет больший коэффициент, чем второй (3 к 1)
        bit = (bit1 * 3 + bit2 * 1) / 4

        return bit


def quantization(num: np.float32, d: int, bit: bool) -> np.float32:
    """Квантовать сигнал в число с заданным коэффициентом"""
    return (num // d + 0.25 + 0.5 * bit) * d


def reverse_quantization(num: np.float32, d: int) -> bool:
    """Восстановить квантованный сигнал из числа с заданным коэффициентом"""
    return num % d > d / 2


def isvd(u, s, v):
    """Обратное сингулярное разложение"""
    return np.dot(u, np.dot(np.diag(s), v))


def extract_avg(array: np.ndarray[np.float64], wm_size: int) -> np.ndarray[np.float64]:
    """Извлечь массив средних арифметических дубликатов битов водяного знака.
    Каждый элемент массива является средним арифметическим между дубликатами
    квантованного числа от отдельного бита исходного водяного знака из каждого канала(2),
    каждого блока (Кол-во блоков // размер водяного знака)
    и каждого из используемых коэффициентов сингулярного разложения (2)"""

    wm_bits_avg = np.zeros(np.int32(wm_size))

    for i in range(wm_size):
        repeated_wm_bit = array[:, i::wm_size]
        if WaterMarkCore.DUPLICATE_NUM is not None:
            repeated_wm_bit = repeated_wm_bit[:, : WaterMarkCore.DUPLICATE_NUM]
        wm_bits_avg[i] = repeated_wm_bit.mean()

    return wm_bits_avg


def one_dim_kmeans(inputs: np.ndarray[np.float64], iter_num=300) -> np.ndarray[bool]:
    """Кластеризировать входные точки (метод k-средних)"""

    threshold = 0
    e_tol = 10 ** (-6)
    center = [inputs.min(), inputs.max()]  # 1. Инициализация центральной точки

    for _ in range(iter_num):
        threshold = (center[0] + center[1]) / 2

        # 2. Проверка расстояния между всеми точками и этими k точками,
        # каждая из которых классифицирована до ближайшего центра
        is_class01 = inputs > threshold

        # 3. Вычисление новой центральной точки
        center = [inputs[~is_class01].mean(), inputs[is_class01].mean()]

        # 4. Условие остановки
        if np.abs((center[0] + center[1]) / 2 - threshold) < e_tol:
            threshold = (center[0] + center[1]) / 2
            break

    is_class01 = inputs > threshold

    return is_class01