获取图像的风格矩阵

风格矩阵代表图像自身各通道之间的相关性，代码如下：

def gram_matrix(input_tensor):
    """
    获取风格矩阵
    """
    # 爱因斯坦求和约定（Einstein summation convention） b:批次大小，i:高、j:宽、channels:特征通道数，bijc和bijd表示沿着高、宽维度进行，即对同一位置的特征向量相乘
    # 结果矩阵形状为(batch_size, channels, channels)，因为每个通道做点积，再求和得到1个值。 而channel1与channel2..channel_x， channel2与channel1...channel_x、...最终会得到一个矩阵。
    result = tf.linalg.einsum('bijc,bijd->bcd', input_tensor, input_tensor)
    input_shape = tf.shape(input_tensor)
    # 获取height * weight = 位置数
    num_locations = tf.cast(input_shape[1] * input_shape[2], tf.float32)
    # 标准化处理，确保输出的Gram矩阵不会因为图像大小不同而由较大的数值差值。
    return result / num_locations

示例

假设input_tensor有 3 个通道（channels=3），并且为了简化，假设输入张量的形状是 (1, 2, 2, 3)，即：

batch_size=1
height=2
width=2
channels=3

假设输入张量 input_tensor 是：

[
  [
    [[1, 2, 3], [4, 5, 6]],   # 第一个位置 (i=0, j=0) 和 第二个位置 (i=0, j=1)
    [[7, 8, 9], [10, 11, 12]] # 第三个位置 (i=1, j=0) 和 第四个位置 (i=1, j=1)
  ]
]

这个张量形状为 (1, 2, 2, 3)，即 (batch_size=1, height=2, width=2, channels=3)。

每个位置(i, j)的通道 c 上有一个值，如 (1, 2, 3) 表示在(i=0, j=0)位置上的三个通道的值。

Gram矩阵的生成

第一步：通道之间的乘积

einsum 表达式 'bijc,bijd->bcd' 将两个input_tensor相乘，并在i 和 j上进行求和。我们来看具体的操作。

在每个 (i, j) 位置上，我们取出通道 c 和 d 的值相乘，并对所有 (i, j) 位置的乘积求和。

通道 0 和 通道 0 (c=0, d=0)

在位置 (0, 0)：1 * 1 = 1
在位置 (0, 1)：4 * 4 = 16
在位置 (1, 0)：7 * 7 = 49
在位置 (1, 1)：10 * 10 = 100
求和：1 + 16 + 49 + 100 = 166

通道 0 和 通道 1 (c=0, d=1)

在位置 (0, 0)：1 * 2 = 2
在位置 (0, 1)：4 * 5 = 20
在位置 (1, 0)：7 * 8 = 56
在位置 (1, 1)：10 * 11 = 110
求和：2 + 20 + 56 + 110 = 188

通道 0 和 通道 2 (c=0, d=2)

在位置 (0, 0)：1 * 3 = 3
在位置 (0, 1)：4 * 6 = 24
在位置 (1, 0)：7 * 9 = 63
在位置 (1, 1)：10 * 12 = 120
求和：3 + 24 + 63 + 120 = 210

通道 1 和 通道 1 (c=1, d=1)

在位置 (0, 0)：2 * 2 = 4
在位置 (0, 1)：5 * 5 = 25
在位置 (1, 0)：8 * 8 = 64
在位置 (1, 1)：11 * 11 = 121
求和：4 + 25 + 64 + 121 = 214

通道 1 和 通道 2 (c=1, d=2)

在位置 (0, 0)：2 * 3 = 6
在位置 (0, 1)：5 * 6 = 30
在位置 (1, 0)：8 * 9 = 72
在位置 (1, 1)：11 * 12 = 132
求和：6 + 30 + 72 + 132 = 240

通道 2 和 通道 2 (c=2, d=2)

在位置 (0, 0)：3 * 3 = 9
在位置 (0, 1)：6 * 6 = 36
在位置 (1, 0)：9 * 9 = 81
在位置 (1, 1)：12 * 12 = 144
求和：9 + 36 + 81 + 144 = 270

第二步：构建 Gram 矩阵

根据上面的计算结果，我们可以构建最终的 Gram 矩阵。矩阵的形状是 (channels, channels)，在本例中是 (3, 3)：

[
  [166, 188, 210],  # 通道 0 和其他通道的相关性
  [188, 214, 240],  # 通道 1 和其他通道的相关性
  [210, 240, 270]   # 通道 2 和其他通道的相关性
]

最后一步：除以位置数

根据代码中的 num_locations = tf.cast(input_shape[1] * input_shape[2], tf.float32)，位置数为 height * width = 2 * 2 = 4。

