【Keras】基于SegNet和U-Net的遥感图像语义分割

上两个月参加了个比赛，做的是对遥感高清图像做语义分割，美其名曰“天空之眼”。这两周数据挖掘课期末project我们组选的课题也是遥感图像的语义分割，所以刚好又把前段时间做的成果重新整理和加强了一下，故写了这篇文章，记录一下用深度学习做遥感图像语义分割的完整流程以及一些好的思路和技巧。

数据集

首先介绍一下数据，我们这次采用的数据集是CCF大数据比赛提供的数据（2015年中国南方某城市的高清遥感图像），这是一个小数据集，里面包含了5张带标注的大尺寸RGB遥感图像（尺寸范围从3000×3000到6000×6000），里面一共标注了4类物体，植被（标记1）、建筑（标记2）、水体（标记3）、道路（标记4）以及其他(标记0)。其中，耕地、林地、草地均归为植被类，为了更好地观察标注情况，我们将其中三幅训练图片可视化如下：蓝色-水体，黄色-房屋，绿色-植被，棕色-马路。更多数据介绍可以参看这里。

现在说一说我们的数据处理的步骤。我们现在拥有的是5张大尺寸的遥感图像，我们不能直接把这些图像送入网络进行训练，因为内存承受不了而且他们的尺寸也各不相同。因此，我们首先将他们做随机切割，即随机生成x,y坐标，然后抠出该坐标下256*256的小图，并做以下数据增强操作：

1. 原图和label图都需要旋转：90度，180度，270度
2. 原图和label图都需要做沿y轴的镜像操作
3. 原图做模糊操作
4. 原图做光照调整操作
5. 原图做增加噪声操作（高斯噪声，椒盐噪声）

这里我没有采用Keras自带的数据增广函数，而是自己使用opencv编写了相应的增强函数。

img_w = 256  
img_h = 256  

image_sets = ['1.png','2.png','3.png','4.png','5.png']

def gamma_transform(img, gamma):
    gamma_table = [np.power(x / 255.0, gamma) * 255.0 for x in range(256)]
    gamma_table = np.round(np.array(gamma_table)).astype(np.uint8)
    return cv2.LUT(img, gamma_table)

def random_gamma_transform(img, gamma_vari):
    log_gamma_vari = np.log(gamma_vari)
    alpha = np.random.uniform(-log_gamma_vari, log_gamma_vari)
    gamma = np.exp(alpha)
    return gamma_transform(img, gamma)


def rotate(xb,yb,angle):
    M_rotate = cv2.getRotationMatrix2D((img_w/2, img_h/2), angle, 1)
    xb = cv2.warpAffine(xb, M_rotate, (img_w, img_h))
    yb = cv2.warpAffine(yb, M_rotate, (img_w, img_h))
    return xb,yb

def blur(img):
    img = cv2.blur(img, (3, 3));
    return img

def add_noise(img):
    for i in range(200): #添加点噪声
        temp_x = np.random.randint(0,img.shape[0])
        temp_y = np.random.randint(0,img.shape[1])
        img[temp_x][temp_y] = 255
    return img


def data_augment(xb,yb):
    if np.random.random() < 0.25:
        xb,yb = rotate(xb,yb,90)
    if np.random.random() < 0.25:
        xb,yb = rotate(xb,yb,180)
    if np.random.random() < 0.25:
        xb,yb = rotate(xb,yb,270)
    if np.random.random() < 0.25:
        xb = cv2.flip(xb, 1)  # flipcode > 0：沿y轴翻转
        yb = cv2.flip(yb, 1)

    if np.random.random() < 0.25:
        xb = random_gamma_transform(xb,1.0)

    if np.random.random() < 0.25:
        xb = blur(xb)

    if np.random.random() < 0.2:
        xb = add_noise(xb)

    return xb,yb

def creat_dataset(image_num = 100000, mode = 'original'):
    print('creating dataset...')
    image_each = image_num / len(image_sets)
    g_count = 0
    for i in tqdm(range(len(image_sets))):
        count = 0
        src_img = cv2.imread('./data/src/' + image_sets[i])  # 3 channels
        label_img = cv2.imread('./data/label/' + image_sets[i],cv2.IMREAD_GRAYSCALE)  # single channel
        X_height,X_width,_ = src_img.shape
        while count < image_each:
            random_width = random.randint(0, X_width - img_w - 1)
            random_height = random.randint(0, X_height - img_h - 1)
            src_roi = src_img[random_height: random_height + img_h, random_width: random_width + img_w,:]
            label_roi = label_img[random_height: random_height + img_h, random_width: random_width + img_w]
            if mode == 'augment':
                src_roi,label_roi = data_augment(src_roi,label_roi)

            visualize = np.zeros((256,256)).astype(np.uint8)
            visualize = label_roi *50

            cv2.imwrite(('./aug/train/visualize/%d.png' % g_count),visualize)
            cv2.imwrite(('./aug/train/src/%d.png' % g_count),src_roi)
            cv2.imwrite(('./aug/train/label/%d.png' % g_count),label_roi)
            count += 1 
            g_count += 1

经过上面数据增强操作后，我们得到了较大的训练集：100000张256*256的图片。

卷积神经网络

面对这类图像语义分割的任务，我们可以选取的经典网络有很多，比如FCN,U-Net,SegNet,DeepLab,RefineNet,Mask Rcnn,Hed Net这些都是非常经典而且在很多比赛都广泛采用的网络架构。所以我们就可以从中选取一两个经典网络作为我们这个分割任务的解决方案。我们根据我们小组的情况，选取了U-Net和SegNet作为我们的主体网络进行实验。

SegNet

SegNet已经出来好几年了，这不是一个最新、效果最好的语义分割网络，但是它胜在网络结构清晰易懂，训练快速坑少，所以我们也采取它来做同样的任务。SegNet网络结构是编码器-解码器的结构，非常优雅，值得注意的是，SegNet做语义分割时通常在末端加入CRF模块做后处理，旨在进一步精修边缘的分割结果。有兴趣深究的可以看看这里

现在讲解代码部分，首先我们先定义好SegNet的网络结构。

def SegNet():  
    model = Sequential()  
    #encoder  
    model.add(Conv2D(64,(3,3),strides=(1,1),input_shape=(3,img_w,img_h),padding='same',activation='relu'))  
    model.add(BatchNormalization())  
    model.add(Conv2D(64,(3,3),strides=(1,1),padding='same',activation='relu'))  
    model.add(BatchNormalization())  
    model.add(MaxPooling2D(pool_size=(2,2)))  
    #(128,128)  
    model.add(Conv2D(128, (3, 3), strides=(1, 1), padding='same', activation='relu'))  
    model.add(BatchNormalization())  
    model.add(Conv2D(128, (3, 3), strides=(1, 1), padding='same', activation='relu'))  
    model.add(BatchNormalization())  
    model.add(MaxPooling2D(pool_size=(2,