202003_models/CBDNet_CVPR2019.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "CBDNet_CVPR2019.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyNe66Pvzdd/4AAN/HOlEKmw",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/OUCTheoryGroup/colab_demo/blob/master/CBDNet_CVPR2019.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RNN7mH7Stqrl",
        "colab_type": "text"
      },
      "source": [
        "## 图像去噪网络 CBDNet\n",
        "\n",
        "> 本代码解析使用了 浙江大学  IDKiro 复现的代码，对他的辛苦工作表示感谢！\n",
        "> 在原代码的基础上，我对代码进行了简化，因此效果应该略有不同。\n",
        "> 原代码地址：https://github.com/IDKiro/CBDNet-pytorch\n",
        "\n",
        "《Toward Convolutional Blind Denoising of Real Photographs》， *CVPR* 2019\n",
        "\n",
        "Shi Guo, Zhifei Yan, Kai Zhang, Wangmeng Zuo, Lei Zhang.\n",
        "\n",
        "这是哈工大与香港理工大Lei Zhang老师课题组合作完成的论文，这两个团队在图像去噪方面一直走在前沿，许多经典工作都是他们提出的，如WNNM、DnCNN等。这一篇也是其在深度图像去噪方面的新的文章。与其前面的工作不同的是，以前的图像去噪大多使用合成数据，这篇文章研究了CNN在真实图像上的去噪效果，其主要贡献在于以下几点：\n",
        "\n",
        "- 提出了一个更加真实的噪声模型，其考虑了信号依赖噪声和ISP流程对噪声的影响，展示了图像噪声模型在真实噪声图像中起着关键作用。\n",
        "- 提出了CBDNet模型，其包括了一个噪声估计子网络和一个非盲去噪子网络，可以实现图像的盲去噪（即未知噪声水平）。\n",
        "- 提出了非对称学习（asymmetric learning）的损失函数，并允许用户交互式调整去噪结果，增强了去噪结果的鲁棒性。\n",
        "- 将合成噪声图像与真实噪声图像一起用于网络的训练，提升网络的去噪效果和泛化能力。\n",
        "\n",
        "下面提供详细的代码解析。首先下载数据，论文中使用的数据集非常大，为了方便colab平台运行，制作了一个mini_denoise_dataset。\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2XFvF0f3ifgm",
        "colab_type": "code",
        "outputId": "94111fb7-af57-49db-a7e7-e67416c772cd",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 773
        }
      },
      "source": [
        "! wget https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/mini_denoise_dataset.zip\n",
        "! unzip mini_denoise_dataset"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "--2020-03-03 01:13:35--  https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/mini_denoise_dataset.zip\n",
            "Resolving q6dz4bbgt.bkt.clouddn.com (q6dz4bbgt.bkt.clouddn.com)... 101.89.125.237, 101.89.125.227, 101.89.125.229, ...\n",
            "Connecting to q6dz4bbgt.bkt.clouddn.com (q6dz4bbgt.bkt.clouddn.com)|101.89.125.237|:80... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 286939017 (274M) [application/zip]\n",
            "Saving to: ‘mini_denoise_dataset.zip’\n",
            "\n",
            "mini_denoise_datase 100%[===================>] 273.65M  2.56MB/s    in 1m 57s  \n",
            "\n",
            "2020-03-03 01:15:33 (2.35 MB/s) - ‘mini_denoise_dataset.zip’ saved [286939017/286939017]\n",
            "\n",
            "Archive:  mini_denoise_dataset.zip\n",
            "   creating: mini_denoise_dataset/\n",
            "   creating: mini_denoise_dataset/test/\n",
            "  inflating: mini_denoise_dataset/test/0001.bmp  \n",
            "  inflating: mini_denoise_dataset/test/0002.bmp  \n",
            "  inflating: mini_denoise_dataset/test/0003.bmp  \n",
            "  inflating: mini_denoise_dataset/test/0004.bmp  \n",
            "  inflating: mini_denoise_dataset/test/0005.bmp  \n",
            "  inflating: mini_denoise_dataset/test/0006.bmp  \n",
            "  inflating: mini_denoise_dataset/test/0007.bmp  \n",
            "  inflating: mini_denoise_dataset/test/0008.bmp  \n",
            "  inflating: mini_denoise_dataset/test/0009.bmp  \n",
            "  inflating: mini_denoise_dataset/test/0010.bmp  \n",
            "   creating: mini_denoise_dataset/train/\n",
            "   creating: mini_denoise_dataset/train/Batch_001/\n",
            "  inflating: mini_denoise_dataset/train/Batch_001/IMG_20160202_015216Reference.bmp  \n",
            "  inflating: mini_denoise_dataset/train/Batch_001/IMG_20160202_015247Noisy.bmp  \n",
            "  inflating: mini_denoise_dataset/train/Batch_001/IMG_20160202_015252Noisy.bmp  \n",
            "   creating: mini_denoise_dataset/train/Batch_002/\n",
            "  inflating: mini_denoise_dataset/train/Batch_002/IMG_20160202_021953Reference.bmp  \n",
            "  inflating: mini_denoise_dataset/train/Batch_002/IMG_20160202_022024Noisy.bmp  \n",
            "  inflating: mini_denoise_dataset/train/Batch_002/IMG_20160202_022030Noisy.bmp  \n",
            "   creating: mini_denoise_dataset/train/Batch_003/\n",
            "  inflating: mini_denoise_dataset/train/Batch_003/IMG_20160204_035200Reference.bmp  \n",
            "  inflating: mini_denoise_dataset/train/Batch_003/IMG_20160204_035217Noisy.bmp  \n",
            "  inflating: mini_denoise_dataset/train/Batch_003/IMG_20160204_035222Noisy.bmp  \n",
            "   creating: mini_denoise_dataset/train/Batch_004/\n",
            "  inflating: mini_denoise_dataset/train/Batch_004/IMG_20160204_042548Reference.bmp  \n",
            "  inflating: mini_denoise_dataset/train/Batch_004/IMG_20160204_042603Noisy.bmp  \n",
            "  inflating: mini_denoise_dataset/train/Batch_004/IMG_20160204_042608Noisy.bmp  \n",
            "   creating: mini_denoise_dataset/train/Batch_005/\n",
            "  inflating: mini_denoise_dataset/train/Batch_005/IMG_20160204_052934Reference.bmp  \n",
            "  inflating: mini_denoise_dataset/train/Batch_005/IMG_20160204_052950Noisy.bmp  \n",
            "  inflating: mini_denoise_dataset/train/Batch_005/IMG_20160204_052954Noisy.bmp  \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QO5OkcnczqzJ",
        "colab_type": "text"
      },
      "source": [
        "引入基本的库："
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "93YfWRnfdh-r",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os, time, scipy.io, shutil\n",
        "from PIL import Image\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import numpy as np\n",
        "import glob\n",
        "import re\n",
        "import cv2"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bav526pkzt_i",
        "colab_type": "text"
      },
      "source": [
        "## 2. 定义 CBDNet 网络\n",
        "\n",
        "网络结构如下所示：\n",
        "\n",
        "![CBDNet网络结构](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200303081046.jpg)\n",
        "\n",
        "整体架构可以看到，网络由一个 全卷积网络FCN，和一个 UNet 组成。\n",
        "\n",
        "### 2.1 FCN部分\n",
        "\n",
        "包括 5 次 conv 操作，使用 3x3 的卷积核，使用了1个像素的padding来保证尺寸一致，feature map 数量依次为：3 ==> 32 ==> 32 ==> 32 ==> 32 ==> 3\n",
        "\n",
        "这部分的代码如下 （和论文架构图中展示的完全一致）：\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EZ4YdqzZ_hMQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class FCN(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(FCN, self).__init__()\n",
        "\n",
        "        # 3 ==> 32 的输入卷积\n",
        "        self.inc = nn.Sequential(\n",
        "            nn.Conv2d(3, 32, 3, padding=1),\n",
        "            nn.ReLU(inplace=True))\n",
        "        \n",
        "        # 32 ==> 32 的中间卷积\n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv2d(32, 32, 3, padding=1),\n",
        "            nn.ReLU(inplace=True)\n",
        "        )\n",
        "        \n",
        "        # 32 ==> 3 的输出卷积 \n",
        "        self.outc = nn.Sequential(\n",
        "            nn.Conv2d(32, 3, 3, padding=1),\n",
        "            nn.ReLU(inplace=True)\n",
        "        )\n",
        "    \n",
        "    def forward(self, x):\n",
        "        # 第 1 次卷积\n",
        "        conv1 = self.inc(x)\n",
        "        # 第 2 次卷积\n",
        "        conv2 = self.conv(conv1)\n",
        "        # 第 3 次卷积\n",
        "        conv3 = self.conv(conv2)\n",
        "        # 第 4 次卷积\n",
        "        conv4 = self.conv(conv3)\n",
        "        # 第 5 次卷积\n",
        "        conv5 = self.outc(conv4)\n",
        "        return conv5"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "31BW970pA28x",
        "colab_type": "text"
      },
      "source": [
        "### 2.2 UNet部分\n",
        "\n",
        "如上面架构图所示，UNet 使用到了大量卷积，全部是 3x3 ，加了 1 个像素的padding保证尺寸，这里编写了一个 single_conv 类用于卷积操作，包括卷积 和 ReLU 函数。这个类创建需要两个参数：输入的通道数 in_ch，输出的通道数 out_ch。具体如下：\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JYY6wJA0EatC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class single_conv(nn.Module):\n",
        "    def __init__(self, in_ch, out_ch):\n",
        "        super(single_conv, self).__init__()\n",
        "        self.conv = nn.Sequential(\n",
        "            nn.Conv2d(in_ch, out_ch, 3, padding=1),\n",
        "            nn.ReLU(inplace=True))\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv(x)\n",
        "        return x"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6jD_J8bJEiHc",
        "colab_type": "text"
      },
      "source": [
        "从网络架构图可以看到，网络有两次下采样，还有两次上采样。下采样使用 2x2 的均值 pooling 就可以了，如何处理上采样呢？\n",
        "\n",
        "上采样采用的是反卷积，这里编写了一个 up 类。输入的通道数是 in_ch，输出的通道数是 in_ch//2。因为是反卷积，所以使用 nn.ConvTranspose2d 函数，卷积核大小为2。即原图像中的 1 个像素经过卷积会变成 2*2 的区域。同时，卷积的步长为 2，卷积结果紧密的拼接为一张大图。\n",
        "\n",
        "同时，也有一些需要特殊考虑的地方，从图中可以看出，这里有一个特征融合的步骤，这时可能会产生一定的问题，那就是两个feature map 尺寸可能不一样，比如：之前尺寸是 7，下采样再上采样的话，尺寸变化为： 7 ==> 3 ==> 6 ，因为 2x2 pooling 的时候，不会考虑最边上的像素！\n",
        "\n",
        "所以，在 forward 函数中，加入了一个 padding 操作，代码如下："
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TfTlKE8nEmEi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class up(nn.Module):\n",
        "    def __init__(self, in_ch):\n",
        "        super(up, self).__init__()\n",
        "        self.up = nn.ConvTranspose2d(in_ch, in_ch//2, 2, stride=2)\n",
        "\n",
        "    # forward 需要两个输入，x1 是需要上采样的小尺寸 feature map\n",
        "    # x2 是以前的大尺寸 feature map，因为中间的 pooling 可能损失了边缘像素，\n",
        "    # 所以上采样以后的 x1 可能会比 x2 尺寸小\n",
        "    def forward(self, x1, x2):\n",
        "        # x1 上采样\n",
        "        x1 = self.up(x1)\n",
        "        \n",
        "        # 输入数据是四维的，第一个维度是样本数，剩下的三个维度是 CHW\n",
        "        # 所以 Y 方向上的悄寸差别在 [2],  X 方向上的尺寸差别在 [3] \n",
        "        diffY = x2.size()[2] - x1.size()[2]\n",
        "        diffX = x2.size()[3] - x1.size()[3]\n",
        "        # 给 x1 进行 padding 操作\n",
        "        x1 = F.pad(x1, (diffX // 2, diffX - diffX//2,\n",
        "                        diffY // 2, diffY - diffY//2))\n",
        "        # 把 x2 加到反卷积后的 feature map\n",
        "        x = x2 + x1\n",
        "        return x"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cjglQG5VLsAK",
        "colab_type": "text"
      },
      "source": [
        "需要注意的是，输出层也写了一个类，输出部分是将 64 个 feature map，利用 1x1 的卷积变成 3 个 feature map。 教程里介绍 GoogLetNet，ResNet 的时候也有写，1x1 的卷积可以较好的起到降维作用。**大家注意这里，最后一层不使用激活函数。**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wMQNmSbCLrKt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class outconv(nn.Module):\n",
        "    def __init__(self, in_ch, out_ch):\n",
        "        super(outconv, self).__init__()\n",
        "        self.conv = nn.Conv2d(in_ch, out_ch, 1)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv(x)\n",
        "        return x"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PdJAULTLNY4b",
        "colab_type": "text"
      },
      "source": [
        "下面是 UNet 部分的完整代码，和架构图示完全一致，包括：\n",
        "- input conv : 6 ==> 64 ==> 64\n",
        "- down1 : 2x2 的均值 pooling\n",
        "- conv1 : 64 ==> 128 ==> 128 ==> 128\n",
        "- down2 : 2x2 的均值 pooling\n",
        "- conv2 : 128 ==> 256 ==> 256 ==> 256 ==> 256 ==> 256 ==> 256\n",
        "- up1 : conv2 反卷积，和 conv1 的结果相加，输入256，输出128\n",
        "- conv3 :  128 ==> 128 ==> 128 ==> 128\n",
        "- up2 : conv3 反卷积，和 input conv 的结果相加，输入128，输出64\n",
        "- conv4 : 64 ==> 64 ==> 64\n",
        "- output conv:  65 ==> 3，用1x1的卷积降维，得到降噪结果\n",
        "\n",
        "下面结合代码详细解释："
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mJnLGGyYNXHh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class UNet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(UNet, self).__init__()\n",
        "\n",
        "        self.inc = nn.Sequential(\n",
        "            single_conv(6, 64),\n",
        "            single_conv(64, 64))\n",
        "\n",
        "        self.down1 = nn.AvgPool2d(2)\n",
        "        self.conv1 = nn.Sequential(\n",
        "            single_conv(64, 128),\n",
        "            single_conv(128, 128),\n",
        "            single_conv(128, 128))\n",
        "\n",
        "        self.down2 = nn.AvgPool2d(2)\n",
        "        self.conv2 = nn.Sequential(\n",
        "            single_conv(128, 256),\n",
        "            single_conv(256, 256),\n",
        "            single_conv(256, 256),\n",
        "            single_conv(256, 256),\n",
        "            single_conv(256, 256),\n",
        "            single_conv(256, 256))\n",
        "\n",
        "        self.up1 = up(256)\n",
        "        self.conv3 = nn.Sequential(\n",
        "            single_conv(128, 128),\n",
        "            single_conv(128, 128),\n",
        "            single_conv(128, 128))\n",
        "\n",
        "        self.up2 = up(128)\n",
        "        self.conv4 = nn.Sequential(\n",
        "            single_conv(64, 64),\n",
        "            single_conv(64, 64))\n",
        "\n",
        "        self.outc = outconv(64, 3)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # input conv : 6 ==> 64 ==> 64\n",
        "        inx = self.inc(x)\n",
        "\n",
        "        # 均值 pooling, 然后 conv1 : 64 ==> 128 ==> 128 ==> 128\n",
        "        down1 = self.down1(inx)\n",
        "        conv1 = self.conv1(down1)\n",
        "\n",
        "        # 均值 pooling，然后 conv2 : 128 ==> 256 ==> 256 ==> 256 ==> 256 ==> 256 ==> 256\n",
        "        down2 = self.down2(conv1)\n",
        "        conv2 = self.conv2(down2)\n",
        "\n",
        "        # up1 : conv2 反卷积，和 conv1 的结果相加，输入256，输出128\n",
        "        up1 = self.up1(conv2, conv1)\n",
        "        # conv3 : 128 ==> 128 ==> 128 ==> 128\n",
        "        conv3 = self.conv3(up1)\n",
        "\n",
        "        # up2 : conv3 反卷积，和 input conv 的结果相加，输入128，输出64\n",
        "        up2 = self.up2(conv3, inx)\n",
        "        # conv4 : 64 ==> 64 ==> 64\n",
        "        conv4 = self.conv4(up2)\n",
        "\n",
        "        # output conv: 65 ==> 3，用1x1的卷积降维，得到降噪结果\n",
        "        out = self.outc(conv4)\n",
        "        return out"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nRRAm960TiQF",
        "colab_type": "text"
      },
      "source": [
        "### 2.3 CBDNet 分析\n",
        "\n",
        "下面是 CBDNet 整个网络的代码，先将数据输入 FCN，得到估计的噪声强度： noise_level，为 3 通道。然后将 3通道的原图像，和 noise_level 拼接在一起，作为 UNet 的输入。\n",
        "\n",
        "UNet 经过一系列操作，得到 out ，这里的 out 被认为是噪声的 residual mapping，和 输入图像加在一起，输出最终的去噪图像。\n",
        "\n",
        "可以看出，这里也采用了一 residual learning 的思想，认为 噪声的 residual mapping 学习起来更加容易。\n",
        "\n",
        "整体网络的代码如下："
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xkQt4elWTni3",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class CBDNet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(CBDNet, self).__init__()\n",
        "        self.fcn = FCN()\n",
        "        self.unet = UNet()\n",
        "    \n",
        "    def forward(self, x):\n",
        "        noise_level = self.fcn(x)\n",
        "        concat_img = torch.cat([x, noise_level], dim=1)\n",
        "        out = self.unet(concat_img) + x\n",
        "        return noise_level, out"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ef45SuXPY2Q-",
        "colab_type": "text"
      },
      "source": [
        "### 2.4 损失函数设计\n",
        "\n",
        "损失函数包括三部分：\n",
        "\n",
        "![损失函数说明](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200303132845.jpg)\n",
        "\n",
        "三个部分加在一起，就是最终的损失函数，下面结合代码来看。"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YV21RySPdw8Q",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class fixed_loss(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        \n",
        "    def forward(self, out_image, gt_image, est_noise, gt_noise, if_asym):\n",
        "        # 分别得到图像的高度和宽度\n",
        "        h_x = est_noise.size()[2]\n",
        "        w_x = est_noise.size()[3]\n",
        "        # 每个样本为 CHW ，把 H 方向第一行的数据去掉，统计一下一共多少元素\n",
        "        count_h = self._tensor_size(est_noise[:, :, 1:, :])\n",
        "        # 每个样本为 CHW ，把 W 方向第一列的数据去掉，统计一下一共多少元素\n",
        "        count_w = self._tensor_size(est_noise[:, :, : ,1:])\n",
        "        # H 方向，第一行去掉得后的矩阵，减去最后一行去掉后的矩阵，即下方像素减去上方像素，平方，然后求和\n",
        "        h_tv = torch.pow((est_noise[:, :, 1:, :] - est_noise[:, :, :h_x-1, :]), 2).sum()\n",
        "        # W 方向，第一列去掉得后的矩阵，减去最后一列去掉后的矩阵，即右方像素减去左方像素，平方，然后求和\n",
        "        w_tv = torch.pow((est_noise[:, :, :, 1:] - est_noise[:, :, :, :w_x-1]), 2).sum()\n",
        "        # 求平均，得到平均每个像素上的 tvloss\n",
        "        tvloss = h_tv / count_h + w_tv / count_w\n",
        "\n",
        "        loss = torch.mean( \\\n",
        "                # 第三部分：重建损失\n",
        "                torch.pow((out_image - gt_image), 2)) + \\\n",
        "                # 第一部分：对比损失\n",
        "                if_asym * 0.5 * torch.mean(torch.mul(torch.abs(0.3 - F.relu(gt_noise - est_noise)), torch.pow(est_noise - gt_noise, 2))) + \\\n",
        "                # 第二部分：起平滑作用的 tvloss\n",
        "                0.05 * tvloss\n",
        "        return loss\n",
        "\n",
        "    def _tensor_size(self,t):\n",
        "        return t.size()[1]*t.size()[2]*t.size()[3]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vjMhEpNLz5kL",
        "colab_type": "text"
      },
      "source": [
        "从上面的代码中可以看到，对比损失前系数为 0.5， alpha 取值为 0.3，tvloss 系数为 0.05，和论文里的默认参数一致。\n",
        "\n",
        "**【划重点】这里需要专门指出的是：**\n",
        "\n",
        "对于 gt_noise，只有在使用合成数据进行训练时才会用到；以前的图像去噪，大多在真实图像上加一个随机Gauss噪声，得到噪声图像，这时 gt_noise 是已知的，就能够输入。\n",
        "\n",
        "这个教程里处理的是真实图像，因此没有 gt_noise，所以在训练时，gt_noise 一直是0。原来代码里专门有一部分是人工合成噪声来训练，为方便理解代码，我暂时去掉了这部分。\n",
        "\n",
        "\n",
        "下面是两个程序中要用到的两个小函数："
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nHW4KXRqsYfD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# 这个类用于存储 loss，观察结果时使用\n",
        "# 每轮训练一张图像，就计算一下 loss 的均值存储在 self.avg 里，用于输出观察变化\n",
        "# 同时，把当前 loss 的值存储在 self.val 里\n",
        "class AverageMeter(object):\n",
        "\tdef __init__(self):\n",
        "\t\tself.reset()\n",
        "\n",
        "\tdef reset(self):\n",
        "\t\tself.val = 0\n",
        "\t\tself.avg = 0\n",
        "\t\tself.sum = 0\n",
        "\t\tself.count = 0\n",
        "\n",
        "\tdef update(self, val, n=1):\n",
        "\t\tself.val = val\n",
        "\t\tself.sum += val * n\n",
        "\t\tself.count += n\n",
        "\t\tself.avg = self.sum / self.count\n",
        "\n",
        "# 图像矩阵由 hwc 转换为 chw ，这个就不多解释了\n",
        "def hwc_to_chw(img):\n",
        "    return np.transpose(img, axes=[2, 0, 1])\n",
        "# 图像矩阵由 chw 转换为 hwc ，这个也不多解释\n",
        "def chw_to_hwc(img):\n",
        "    return np.transpose(img, axes=[1, 2, 0])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eNv0v8q_1DpR",
        "colab_type": "text"
      },
      "source": [
        "## 3. 初始化基本变量，开始训练"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mEAWejlSixGN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# 训练的时候，输入图像尺寸都是 ps x ps 的\n",
        "ps = 256\n",
        "\n",
        "train_dir = './mini_denoise_dataset/train/'\n",
        "train_fns = glob.glob(train_dir + 'Batch_*')\n",
        "\n",
        "origin_imgs = [None] * len(train_fns)\n",
        "noised_imgs = [None] * len(train_fns)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vKvLjoa-3Akd",
        "colab_type": "text"
      },
      "source": [
        "定义网络模型、优化器、损失函数"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iEM-JOGLzKRk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# 使用GPU训练，可以在菜单 \"代码执行工具\" -> \"更改运行时类型\" 里进行设置\n",
        "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "# 创建 模型 + 优化器 + 损失函数\n",
        "model = CBDNet().to(device)\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)\n",
        "criterion = fixed_loss()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VymLsD7Q3FWu",
        "colab_type": "text"
      },
      "source": [
        "开始训练："
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u3ZBSNM9b0ni",
        "colab_type": "code",
        "outputId": "d1555402-9a82-4a78-875d-4c87c90e6b69",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "cnt = 0\n",
        "total_loss = AverageMeter()\n",
        "# 设置为训练模式，即启用 BatchNormalization 和 Dropout\n",
        "model.train()\n",
        "\n",
        "for epoch in range(200):\n",
        "    # 内存中清空图片\n",
        "    for i in range(len(train_fns)):\n",
        "        origin_imgs[i] = []\n",
        "        noised_imgs[i] = []\n",
        "\n",
        "    # 打乱训练图片的顺序\n",
        "    for idx in np.random.permutation(len(train_fns)):    \n",
        "        # 读入origin image；RGB通道反过来，然后归一化；转化为 float32类型\n",
        "        origin_img = cv2.imread(glob.glob(train_fns[idx] + '/*Reference.bmp')[0])\n",
        "        origin_img = origin_img[:,:,::-1] / 255.0\n",
        "        origin_imgs[idx] = np.array(origin_img).astype('float32')\n",
        "\n",
        "        # 读入noised image；因为一个文件夹里有2张噪声图，这里写了一个循环\n",
        "        train_noised_list = glob.glob(train_fns[idx] + '/*Noisy.bmp')  \n",
        "        for nidx in range(len(train_noised_list)):\n",
        "            noised_img = cv2.imread(train_noised_list[nidx])\n",
        "            noised_img = noised_img[:,:,::-1] / 255.0\n",
        "            noised_img = np.array(noised_img).astype('float32')\n",
        "            noised_imgs[idx].append(noised_img)\n",
        "\n",
        "            H, W, C = origin_img.shape\n",
        "            # 从图像中随机取 256x256 大小的块\n",
        "            xx = np.random.randint(0, W-ps+1)\n",
        "            yy = np.random.randint(0, H-ps+1)\n",
        "            temp_origin_img = origin_imgs[idx][yy:yy+ps, xx:xx+ps, :]\n",
        "            temp_noised_img = noised_imgs[idx][nidx][yy:yy+ps, xx:xx+ps, :]\n",
        "\n",
        "            # 生成 0，1 随机数，随机做图像的左右、上下、通道翻转，增加训练样本的多样性\n",
        "            if np.random.randint(0, 2) == 1:  # 左右翻转\n",
        "                temp_origin_img = np.flip(temp_origin_img, axis=1)\n",
        "                temp_noised_img = np.flip(temp_noised_img, axis=1)\n",
        "            if np.random.randint(0, 2) == 1:  # 上下翻转\n",
        "                temp_origin_img = np.flip(temp_origin_img, axis=0)\n",
        "                temp_noised_img = np.flip(temp_noised_img, axis=0)\n",
        "            if np.random.randint(0, 2) == 1:  # 通道翻转\n",
        "                temp_origin_img = np.transpose(temp_origin_img, (1, 0, 2))\n",
        "                temp_noised_img = np.transpose(temp_noised_img, (1, 0, 2))\n",
        "\n",
        "            temp_noised_img_chw = hwc_to_chw(temp_noised_img)\n",
        "            temp_origin_img_chw = hwc_to_chw(temp_origin_img)\n",
        "\n",
        "            cnt += 1\n",
        "\n",
        "            # 这里给输入数据增加一个维度，即原来是三维的，现在是四维的，方便CNN处理\n",
        "            input_var  = torch.from_numpy(temp_noised_img_chw.copy()).type(torch.FloatTensor).unsqueeze(0).to(device)\n",
        "            target_var = torch.from_numpy(temp_origin_img_chw.copy()).type(torch.FloatTensor).unsqueeze(0).to(device)\n",
        "\n",
        "            # 噪声图像输入网络处理\n",
        "            noise_level_est, output = model(input_var)\n",
        "            # 计算损失\n",
        "            loss = criterion(output, target_var, noise_level_est, 0, 0)\n",
        "            total_loss.update(loss.item())\n",
        "            # 常规操作： 梯度归零 + 反向传播 + 优化\n",
        "            optimizer.zero_grad()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "    print('[Epoch %d] [Img count %d] [Loss.val: %.4f] ([loss.avg: %.4f])\\t' % (epoch, cnt, total_loss.val, total_loss.avg))"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[Epoch 0] [Img count 10] [Loss.val: 0.0002] ([loss.avg: 0.0012])\t\n",
            "[Epoch 1] [Img count 20] [Loss.val: 0.0004] ([loss.avg: 0.0012])\t\n",
            "[Epoch 2] [Img count 30] [Loss.val: 0.0016] ([loss.avg: 0.0013])\t\n",
            "[Epoch 3] [Img count 40] [Loss.val: 0.0001] ([loss.avg: 0.0012])\t\n",
            "[Epoch 4] [Img count 50] [Loss.val: 0.0015] ([loss.avg: 0.0012])\t\n",
            "[Epoch 5] [Img count 60] [Loss.val: 0.0002] ([loss.avg: 0.0012])\t\n",
            "[Epoch 6] [Img count 70] [Loss.val: 0.0012] ([loss.avg: 0.0012])\t\n",
            "[Epoch 7] [Img count 80] [Loss.val: 0.0003] ([loss.avg: 0.0012])\t\n",
            "[Epoch 8] [Img count 90] [Loss.val: 0.0013] ([loss.avg: 0.0012])\t\n",
            "[Epoch 9] [Img count 100] [Loss.val: 0.0002] ([loss.avg: 0.0012])\t\n",
            "[Epoch 10] [Img count 110] [Loss.val: 0.0003] ([loss.avg: 0.0011])\t\n",
            "[Epoch 11] [Img count 120] [Loss.val: 0.0010] ([loss.avg: 0.0011])\t\n",
            "[Epoch 12] [Img count 130] [Loss.val: 0.0002] ([loss.avg: 0.0011])\t\n",
            "[Epoch 13] [Img count 140] [Loss.val: 0.0010] ([loss.avg: 0.0011])\t\n",
            "[Epoch 14] [Img count 150] [Loss.val: 0.0013] ([loss.avg: 0.0011])\t\n",
            "[Epoch 15] [Img count 160] [Loss.val: 0.0013] ([loss.avg: 0.0011])\t\n",
            "[Epoch 16] [Img count 170] [Loss.val: 0.0023] ([loss.avg: 0.0011])\t\n",
            "[Epoch 17] [Img count 180] [Loss.val: 0.0004] ([loss.avg: 0.0011])\t\n",
            "[Epoch 18] [Img count 190] [Loss.val: 0.0012] ([loss.avg: 0.0011])\t\n",
            "[Epoch 19] [Img count 200] [Loss.val: 0.0018] ([loss.avg: 0.0011])\t\n",
            "[Epoch 20] [Img count 210] [Loss.val: 0.0008] ([loss.avg: 0.0011])\t\n",
            "[Epoch 21] [Img count 220] [Loss.val: 0.0001] ([loss.avg: 0.0011])\t\n",
            "[Epoch 22] [Img count 230] [Loss.val: 0.0003] ([loss.avg: 0.0011])\t\n",
            "[Epoch 23] [Img count 240] [Loss.val: 0.0004] ([loss.avg: 0.0011])\t\n",
            "[Epoch 24] [Img count 250] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 25] [Img count 260] [Loss.val: 0.0014] ([loss.avg: 0.0010])\t\n",
            "[Epoch 26] [Img count 270] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 27] [Img count 280] [Loss.val: 0.0018] ([loss.avg: 0.0010])\t\n",
            "[Epoch 28] [Img count 290] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 29] [Img count 300] [Loss.val: 0.0009] ([loss.avg: 0.0010])\t\n",
            "[Epoch 30] [Img count 310] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 31] [Img count 320] [Loss.val: 0.0016] ([loss.avg: 0.0010])\t\n",
            "[Epoch 32] [Img count 330] [Loss.val: 0.0004] ([loss.avg: 0.0010])\t\n",
            "[Epoch 33] [Img count 340] [Loss.val: 0.0005] ([loss.avg: 0.0010])\t\n",
            "[Epoch 34] [Img count 350] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 35] [Img count 360] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 36] [Img count 370] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 37] [Img count 380] [Loss.val: 0.0004] ([loss.avg: 0.0010])\t\n",
            "[Epoch 38] [Img count 390] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 39] [Img count 400] [Loss.val: 0.0004] ([loss.avg: 0.0010])\t\n",
            "[Epoch 40] [Img count 410] [Loss.val: 0.0019] ([loss.avg: 0.0010])\t\n",
            "[Epoch 41] [Img count 420] [Loss.val: 0.0014] ([loss.avg: 0.0010])\t\n",
            "[Epoch 42] [Img count 430] [Loss.val: 0.0011] ([loss.avg: 0.0010])\t\n",
            "[Epoch 43] [Img count 440] [Loss.val: 0.0007] ([loss.avg: 0.0010])\t\n",
            "[Epoch 44] [Img count 450] [Loss.val: 0.0007] ([loss.avg: 0.0010])\t\n",
            "[Epoch 45] [Img count 460] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 46] [Img count 470] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 47] [Img count 480] [Loss.val: 0.0009] ([loss.avg: 0.0010])\t\n",
            "[Epoch 48] [Img count 490] [Loss.val: 0.0005] ([loss.avg: 0.0010])\t\n",
            "[Epoch 49] [Img count 500] [Loss.val: 0.0017] ([loss.avg: 0.0010])\t\n",
            "[Epoch 50] [Img count 510] [Loss.val: 0.0014] ([loss.avg: 0.0010])\t\n",
            "[Epoch 51] [Img count 520] [Loss.val: 0.0025] ([loss.avg: 0.0010])\t\n",
            "[Epoch 52] [Img count 530] [Loss.val: 0.0017] ([loss.avg: 0.0010])\t\n",
            "[Epoch 53] [Img count 540] [Loss.val: 0.0014] ([loss.avg: 0.0010])\t\n",
            "[Epoch 54] [Img count 550] [Loss.val: 0.0016] ([loss.avg: 0.0010])\t\n",
            "[Epoch 55] [Img count 560] [Loss.val: 0.0028] ([loss.avg: 0.0010])\t\n",
            "[Epoch 56] [Img count 570] [Loss.val: 0.0030] ([loss.avg: 0.0010])\t\n",
            "[Epoch 57] [Img count 580] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 58] [Img count 590] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 59] [Img count 600] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 60] [Img count 610] [Loss.val: 0.0014] ([loss.avg: 0.0010])\t\n",
            "[Epoch 61] [Img count 620] [Loss.val: 0.0006] ([loss.avg: 0.0010])\t\n",
            "[Epoch 62] [Img count 630] [Loss.val: 0.0037] ([loss.avg: 0.0010])\t\n",
            "[Epoch 63] [Img count 640] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 64] [Img count 650] [Loss.val: 0.0013] ([loss.avg: 0.0010])\t\n",
            "[Epoch 65] [Img count 660] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 66] [Img count 670] [Loss.val: 0.0034] ([loss.avg: 0.0010])\t\n",
            "[Epoch 67] [Img count 680] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 68] [Img count 690] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 69] [Img count 700] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 70] [Img count 710] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 71] [Img count 720] [Loss.val: 0.0016] ([loss.avg: 0.0010])\t\n",
            "[Epoch 72] [Img count 730] [Loss.val: 0.0008] ([loss.avg: 0.0010])\t\n",
            "[Epoch 73] [Img count 740] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 74] [Img count 750] [Loss.val: 0.0014] ([loss.avg: 0.0010])\t\n",
            "[Epoch 75] [Img count 760] [Loss.val: 0.0020] ([loss.avg: 0.0010])\t\n",
            "[Epoch 76] [Img count 770] [Loss.val: 0.0007] ([loss.avg: 0.0010])\t\n",
            "[Epoch 77] [Img count 780] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 78] [Img count 790] [Loss.val: 0.0034] ([loss.avg: 0.0010])\t\n",
            "[Epoch 79] [Img count 800] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 80] [Img count 810] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 81] [Img count 820] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 82] [Img count 830] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 83] [Img count 840] [Loss.val: 0.0004] ([loss.avg: 0.0010])\t\n",
            "[Epoch 84] [Img count 850] [Loss.val: 0.0017] ([loss.avg: 0.0010])\t\n",
            "[Epoch 85] [Img count 860] [Loss.val: 0.0011] ([loss.avg: 0.0010])\t\n",
            "[Epoch 86] [Img count 870] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 87] [Img count 880] [Loss.val: 0.0008] ([loss.avg: 0.0010])\t\n",
            "[Epoch 88] [Img count 890] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 89] [Img count 900] [Loss.val: 0.0018] ([loss.avg: 0.0010])\t\n",
            "[Epoch 90] [Img count 910] [Loss.val: 0.0017] ([loss.avg: 0.0010])\t\n",
            "[Epoch 91] [Img count 920] [Loss.val: 0.0011] ([loss.avg: 0.0010])\t\n",
            "[Epoch 92] [Img count 930] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 93] [Img count 940] [Loss.val: 0.0018] ([loss.avg: 0.0010])\t\n",
            "[Epoch 94] [Img count 950] [Loss.val: 0.0004] ([loss.avg: 0.0010])\t\n",
            "[Epoch 95] [Img count 960] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 96] [Img count 970] [Loss.val: 0.0004] ([loss.avg: 0.0010])\t\n",
            "[Epoch 97] [Img count 980] [Loss.val: 0.0009] ([loss.avg: 0.0010])\t\n",
            "[Epoch 98] [Img count 990] [Loss.val: 0.0022] ([loss.avg: 0.0010])\t\n",
            "[Epoch 99] [Img count 1000] [Loss.val: 0.0011] ([loss.avg: 0.0010])\t\n",
            "[Epoch 100] [Img count 1010] [Loss.val: 0.0009] ([loss.avg: 0.0010])\t\n",
            "[Epoch 101] [Img count 1020] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 102] [Img count 1030] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 103] [Img count 1040] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 104] [Img count 1050] [Loss.val: 0.0015] ([loss.avg: 0.0010])\t\n",
            "[Epoch 105] [Img count 1060] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 106] [Img count 1070] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 107] [Img count 1080] [Loss.val: 0.0013] ([loss.avg: 0.0010])\t\n",
            "[Epoch 108] [Img count 1090] [Loss.val: 0.0020] ([loss.avg: 0.0010])\t\n",
            "[Epoch 109] [Img count 1100] [Loss.val: 0.0003] ([loss.avg: 0.0010])\t\n",
            "[Epoch 110] [Img count 1110] [Loss.val: 0.0003] ([loss.avg: 0.0010])\t\n",
            "[Epoch 111] [Img count 1120] [Loss.val: 0.0011] ([loss.avg: 0.0010])\t\n",
            "[Epoch 112] [Img count 1130] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 113] [Img count 1140] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 114] [Img count 1150] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 115] [Img count 1160] [Loss.val: 0.0006] ([loss.avg: 0.0010])\t\n",
            "[Epoch 116] [Img count 1170] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 117] [Img count 1180] [Loss.val: 0.0015] ([loss.avg: 0.0010])\t\n",
            "[Epoch 118] [Img count 1190] [Loss.val: 0.0008] ([loss.avg: 0.0010])\t\n",
            "[Epoch 119] [Img count 1200] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 120] [Img count 1210] [Loss.val: 0.0013] ([loss.avg: 0.0010])\t\n",
            "[Epoch 121] [Img count 1220] [Loss.val: 0.0014] ([loss.avg: 0.0010])\t\n",
            "[Epoch 122] [Img count 1230] [Loss.val: 0.0003] ([loss.avg: 0.0010])\t\n",
            "[Epoch 123] [Img count 1240] [Loss.val: 0.0009] ([loss.avg: 0.0010])\t\n",
            "[Epoch 124] [Img count 1250] [Loss.val: 0.0004] ([loss.avg: 0.0010])\t\n",
            "[Epoch 125] [Img count 1260] [Loss.val: 0.0004] ([loss.avg: 0.0010])\t\n",
            "[Epoch 126] [Img count 1270] [Loss.val: 0.0024] ([loss.avg: 0.0010])\t\n",
            "[Epoch 127] [Img count 1280] [Loss.val: 0.0008] ([loss.avg: 0.0010])\t\n",
            "[Epoch 128] [Img count 1290] [Loss.val: 0.0011] ([loss.avg: 0.0010])\t\n",
            "[Epoch 129] [Img count 1300] [Loss.val: 0.0008] ([loss.avg: 0.0010])\t\n",
            "[Epoch 130] [Img count 1310] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 131] [Img count 1320] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 132] [Img count 1330] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 133] [Img count 1340] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 134] [Img count 1350] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 135] [Img count 1360] [Loss.val: 0.0009] ([loss.avg: 0.0010])\t\n",
            "[Epoch 136] [Img count 1370] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 137] [Img count 1380] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 138] [Img count 1390] [Loss.val: 0.0030] ([loss.avg: 0.0010])\t\n",
            "[Epoch 139] [Img count 1400] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 140] [Img count 1410] [Loss.val: 0.0015] ([loss.avg: 0.0010])\t\n",
            "[Epoch 141] [Img count 1420] [Loss.val: 0.0005] ([loss.avg: 0.0010])\t\n",
            "[Epoch 142] [Img count 1430] [Loss.val: 0.0011] ([loss.avg: 0.0010])\t\n",
            "[Epoch 143] [Img count 1440] [Loss.val: 0.0009] ([loss.avg: 0.0010])\t\n",
            "[Epoch 144] [Img count 1450] [Loss.val: 0.0004] ([loss.avg: 0.0010])\t\n",
            "[Epoch 145] [Img count 1460] [Loss.val: 0.0016] ([loss.avg: 0.0010])\t\n",
            "[Epoch 146] [Img count 1470] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 147] [Img count 1480] [Loss.val: 0.0006] ([loss.avg: 0.0010])\t\n",
            "[Epoch 148] [Img count 1490] [Loss.val: 0.0006] ([loss.avg: 0.0010])\t\n",
            "[Epoch 149] [Img count 1500] [Loss.val: 0.0015] ([loss.avg: 0.0010])\t\n",
            "[Epoch 150] [Img count 1510] [Loss.val: 0.0006] ([loss.avg: 0.0010])\t\n",
            "[Epoch 151] [Img count 1520] [Loss.val: 0.0006] ([loss.avg: 0.0010])\t\n",
            "[Epoch 152] [Img count 1530] [Loss.val: 0.0022] ([loss.avg: 0.0010])\t\n",
            "[Epoch 153] [Img count 1540] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 154] [Img count 1550] [Loss.val: 0.0015] ([loss.avg: 0.0010])\t\n",
            "[Epoch 155] [Img count 1560] [Loss.val: 0.0024] ([loss.avg: 0.0010])\t\n",
            "[Epoch 156] [Img count 1570] [Loss.val: 0.0011] ([loss.avg: 0.0010])\t\n",
            "[Epoch 157] [Img count 1580] [Loss.val: 0.0019] ([loss.avg: 0.0010])\t\n",
            "[Epoch 158] [Img count 1590] [Loss.val: 0.0009] ([loss.avg: 0.0010])\t\n",
            "[Epoch 159] [Img count 1600] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 160] [Img count 1610] [Loss.val: 0.0014] ([loss.avg: 0.0010])\t\n",
            "[Epoch 161] [Img count 1620] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 162] [Img count 1630] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 163] [Img count 1640] [Loss.val: 0.0006] ([loss.avg: 0.0010])\t\n",
            "[Epoch 164] [Img count 1650] [Loss.val: 0.0013] ([loss.avg: 0.0010])\t\n",
            "[Epoch 165] [Img count 1660] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 166] [Img count 1670] [Loss.val: 0.0011] ([loss.avg: 0.0010])\t\n",
            "[Epoch 167] [Img count 1680] [Loss.val: 0.0008] ([loss.avg: 0.0010])\t\n",
            "[Epoch 168] [Img count 1690] [Loss.val: 0.0014] ([loss.avg: 0.0010])\t\n",
            "[Epoch 169] [Img count 1700] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 170] [Img count 1710] [Loss.val: 0.0023] ([loss.avg: 0.0010])\t\n",
            "[Epoch 171] [Img count 1720] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 172] [Img count 1730] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 173] [Img count 1740] [Loss.val: 0.0022] ([loss.avg: 0.0010])\t\n",
            "[Epoch 174] [Img count 1750] [Loss.val: 0.0003] ([loss.avg: 0.0010])\t\n",
            "[Epoch 175] [Img count 1760] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 176] [Img count 1770] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 177] [Img count 1780] [Loss.val: 0.0002] ([loss.avg: 0.0010])\t\n",
            "[Epoch 178] [Img count 1790] [Loss.val: 0.0018] ([loss.avg: 0.0010])\t\n",
            "[Epoch 179] [Img count 1800] [Loss.val: 0.0040] ([loss.avg: 0.0010])\t\n",
            "[Epoch 180] [Img count 1810] [Loss.val: 0.0021] ([loss.avg: 0.0010])\t\n",
            "[Epoch 181] [Img count 1820] [Loss.val: 0.0019] ([loss.avg: 0.0010])\t\n",
            "[Epoch 182] [Img count 1830] [Loss.val: 0.0010] ([loss.avg: 0.0010])\t\n",
            "[Epoch 183] [Img count 1840] [Loss.val: 0.0029] ([loss.avg: 0.0010])\t\n",
            "[Epoch 184] [Img count 1850] [Loss.val: 0.0015] ([loss.avg: 0.0010])\t\n",
            "[Epoch 185] [Img count 1860] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 186] [Img count 1870] [Loss.val: 0.0011] ([loss.avg: 0.0010])\t\n",
            "[Epoch 187] [Img count 1880] [Loss.val: 0.0011] ([loss.avg: 0.0010])\t\n",
            "[Epoch 188] [Img count 1890] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 189] [Img count 1900] [Loss.val: 0.0017] ([loss.avg: 0.0010])\t\n",
            "[Epoch 190] [Img count 1910] [Loss.val: 0.0006] ([loss.avg: 0.0010])\t\n",
            "[Epoch 191] [Img count 1920] [Loss.val: 0.0014] ([loss.avg: 0.0010])\t\n",
            "[Epoch 192] [Img count 1930] [Loss.val: 0.0012] ([loss.avg: 0.0010])\t\n",
            "[Epoch 193] [Img count 1940] [Loss.val: 0.0003] ([loss.avg: 0.0010])\t\n",
            "[Epoch 194] [Img count 1950] [Loss.val: 0.0031] ([loss.avg: 0.0010])\t\n",
            "[Epoch 195] [Img count 1960] [Loss.val: 0.0007] ([loss.avg: 0.0010])\t\n",
            "[Epoch 196] [Img count 1970] [Loss.val: 0.0003] ([loss.avg: 0.0010])\t\n",
            "[Epoch 197] [Img count 1980] [Loss.val: 0.0005] ([loss.avg: 0.0010])\t\n",
            "[Epoch 198] [Img count 1990] [Loss.val: 0.0001] ([loss.avg: 0.0010])\t\n",
            "[Epoch 199] [Img count 2000] [Loss.val: 0.0017] ([loss.avg: 0.0010])\t\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C8jsKcL83V6F",
        "colab_type": "text"
      },
      "source": [
        "可以看到，loss 一直在下降，受限于colab的硬件资源，如果在本地服务器用更多的训练数据，效果改进可能会更加明显。\n",
        "\n",
        "## 4. 模型测试\n",
        "\n",
        "一些基本参数："
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QOx6E7pQzoAl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "test_dir = './mini_denoise_dataset/test/'\n",
        "test_fns = glob.glob(test_dir + '*.bmp')\n",
        "\n",
        "# 建立 result 目录，保存图片处理结果\n",
        "result_dir = './result/'\n",
        "if not os.path.exists( result_dir ):\n",
        "    os.mkdir( result_dir )"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "08AZTF6u4Rc2",
        "colab_type": "text"
      },
      "source": [
        "开始测试："
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5_oa6jN84Q25",
        "colab_type": "code",
        "outputId": "0d608368-57e3-48d5-82f7-24636cc40136",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 187
        }
      },
      "source": [
        "for ind, test_img_path in enumerate(test_fns):\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "        print(test_img_path)\n",
        "        # 读入图像，切换RGB通道并归一化，转化为 numpy float32格式\n",
        "        noisy_img = cv2.imread(test_img_path)\n",
        "        noisy_img = noisy_img[:,:,::-1] / 255.0\n",
        "        noisy_img = np.array(noisy_img).astype('float32')\n",
        "\n",
        "        # 转化为 chw 才符合 pytorch 网络的输入格式\n",
        "        temp_noisy_img_chw = hwc_to_chw(noisy_img)\n",
        "        # 图像放到 gpu 上\n",
        "        input_var = torch.from_numpy(temp_noisy_img_chw.copy()).type(torch.FloatTensor).unsqueeze(0).to(device)\n",
        "        # 输入模型得到结果\n",
        "        _, output = model(input_var)\n",
        "\n",
        "        # 输出结果转化为 numpy ，同时，把数据转到 0，1 之间（因为可能会有一些异常值）\n",
        "        output_np = output.squeeze().cpu().detach().numpy()\n",
        "        output_np = chw_to_hwc(np.clip(output_np, 0, 1))\n",
        "        # 把噪声图像，和降噪后的图像拼接在一起，然后保存图像\n",
        "        tempImg = np.concatenate((noisy_img, output_np), axis=1)*255.0\n",
        "        \n",
        "        Image.fromarray(np.uint8(tempImg)).save(fp=result_dir + 'test_%d.jpg'%(ind), format='JPEG')"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "./mini_denoise_dataset/test/0006.bmp\n",
            "./mini_denoise_dataset/test/0010.bmp\n",
            "./mini_denoise_dataset/test/0003.bmp\n",
            "./mini_denoise_dataset/test/0002.bmp\n",
            "./mini_denoise_dataset/test/0005.bmp\n",
            "./mini_denoise_dataset/test/0007.bmp\n",
            "./mini_denoise_dataset/test/0001.bmp\n",
            "./mini_denoise_dataset/test/0004.bmp\n",
            "./mini_denoise_dataset/test/0009.bmp\n",
            "./mini_denoise_dataset/test/0008.bmp\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ElnAIC2CZouE",
        "colab_type": "text"
      },
      "source": [
        "代码运行结束，下面是10张测试图片的效果：\n",
        "\n",
        "\n",
        "![替代文字](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200302212300.jpg)\n",
        "\n",
        "![替代文字](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200302212301.jpg)\n",
        "\n",
        "![替代文字](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200302212302.jpg)\n",
        "\n",
        "![替代文字](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200302212303.jpg)\n",
        "\n",
        "![替代文字](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200302212304.jpg)\n",
        "\n",
        "![替代文字](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200302212305.jpg)\n",
        "\n",
        "![替代文字](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200302212306.jpg)\n",
        "\n",
        "![替代文字](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200302212307.jpg)\n",
        "\n",
        "![替代文字](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200302212308.jpg)\n",
        "\n",
        "![替代文字](https://gaopursuit.oss-cn-beijing.aliyuncs.com/202003/20200302212309.jpg)\n",
        "\n",
        "\n",
        "可以看到，去噪的效果有了，但是感觉有些模糊。也许是训练数据用的不够，也许是训练的 epoch 还不够多 ~~~ 有机会会在本地服务器，用更多的训练数据实验一下。"
      ]
    }
  ]
}