最终的 Gram 矩阵需要除以这个位置数：

[
  [166/4, 188/4, 210/4],  # => [41.5, 47.0, 52.5]
  [188/4, 214/4, 240/4],  # => [47.0, 53.5, 60.0]
  [210/4, 240/4, 270/4]   # => [52.5, 60.0, 67.5]
]

最终 Gram 矩阵为：

[
  [41.5, 47.0, 52.5],
  [47.0, 53.5, 60.0],
  [52.5, 60.0, 67.5]
]

这个 Gram 矩阵显示了每个通道之间的相关性。例如：

(0, 1) 位置的47.0表示通道 0 和通道 1 在所有位置上的特征值相关性。
对角线上的元素（如 41.5、53.5、67.5）表示通道自身的相关性，即通道内的特征强度。
。

应用举例

风格转化：

import matplotlib as mpl
import matplotlib.pyplot as plt
import tensorflow as tf

mpl.rcParams['figure.figsize'] = (12, 12)
mpl.rcParams['axes.grid'] = False
import numpy as np
import PIL.Image
import tensorflow_hub as hub


class StyleContentModel(tf.keras.models.Model):
    def __init__(self, style_layers, content_layers):
        super(StyleContentModel, self).__init__()
        self.vgg = vgg_layers(style_layers + content_layers)  # 这里是list的元素拼接，组成一个新的大list，item为vgg激活层的名称
        self.style_layers = style_layers
        self.content_layers = content_layers
        self.num_style_layers = len(style_layers)
        self.vgg.trainable = False

    def call(self, inputs):
        "Expects float input in [0,1]"
        inputs = inputs * 255.0
        preprocessed_input = tf.keras.applications.vgg19.preprocess_input(inputs)  # 标准化处理，图像素值缩放，均值减除
        outputs = self.vgg(preprocessed_input)
        style_outputs, content_outputs = (outputs[:self.num_style_layers], outputs[self.num_style_layers:])
        style_outputs = [gram_matrix(style_output) for style_output in style_outputs]
        # zip(self.style_layers, style_outputs)相当于 [
        # ('conv1', output1),
        # ('conv2', output2),
        # ('conv3', output3)
        # ]
        content_dict = {content_name: value for content_name, value in zip(self.content_layers, content_outputs)}
        style_dict = {style_name: value for style_name, value in zip(self.style_layers, style_outputs)}
        return {'content': content_dict, 'style': style_dict}


def load_img(path_to_img):
    """
    从给定地址加载图片
    """
    max_dim = 512
    img = tf.io.read_file(path_to_img)
    img = tf.image.decode_image(img, channels=3)
    # 用于处理图像，会自动 /255.进行缩放
    img = tf.image.convert_image_dtype(img, tf.float32)

    # tf.cast不会缩放，取到高、宽
    shape = tf.cast(tf.shape(img)[:-1], tf.float32)
    long_dim = max(shape)
    scale = max_dim / long_dim

    new_shape = tf.cast(shape * scale, tf.int32)
    img = tf.image.resize(img, new_shape)
    img = img[tf.newaxis, :]
    return img


def imshow(image, title=None):
    if len(image.shape) > 3:
        image = tf.squeeze(image, axis=0)
    plt.imshow(image)
    if title:
        plt.title(title)


def tensor_to_image(tensor):
    tensor = - tensor * 255
    tensor = np.array(tensor, dtype=np.uint8)
    if np.ndim(tensor) > 3:
        assert tensor.shape[0] == 1
        tensor = tensor[0]
    return PIL.Image.fromarray(tensor)


def vgg_layers(layer_names):
    # 创建一个自定义的模型
    # 在输入vgg.input后（也就是一张图片后），这个函数会返回上面定义的那些网络层的激活值。
    vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
    vgg.trainable = False
    outputs = [vgg.get_layer(name).output for name in layer_names]
    model = tf.keras.Model([vgg.input], outputs)
    return model


def gram_matrix(input_tensor):
    """
    获取风格矩阵
    """
    # 爱因斯坦求和约定（Einstein summation convention） b:批次大小，i:高、j:宽、channels:特征通道数，bijc和bijd表示沿着高、宽维度进行，即对同一位置的特征向量相乘
    # 结果矩阵形状为(batch_size, channels, channels)，因为每个通道做点积，再求和得到1个值。 而channel1与channel2..channel_x， channel2与channel1...channel_x、...最终会得到一个矩阵。
    result = tf.linalg.einsum('bijc,bijd->bcd', input_tensor, input_tensor)
    input_shape = tf.shape(input_tensor)
    # 获取height * weight = 位置数
    num_locations = tf.cast(input_shape[1] * input_shape[2], tf.float32)
    # 标准化处理，确保输出的Gram矩阵不会因为图像大小不同而由较大的数值差值。
    return result / num_locations


