{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 代码实现"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们来看FCN-8s的代码实现。这里，我们选择ResNet101作为模型的主干网络."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import models\n",
    "\n",
    "# 使用 ResNet101 作为主干网络，模型使用ImageNet预训练\n",
    "pretrained_net = models.resnet101(pretrained='imagenet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn as nn\n",
    "\n",
    "# FCN-8s模型\n",
    "class FCN8s(nn.Module):\n",
    "    def __init__(self, num_classes):\n",
    "        super(FCN8s, self).__init__()\n",
    "        # 这里使用children()调用ResNet101的部分网络\n",
    "        # 该深度特征图的大小为输入图像的1/8\n",
    "        self.stage1 = nn.Sequential(*list(pretrained_net.children())[:-4])\n",
    "        # 该深度特征图的大小为输入图像的1/16\n",
    "        self.stage2 = list(pretrained_net.children())[-4]\n",
    "        # 该深度特征图的大小为输入图像的1/32\n",
    "        self.stage3 = list(pretrained_net.children())[-3]\n",
    "        \n",
    "        # 调整stage3输出的通道数\n",
    "        self.scores1 = nn.Conv2d(2048, num_classes, 1)\n",
    "        # 调整stage2输出的通道数\n",
    "        self.scores2 = nn.Conv2d(1024, num_classes, 1)\n",
    "        # 调整stage1输出的通道数\n",
    "        self.scores3 = nn.Conv2d(512, num_classes, 1)\n",
    "        \n",
    "        # 对pool3与pool4、pool5的输出进行8倍上采样\n",
    "        self.upsample_8x = nn.ConvTranspose2d(\n",
    "            num_classes, num_classes, 16, 8, 4, bias=False)\n",
    "        self.upsample_8x.weight.data = bilinear_kernel(\n",
    "            num_classes, num_classes, 16) # 使用双线性 kernel\n",
    "        \n",
    "        # 对pool4和pool5融合的特征2倍上采样\n",
    "        self.upsample_4x = nn.ConvTranspose2d(\n",
    "            num_classes, num_classes, 4, 2, 1, bias=False)\n",
    "        self.upsample_4x.weight.data = bilinear_kernel(\n",
    "            num_classes, num_classes, 4) # 使用双线性 kernel\n",
    "        \n",
    "        # 对pool5的输出2倍上采样\n",
    "        self.upsample_2x = nn.ConvTranspose2d(\n",
    "            num_classes, num_classes, 4, 2, 1, bias=False)   \n",
    "        self.upsample_2x.weight.data = bilinear_kernel(\n",
    "            num_classes, num_classes, 4) # 使用双线性 kernel\n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.stage1(x)\n",
    "        # s1 1/8\n",
    "        s1 = x\n",
    "        \n",
    "        x = self.stage2(x)\n",
    "        # s2 1/16\n",
    "        s2 = x\n",
    "\n",
    "        x = self.stage3(x)\n",
    "        # s3 1/32\n",
    "        s3 = x\n",
    "                \n",
    "        # 调整pool5输出特征图的通道数\n",
    "        s3 = self.scores1(s3)\n",
    "        # 进行两倍上采样\n",
    "        s3 = self.upsample_2x(s3)\n",
    "        \n",
    "        # 调整pool4输出特征图的通道数\n",
    "        s2 = self.scores2(s2)\n",
    "        # 融合pool5、pool4的特征图\n",
    "        s2 = s2 + s3\n",
    "        \n",
    "        # 调整pool3输出特征图的通道数\n",
    "        s1 = self.scores3(s1)\n",
    "        # 将s2两倍上采样\n",
    "        s2 = self.upsample_4x(s2)\n",
    "        # 融合特征图\n",
    "        s = s1 + s2\n",
    "        \n",
    "        # 8倍上采样得到与原图像大小一致的特征图\n",
    "        s = self.upsample_8x(s2)\n",
    "        return s\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-22T06:50:45.953022Z",
     "iopub.status.busy": "2022-11-22T06:50:45.952484Z",
     "iopub.status.idle": "2022-11-22T06:50:48.729717Z",
     "shell.execute_reply": "2022-11-22T06:50:48.727912Z",
     "shell.execute_reply.started": "2022-11-22T06:50:45.952908Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import random\n",
    "import cv2\n",
    "import d2l.torch as d2l\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import Dataset\n",
    "from torch.autograd import Variable\n",
    "from torchvision.transforms import transforms as tfs\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "根据原论文，使用Pascal Voc作为实验的数据集。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据集地址\n",
    "# 请读者自行下载数据集\n",
    "voc_root = './voc/VOCdevkit/VOC2012'\n",
    "\n",
    "\n",
    "# 定义Pascal Voc类\n",
    "class VOCSegDataset(Dataset):\n",
    "    def __init__(self, train, crop_size, transforms):\n",
    "        # 定义数据大小\n",
    "        self.crop_size = crop_size\n",
    "        # 定义数据增强类型\n",
    "        self.transforms = transforms\n",
    "        \n",
    "        # 数据以及对应标签的读取\n",
    "        data_list, label_list = read_images(voc_root, train=train)\n",
    "        self.data_list = self._filter(data_list)\n",
    "        self.label_list = self._filter(label_list)\n",
    "        print('Read'+str(len(self.data_list))+'images')\n",
    "        \n",
    "    def _filter(self, images):\n",
    "        return [im for im in images if \n",
    "                (Image.open(im).size[1] >= self.crop_size[0] and\n",
    "                Image.open(im).size[0] >= self.crop_size[1])]\n",
    "    # 定义数据如何进行传输\n",
    "    def __getitem__(self,idx):\n",
    "        img = self.data_list[idx]\n",
    "        label = self.label_list[idx]\n",
    "        # 读取图像\n",
    "        img = cv2.imread(img)\n",
    "        # 将图像通道从BGR转换成RGB\n",
    "        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n",
    "        \n",
    "        label = cv2.imread(label)\n",
    "        # 因为标签与像素点一一对应，因此也需要将其进行转换\n",
    "        label = cv2.cvtColor(label, cv2.COLOR_BGR2RGB)\n",
    "        # 进行数据增强\n",
    "        img,label = self.transforms(img, label, self.crop_size)\n",
    "        return img, label\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.data_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们对上述类中出现的部分函数进行实现。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-22T06:51:17.650295Z",
     "iopub.status.busy": "2022-11-22T06:51:17.649804Z",
     "iopub.status.idle": "2022-11-22T06:51:17.695748Z",
     "shell.execute_reply": "2022-11-22T06:51:17.694813Z",
     "shell.execute_reply.started": "2022-11-22T06:51:17.650255Z"
    }
   },
   "outputs": [],
   "source": [
    "# 对上述数据集类的方法进行实现\n",
    "\n",
    "# 读取数据集\n",
    "def read_images(root, train=True):\n",
    "    txt_filename = root + \"/ImageSets/Segmentation/\" \\\n",
    "        + ('train.txt' if train else 'val.txt')\n",
    "    with open(txt_filename, 'r') as f:\n",
    "        images = f.read().split()\n",
    "    data = [os.path.join(root, 'JPEGImages', i + '.jpg') \n",
    "                                    for i in images]\n",
    "    label = [os.path.join(root, 'SegmentationClass', i+'.png') \n",
    "                                    for i in images]\n",
    "    return data, label\n",
    "\n",
    "\n",
    "# 进行随机裁剪\n",
    "def rand_crop(data, label, height,width):\n",
    "    h, w, _ = data.shape\n",
    "    top = random.randint(0, h - height)\n",
    "    left = random.randint(0, w - width)\n",
    "    # 裁剪数据\n",
    "    data = data[top:top + height, left:left + width]\n",
    "    # 裁剪标签\n",
    "    label = label[top:top + height, left:left + width]\n",
    "    return data, label\n",
    "\n",
    "\n",
    "# 为图像增强时的像素点建立对应的标签\n",
    "def image2label(im):\n",
    "    data = np.array(im, dtype='int32')\n",
    "    idx = (data[:,:,0] * 256 + data[:,:,1]) * 256 + data[:,:,2]\n",
    "    return np.array(cm2lbl[idx], dtype='int64')\n",
    "\n",
    "\n",
    "# 定义数据增强方法\n",
    "def img_transforms(im, label, crop_size):\n",
    "    im,label = rand_crop(im, label, *crop_size)\n",
    "    im_tfs = tfs.Compose([\n",
    "        tfs.ToTensor(),\n",
    "        tfs.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])\n",
    "    ])\n",
    "    im = im_tfs(im)\n",
    "    label = image2label(label)\n",
    "    label = torch.from_numpy(label)\n",
    "    return im, label\n",
    "\n",
    "\n",
    "\n",
    "# 删除双线性插值函数中的输出\n",
    "def bilinear_kernel(in_channels, out_channels, kernel_size):\n",
    "    factor = (kernel_size + 1) // 2\n",
    "    if kernel_size % 2 == 1:\n",
    "        center = factor - 1\n",
    "    else:\n",
    "        center = factor - 0.5\n",
    "    og = np.ogrid[:kernel_size, :kernel_size]\n",
    "    filt = (1 - abs(og[0] - center) / factor) * \\\n",
    "        (1 - abs(og[1] - center) / factor)\n",
    "    weight = np.zeros((in_channels, out_channels,\n",
    "                       kernel_size, kernel_size), dtype='float32')\n",
    "    weight[range(in_channels), range(out_channels),:,:] = filt\n",
    "    return torch.from_numpy(weight)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接着，定义一些需要用到的全局变量，如类别信息等，并开始加载数据集，构建数据加载器。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-22T06:51:25.202586Z",
     "iopub.status.busy": "2022-11-22T06:51:25.202182Z",
     "iopub.status.idle": "2022-11-22T06:52:07.293551Z",
     "shell.execute_reply": "2022-11-22T06:52:07.292431Z",
     "shell.execute_reply.started": "2022-11-22T06:51:25.202553Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Read1114images\n",
      "Read1078images\n"
     ]
    }
   ],
   "source": [
    "# Pascal Voc的类别\n",
    "classes = ['background','aeroplane','bicycle','bird','boat',\n",
    "           'bottle','bus','car','cat','chair','cow','diningtable',\n",
    "           'dog','horse','motorbike','person','potted plant',\n",
    "           'sheep','sofa','train','tv/monitor']\n",
    "\n",
    "# colormap的数量与类别数对应，每一种类别都有其独一无二的颜色，便于绘制图像时进行观察\n",
    "colormap = [[0,0,0], [128,0,0], [0,128,0], [128,128,0], [0,0,128],\n",
    "            [128,0,128], [0,128,128], [128,128,128], [64,0,0], \n",
    "            [192,0,0], [64,128,0], [192,128,0], [64,0,128], \n",
    "            [192,0,128], [64,128,128], [192,128,128], [0,64,0], \n",
    "            [128,64,0], [0,192,0], [128,192,0], [0,64,128]]\n",
    "\n",
    "# 读取数据和标签\n",
    "data, label = read_images(voc_root)\n",
    "num_classes = len(classes)\n",
    "\n",
    "# 将colormap转换成类别\n",
    "cm2lbl = np.zeros(256**3)\n",
    "for i, cm in enumerate(colormap):\n",
    "    cm2lbl[(cm[0]*256+cm[1])*256+cm[2]]=i\n",
    "\n",
    "# 设置输入图像大小\n",
    "input_shape = (320,480)\n",
    "voc_train = VOCSegDataset(True, input_shape, img_transforms)\n",
    "voc_test = VOCSegDataset(False, input_shape, img_transforms)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#设置batch_size，生成数据加载器\n",
    "BATCH_SIZE = 64\n",
    "train_data = DataLoader(voc_train, batch_size=BATCH_SIZE, shuffle=True)\n",
    "test_data = DataLoader(voc_test, batch_size=BATCH_SIZE)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "先展示数据集中的部分图像与其对应的标签。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2022-11-22T06:52:11.770267Z",
     "iopub.status.busy": "2022-11-22T06:52:11.769852Z",
     "iopub.status.idle": "2022-11-22T06:52:12.587582Z",
     "shell.execute_reply": "2022-11-22T06:52:12.586360Z",
     "shell.execute_reply.started": "2022-11-22T06:52:11.770231Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 185268 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 198820 characters]",
      "text/plain": [
       "<Figure size 1000x1000 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "for i, (img, label) in enumerate(voc_train):\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.subplot(221)\n",
    "    plt.imshow(img.moveaxis(0, 2))\n",
    "    plt.subplot(222)\n",
    "    plt.imshow(label)  \n",
    "    plt.show()\n",
    "    plt.close()\n",
    "    if i ==1:\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接着，编写相关函数实现计算mIoU。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _fast_hist(label_true, label_pred, n_class):\n",
    "    mask = (label_true >= 0) & (label_true < n_class)\n",
    "    hist = np.bincount(\n",
    "        n_class * label_true[mask].astype(int) +\n",
    "        label_pred[mask], minlength=n_class ** 2).reshape(\n",
    "                                                n_class, n_class)\n",
    "    return hist\n",
    "\n",
    "\n",
    "def mIoU(label_trues, label_preds, n_class):\n",
    "    # 计算mIoU\n",
    "    hist = np.zeros((n_class, n_class))\n",
    "    \n",
    "    for lt, lp in zip(label_trues, label_preds):\n",
    "        hist += _fast_hist(lt.flatten(), lp.flatten(), n_class)\n",
    "        \n",
    "    iu = np.diag(hist) / (hist.sum(axis=1) + \n",
    "                          hist.sum(axis=0) - np.diag(hist))\n",
    "    mean_iu = np.nanmean(iu)\n",
    "\n",
    "    return mean_iu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "加载FCN-8s模型，并对其进行训练。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = device = torch.device('cuda' if torch.cuda.is_available() \n",
    "                               else 'cpu')\n",
    "\n",
    "# 加载模型 \n",
    "model = FCN8s(num_classes)\n",
    "model.to(device)\n",
    "model = nn.DataParallel(model)\n",
    "\n",
    "# 设置优化器\n",
    "optimizer = torch.optim.SGD(model.parameters(), lr=1e-2, \n",
    "                            weight_decay=1e-4)\n",
    "\n",
    "# 设置损失函数\n",
    "criterion = nn.NLLLoss()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# 由于没有验证集，这里每训练完一次模型便在测试集上进行测试\n",
    "\n",
    "# 不显示无关信息\n",
    "np.seterr(divide='ignore',invalid='ignore')\n",
    "\n",
    "\n",
    "Epoch = 120\n",
    "for epoch in range(Epoch):\n",
    "    \n",
    "    # 训练损失\n",
    "    train_loss = 0\n",
    "    # 训练mIoU\n",
    "    train_mean_iu = 0\n",
    "    \n",
    "    # 开始训练\n",
    "    model = model.train()\n",
    "\n",
    "    for data in train_data:\n",
    "        x = data[0].to(device)\n",
    "        y = data[1].to(device)\n",
    "\n",
    "        # 梯度清零\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # 获得模型预测概率分布\n",
    "        outputs = model(x)\n",
    "        # 获得模型预测的标签，大小为(b,n,h,w)\n",
    "        outputs = F.log_softmax(outputs, dim=1)\n",
    "        \n",
    "        # 计算损失\n",
    "        loss = criterion(outputs, y)\n",
    "        \n",
    "        # 反向传播\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        # 记录损失\n",
    "        train_loss += loss.item()\n",
    "        \n",
    "        # 获得预测标签\n",
    "        label_pred = outputs.max(dim=1)[1].data.cpu().numpy()\n",
    "        label_true = y.data.cpu().numpy()\n",
    "        \n",
    "        for lbt, lbp in zip(label_true, label_pred):\n",
    "            # 返回每张图的pred和gt标签计算mIoU\n",
    "            mean_iu = mIoU(lbt, lbp, num_classes)\n",
    "            train_mean_iu += mean_iu\n",
    "    \n",
    "\n",
    "    # 测试损失\n",
    "    eval_loss = 0\n",
    "    # 测试mIoU\n",
    "    eval_mean_iu = 0\n",
    "    \n",
    "    # 进行模型测试\n",
    "    model = model.eval()\n",
    "    \n",
    "    for data in test_data:\n",
    "        x_test = data[0].to(device)\n",
    "        y_test = data[1].to(device)\n",
    "        outputs_test = model(x_test)\n",
    "        outputs_test = F.log_softmax(outputs_test, dim=1)\n",
    "        \n",
    "        loss = criterion(outputs_test, y_test)\n",
    "        eval_loss += loss.item()\n",
    "        \n",
    "        label_pred = outputs_test.max(dim=1)[1].data.cpu().numpy()\n",
    "        label_true = y_test.data.cpu().numpy()\n",
    "        \n",
    "        for lbt, lbp in zip(label_true, label_pred):\n",
    "            mean_iu = mIoU(lbt, lbp, num_classes)\n",
    "            eval_mean_iu += mean_iu\n",
    "        \n",
    "    epoch_str = ('Epoch: {}, Train Loss: {:.5f}, Train Mean IU: {:.5f},\\\n",
    "                 Valid Loss: {:.5f}, Valid Mean IU: {:.5f} '.format(\n",
    "                epoch, train_loss / len(train_data), \n",
    "                train_mean_iu / len(voc_train), \n",
    "                eval_loss / len(test_data), \n",
    "                eval_mean_iu / len(voc_test)))\n",
    "   \n",
    "    print(epoch_str)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "为了方便，在这里不展示具体的训练信息输出。接着，随机选择一些图像，并展示其真实标签与模型的预测标签。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 4159548 characters]",
      "text/plain": [
       "<Figure size 6000x5000 with 30 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 将预测的标签映射到colormap上\n",
    "cm = np.array(colormap).astype('uint8')\n",
    "\n",
    "def predict(im, label):\n",
    "    im = Variable(im.unsqueeze(0)).cuda()\n",
    "    out = model(im)\n",
    "    pred = out.max(1)[1].squeeze().cpu().data.numpy()\n",
    "    pred = cm[pred]\n",
    "    return pred, cm[label.numpy()]\n",
    "\n",
    "\n",
    "_, figs = plt.subplots(10, 3, figsize=(12, 10))\n",
    "\n",
    "for i in range(10):\n",
    "    x, y = voc_test[i]\n",
    "    x.to(device)\n",
    "    y.to(device)\n",
    "    pred, label = predict(x, y)\n",
    "    figs[i, 0].imshow(Image.open(voc_test.data_list[i]))\n",
    "    figs[i, 0].axes.get_xaxis().set_visible(False)\n",
    "    figs[i, 0].axes.get_yaxis().set_visible(False)\n",
    "    figs[i, 1].imshow(label)\n",
    "    figs[i, 1].axes.get_xaxis().set_visible(False)\n",
    "    figs[i, 1].axes.get_yaxis().set_visible(False)\n",
    "    figs[i, 2].imshow(pred)\n",
    "    figs[i, 2].axes.get_xaxis().set_visible(False)\n",
    "    figs[i, 2].axes.get_yaxis().set_visible(False)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.10.7"
  },
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
