watermark_generate_pkg/watermark_generate/tools/picture_watermark.py

"""
图片嵌入水印工具类
"""
from watermark_generate.blind_watermark import WaterMark
from watermark_generate.tools import logger_tool

logger = logger_tool.logger


class PictureWatermarkEmbeder:
    def __init__(self, secret):
        """
        初始化图片盲水印嵌入器
        :param secret: 密码标签
        """
        self.bwm = WaterMark(password_img=1, password_wm=1)
        self.bwm.read_wm(secret, mode='str')
        self.secret = secret
        self.dest_img = None

    def embed(self, src_img, dest_img):
        """
        盲水印嵌入方法
        :param src_img: 原始图片位置
        :param dest_img: 嵌入图片存放位置
        """
        self.bwm.read_img(src_img)
        self.bwm.embed(dest_img)
        self.dest_img = dest_img
        logger.debug(f'嵌入图片位置{self.dest_img}')

    def verify(self) -> bool:
        """
        验证嵌入结果，检查嵌入的字符串与提取字符串是否一致
        :return: 嵌入结果
        """
        wm_extract = self.bwm.extract(self.dest_img, wm_shape=self.bwm.wm_size, mode='str')
        logger.debug(f'初始水印{self.secret}，水印长度{self.bwm.wm_size}，水印提取{wm_extract}')
        return wm_extract == self.secret


def extract(embed_img, secret_len=512):
    bwm = WaterMark(password_img=1, password_wm=1)
    # todo 根据生成的密码标签长度修改此处计算水印长度函数
    secret = bwm.extract(embed_img, wm_shape=get_wm_bit(secret_len), mode='str')
    logger.debug(f'水印长度{get_wm_bit(secret_len)}，水印提取{secret}')
    return secret


def get_wm_bit(len):
    return len * 8 - 2