全部代码：
"""-------------------------
Project: DWT数字水印嵌入、提取
Author: Tysay
Environment: Python3.8 + VsCode
History: 2022/12/01
-------------------------"""
import cv2
import pywt
import numpy as np
import matplotlib.pyplot as pltplt.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 = keyself.img = cv2.imread(origin)self.mark = cv2.imread(watermark)self.coef = weightdef 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, 1for k in range(self.key):for i in range(r):for j in range(c): # arnold置换算法x = (i + b * j) % ry = (a * i + (a * b + 1) * j) % cp[x, y] = img[i, j]return pdef deArnold(self, img):r, c = img.shapep = np.zeros((r, c), np.uint8)a, b = 1, 1for k in range(self.key):for i in range(r):for j in range(c): # 逆arnold置换算法x = ((a * b + 1) * i - b * j) % ry = (-a * i + j) % cp[x, y] = img[i, j]return pdef 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)] = cwaterTmg1 = 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 * a1cH3 = cH3 + ch1 * a2cV3 = cV3 + cv1 * a3cD3 = 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 newImgdef 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.coefca1 = (cl - dl) * a1ch1 = (cH3 - dH3) * a2cv1 = (cV3 - dV3) * a3cd1 = (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 WaterMarkDWTif __name__ == '__main__':# 读取载体图像、水印图像Img = 'DWTimage/Vector.bmp'waterImg = 'DWTimage/mark1.bmp'key = 10coef = [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, '获取的水印图')
