地表建筑物识别——Task01赛题理解

前言：这次参加Datawhale与天池联合发起的零基础入门语义分割之地表建筑物识别挑战赛。自己在大四上（2018下半年）也简单接触过语义分割，当时是基于FCN做CT图像乳腺肿瘤的分割。现在参加这个入门级赛事，主要是看看大佬们是怎么走一个完整的深度学习项目的，提升用深度学习在自己的科研方向（波前整形、散斑成像）等方面的应用能力。当然，学习常用语义分割模型（FCN、Unet、DeepLab、SegNet、PSPNet等），模型集成方法以及各种评价 / 损失函数是其中的核心内容。

Task1：赛题理解与baseline（3 天）
– 学习主题：理解赛题内容解题流程
– 学习内容：赛题理解、数据读取、比赛baseline 构建
– 学习成果：比赛baseline 提交

赛题理解

本赛题使用航拍数据 (来源于Inria Aerial Image Labeling)，需要参赛选手完成地表建筑物识别，将地表航拍图像素划分为有建筑物和无建筑物两类。如下图，左边为原始航拍图，右边为对应的建筑物标注。

数据说明

FileName	Size	含义
test_a.zip	314.49MB	测试集A榜图片
test_a_samplesubmit.csv	46.39KB	测试集A榜提交样例
train.zip	3.68GB	训练集图片
train_mask.csv.zip	97.52MB	训练集图片RLE标注

读取数据

赛题为语义分割任务，因此具体的标签为图像像素类别。在赛题数据中像素属于2 类（无建筑物和有建筑物），因此标签为有建筑物的像素。赛题原始图片为jpg 格式，标签为RLE 编码的字符串。

RLE 全称（run-length encoding），翻译为游程编码或行程长度编码，对连续的黑、白像素数以不同的码字进行编码。RLE 是一种简单的非破坏性资料压缩法，经常用在在语义分割比赛中对标签进行编码。RLE 与图片之间的转换如下，关于代码详细解释可参考博客[1]

import numpy as np
import pandas as pd
import cv2

# 将图片编码为rle格式
def rle_encode(im):
    '''
    im: numpy array, 1 - mask, 0 - background
    Returns run length as string formated
    '''
    pixels = im.flatten(order = 'F')
    pixels = np.concatenate([[0], pixels, [0]])
    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
    runs[1::2] -= runs[::2]
    return ' '.join(str(x) for x in runs)


# 将rle格式进行解码为图片
def rle_decode(mask_rle: str = '', shape=(512, 512)):
    '''
    mask_rle: run-length as string formated (start length)
    shape: (height,width) of array to return 
    Returns numpy array, 1 - mask, 0 - background

    '''
    s = mask_rle.split()
    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0:][::2], s[1:][::2])]
    starts -= 1
    ends = starts + lengths
    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)
    for lo, hi in zip(starts, ends):
        img[lo:hi] = 1
    return img.reshape(shape, order='F')

解题思路

由于本次赛题是一个典型的语义分割任务，因此可以直接使用语义分割的模型来完成：

• 步骤1：使用FCN 模型模型跑通具体模型训练过程，并对结果进行预测提交；

• 步骤2：在现有基础上加入数据扩增方法，并划分验证集以监督模型精度；

• 步骤3：使用更加强大模型结构（如Unet 和PSPNet）或尺寸更大的输入完成训练；

• 步骤4：训练多个模型完成模型集成操作；

课后作业

1. 理解RLE 编码过程，并完成赛题数据读取并可视化；
2. 统计所有图片整图中没有任何建筑物像素占所有训练集图片的比例；
3. 统计所有图片中建筑物像素占所有像素的比例；
4. 统计所有图片中建筑物区域平均区域大小；

作业主要涉及到读取训练集图片以及将对应的RLE编码解码转换为mask图片后进行简单判断即可，代码如下：

import numpy as np
import pandas as pd
import cv2
import matplotlib.pyplot as plt
from rle import rle_encode, rle_decode


train_mask = pd.read_csv('./datasets/train_mask.csv', sep = '\t', names = ['name', 'mask'])

