TensorFlow 中的 DeepLab-v3-plus 语义分割

本仓库尝试在 TensorFlow 中复现 用于语义图像分割的带空洞可分离卷积的编码器-解码器模型（DeepLabv3+），以在 PASCAL VOC 数据集 和 Cityscapes 数据集 上进行语义图像分割。该实现主要基于 我的 DeepLabv3 实现，而后者最初则参考了 DrSleep 的 DeepLab v2 实现 以及 TensorFlow 模型库中的 ResNet 实现。

环境搭建

需求：

• tensorflow >=1.6
• numpy
• matplotlib
• pillow
• opencv-python

您可以通过运行 pip install -r requirements.txt 来安装这些依赖项。

数据集准备

该项目使用 TFRecord 格式 来处理训练和评估过程中的数据。从原始图像文件创建 TFRecord 文件非常简单，我们将在下面介绍具体步骤。

Cityscapes

注意： 本项目包含用于生成 Cityscapes 和 Pascal VOC 数据集 TFRecord 文件的脚本，但不支持其他数据集。

为 Cityscapes 创建 TFRecord 文件

要下载 Cityscapes 数据集，您首先需要在其 官网 上注册。注册后，请确保同时下载 leftImg8bit 和 gtFine 数据集。最终的文件夹结构应如下所示：

+ cityscapes
 + leftImg8bit
 + gtFine

接下来，为了生成数据集的训练标签，克隆 cityScapesScripts 项目：

git clone https://github.com/mcordts/cityscapesScripts.git
cd cityscapesScripts

然后，在您的 Cityscapes 数据集根目录下运行：

# 必须定义 $CITYSCAPES_ROOT 变量
python cityscapesscripts/preparation/createTrainIdLabelImgs.py

最后，您可以运行本仓库提供的转换脚本 create_cityscapes_tf_record.py。

Pascal VOC

• 下载并解压 PASCAL VOC 训练/验证数据（2GB 的 tar 文件），通过 --data_dir 参数指定保存路径。
• 下载并解压 增强版分割数据（感谢 DrSleep），同样通过 --data_dir 和 --label_data_dir 参数指定保存路径（即 $data_dir/$label_data_dir）。

为 Pascal VOC 创建 TFRecord 文件

准备好数据集后，您可以通过以下命令为 Pascal VOC 创建 TFRecord 文件：

python create_pascal_tf_record.py --data_dir DATA_DIR \
                                  --image_data_dir IMAGE_DATA_DIR \
                                  --label_data_dir LABEL_DATA_DIR 

训练

在训练之前，您需要从 slim 下载并解压 预训练的 Resnet v2 101 模型，并通过 --pre_trained_model 参数指定其保存路径。此外，您还需要将原始数据转换为 TensorFlow 的 TFRecord 格式。完成数据集准备的所有步骤并创建好训练和验证数据的 TFRecord 文件后，即可开始训练模型：

python train.py --model_dir MODEL_DIR --pre_trained_model PRE_TRAINED_MODEL

其中，--pre_trained_model 指的是预训练的 Resnet 模型，而 --model_dir 则用于保存训练好的 DeepLabv3+ 检查点。如果 --model_dir 中已存在有效的检查点，则模型将从该检查点继续训练。

您还可以通过以下命令查看其他可用选项：

python train.py --help

对于在 Pascal VOC 2012 验证集上达到 77.31% mIoU 的训练模型推理结果，您可以从 这里 下载并解压到 --model_dir 目录中。

训练过程可以通过 TensorBoard 进行可视化，命令如下：

tensorboard --logdir MODEL_DIR

评估

要评估模型性能，可以使用以下命令：

python evaluate.py --help

目前，本实现所构建的最佳模型在 Pascal VOC 2012 验证集上的 mIoU 达到 77.31%。

网络主干	训练操作系统	评估操作系统	SC	论文中的 mIoU	仓库中的 mIoU
Resnet101	16	16		78.85%	77.31%

上述模型是在 Tesla V100 显卡和 TensorFlow 1.6 环境下，使用以下参数训练约 9.5 小时得到的：

python train.py --train_epochs 43 --batch_size 15 --weight_decay 2e-4 --model_dir models/ba=15,wd=2e-4,max_iter=30k --max_iter 30000

