【AI特训营】：基于BiseNetV2的洪水区域分割

★★★ 本文源自AlStudio社区精品项目，【点击此处】查看更多精品内容 >>>

基于BiseNetv2的洪水分割任务

BiseNet v2:

基于轻量化网络模型的设计作为一个热门的研究方法，许多研究者都在运算量、参数量和精度之间寻找平衡，希望使用尽量少的运算量和参数量的同时获得较高的模型精度。目前，轻量级模型主要有SqueezeNet、MobileNet系列和ShuffleNet系列等，这些模型在图像分类领域取得了不错的效果，可以作为基本的主干网络应用于语义分割任务当中。
然而，在语义分割领域，由于需要对输入图片进行逐像素的分类，运算量很大。通常，为了减少语义分割所产生的计算量，通常而言有两种方式：减小图片大小和降低模型复杂度。减小图片大小可以最直接地减少运算量，但是图像会丢失掉大量的细节从而影响精度。降低模型复杂度则会导致模型的特征提取能力减弱，从而影响分割精度。所以，如何在语义分割任务中应用轻量级模型，兼顾实时性和精度性能具有相当大的挑战性。
与几种最新的实时语义分割方法相比，BiSeNet V2具有良好的性能。具体来说，对于2048x1,024的输入，BiseNet2在Cityscapes测试集中的平均IoU达到72.6％，在一张NVIDIA GeForce GTX 1080 Ti卡上的速度为156 FPS，这比现有方法要快得多，而且可以实现更好的分割精度。

在本次任务中，我们使用 **BiSeNetV2** 来完成任务。

数据集介绍

**数据集地址：**https://www.kaggle.com/datasets/faizalkarim/flood-area-segmentation

数据集包含洪水灾区的图像和显示水域的相应遮罩图像。
有290张图片和自动标注的面具。掩模图像是使用开源数据标签软件Label Studio创建的。任务是创建一个分割模型，该模型可以准确地分割洪水灾区的给定图片中的水域。

此类模型可用于洪水调查、更好的决策和规划。由于数据较少，可以使用预先训练的模型和数据扩充。

数据集可视化：

1.解压数据集，配置环境

数据集只有一个zip文件，将其解压至data/dataset文件夹下，这个文件夹下的文件在每次关闭环境后会清空，因此每次运行环境时均需重新解压。
本项目使用PaddleSeg进行训练，地址：https://github.com/PaddlePaddle/PaddleSeg ，在这里，我们使用pypi上打包好的包进行安装，和数据集一样，每次打开环境的时候均需重新安装

!unzip data/data198914/archive.zip -d data/dataset
!pip install paddleseg

2.数据预处理

原数据集的mask图像为：

不是正常mask的形式，因此，首先需要将mask转为只有0和1的值的图片，放入data/dataset/gt_mask文件夹下，处理后的图片应该在肉眼观察下为一片黑

根据PaddleSeg的数据集格式，我们将训练图片和mask图按照img mask\n的形式写入数据集文件，按照8：2的比例划分训练集与测试集。

在处理数据的过程中，有部分图像的大小与mask的大小不同，因此在写入数据集文件的时候忽略掉这部分数据

import os
import random
import cv2
import numpy as np

%cd /home/aistudio
imgs = os.listdir('data/dataset/Image')
root_path = '/home/aistudio/data/dataset/'
os.system('mkdir /home/aistudio/data/dataset/gt_mask')
ignore = ['14.jpg',
          '15.jpg',
          '0.jpg',
          '2053.jpg',
          '1061.jpg',
          '2052.jpg',
          '1079.jpg',
          '3059.jpg']
with open('train.list', 'w') as f1:
    with open('val.list', 'w') as f2:
        for img in imgs:
            if img[-4:] != '.jpg' or ignore.count(img) != 0 :
                continue
            img_path = root_path + 'Image/'
            mask_path = root_path + 'Mask/'
            n = random.randint(0, 10)

            mask = cv2.imread(mask_path+img[:-4]+'.png')
            imgs = cv2.imread(img_path+img)
            if imgs.shape != mask.shape:
                print(img_path+img)
                print(imgs.shape)
                print(mask.shape)
            gt_mask = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.uint8)
            gt_mask[mask[:,:,0]!=0] = 1
            cv2.imwrite(root_path+'gt_mask/'+img[:-4]+'.png', gt_mask)

            if n > 8:
                f2.writelines(img_path+img+' '+root_path+'gt_mask/'+img[:-4]+'.png\n')
            else:
                f1.writelines(img_path+img+' '+root_path+'gt_mask/'+img[:-4]+'.png\n')
    

3.使用PaddleSeg进行训练

在bisenet_cityscapes_1024x1024_160k.yml的基础上进行更改，只需要将数据集的读取更改即可，这部分配置文件在configs/_base_/cityscapes.yml中，更改为：

batch_size: 1
iters: 80000

train_dataset:
  type: Dataset
  dataset_root: ./
  train_path: /home/aistudio/train.list
  num_classes: 2
  transforms:
    - type: ResizeStepScaling
      min_scale_factor: 0.5
      max_scale_factor: 2.0
      scale_step_size: 0.25
    - type: RandomPaddingCrop
      crop_size: [1024, 512]
    - type: RandomHorizontalFlip
    - type: RandomDistort
      brightness_range: 0.4
      contrast_range: 0.4
      saturation_range: 0.4
    - type: Normalize
  mode: train

val_dataset:
  type: Dataset
  dataset_root: ./
  val_path: /home/aistudio/val.list
  num_classes: 2
  transforms:
    - type: Normalize
  mode: val


optimizer:
  type: sgd
  momentum: 0.9
  weight_decay: 4.0e-5

lr_scheduler:
  type: PolynomialDecay
  learning_rate: 0.01
  end_lr: 0
  power: 0.9

loss:
  types:
    - type: CrossEntropyLoss
  coef: [1]

其中train_path和val_path为生成的数据集文件，是两个.list文件

%cd PaddleSeg
!python tools/train.py \
        --config configs/bisenet/bisenet_cityscapes_1024x1024_160k.yml \
        --save_interval 500 \
        --do_eval \
        --use_vdl \
        --save_dir output

经过8个小时的训练，共计训练了160000个iter，已经能够产生较好的效果，训练的评测图和Loss如图所示：

Acc是指预测的mask与真实mask的准确度 Acc=预测正确的像素数量/总像素数量
Iou是指mask与真实mask的重叠率 Iou=预测的mask与真实mask的交集/预测的mask与真实mask的并集
一般来说，两者越高证明效果越好

4.验证集测试

使用best_model测试模型效果
结果为：
Class IoU: [0.7836 0.7758]
Class Precision: [0.8853 0.867 ]
Class Recall: [0.8722 0.8806]
第一列为背景类，第二列为目标类(洪水覆盖区域)

!python tools/val.py \
        --config configs/bisenet/bisenet_cityscapes_1024x1024_160k.yml \
        --model_path output/best_model/model.pdparams

5.模型预测

对于单张图像，可以使用以下代码预测

!python tools/predict.py \
       --config configs/bisenet/bisenet_cityscapes_1024x1024_160k.yml \
       --model_path output/best_model/model.pdparams \
       --image_path /home/aistudio/data/dataset/Image/1038.jpg \
       --save_dir output/result

预测结果保存在output/result目录下：
可视化图像为：

mask为：

6.模型导出

对于训练好的模型，使用predict.py只是为了测试的可视化，对于模型的部署，首先需要导出模型，将用于训练的结构裁剪掉，使得模型更加轻便快速。

导出的推理模型会保存在output/infer_model下

!python tools/export.py \
       --config configs/bisenet/bisenet_cityscapes_1024x1024_160k.yml \
       --model_path output/best_model/model.pdparams \
       --save_dir output/infer_model

7.模型推理

使用导出的模型推理图片，PaddleSeg的推理代码可以输入三种类型的数据: 单张图片、数据集文件夹、list清单

!python deploy/python/infer.py \
        --config output/infer_model/deploy.yaml \
        --image_path /home/aistudio/data/dataset/Image/1038.jpg \
        --save_dir output/result

总结

本项目是AI训练营的项目之一，通过PaddleSeg，可以使用大量优秀的网络进行分割训练、部署、推理工作，这也大大降低了该领域的入门门槛，其次，相比与mmlab开源的mmsegmentation框架，PaddleSeg能够实现多模型的训练而只需更改一次数据路径，对于工业界而言更加方便。