#读取第一张图，并将对应的rle解码为mask矩阵
img = cv2.imread("./datasets/train/" + train_mask['name'].iloc[0])
mask = rle_decode(train_mask['mask'].iloc[0])
print(rle_encode(mask) == train_mask['mask'].iloc[0])  #true

plt.figure()
plt.subplot(1,2,1)
plt.imshow(img)
plt.subplot(1,2,2)
plt.imshow(mask)
plt.show()

noBuild = 0     
BuildArea = 0    
for idx in np.arange(train_mask.shape[0]):   
    if pd.isnull(train_mask['mask'].iloc[idx]):
        noBuild += 1
    else:
        mask = rle_decode(train_mask['mask'].iloc[idx])
        BuildArea += np.sum(mask)

#统计所有图片中建筑物区域平均区域大小
meanBuildArea = BuildArea / (train_mask.shape[0]-noBuild)

#统计所有图片中建筑物像素占所有像素的比例
buildPixPerc = BuildArea / (train_mask.shape[0]*mask.shape[0]*mask.shape[1])

#统计所有图片整图中没有任何建筑物像素占所有训练集图片的比例
noBuildPerc = noBuild/train_mask.shape[0] 

print("The percentage of image containing no buildings: %.4f" % noBuildPerc)
print("The percentage of pixels of building: %.4f" % buildPixPerc)
print("The mean area of buildings in an image: %d" % meanBuildArea)

运行结果为：

The percentage of image containing no buildings: 0.1735
The percentage of pixels of building: 0.1571
The mean area of buildings in an image:49820

全卷积网络FCN

关于全卷积网络FCN(Fully Convolutional Network)， 网上已有不少优秀博客[2-3]对其解读，FCN的要点包括：

• FCN将经典CNN (VGG)的后几个全连接层表示成卷积层，卷积核的大小(通道数，宽，高)分别为（4096,7,7）、（4096,1,1）、（1000,1,1），卷积跟全连接是不一样的概念和计算过程，使用的是之前CNN已经训练好的权值和偏置。FCN所有的层都是卷积层，故称为全卷积网络。与经典的CNN在卷积层使用全连接层得到固定长度的特征向量进行分类不同，FCN可以接受任意尺寸的输入图像。
• FCN采用跳级结构(skip layer)将不同深度层的feature map融合，因此有FCN-32s , FCN-16x , FCN-8s。 采用反卷积层(Transposed Conv)对融合后的特征图进行上采样，使它恢复到与输入图像相同的尺寸，从而在保留了原始输入图像中的空间信息的同时，可以对每一个像素都产生一个预测，这种逐像素的分类解决了语义级别的图像分割问题。
  

基于FCN Baseline 跑通语义分割训练过程

基于FCN Baseline训练的代码架构为：

-----utils
       |–utils
       |–rle
       |–readData
       |–dataloader
       |–loss
-----main_baseline

加载数据集

我们使用Pytorch读取赛题数据。通过Dataset 对数据进行读取并进行数据扩增，DataLoder 对Dataset进行封装并进行批量读取。定义自己的Dataset 类需要重载__getitem__()和__len__函数，注意这里使用albumentations进行数据扩增，self.as_tensor 能对一批大小为(H, W, 3)的RGB图片转换为标准化的(3, IMAGE_SIZE, IMAGE_SIZE) 的tensor。例如原来的航拍图片是512x512的，可能是为了节省内存，转换为256x256的图片进行训练，由于FCN对输入图片大小无要求，测试时可以直接输入512x512图片评估分割效果。

import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader, Subset
from torchvision import transforms as T
import albumentations as A
import cv2
from .rle import rle_encode, rle_decode