推理

要对您的图像应用语义分割，可以使用以下命令：

python inference.py --data_dir DATA_DIR --infer_data_list INFER_DATA_LIST --model_dir MODEL_DIR 

训练好的模型可以从 这里 下载。关于掩码的详细说明，例如颜色含义等信息，可以在 DrSleep 的仓库 中找到。

待办事项：

欢迎提交 Pull Request。

• 实现解码器
• 使用 ResNet 作为网络主干
• 在 Cityscapes 数据集上进行训练
• 使用 Xception 作为网络主干
• 实现深度可分离卷积
• 提高网络对 GPU 内存的利用率（例如支持更大的批量大小）
• 支持多 GPU 训练
• 支持通道优先的内存布局（据称可在 GPU 上显著提升性能）
• 使用 MS-COCO 数据集预训练模型
• 编写单元测试

致谢

本仓库大量借鉴了以下项目的代码：

• DrSleep 的 DeepLab-ResNet（DeepLabv2）
• TensorFlow 官方模型库
• TensorFlow 物体检测 API
• TensorFlow-Slim
• TensorFlow

tensorflow-deeplab-v3-plus 快速上手指南

本指南帮助中国开发者快速部署并使用基于 TensorFlow 的 DeepLabv3+ 语义分割模型。

环境准备

系统要求

• Python: 建议 Python 3.6+
• TensorFlow: 版本 >= 1.6 (注意：此项目基于 TF 1.x，不直接兼容 TF 2.x)
• GPU: 推荐使用 NVIDIA GPU 以加速训练和推理（可选，但强烈推荐）

前置依赖

请确保已安装以下 Python 库：

• numpy
• matplotlib
• pillow
• opencv-python

国内加速建议：在安装依赖时，建议使用清华或阿里镜像源以提升下载速度。 例如：pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆仓库

   git clone https://github.com/rishizek/tensorflow-deeplab-v3-plus.git
   cd tensorflow-deeplab-v3-plus
   

2. 安装依赖

   pip install -r requirements.txt
   

   (国内用户推荐添加 -i https://pypi.tuna.tsinghua.edu.cn/simple)

3. 下载预训练模型 (可选，用于直接推理) 如果仅需进行推理测试，可下载作者在 Pascal VOC 上训练的模型 (mIoU 77.31%)：

   # 手动下载链接: https://www.dropbox.com/s/1xrd4c5atyrkb6z/deeplabv3plus_ver1.tar.gz?dl=0
   # 下载后解压到指定目录，例如 models/
   

基本使用

1. 数据集准备 (以 Pascal VOC 为例)

本项目使用 TFRecord 格式数据。你需要先下载原始数据并转换为 TFRecord。

• 下载数据：

  • PASCAL VOC train/val
  • 增强分割数据 (SegmentationClassAug)

• 生成 TFRecord： 将数据解压后，运行以下脚本生成训练所需的 .tfrecord 文件：

  python create_pascal_tf_record.py --data_dir DATA_DIR \
                                    --image_data_dir IMAGE_DATA_DIR \
                                    --label_data_dir LABEL_DATA_DIR
  

  注：请将 DATA_DIR 等替换为你本地的实际路径。

2. 模型训练

在开始训练前，需下载预训练的 ResNet v2 101 backbone 模型：

• 下载地址：resnet_v2_101_2017_04_14.tar.gz

执行训练命令：

python train.py --model_dir models/my_experiment \
                --pre_trained_model /path/to/resnet_v2_101.ckpt

• --model_dir: 保存训练检查点的目录。
• --pre_trained_model: 预训练 ResNet 模型的路径。

可视化训练过程：

tensorboard --logdir models/my_experiment

3. 模型推理 (Inference)

使用训练好的模型对你的图片进行语义分割：

python inference.py --data_dir /path/to/images \
                    --infer_data_list image_list.txt \
                    --model_dir models/my_experiment

• --data_dir: 包含待处理图片的文件夹。
• --infer_data_list: 包含待处理图片文件名的文本列表。
• --model_dir: 存放训练好模型检查点的目录。

生成的分割掩码（Mask）颜色含义可参考原项目引用的 DrSleep 仓库文档。