def clip_0_1(image):
    """修剪图片"""
    return tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0)


def style_content_loss(outputs, style_targets, content_targets, num_style_layers, num_content_layers):
    """
    对比最终图片image与内容图片content_image及风格图片style_image的差别，也就是损失函数，差别越大损失越大。
    """
    style_weight = 1e-2  # 控制风格化到什么程度0.02
    content_weight = 1e4  # 控制内容保留的程度10000

    style_outputs = outputs['style']  # image当前的风格矩阵
    content_outputs = outputs['content']  # image当前的内容激活矩阵

    # 计算风格损失
    style_loss = tf.add_n(
        [tf.reduce_mean((style_outputs[name] - style_targets[name]) ** 2) for name in style_outputs.keys()]
    )
    style_loss *= style_weight / num_style_layers

    # 计算内容损失
    content_loss = tf.add_n(
        [tf.reduce_mean((content_outputs[name] - content_targets[name]) ** 2) for name in content_outputs.keys()]
    )
    content_loss *= content_weight / num_content_layers
    loss = style_loss + content_loss
    return loss

def train_step(image, extractor, style_targets, content_targets, num_style_layers, num_content_layers):
    opt = tf.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1)
    # tape会记录下前向传播的每个步骤，后面好自动执行反向传播
    with tf.GradientTape() as tape:
        outputs = extractor(image)  # 获取当前image的内容激活值矩阵和风格矩阵
        loss = style_content_loss(outputs, style_targets, content_targets, num_style_layers, num_content_layers)  # 计算损失
    # 获取image相对于loss的梯度，image就相当于w和参数b一样
    grad = tape.gradient(loss, image)
    # 使用梯度来改变image，也即是说image会变得越来越像content_image风格图片style_image
    opt.apply_gradients([(grad, image)])
    image.assign(clip_0_1(image))


if __name__ == "__main__":
    content_path = tf.keras.utils.get_file('YellowLabradorLooking_new.jpg',
                                           'https://storage.googleapis.com/download.tensorflow.org/example_images/YellowLabradorLooking_new.jpg')
    style_path = tf.keras.utils.get_file('The_Great_Wave_off_Kanagawa.jpg',
                                         'https://upload.wikimedia.org/wikipedia/commons/0/0a/The_Great_Wave_off_Kanagawa.jpg')
    # 全局
    content_image = load_img(content_path)
    style_image = load_img(style_path)
    plt.subplot(1, 2, 1)
    imshow(content_image, 'Content Image')
    plt.subplot(1, 2, 2)
    imshow(style_image, 'Style Image')
    hub_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
    stylized_image = hub_model(tf.constant(content_image), tf.constant(style_image))[0]
    result = tensor_to_image(stylized_image)
    plt.show()
    plt.imshow(tf.squeeze(stylized_image, 0))
    plt.show()

    # vgg就代表了一个训练好了的VGG19模型，include_top=False接代表不需要最后一层。 因为只用它来风格转换，不需要最后一层，最后一层是用来识别图片的。
    # vgg = tf.keras.applications.VGG19(include_top=False, weights="imagenet")
    # for layer in vgg.layers:
    #     print(layer.name)
    # 风格转换产生的图片既要有内容图片的内容又要有风格图片的风格。下面我们将使用内容图片的VGG的block5_conv2层来生成最终图像的内容，
    # 同时用风格图片的VGG的block1_conv1，block2_conv1，block3_conv1，block4_conv1，block5_conv1来生成最终图片的风格。

    content_layers = ['block5_conv2']
    style_layers = ['block1_conv1', 'block2_conv1', 'block3_conv1', 'block4_conv1', 'block5_conv1']
    num_content_layers = len(content_layers)
    num_style_layers = len(style_layers)
    extractor = StyleContentModel(style_layers, content_layers)
    style_targets = extractor(style_image)['style']  # 获取风格图片的风格矩阵
    content_targets = extractor(content_image)['content']  # 获取内容图片的内容激活值矩阵

    # 复制内容图片到image
    # 后面会不断的根据content_image和style_image来改变image，使得image的风格越来越像style_image

    image = tf.Variable(content_image)
    for i in range(100):
        print(f"训练第{i}次..........")
        train_step(image, extractor, style_targets, content_targets, num_style_layers, num_content_layers)
    plt.imshow(tf.squeeze(image, 0))
    plt.show()