From e026f61c65084549ae364572fd117e80a9a9e1cd Mon Sep 17 00:00:00 2001
From: guofei9987 <guofei9987@foxmail.com>
Date: Sat, 11 Nov 2023 18:44:18 +0800
Subject: [PATCH] make block to be 2-D, not 1-D

---
 blind_watermark/bwm_core.py | 97 ++++++++++++++++++++++++++++++++++++-
 blind_watermark/version.py  |  2 +-
 2 files changed, 97 insertions(+), 2 deletions(-)

diff --git a/blind_watermark/bwm_core.py b/blind_watermark/bwm_core.py
index 4f4ee90..a6d344c 100644
--- a/blind_watermark/bwm_core.py
+++ b/blind_watermark/bwm_core.py
@@ -49,6 +49,76 @@ def read_img_arr(self, img):
         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]
+
+        # 分块的数量
+        self.ca_num_h, self.ca_num_w = self.img_h // 2 // self.block_h, self.img_w // 2 // self.block_w
+        self.block_num = self.ca_num_h * self.ca_num_w
+
+        # ca 水印作用的像素区域
+        self.ca_h, self.ca_w = self.ca_num_h * self.block_h, self.ca_num_w * self.block_w
+        # 像素区域
+        self.img_h_new, self.img_w_new = self.ca_h * 2, self.ca_w * 2
+
+        self.img_YUV = cv2.cvtColor(self.img, cv2.COLOR_BGR2YUV)
+        for channel in range(3):
+            self.ca[channel], self.hvd[channel] = dwt2(self.img_YUV[:self.img_h_new, :self.img_w_new, channel], 'haar')
+
+    def embed(self):
+        embed_ca = np.zeros(shape=(self.ca_h, self.ca_w, 3))
+        embed_YUV = [np.array([]), np.array([]), np.array([])]
+
+        self.idx_shuffle = random_strategy1(self.password_img, self.block_num,
+                                            self.block_h * self.block_w)
+        for channel in range(3):
+            tmp_blocks = list()
+            for i in range(self.block_num):
+                idx1, idx2 = i // self.ca_num_w, i % self.ca_num_w
+                tmp_block = self.ca[channel][
+                            idx1 * self.block_h:idx1 * self.block_h + self.block_h,
+                            idx2 * self.block_w:idx2 * self.block_w + self.block_w]
+                tmp_blocks.append(tmp_block)
+
+            # 为了可以并行，写成这样
+            tmp = self.pool.map(self.block_add_wm,
+                                [(tmp_blocks[i], self.idx_shuffle[i], i) for i in range(self.block_num)])
+
+            for i in range(self.block_num):
+                idx1, idx2 = i // self.ca_num_w, i % self.ca_num_w
+                embed_ca[idx1 * self.block_h:idx1 * self.block_h + self.block_h,
+                idx2 * self.block_w:idx2 * self.block_w + self.block_w, channel] = tmp[i]
+
+            # 逆变换回去
+            embed_YUV[channel] = idwt2((embed_ca[:, :, channel], self.hvd[channel]), "haar")
+
+        # 合并3通道
+        embed_img_YUV = np.stack(embed_YUV, axis=2)
+
+        embed_img1 = cv2.cvtColor(embed_img_YUV.astype(np.float32), cv2.COLOR_YUV2BGR)
+        # 之前由于不能整除导致留了一些边角料，拼回去
+        embed_img = self.img.copy()
+
+        embed_img[:self.img_h_new, :self.img_w_new, :] = embed_img1
+        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 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,
@@ -103,7 +173,7 @@ def block_add_wm_fast(self, arg):
 
         return idct(np.dot(u, np.dot(np.diag(s), v)))
 
-    def embed(self):
+    def embed2(self):
         self.init_block_index()
 
         embed_ca = copy.deepcopy(self.ca)
@@ -166,6 +236,31 @@ def block_get_wm_fast(self, args):
     def extract_raw(self, img):
         # 每个分块提取 1 bit 信息
         self.read_img_arr(img=img)
+
+        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_h * self.block_w,  # 16
+                                            )
+        for channel in range(3):
+            tmp_blocks = list()
+            for i in range(self.block_num):
+                idx1, idx2 = i // self.ca_num_w, i % self.ca_num_w
+                tmp_block = self.ca[channel][
+                            idx1 * self.block_h:idx1 * self.block_h + self.block_h,
+                            idx2 * self.block_w:idx2 * self.block_w + self.block_w]
+                tmp_blocks.append(tmp_block)
+
+            wm_block_bit[channel, :] = self.pool.map(
+                self.block_get_wm,
+                [(tmp_blocks[i], self.idx_shuffle[i])
+                 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 个分块提取的水印，全都记录下来
diff --git a/blind_watermark/version.py b/blind_watermark/version.py
index 04e6017..b7bca8c 100644
--- a/blind_watermark/version.py
+++ b/blind_watermark/version.py
@@ -1,4 +1,4 @@
-__version__ = '0.4.4'
+__version__ = '0.5.1'
 
 
 class Notes: