数字水印(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, '获取的水印图')
