import cv2
import numpy as np
import pywt
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 arnold(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):
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):
x = ((a * b + 1) * i - b * j) % r
y = (-a * i + j) % c
p[x, y] = img[i, j]
return p
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)
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)
return waterImg
if __name__ == '__main__':
img = 'a.png'
k = 20
xs = [0.2, 0.2, 0.5, 0.4]
watermark_img = 'newImg.png'
W1 = WaterMarkDWT(img, watermark_img, k, xs)
decrypted_w
atermark = W1.get()
cv2.imwrite('shuiying.png', decrypted_watermark)
|
