|
|
"""Файл, хранящий класс 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) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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] |
|
|
|