# Custom Dataset class
class TianChiDataset(Dataset):
    def __init__(self, imgPaths, rles, IMAGE_SIZE, test_mode=False):
        super(TianChiDataset, self).__init__()
        self.imgPaths = imgPaths
        self.rles = rles
        self.test_mode = test_mode
 
        self.transform = A.Compose([
            A.Resize(IMAGE_SIZE, IMAGE_SIZE),
            A.HorizontalFlip(p=0.5),
            A.VerticalFlip(p=0.5),
            A.RandomRotate90(),
        ])
        self.as_tensor = T.Compose([
            T.ToPILImage(),
            T.Resize(IMAGE_SIZE),
            T.ToTensor(),
            T.Normalize([0.625, 0.448, 0.688],  #normMean
                        [0.131, 0.177, 0.101]), #normStd
        ])

        self.len = len(imgPaths)

    # get data operation
    def __getitem__(self, index):
        img = cv2.imread(self.imgPaths[index])
        if not self.test_mode:
            mask = rle_decode(self.rles[index])

            #data augmentation via albumentations
            augments = self.transform(image = img, mask = mask) 
            return self.as_tensor(augments['image']), augments['mask'][None] #增加第1个维度
        else:
            return self.as_tensor(img), ''

    def __len__(self):
        return self.len

函数 get_dataloader 以自定义的Dataset类为输入，划分训练集和验证集，得到封装的DataLoder类。预训练时训练集有 $30000÷5=6000$ 张图，验证集有$30000÷300=100$ 张图，代码如下：

def get_dataloader(dataset, BATCH_SIZE):

    train_idx, valid_idx = [], []
    for i in range(len(dataset)):
        if i % 300 == 0:
            valid_idx.append(i)
        # else:
        elif i % 5 == 1:  #pretrain
            train_idx.append(i)
    train_ds = Subset(dataset, train_idx)
    valid_ds = Subset(dataset, valid_idx)

    #define training and validation data loaders
    train_loader = DataLoader(
        train_ds, batch_size = BATCH_SIZE, shuffle=True, num_workers=0)
    val_loader = DataLoader(
        valid_ds, batch_size = BATCH_SIZE, shuffle=False, num_workers=0)
    return train_loader, val_loader

loss函数

语义分割任务，常用Dice coefficient来衡量选手结果与真实标签的差异性，Dice coefficient可以按像素差异性来比较结果的差异性，具体计算方式为：
$$\frac{2\ast |X\cap Y|}{|X|+|Y|}$$
其中 $X$ 是预测结果，$Y$ 为真实标签的结果。当$X$与$Y$完全相同时Dice coefficient为1，排行榜使用所有测试集图片的平均Dice coefficient来衡量，分数值越大越好。

这里定义了SoftDiceLoss，除此之外，还使用二分类常用的二值交叉熵误差 (BCE)，nn.BCEWithLogitsLoss()是以sigmoid形式为输入时计算BCE的数值稳定版本。

import torch.nn as nn

class SoftDiceLoss(nn.Module):
    def __init__(self, smooth=1., dims=(-2,-1)):

        super(SoftDiceLoss, self).__init__()
        self.smooth = smooth
        self.dims = dims
    
    def forward(self, x, y):
        tp = (x * y).sum(self.dims)
        fp = (x * (1 - y)).sum(self.dims)
        fn = ((1 - x) * y).sum(self.dims)
        
        dc = (2 * tp + self.smooth) / (2 * tp + fp + fn + self.smooth)
        dc = dc.mean()
        return 1 - dc

# Numerically stable version of the binary cross-entropy loss function with sigmoid input.
# z * -log(sigmoid(x)) + (1 - z) * -log(1 - sigmoid(x))    
bce_fn = nn.BCEWithLogitsLoss()
dice_fn = SoftDiceLoss()

def loss_fn(y_pred, y_true):
    bce = bce_fn(y_pred, y_true)
    dice = dice_fn(y_pred.sigmoid(), y_true)
    return 0.8*bce+ 0.2*dice

加载预训练FCN模型，定义optimizer

# Define model, optimizer
model = torchvision.models.segmentation.fcn_resnet50(True)
    
# pth = torch.load("../input/pretrain-coco-weights-pytorch/fcn_resnet50_coco-1167a1af.pth")
# for key in ["aux_classifier.0.weight", "aux_classifier.1.weight", "aux_classifier.1.bias", "aux_classifier.1.running_mean", "aux_classifier.1.running_var", "aux_classifier.1.num_batches_tracked", "aux_classifier.4.weight", "aux_classifier.4.bias"]:
#     del pth[key]
    
