RedMark Noise Layer 之 JPEG压缩——基于深度学习的水印生成网络“JPEG压缩模块”详解

一、简介

论文链接：点击此链接查看RedMark的文献

二、RedMark的JPEG压缩模块

2.1 JPEG压缩流程

JPEG压缩流程，如下所示：

具体分析见：点击此连接查看JPEG压缩具体过程

2.2 RedMark中JPEG压缩实现

论文主要思想： 作者认为，在JPEG压缩的实现过程中，舍入操作是不可微的，需要通过可微操作近似模拟，以促进梯度训练过程中的反向传播。该文通过在 [-0.5，0.5] 范围内的均匀噪声来模拟舍入操作。

其中，Iw DCT是DCT域中带有水印的图像，Iw ∗ DCT是其近似量化版本。Q表示量化矩阵，σ 是均匀噪声 [-0.5，0.5]。

RedMark的JPEG压缩模块，如下图所示：

结合JPEG压缩的步骤，分析得RedMark的实现如下：

1. 颜色模式转换，无（因为它处理的是灰度图像）。

2. 向下采样，无。

3. 8×8分块，通过reshape实现。

4. DCT变换，将64×64的离散余弦变换系数作为固定的卷积权重，与图片进行卷积，从而实现DCT变换。

5. 量化，DCT变换的结果 ÷ 量化矩阵 得到 量化结果。

6. 舍入，量化结果 + 随机噪声；该噪声为[ -0.5, 0.5 ]上均匀分布的随机噪声，通过该随机噪声来模拟舍入操作。

7. 反量化，模拟舍入结果 × 量化矩阵 得到 反量化结果。

8. 反DCT变换，将64×64的IDCT系数作为固定的卷积权重，与图片进行卷积，从而实现IDCT变换。

9. 块还原，将图像reshape回去。

综上，1～6步实现JPEG压缩的过程，7～9实现JPEG解压缩的过程。

2.3 JPEG压缩具体代码

JPEG 模块的代码，如下所示：

import matplotlib.pyplot as plt
from tensorflow.keras import layers
from PIL import Image
import numpy as np
import tensorflow as tf
import scipy.io as sio
from skimage.metrics import peak_signal_noise_ratio
from skimage.metrics import structural_similarity


def scalar_output_shape(input_shape):
    return input_shape


def UniformNoise(x, val):
    noise = tf.random_uniform(shape=[1, 4, 4, 64], minval=-val, maxval=val, dtype=tf.float32, seed=None)
    return x + noise


Q = 50
jpeg_noise = 0.55
q_mtx = sio.loadmat('transforms/jpeg_qm.mat')['qm']
q_mtx = q_mtx.astype('float32')
if (Q < 50):
    S = 5000 / Q
else:
    S = 200 - 2 * Q
q_mtx = np.floor((S * q_mtx + 50.0) / 100.0)
print("q_mtx:", q_mtx)
q_mtx = np.reshape(q_mtx, (64, 1)) * 1.0
q_mtx = np.repeat(q_mtx[np.newaxis, ...], 4, axis=0)
q_mtx = np.repeat(q_mtx[np.newaxis, ...], 4, axis=0)
q_mtx = np.squeeze(q_mtx)
q_mtx[q_mtx == 0] = 1

# Load input image
img = Image.open('./images/img/img (1).jpg').convert('L')
img = img.resize((32, 32))
img = np.array(img, dtype=np.float32)
input_img = img.reshape((1, 32, 32, 1))

rearranged_img = l1 = layers.Lambda(tf.space_to_depth, arguments={'block_size': 8}, name='rearrange_img')(
    input_img)

#####################  Jpeg_attake  model ############################
model = tf.keras.models.Sequential()
model.add(layers.Conv2D(64, (1, 1), activation='linear', input_shape=(4, 4, 64), padding='same', use_bias=False,
                        trainable=False, name='dct'))
model.add(layers.Lambda(lambda x: (x * 255) / q_mtx, output_shape=scalar_output_shape, name='jpg1'))
model.add(layers.Lambda(UniformNoise, arguments={'val': jpeg_noise}))
model.add(layers.Lambda(lambda x: (x / 255) * q_mtx, output_shape=scalar_output_shape, name='jpg2'))
model.add(layers.Conv2D(64, (1, 1), activation='linear', input_shape=(4, 4, 64), padding='same', use_bias=False,
                        trainable=False, name='idct'))

# Set IDCT coefficients
idct_mtx = sio.loadmat('transforms/IDCT_coef.mat')['IDCT_coef']
idct_mtx = np.reshape(idct_mtx, [1, 1, 64, 64])
model.get_layer('idct').set_weights(np.array([idct_mtx]))

# Set DCT coefficients
dct_mtx = sio.loadmat('transforms/DCT_coef.mat')['DCT_coef']
dct_mtx = np.reshape(dct_mtx, [1, 1, 64, 64])
model.get_layer('dct').set_weights(np.array([dct_mtx]))

model.compile(optimizer='adam', loss='mse', metrics=['accuracy'])
model.summary()
output_img = model.predict(rearranged_img, steps=1)

# Reshape the output image
rearrange_output_img = layers.Lambda(tf.depth_to_space, arguments={'block_size': 8}, name='rearrange_output_img')(
    output_img)

# Show the output image
sess = tf.Session()
rearrange_output_img_arr = sess.run(rearrange_output_img)
sess.close()
plt.figure(figsize=(10, 10))
plt.subplot(1, 2, 1)
plt.imshow(input_img[0, :, :, 0], cmap='gray')
plt.title('Input image')
plt.subplot(1, 2, 2)
plt.imshow(rearrange_output_img_arr[0, :, :, 0], cmap='gray')
plt.title('Output image')
plt.show()

运行结果，如下所示：

其中，input image表示输入的原始图像，output image表示利用RedMark的JPEG压缩与解压缩之后的结果。