Skip to content

Commit

Permalink
del useless codes
Browse files Browse the repository at this point in the history
  • Loading branch information
@guofei9987
guofei9987 committed Nov 11, 2023
1 parent e026f61 commit 50dbc7f
Showing 1 changed file with 2 additions and 94 deletions.
96 changes: 2 additions & 94 deletions blind_watermark/bwm_core.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
# @Author : github.com/guofei9987
import numpy as np
from numpy.linalg import svd
import copy
import cv2
from cv2 import dct, idct
from pywt import dwt2, idwt2
Expand All @@ -17,26 +16,15 @@ def __init__(self, password_img=1, mode='common', processes=None):
self.password_img = password_img
self.d1, self.d2 = 36, 20 # d1/d2 越大鲁棒性越强,但输出图片的失真越大

# init data
self.img, self.img_YUV = None, None # self.img 是原图，self.img_YUV 对像素做了加白偶数化
self.ca, self.hvd, = [np.array([])] * 3, [np.array([])] * 3 # 每个通道 dct 的结果
self.ca_block = [np.array([])] * 3 # 每个 channel 存一个四维 array，代表四维分块后的结果
self.ca_part = [np.array([])] * 3 # 四维分块后，有时因不整除而少一部分，self.ca_part 是少这一部分的 self.ca
self.ca, self.hvd, = [None] * 3, [None] * 3 # 每个通道 dct 的结果

self.wm_size, self.block_num = 0, 0 # 水印的长度，原图片可插入信息的个数
self.pool = AutoPool(mode=mode, processes=processes)

self.fast_mode = False
self.alpha = None # 用于处理透明图

def init_block_index(self):
self.block_num = self.ca_block_shape[0] * self.ca_block_shape[1]
assert self.wm_size < self.block_num, IndexError(
'最多可嵌入{}kb信息，多于水印的{}kb信息，溢出'.format(self.block_num / 1000, self.wm_size / 1000))
# self.part_shape 是取整后的ca二维大小,用于嵌入时忽略右边和下面对不齐的细条部分。
self.part_shape = self.ca_block_shape[:2] * self.block_shape
self.block_index = [(i, j) for i in range(self.ca_block_shape[0]) for j in range(self.ca_block_shape[1])]

def read_img_arr(self, img):
# 处理透明图
self.alpha = None
Expand All @@ -47,7 +35,6 @@ def read_img_arr(self, img):

# 读入图片->YUV化->加白边使像素变偶数->四维分块
self.img = img.astype(np.float32)
self.img_shape = self.img.shape[:2]

self.block_h, self.block_w = self.block_shape
self.img_h, self.img_w = self.img.shape[:2]
Expand Down Expand Up @@ -107,35 +94,6 @@ def embed(self):
embed_img = cv2.merge([embed_img.astype(np.uint8), self.alpha])
return embed_img

def read_img_arr2(self, img):
# 处理透明图
self.alpha = None
if img.shape[2] == 4:
if img[:, :, 3].min() < 255:
self.alpha = img[:, :, 3]
img = img[:, :, :3]

# 读入图片->YUV化->加白边使像素变偶数->四维分块
self.img = img.astype(np.float32)
self.img_shape = self.img.shape[:2]

# 如果不是偶数，那么补上白边，Y（明亮度）UV（颜色）
self.img_YUV = cv2.copyMakeBorder(cv2.cvtColor(self.img, cv2.COLOR_BGR2YUV),
0, self.img.shape[0] % 2, 0, self.img.shape[1] % 2,
cv2.BORDER_CONSTANT, value=(0, 0, 0))

self.ca_shape = [(i + 1) // 2 for i in self.img_shape]

self.ca_block_shape = (self.ca_shape[0] // self.block_shape[0], self.ca_shape[1] // self.block_shape[1],
self.block_shape[0], self.block_shape[1])
strides = 4 * np.array([self.ca_shape[1] * self.block_shape[0], self.block_shape[1], self.ca_shape[1], 1])

for channel in range(3):
self.ca[channel], self.hvd[channel] = dwt2(self.img_YUV[:, :, channel], 'haar')
# 转为4维度
self.ca_block[channel] = np.lib.stride_tricks.as_strided(self.ca[channel].astype(np.float32),
self.ca_block_shape, strides)

def read_wm(self, wm_bit):
self.wm_bit = wm_bit
self.wm_size = wm_bit.size
Expand Down Expand Up @@ -173,40 +131,6 @@ def block_add_wm_fast(self, arg):

return idct(np.dot(u, np.dot(np.diag(s), v)))

def embed2(self):
self.init_block_index()

embed_ca = copy.deepcopy(self.ca)
embed_YUV = [np.array([])] * 3

self.idx_shuffle = random_strategy1(self.password_img, self.block_num,
self.block_shape[0] * self.block_shape[1])
for channel in range(3):
tmp = self.pool.map(self.block_add_wm,
[(self.ca_block[channel][self.block_index[i]], self.idx_shuffle[i], i)
for i in range(self.block_num)])

for i in range(self.block_num):
self.ca_block[channel][self.block_index[i]] = tmp[i]

# 4维分块变回2维
self.ca_part[channel] = np.concatenate(np.concatenate(self.ca_block[channel], 1), 1)
# 4维分块时右边和下边不能整除的长条保留，其余是主体部分，换成 embed 之后的频域的数据
embed_ca[channel][:self.part_shape[0], :self.part_shape[1]] = self.ca_part[channel]
# 逆变换回去
embed_YUV[channel] = idwt2((embed_ca[channel], self.hvd[channel]), "haar")

# 合并3通道
embed_img_YUV = np.stack(embed_YUV, axis=2)
# 之前如果不是2的整数，增加了白边，这里去除掉
embed_img_YUV = embed_img_YUV[:self.img_shape[0], :self.img_shape[1]]
embed_img = cv2.cvtColor(embed_img_YUV, cv2.COLOR_YUV2BGR)
embed_img = np.clip(embed_img, a_min=0, a_max=255)

if self.alpha is not None:
embed_img = cv2.merge([embed_img.astype(np.uint8), self.alpha])
return embed_img

def block_get_wm(self, args):
if self.fast_mode:
return self.block_get_wm_fast(args)
Expand Down Expand Up @@ -258,23 +182,6 @@ def extract_raw(self, img):
for i in range(self.block_num)])
return wm_block_bit

def extract_raw2(self, img):
# 每个分块提取 1 bit 信息
self.read_img_arr(img=img)
self.init_block_index()

wm_block_bit = np.zeros(shape=(3, self.block_num)) # 3个channel，length 个分块提取的水印，全都记录下来

self.idx_shuffle = random_strategy1(seed=self.password_img,
size=self.block_num,
block_shape=self.block_shape[0] * self.block_shape[1], # 16
)
for channel in range(3):
wm_block_bit[channel, :] = self.pool.map(self.block_get_wm,
[(self.ca_block[channel][self.block_index[i]], self.idx_shuffle[i])
for i in range(self.block_num)])
return wm_block_bit

def extract_avg(self, wm_block_bit):
# 对循环嵌入+3个 channel 求平均
wm_avg = np.zeros(shape=self.wm_size)
Expand Down Expand Up @@ -320,6 +227,7 @@ def random_strategy1(seed, size, block_shape):


def random_strategy2(seed, size, block_shape):
# same with all blocks
one_line = np.random.RandomState(seed) \
.random(size=(1, block_shape)) \
.argsort(axis=1)
Expand Down

0 comments on commit 50dbc7f

Please sign in to comment.