model.classifier[4] = nn.Conv2d(512, 1, kernel_size=(1, 1), stride=(1, 1))
model.to(DEVICE)
#print(model.buffers)  #visualize the network architature

optimizer = torch.optim.AdamW(model.parameters(),
                  lr=1e-4, weight_decay=1e-3)

网络训练与验证

header = r'''
        Train | Valid
Epoch |  Loss |  Loss | Time, m
'''
#          Epoch         metrics            time
raw_line = '{:6d}' + '\u2502{:7.3f}'*2 + '\u2502{:6.2f}'
print(header)

best_loss = 10
train_result = {'iters':[], 'train_losses': []}

iters = 0
for epoch in range(1, EPOCHES+1):
    train_losses = []
    start_time = time.time()
    model.train()
    for image, target in progressbar(train_loader):
        
        image, target = image.to(DEVICE), target.float().to(DEVICE)
        optimizer.zero_grad()
        output = model(image)['out']   #orderedDict
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
        train_losses.append(loss.item())
        iters += 1

        if iters % ITER_PER_EPOCH < 20 or iters % 100 == 0:
            train_result['iters'].append(iters)
            train_result['train_losses'].append(loss.item())

    #Valid per epoch
    with torch.no_grad():
        val_losses = []
        model.eval() 
        for image, target in val_loader:
            image, target = image.to(DEVICE), target.float().to(DEVICE)
            output = model(image)['out']
            loss = loss_fn(output, target)
            val_losses.append(loss.item()) 

    vloss = np.array(val_losses).mean()
    print(raw_line.format(epoch, np.array(train_losses).mean(), 
            vloss, (time.time()-start_time)/60**1))

    if vloss < best_loss:
        best_loss = vloss
        torch.save(model.state_dict(), './checkpoints/model_best.pth')

plot_lossCurve(train_result)

说明：这里对每个epoch的前20次迭代以及每隔100次迭代的训练误差进行保存，绘制误差下降曲线；每次epoch结束时进行验证，将最小验证误差对应的模型权重参数保存为checkpoint。

Baseline模型测试结果

• 加载训练时保存的权重参数用于预测，测试集有2500张图片。将网络预测的mask进行RLE编码，与图片名称一起保存为csv文件。

model.load_state_dict(torch.load('./checkpoints/model_best.pth'))
model.eval()

subm = []
test_mask = pd.read_csv('./datasets/test_a_samplesubmit.csv', sep='\t', names=['name', 'mask'])
test_mask['name'] = test_mask['name'].apply(lambda x: './datasets/test_a/' + x)

for idx, name in enumerate(progressbar(test_mask['name'].iloc[:])):
    image = cv2.imread(name)  #ndarray(512,512,3)
    with torch.no_grad():  
        image = as_tensor(image).unsqueeze(0) #tensor(1, 3, 512, 512)
        score = model(image.to(DEVICE))['out'][0][0] #tensor(1, 1, 512, 512)->(512, 512)
        score_sigmoid = score.sigmoid().cpu().numpy()
        mask = (score_sigmoid > 0.5).astype(np.uint8)          
        # break
    subm.append([name.split('/')[-1], rle_encode(mask)])

#save predicted rle labels of test set into csv
subm = pd.DataFrame(subm)
subm.to_csv('./datasets/test_mask_temp.csv', index=None, header=None, sep='\t')

• 第2000张测试集图片的语义分割预测效果如下，对于仅有6000张训练集图片以及训练10个epoch，花费不到1小时来说，这样的预测效果还算挺好的！
  

如何进一步提升语义分割预测精度呢？

1. 使⽤更强的数据增强⽅法；
2. 模型调参，如学习率、图像尺⼨等；
3. 调整优化算法、损失函数，考虑正则⽅法；
4. 更换更强模型，如UNet、DeepLab等；
5. 考虑集成⽅法；
6. . ……

注意样本的分配，绘制学习曲线，基于模型的学习效果（是否过拟合、⽋拟合）来确定合适的优化策略

参考文献

[1]对mask进行rle编码然后进行解码-详细注释
[2] 全卷积网络 FCN 详解
[3] FCN的学习及理解（Fully Convolutional Networks for Semantic Segmentation）