Bootstrap

数字水印(DWT算法Python实现)

 全部代码:

"""-------------------------
Project:     DWT数字水印嵌入、提取
Author:      Tysay
Environment: Python3.8 + VsCode
History:     2022/12/01
-------------------------"""
import cv2
import pywt
import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']  # 替换sans-serif字体
plt.rcParams['axes.unicode_minus'] = False  # 解决坐标轴负数的负号显示问题


class WaterMarkDWT:
    def __init__(self, origin: str, watermark: str, key: int, weight: list):
        self.key = key
        self.img = cv2.imread(origin)
        self.mark = cv2.imread(watermark)
        self.coef = weight

    def imshow(self, img, title='Img', size: tuple = None):
        if isinstance(img, str):
            img = cv2.imread(img)
        if size:
            img = cv2.resize(img, size)
        plt.imshow(img, cmap='gray')
        plt.title(title)
        plt.axis('off')
        plt.show()

    def arnold(self, img):
        # 获取图片的行、列
        r, c = img.shape
        # 构建一个全为0的矩阵,size和img一样
        p = np.zeros((r, c), np.uint8)

        a, b = 1, 1
        for k in range(self.key):
            for i in range(r):
                for j in range(c):  # arnold置换算法
                    x = (i + b * j) % r
                    y = (a * i + (a * b + 1) * j) % c
                    p[x, y] = img[i, j]
        return p

    def deArnold(self, img):
        r, c = img.shape
        p = np.zeros((r, c), np.uint8)

        a, b = 1, 1
        for k in range(self.key):
            for i in range(r):
                for j in range(c):  # 逆arnold置换算法
                    x = ((a * b + 1) * i - b * j) % r
                    y = (-a * i + j) % c
                    p[x, y] = img[i, j]
        return p

    def set(self, size: tuple = (1200, 1200)):
        # 将载体图片和水印图片都做resize处理
        Img = cv2.resize(self.img, size)
        waterImg = cv2.resize(self.mark, (size[0] // 2 + 1, size[1] // 2 + 1))

        # 载体图像灰度处理
        Img1 = cv2.cvtColor(Img, cv2.COLOR_RGB2GRAY)
        waterTmg1 = cv2.cvtColor(waterImg, cv2.COLOR_RGB2GRAY)

        # 对水印图像进行Arnold变换
        waterTmg1 = self.arnold(waterTmg1)

        # 载体图像三级小波变换
        c = pywt.wavedec2(Img1, 'db2', level=3)
        # 采用Haar小波进行三级小波分解,得到不同分辨率级下的多个细节子图和一个逼近子图
        [cl, (cH3, cV3, cD3), (cH2, cV2, cD2), (cH1, cV1, cD1)] = c

        waterTmg1 = cv2.resize(waterTmg1, (size[0] // 4 + 1, size[1] // 4 + 1))

        # 水印图像一级小波变换
        d = pywt.wavedec2(waterTmg1, 'db2', level=1)
        [ca1, (ch1, cv1, cd1)] = d

        # 自定义嵌入系数(a1,a2,a3,a4 是对应的加权因子)
        a1, a2, a3, a4 = self.coef

        # 嵌入
        cl = cl + ca1 * a1
        cH3 = cH3 + ch1 * a2
        cV3 = cV3 + cv1 * a3
        cD3 = cD3 + cd1 * a4

        # 图像重构
        newImg = pywt.waverec2(
            [cl, (cH3, cV3, cD3), (cH2, cV2, cD2), (cH1, cV1, cD1)], 'db2')
        newImg = np.array(newImg, np.uint8)

        # 保存图像
        cv2.imwrite('newImg.bmp', newImg)
        return newImg

    def get(self, size: tuple = (1200, 1200), flag: int = None):
        # 原始图像灰度处理
        img = cv2.resize(self.img, size)

        img1 = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        img2 = cv2.cvtColor(self.mark, cv2.COLOR_RGB2GRAY)

        # 载体图像三级小波变换
        c = pywt.wavedec2(img2, 'db2', level=3)
        [cl, (cH3, cV3, cD3), (cH2, cV2, cD2), (cH1, cV1, cD1)] = c

        # 原始图像三级小波变换
        d = pywt.wavedec2(img1, 'db2', level=3)
        [dl, (dH3, dV3, dD3), (dH2, dV2, dD2), (dH1, dV1, dD1)] = d

        # 嵌入算法逆运算
        a1, a2, a3, a4 = self.coef

        ca1 = (cl - dl) * a1
        ch1 = (cH3 - dH3) * a2
        cv1 = (cV3 - dV3) * a3
        cd1 = (cD3 - dD3) * a4

        # 水印图像重构
        waterImg = pywt.waverec2([ca1, (ch1, cv1, cd1)], 'db2')
        waterImg = np.array(waterImg, np.uint8)

        # 对提取的水印图像进行逆Arnold变换
        waterImg = self.deArnold(waterImg)

        # 设置卷积核
        kernel = np.ones((3, 3), np.uint8)
        if flag == 0:
            # 图像腐蚀处理
            waterImg = cv2.erode(waterImg, kernel)
        elif flag == 1:
            # 图像膨胀处理
            waterImg = cv2.dilate(waterImg, kernel)

        # 保存
        cv2.imwrite('waterImg.bmp', waterImg)
        return waterImg

实验结果:

from WaterMarkClass import WaterMarkDWT

if __name__ == '__main__':
    # 读取载体图像、水印图像
    Img = 'DWTimage/Vector.bmp'
    waterImg = 'DWTimage/mark1.bmp'
    key = 10
    coef = [0.1, 0.2, 0.1, 0.1]
    # 水印嵌入
    W1 = WaterMarkDWT(Img, waterImg, key, coef)
    newimg = W1.set()
    W1.imshow(Img, '原图')
    W1.imshow(newimg, '带有水印的原图', (5285,2973))

    # 读取嵌入水印图像
    newImg = 'newImg.bmp'
    _coef = [10, 5, 10, 10]
    # 水印提取
    W2 = WaterMarkDWT(Img, newImg, key, _coef)
    wmark = W2.get()
    W2.imshow(waterImg, '原水印图')
    W2.imshow(wmark, '获取的水印图')

 

 

 

 

;