11_fast_style_transfer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 第11天：TensorFlow2项目实战—快速风格迁移"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "python版本： 3.7.9 (default, Aug 31 2020, 17:10:11) [MSC v.1916 64 bit (AMD64)]\n",
      "tensorlfow版本: 2.1.0\n",
      "可用GPU数量: 1\n",
      "GPU显卡信息:\n",
      "Device mapping:\n",
      "/job:localhost/replica:0/task:0/device:GPU:0 -> device: 0, name: GeForce RTX 2080 Ti, pci bus id: 0000:01:00.0, compute capability: 7.5\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras.layers import Conv2D, ReLU, UpSampling2D\n",
    "# 需要根据python版本，tf版本安装对应的tensorflow_addons\n",
    "from tensorflow_addons.layers import InstanceNormalization\n",
    "\n",
    "# 计算损失的python文件\n",
    "from lossnetwork import loadimg, grammatrix, stylelossfunc, contentlossfunc,totalvariationloss\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import sys\n",
    "import os\n",
    "import shutil\n",
    "plt.rcParams['font.family'] = 'SimHei'  # 绘图显示中文 \n",
    "plt.rcParams['axes.unicode_minus']=False  # 绘图显示负号\n",
    "\n",
    "print('python版本：', sys.version)\n",
    "print('tensorlfow版本:',tf.__version__)\n",
    "print('可用GPU数量:', len(tf.config.experimental.list_physical_devices('GPU')))\n",
    "print('GPU显卡信息:')\n",
    "sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(log_device_placement=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "基于迭代生成风格迁移图片，是将迁移图片作为参数，通过模型的训练不断更新参数，这是一个模型训练的过程，因此该方法速度缓慢。试想如果可以构建一个带有图像风格的转换网络，将内容图像作为这个网络的输入，输出的是风格迁移后图像，这就会很快。以上就是快速风格迁移的思路，参考论文[Perceptual Losses for Real-Time Style Transferand Super-Resolution](https://arxiv.org/pdf/1603.08155v1.pdf)以及[补充材料](https://static-content.springer.com/esm/chp%3A10.1007%2F978-3-319-46475-6_43/MediaObjects/419974_1_En_43_MOESM1_ESM.pdf)。\n",
    "\n",
    "![11.0.png](11.0.png)\n",
    "### 1、快速风格迁移原理\n",
    "\n",
    "\n",
    "原理：对于给定的一个风格图片，将大量的内容图片作为输入数据集，并按照基于迭代的风格迁移的方式计算损失，更新模型参数训练模型。模型训练完毕后，将任意一个内容图像输入进这个训练好的模型后，就可以直接输出带有给定风格的风格迁移后的图像。 \n",
    "\n",
    "快速风格迁移的网络结构包含两部分：Image Transform Net（图像转换模型）和Loss Network（计算损失网络），如下图所示：\n",
    "\n",
    "![11.1.png](11.1.png)\n",
    "\n",
    "其中上图中的图像转换模型就是一个需要训练的模型，该模型训练完成后，就可以对于任意给定的图像，直接输出带有某种风格的图像；计算损失网络，其实也就是特征提取网络，通过计算图像转换网络的输出图像和给定的风格图片之间的风格损失、和输入图片之间的内容损失，从而为图像转换网络的更新参数提供梯度，一般选择已经训练好的模型。\n",
    "![11.2.png](11.2.png)\n",
    "\n",
    "\n",
    "#### 1.1 图像转换模型(Image Transform Net)\n",
    "\n",
    "该模型一般使用下面的深度残差神经网络：\n",
    "\n",
    "![11.3.png](11.3.png)\n",
    "\n",
    "图像转换模型包括1个反射填充层，3个下采样卷积层、5个残差块、3个上采样卷积层。除了最末的输出层以外，所有的卷积层(残差块内除外)后都连接一个Batch Normalization和Relu层。每个残差块的结构如下：\n",
    "\n",
    "![11.4.png](11.4.png)\n",
    "\n",
    "#### 1.2 计算损失网络(Loss Network)\n",
    "\n",
    "和迭代方法的计算是一样的，本文以VGG16为例。直接运行下面的语句下载较慢。模型百度网盘链接，密码。下载完成后，在文件夹地址栏输入%userprofile%并运行，将下载的*.h5文件存放到该文件夹下的.keras/models文件夹中，再次运行下面的下载语句即可。 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 语句直接下载比较慢\n",
    "VGG16model = tf.keras.applications.VGG16(include_top=False, weights='imagenet') # 不含全连接层的预训练模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"vgg16\"\n",
      "_________________________________________________________________\n",
      "Layer (type)                 Output Shape              Param #   \n",
      "=================================================================\n",
      "input_1 (InputLayer)         [(None, None, None, 3)]   0         \n",
      "_________________________________________________________________\n",
      "block1_conv1 (Conv2D)        (None, None, None, 64)    1792      \n",
      "_________________________________________________________________\n",
      "block1_conv2 (Conv2D)        (None, None, None, 64)    36928     \n",
      "_________________________________________________________________\n",
      "block1_pool (MaxPooling2D)   (None, None, None, 64)    0         \n",
      "_________________________________________________________________\n",
      "block2_conv1 (Conv2D)        (None, None, None, 128)   73856     \n",
      "_________________________________________________________________\n",
      "block2_conv2 (Conv2D)        (None, None, None, 128)   147584    \n",
      "_________________________________________________________________\n",
      "block2_pool (MaxPooling2D)   (None, None, None, 128)   0         \n",
      "_________________________________________________________________\n",
      "block3_conv1 (Conv2D)        (None, None, None, 256)   295168    \n",
      "_________________________________________________________________\n",
      "block3_conv2 (Conv2D)        (None, None, None, 256)   590080    \n",
      "_________________________________________________________________\n",
      "block3_conv3 (Conv2D)        (None, None, None, 256)   590080    \n",
      "_________________________________________________________________\n",
      "block3_pool (MaxPooling2D)   (None, None, None, 256)   0         \n",
      "_________________________________________________________________\n",
      "block4_conv1 (Conv2D)        (None, None, None, 512)   1180160   \n",
      "_________________________________________________________________\n",
      "block4_conv2 (Conv2D)        (None, None, None, 512)   2359808   \n",
      "_________________________________________________________________\n",
      "block4_conv3 (Conv2D)        (None, None, None, 512)   2359808   \n",
      "_________________________________________________________________\n",
      "block4_pool (MaxPooling2D)   (None, None, None, 512)   0         \n",
      "_________________________________________________________________\n",
      "block5_conv1 (Conv2D)        (None, None, None, 512)   2359808   \n",
      "_________________________________________________________________\n",
      "block5_conv2 (Conv2D)        (None, None, None, 512)   2359808   \n",
      "_________________________________________________________________\n",
      "block5_conv3 (Conv2D)        (None, None, None, 512)   2359808   \n",
      "_________________________________________________________________\n",
      "block5_pool (MaxPooling2D)   (None, None, None, 512)   0         \n",
      "=================================================================\n",
      "Total params: 14,714,688\n",
      "Trainable params: 14,714,688\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# 查看模型结构\n",
    "VGG16model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2、COCO数据集\n",
    "\n",
    "因为是训练图像转换模型，任意的图片数据集都是可以的，本文中用的是[COCO验证数据集](http://images.cocodataset.org/zips/val2014.zip)，大小6.2GB，共包括40504张图片。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3、运行配置"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "激活虚拟环境后，安装对应版本的tensorflow-addons。\n",
    "```\n",
    "pip install tensorflow-addons==0.9.1\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ VGG16模型计算风格损失的层以及权重"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "StyleLossLayer =  ['block1_conv2', 'block2_conv2', 'block3_conv3', 'block4_conv3']\n",
    "StyleWeight = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ VGG16模型计算内容损失的层以及权重"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "ContentLossLayer = ['block3_conv3']\n",
    "ContentWeight = 12"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ 总变分损失权重"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "TotalWeight = 1e-5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ COCO数据集路径"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "COCODataset = r'C:\\Users\\Administrator\\Desktop\\28tensorflow\\val2014'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ 存储训练的图像转换模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "SaveModelPath = r'C:\\Users\\Administrator\\Desktop\\28tensorflow\\style_model'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4、构建图像转换模型\n",
    "\n",
    "提升风格迁移图片质量的方法：\n",
    " + 用 Instance Normalization来代替通常的Batch Normalization，可以改善风格迁移的质量；风格转移的过程，就是要把风格图像的对比度转移到内容图像上，因此需要先去除内容图像本身的对比度。归一化操作其实就是在去除这种对比度，scale+shift则是将风格图像的对比度附加给内容图像。但是如果采用BN，计算出来的平均值和方差是整个batch内所有图像的均值和方差，而对于单张图像而言，其本身的均值和方差才是最能反映它的对比度的，因此BN会导致对比度去除得不彻底。因此Instance Normalization更加有效。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 反射填充\n",
    "class ReflectionPadding2D(tf.keras.layers.Layer):\n",
    "    def __init__(self, padding=1, **kwargs):\n",
    "        super(ReflectionPadding2D, self).__init__(**kwargs)\n",
    "        self.padding = padding\n",
    "\n",
    "    def compute_output_shape(self, s):\n",
    "        return s[0], s[1] + 2 * self.padding, s[2] + 2 * self.padding, s[3]\n",
    "\n",
    "    def call(self, x):\n",
    "        return tf.pad(x,[[0, 0],[self.padding, self.padding],[self.padding, self.padding],[0, 0],],'REFLECT')\n",
    "\n",
    "# 卷积，\n",
    "class ConvLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, channels, kernel_size=3, strides=1):\n",
    "        super(ConvLayer, self).__init__()\n",
    "        reflection_padding = kernel_size // 2\n",
    "        self.reflection_pad = ReflectionPadding2D(reflection_padding)\n",
    "        self.conv2d = Conv2D(channels, kernel_size, strides=strides)\n",
    "\n",
    "    def call(self, x):\n",
    "        x = self.reflection_pad(x)\n",
    "        x = self.conv2d(x)\n",
    "        return x\n",
    "\n",
    "# 上采样卷积\n",
    "class UpsampleConvLayer(tf.keras.layers.Layer):\n",
    "    def __init__(self, channels, kernel_size=3, strides=1, upsample=2):\n",
    "        super(UpsampleConvLayer, self).__init__()\n",
    "        reflection_padding = kernel_size // 2\n",
    "        self.reflection_pad = ReflectionPadding2D(reflection_padding)\n",
    "        self.conv2d = Conv2D(channels, kernel_size, strides=strides)\n",
    "        self.up2d = UpSampling2D(size=upsample)\n",
    "\n",
    "    def call(self, x):\n",
    "        x = self.up2d(x)\n",
    "        x = self.reflection_pad(x)\n",
    "        x = self.conv2d(x)\n",
    "        return x\n",
    "\n",
    "# 残差块\n",
    "class ResidualBlock(tf.keras.Model):\n",
    "    def __init__(self, channels, strides=1):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        self.conv1 = ConvLayer(channels, kernel_size=3, strides=strides)\n",
    "        self.in1 = InstanceNormalization()\n",
    "        self.conv2 = ConvLayer(channels, kernel_size=3, strides=strides)\n",
    "        self.in2 = InstanceNormalization()\n",
    "\n",
    "    def call(self, inputs):\n",
    "        residual = inputs\n",
    "\n",
    "        x = self.in1(self.conv1(inputs))\n",
    "        x = tf.nn.relu(x)\n",
    "\n",
    "        x = self.in2(self.conv2(x))\n",
    "        x = x + residual\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 按照1.1节中给出的模型结构图构建模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 图像转换模型\n",
    "class TransformerNet(tf.keras.Model):\n",
    "    def __init__(self):\n",
    "        super(TransformerNet, self).__init__()\n",
    "        \n",
    "        self.conv1 = ConvLayer(32, kernel_size=9, strides=1)\n",
    "        self.in1 = InstanceNormalization()\n",
    "        self.conv2 = ConvLayer(64, kernel_size=3, strides=2)\n",
    "        self.in2 = InstanceNormalization()\n",
    "        self.conv3 = ConvLayer(128, kernel_size=3, strides=2)\n",
    "        self.in3 = InstanceNormalization()\n",
    "\n",
    "        self.res1 = ResidualBlock(128)\n",
    "        self.res2 = ResidualBlock(128)\n",
    "        self.res3 = ResidualBlock(128)\n",
    "        self.res4 = ResidualBlock(128)\n",
    "        self.res5 = ResidualBlock(128)\n",
    "\n",
    "        self.deconv1 = UpsampleConvLayer(64, kernel_size=3, strides=1, upsample=2)\n",
    "        self.in4 = InstanceNormalization()\n",
    "        self.deconv2 = UpsampleConvLayer(32, kernel_size=3, strides=1, upsample=2)\n",
    "        self.in5 = InstanceNormalization()\n",
    "        self.deconv3 = ConvLayer(3, kernel_size=9, strides=1)\n",
    "        self.in6 = InstanceNormalization()\n",
    "        \n",
    "        self.relu = ReLU()\n",
    "\n",
    "    def call(self, x):\n",
    "        x = self.relu(self.in1(self.conv1(x)))  \n",
    "        x = self.relu(self.in2(self.conv2(x)))\n",
    "        x = self.relu(self.in3(self.conv3(x)))\n",
    "        x = self.res1(x)\n",
    "        x = self.res2(x)\n",
    "        x = self.res3(x)\n",
    "        x = self.res4(x)\n",
    "        x = self.res5(x)\n",
    "        x = self.relu(self.in4(self.deconv1(x)))\n",
    "        x = self.relu(self.in5(self.deconv2(x)))\n",
    "        x = self.in6(self.deconv3(x))\n",
    "        # 确保输出的数据在[0, 255]之间\n",
    "        x = (tf.nn.tanh(x)+ 1) * 255 / 2\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 选择VGG16预训练模型作为计算损失的网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StyleContentModel(tf.keras.models.Model):\n",
    "    def __init__(self, style_layers, content_layers, vgg):\n",
    "        super(StyleContentModel, self).__init__()\n",
    "        vgg.trainable = False\n",
    "\n",
    "        style_outputs = [vgg.get_layer(name).output for name in style_layers]\n",
    "        content_outputs = [vgg.get_layer(name).output for name in content_layers]\n",
    "\n",
    "        self.vgg = tf.keras.Model([vgg.input], [style_outputs, content_outputs])\n",
    "        self.vgg.trainable = False\n",
    "\n",
    "    def call(self, inputs):\n",
    "        preprocessed_input = tf.keras.applications.vgg16.preprocess_input(inputs)\n",
    "        style_outputs, content_outputs = self.vgg(preprocessed_input)\n",
    "        return style_outputs, content_outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5、构建图像数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def buildfigdataset(figpath, figsize, batchsize):\n",
    "    img_names = os.listdir(figpath)\n",
    "    # 进行路径拼接\n",
    "    img_list = [os.path.join(figpath, img_name) for img_name in img_names]\n",
    "    ds = tf.data.Dataset.from_tensor_slices((img_list))\n",
    "    def load_and_preprocess_from_path_label(path):\n",
    "        image = tf.io.read_file(path)  # 读取图片\n",
    "        # 保持通道数\n",
    "        image = tf.image.decode_jpeg(image, channels=3)\n",
    "        # 图片重新调整大小\n",
    "        image = tf.image.resize(image, [figsize, figsize]) \n",
    "        return image\n",
    "    image_ds  = ds.map(load_and_preprocess_from_path_label)\n",
    "    # 转换批次\n",
    "    batch_ds = image_ds.batch(batchsize)\n",
    "    return batch_ds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6、模型训练与保存\n",
    "\n",
    "模型的唯一性由风格图片名称, 风格、内容、总分损失权重确定，对于同样的参数设置可实现继续训练。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FASTSTMODEL:\n",
    "    \n",
    "    def __init__(self, stylefig, modelsavepath=SaveModelPath, model=VGG16model, \n",
    "                 contentfigpath=COCODataset, stylelayer=StyleLossLayer, contentlayer=ContentLossLayer,\n",
    "                 styleweight=StyleWeight, contentweight=ContentWeight, totalweight=TotalWeight,\n",
    "                 figsize=256, lr=0.003, epochs=10, batchs=4):\n",
    "        # 风格图片\n",
    "        self.stylefig = stylefig\n",
    "        \n",
    "        # 模型存储的路径\n",
    "        self.modelsavepath = modelsavepath\n",
    "        \n",
    "        # 下载的模型\n",
    "        self.model = model\n",
    "        # COCO图片数据集\n",
    "        self.contentfigpath =contentfigpath\n",
    "        \n",
    "        # 模型训练参数\n",
    "        self.figsize =figsize\n",
    "        self.lr= lr\n",
    "        self.epochs = epochs\n",
    "        self.batchs =batchs\n",
    "        \n",
    "        # 计算损失函数的参数\n",
    "        self.stylelayer = stylelayer\n",
    "        self.contentlayer = contentlayer\n",
    "        \n",
    "        # 损失权重\n",
    "        self.styleweight = styleweight\n",
    "        self.contentweight = contentweight\n",
    "        self.totalweight = totalweight\n",
    "        \n",
    "        # 图片数据集\n",
    "        self.figds = buildfigdataset(self.contentfigpath, self.figsize, self.batchs)\n",
    "        # 风格图片\n",
    "        self.style_image = loadimg(self.stylefig)\n",
    "    \n",
    "        # 将所有训练的误差存储为文件\n",
    "        self.styleloss = []\n",
    "        self.contentloss = []\n",
    "        \n",
    "        # 获取风格图片的文件 \n",
    "        self.getname()\n",
    "        \n",
    "        # 储存总损失的txt文件名\n",
    "        self.savetotalloss = '%s_savetotalloss.txt' % self.fignameckpt\n",
    "        \n",
    "    # 获取图片名称\n",
    "    def getname(self):\n",
    "        self.fignameckpt = os.path.basename(self.stylefig).split('.')[0]\n",
    "        \n",
    "    # 读取前一次训练的最终总损失\n",
    "    def getlastloss(self):\n",
    "        with open(r'%s/%s' % (self.modelsavepath, self.savetotalloss), 'r') as d:\n",
    "            f = d.readlines()\n",
    "            if f:\n",
    "                return float(f[-1])\n",
    "            else:\n",
    "                return False\n",
    "            \n",
    "    # 存储最终总损失\n",
    "    def savelastloss(self, loss):\n",
    "        with open(r'%s/%s' % (self.modelsavepath, self.savetotalloss), 'a') as d:\n",
    "            d.write(str(loss)+'\\n')\n",
    "\n",
    "        \n",
    "    # 训练保存模型，\n",
    "    def train_fast_style_model(self):\n",
    "\n",
    "        # 损失网络的输出\n",
    "        extractor = StyleContentModel(self.stylelayer, self.contentlayer, self.model)\n",
    "        \n",
    "        # 图像转换模型\n",
    "        transformer = TransformerNet()\n",
    "\n",
    "        # 计算风格图片的Gram矩阵\n",
    "        style_features, _ = extractor(self.style_image)\n",
    "        gram_style = [grammatrix(x) for x in style_features]\n",
    "        \n",
    "        # 优化器tf.train.Checkpoint\n",
    "        optimizer = tf.optimizers.Adam(learning_rate=self.lr)\n",
    "        # 声明\n",
    "        ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=optimizer, transformer=transformer)\n",
    "        # 管理检查点\n",
    "        log_dir = os.path.join(self.modelsavepath, 'figname=%s_sw=%s_cw=%s_tw=%s'% \n",
    "                               (self.fignameckpt, self.styleweight, self.contentweight, self.totalweight))\n",
    "        manager = tf.train.CheckpointManager(ckpt, log_dir, max_to_keep=1)\n",
    "    \n",
    "        if manager.latest_checkpoint:\n",
    "            # 如果有模型则恢复模型\n",
    "            ckpt.restore(manager.latest_checkpoint)\n",
    "        else:\n",
    "            # 首次训练\n",
    "            # 建立一个储存每次训练最终总损失的txt文件\n",
    "            with open(r'%s/%s' % (self.modelsavepath, self.savetotalloss), 'w') as f:\n",
    "                f.write('')\n",
    "            \n",
    "        def train_step(images):\n",
    "            with tf.GradientTape() as tape:\n",
    "                # 输入的图片数据集的 图像转换模型的输出\n",
    "                transformed_images = transformer(images)\n",
    "                \n",
    "                # 损失网络的内容输出\n",
    "                _, content_features = extractor(images)\n",
    "                \n",
    "                # 输入的图片数据集的 图像转换模型的输出\n",
    "                style_transformed, content_transformed = extractor(transformed_images)\n",
    "                \n",
    "                # 风格损失：风格图片和和经过图像转换模型转换后的图像数据集之间的\n",
    "                style_loss = self.styleweight * stylelossfunc(gram_style, style_transformed)\n",
    "                # 内容损失：内容图片和经过图像转换模型转换后的的图像数据集之间的\n",
    "                content_loss = self.contentweight * contentlossfunc(content_features, content_transformed)\n",
    "                # 总分损失：\n",
    "                total_variation_loss = self.totalweight * totalvariationloss(transformed_images)\n",
    "                # 总的损失\n",
    "                loss = style_loss + content_loss + total_variation_loss\n",
    "             \n",
    "            # 梯度\n",
    "            gradients = tape.gradient(loss, transformer.trainable_variables)\n",
    "            # 更新参数\n",
    "            optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))\n",
    "            \n",
    "            # 存储\n",
    "            self.styleloss.append(style_loss.numpy())\n",
    "            self.contentloss.append(content_loss.numpy())\n",
    "            #print('风格损失:', style_loss.numpy(), '内容损失:', content_loss.numpy(), '总分损失:', total_variation_loss.numpy())\n",
    "            return loss.numpy()\n",
    "        \n",
    "        # 开始训练\n",
    "        for epoch in range(self.epochs):\n",
    "            for images in self.figds:\n",
    "                sumloss = train_step(images)\n",
    "                #print('总损失：', sumloss)\n",
    "                ckpt.step.assign_add(1)\n",
    "                step = int(ckpt.step)\n",
    "                \n",
    "                # 读取存储的损失\n",
    "                savedloss = self.getlastloss()\n",
    "                \n",
    "                if savedloss:\n",
    "                    print('代数%s-%s'%(self.epochs, epoch), 'step数%s'%step, '总损失：', sumloss, )\n",
    "                    print('保存模型', manager.save())\n",
    "                    # 存储损失\n",
    "                    self.savelastloss(sumloss)\n",
    "                else:\n",
    "                    print('保存模型', manager.save())\n",
    "                    self.savelastloss(sumloss)                   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "针对不同的风格图片建立不同的模型。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "StyleFig1 = r'C:\\Users\\Administrator\\Desktop\\28tensorflow\\style1.jpg'\n",
    "# 进行图像转换模型的训练\n",
    "fase_trans_model = FASTSTMODEL(StyleFig1, epochs=3)\n",
    "fase_trans_model.train_fast_style_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "StyleFig2 = r'C:\\Users\\Administrator\\Desktop\\28tensorflow\\style2.jpg'\n",
    "# 进行图像转换模型的训练\n",
    "fase_trans_model = FASTSTMODEL(StyleFig2, epochs=3)\n",
    "fase_trans_model.train_fast_style_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "StyleFig3 = r'C:\\Users\\Administrator\\Desktop\\28tensorflow\\style3.jpg'\n",
    "# 进行图像转换模型的训练\n",
    "fase_trans_model = FASTSTMODEL(StyleFig3, epochs=3)\n",
    "fase_trans_model.train_fast_style_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "StyleFig4 = r'C:\\Users\\Administrator\\Desktop\\28tensorflow\\style4.jpg'\n",
    "# 进行图像转换模型的训练\n",
    "fase_trans_model = FASTSTMODEL(StyleFig4, epochs=3)\n",
    "fase_trans_model.train_fast_style_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7、转换图片"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "+ 需要转换的图片的存储路径"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "InFigPath = r'C:\\Users\\Administrator\\Desktop\\28tensorflow\\infig'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FIGFST(FASTSTMODEL):\n",
    "    # 初始化\n",
    "    def __init__(self, stylefig, transfigpath):\n",
    "        super(FIGFST, self).__init__(stylefig)\n",
    "        self.tsfp = transfigpath\n",
    "        self.otfp = self.tsfp + '_' +self.fignameckpt\n",
    "        \n",
    "        # 限制图片分辨率\n",
    "        self.maxpixel = 800\n",
    "        \n",
    "        self.checkoutpath()\n",
    "    \n",
    "    # 如果存在输出文件夹，则清空，不存在则新建\n",
    "    def checkoutpath(self):\n",
    "        # \n",
    "        if os.path.exists(self.otfp):\n",
    "            shutil.rmtree(self.otfp)\n",
    "        os.mkdir(self.otfp)\n",
    "        \n",
    "        \n",
    "    # 读取图片数据\n",
    "    def loadimgplus(self,figpath):\n",
    "        img = tf.io.read_file(figpath)\n",
    "        img = tf.image.decode_image(img, channels=3)\n",
    "        img = tf.cast(img, tf.float32)\n",
    "        \n",
    "        height, width, _ = img.shape\n",
    "        \n",
    "        maxh = max(height, width)\n",
    "        if height > self.maxpixel or width > self.maxpixel:\n",
    "            # 等比例缩放\n",
    "            img = tf.image.resize(img, [int(self.maxpixel/maxh*height), int(self.maxpixel/maxh*width)]) \n",
    "        imglast = img[tf.newaxis, :]\n",
    "        return imglast\n",
    "\n",
    "    # 加载最新的模型\n",
    "    def faststyletransfer(self):\n",
    "        for imagefp in os.listdir(self.tsfp):\n",
    "            figpath = os.path.join(self.tsfp, imagefp)\n",
    "            image = self.loadimgplus(figpath)\n",
    "\n",
    "            # 引入图像转换模型\n",
    "            transformer = TransformerNet()\n",
    "            # 声明\n",
    "            ckpt = tf.train.Checkpoint(transformer=transformer)\n",
    "            # 恢复模型参数\n",
    "            ckpapath = os.path.join(self.modelsavepath,  'figname=%s_sw=%s_cw=%s_tw=%s'% \n",
    "                                    (self.fignameckpt, self.styleweight, self.contentweight, self.totalweight))\n",
    "        \n",
    "            ckpt.restore(tf.train.latest_checkpoint(ckpapath)).expect_partial()\n",
    "\n",
    "            transformed_image = transformer(image)\n",
    "           \n",
    "            transformed_image = tf.cast(tf.squeeze(transformed_image), tf.uint8).numpy()\n",
    "        \n",
    "            img = Image.fromarray(transformed_image, mode='RGB')\n",
    "\n",
    "            img.save(r'%s/trans_%s' % (self.otfp, imagefp))\n",
    "        print('转换完毕')\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "进行图片的转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "转换完毕\n"
     ]
    }
   ],
   "source": [
    "figstyles = FIGFST(StyleFig4, InFigPath)\n",
    "figstyles.faststyletransfer()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 8、动图、视频转换\n",
    "\n",
    "将本地视频文件或者动图按照帧转化为图片，然后再将风格迁移后的图片连接成动图。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image, ImageSequence  # 动图提取图片\n",
    "import cv2  # 视频转成图片\n",
    "import imageio # 图片链接成gif\n",
    "import shutil # 清空文件夹\n",
    "\n",
    "class FASTST:\n",
    "    \n",
    "    def __init__(self, stylefig, videogifpath, infigpath):\n",
    "        \n",
    "        self.styf = stylefig\n",
    "        \n",
    "        self.vidp = videogifpath\n",
    "        self.infp = infigpath\n",
    "        self.oufp = self.infp  + '_' + os.path.basename(self.styf).split('.')[0]\n",
    "        \n",
    "        # 图片转换\n",
    "        self.figmodel = FIGFST(self.styf, self.infp)\n",
    "        \n",
    "        # 转换\n",
    "        self.judgevideogif()\n",
    "    \n",
    "    # 根据视频还是动图选择\n",
    "    def judgevideogif(self):\n",
    "        name = self.vidp.split('.')[-1]\n",
    "        if name == 'gif':\n",
    "            self.parsegif()\n",
    "        else:\n",
    "            self.video2figure()\n",
    "    \n",
    "    # 视频转为图片\n",
    "    def video2figure(self, time_interval=10):\n",
    "        # 清空文件夹中的文件\n",
    "        shutil.rmtree(self.infp)\n",
    "        os.mkdir(self.infp)\n",
    "        \n",
    "        fig_list = []\n",
    "        vidcap = cv2.VideoCapture(self.vidp)\n",
    "        success, image = vidcap.read()\n",
    "        count = 0\n",
    "        while success:\n",
    "            success, image = vidcap.read()\n",
    "            if count % time_interval == 0:\n",
    "                if image is None:\n",
    "                    break\n",
    "                cv2.imencode('.jpg', image)[1].tofile('%s/fig_%d.jpg' % (self.infp, count))\n",
    "                fig_list.append('fig_%d.jpg' % count)\n",
    "            count += 1\n",
    "        print('视频转换为图片，开始转换')\n",
    "        \n",
    "        # 清空文件夹中的文件\n",
    "        shutil.rmtree(self.oufp)\n",
    "        os.mkdir(self.oufp)\n",
    "        \n",
    "        # 开始转换\n",
    "        self.figmodel.faststyletransfer()\n",
    "        return print('转换完成')\n",
    "    \n",
    "    # 提取gif动图变为图片\n",
    "    def parsegif(self):\n",
    "        # 清空文件夹中的文件\n",
    "        shutil.rmtree(self.infp)\n",
    "        os.mkdir(self.infp)\n",
    "        \n",
    "        # 读取GIF\n",
    "        im = Image.open(self.vidp)\n",
    "        # GIF图片流的迭代器\n",
    "        iterfigs = ImageSequence.Iterator(im)\n",
    "        # 获取文件名\n",
    "        count = 1\n",
    "\n",
    "        # 遍历图片流的每一帧\n",
    "        for frame in iterfigs:\n",
    "            frame.save('%s/fig_%d.png' % (self.infp, count))\n",
    "            count += 1\n",
    "        print('动图转换为图片，开始转换')\n",
    "        \n",
    "        # 清空文件夹中的文件\n",
    "        shutil.rmtree(self.oufp)\n",
    "        os.mkdir(self.oufp)\n",
    "        \n",
    "        # 开始转换\n",
    "        self.figmodel.faststyletransfer()\n",
    "        return print('转换完成')\n",
    "        \n",
    "    # 将文件夹子中的图片链接成动图\n",
    "    def create_gif(self, duration=0.1):\n",
    "        name = os.path.basename(self.styf).split('.')[0] + '.gif'\n",
    "        frames = []\n",
    "        for image_name in os.listdir(self.oufp):\n",
    "            frames.append(imageio.imread('%s/%s' % (self.oufp, image_name)))\n",
    "        imageio.mimsave(name, frames, 'GIF', duration=duration)\n",
    "        return print('动图生成完成')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "动图转换为图片，开始转换\n",
      "转换完毕\n",
      "转换完成\n",
      "动图生成完成\n"
     ]
    }
   ],
   "source": [
    "# 视频路径\n",
    "videopath = r'C:\\Users\\Administrator\\Desktop\\28tensorflow\\3.gif'\n",
    "videoinpath = r'C:\\Users\\Administrator\\Desktop\\28tensorflow\\videoin'\n",
    "videost = FASTST(StyleFig4, videopath, videoinpath)\n",
    "videost.create_gif()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf28",
   "language": "python",
   "name": "an21"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}