pytorch-semseg

用 PyTorch 实现的语义分割算法

本仓库旨在将流行的语义分割架构在 PyTorch 中进行复现。

已实现的网络

• PSPNet - 支持加载预训练模型，无需 Caffe 依赖
• ICNet - 可选批量归一化及预训练模型
• FRRN - A 和 B 两种模型
• FCN - 包括 FCN32s、FCN16s 和 FCN8s 三种变体
• U-Net - 可选反卷积和批量归一化
• Link-Net - 支持多种 ResNet 后端
• Segnet - 使用最大池化索引进行上采样

即将支持

• E-Net
• RefineNet

已实现的数据加载器

• CamVid
• Pascal VOC
• ADE20K
• MIT 场景解析基准
• Cityscapes
• NYUDv2
• Sun-RGBD

需求

• pytorch >=0.4.0
• torchvision ==0.2.0
• scipy
• tqdm
• tensorboardX

一行安装命令

pip install -r requirements.txt

数据

• 从以下 URL 列表下载所需数据集：这里。
• 解压 zip 或 tar 文件，并在 config.yaml 中相应地修改路径。

使用方法

配置文件设置

# 模型配置
model:
    arch: <名称> [选项: 'fcn[8,16,32]s, unet, segnet, pspnet, icnet, icnetBN, linknet, frrn[A,B]']
    <model_keyarg_1>:<value>

# 数据配置
data:
    dataset: <名称> [选项: 'pascal, camvid, ade20k, mit_sceneparsing_benchmark, cityscapes, nyuv2, sunrgbd, vistas'] 
    train_split: <用于训练的划分>
    val_split: <用于验证的划分>
    img_rows: 512
    img_cols: 1024
    path: <数据路径>
    <dataset_keyarg1>:<value>

# 训练配置
training:
    n_workers: 64
    train_iters: 35000
    batch_size: 16
    val_interval: 500
    print_interval: 25
    loss:
        name: <损失类型> [选项: '交叉熵, bootstrapped_cross_entropy, multi_scale_crossentropy']
        <loss_keyarg1>:<value>

    # 优化器配置
    optimizer:
        name: <优化器名称> [选项: 'sgd, adam, adamax, asgd, adadelta, adagrad, rmsprop']
        lr: 1.0e-3
        <optimizer_keyarg1>:<value>

        # 学习率预热配置
        warmup_iters: <学习率预热的迭代次数>
        mode: <预热模式: 'constant' 或 'linear'>
        gamma: <预热的 gamma 值>
       
    # 数据增强配置
    augmentations:
        gamma: x                                     #[gamma 在 1 到 1+x 之间变化]
        hue: x                                       #[色调在 -x 到 x 之间变化]
        brightness: x                                #[亮度在 1-x 到 1+x 之间变化]
        saturation: x                                #[饱和度在 1-x 到 1+x 之间变化]
        contrast: x                                  #[对比度在 1-x 到 1+x 之间变化]
        rcrop: [h, w]                                #[裁剪为 (h,w) 大小]
        translate: [dh, dw]                          #[沿 (dh, dw) 方向反射式平移]
        rotate: d                                    #[旋转 -d 到 d 度]
        scale: [h,w]                                 #[缩放至 (h,w) 大小]
        ccrop: [h,w]                                 #[中心裁剪为 (h,w)]
        hflip: p                                     #[水平翻转，概率为 p]
        vflip: p                                     #[垂直翻转，概率为 p]

    # 学习率调度配置
    lr_schedule:
        name: <调度类型> [选项: 'constant_lr, poly_lr, multi_step, cosine_annealing, exp_lr']
        <scheduler_keyarg1>:<value>

    # 从检查点恢复
    resume: <检查点路径>

训练模型：

python train.py [-h] [--config [CONFIG]] 

--config                要使用的配置文件

验证模型：

usage: validate.py [-h] [--config [CONFIG]] [--model_path [MODEL_PATH]]
                       [--eval_flip] [--measure_time]

  --config              要使用的配置文件
  --model_path          已保存模型的路径
  --eval_flip           启用翻转图像评估 | 默认为 True
  --measure_time        启用时间（fps）测量评估 | 默认为 True

针对自定义图像使用某个数据集测试模型：

python test.py [-h] [--model_path [MODEL_PATH]] [--dataset [DATASET]]
               [--dcrf [DCRF]] [--img_path [IMG_PATH]] [--out_path [OUT_PATH]]
 
  --model_path          已保存模型的路径
  --dataset             要使用的数据集 ['pascal, camvid, ade20k 等']
  --dcrf                启用基于 DenseCRF 的后处理
  --img_path            输入图像的路径
  --out_path            输出分割图的路径

如果您在研究中使用了此代码，请考虑引用：

@article{mshahsemseg,
    Author = {Meet P Shah},
    Title = {Semantic Segmentation Architectures Implemented in PyTorch.},
    Journal = {https://github.com/meetshah1995/pytorch-semseg},
    Year = {2017}
}

pytorch-semseg 快速上手指南

pytorch-semseg 是一个基于 PyTorch 实现的语义分割算法库，集成了 PSPNet、FCN、U-Net、SegNet 等多种主流架构，并支持 Cityscapes、Pascal VOC 等常用数据集。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux / macOS (Windows 需自行配置兼容环境)
• Python: 建议 Python 3.6+
• 核心依赖:
  • pytorch >= 0.4.0
  • torchvision == 0.2.0
  • scipy
  • tqdm
  • tensorboardX

国内加速建议：安装依赖时推荐使用清华或阿里镜像源以提升下载速度。

安装步骤

1. 克隆项目

首先从 GitHub 克隆代码仓库：

git clone https://github.com/meetshah1995/pytorch-semseg.git
cd pytorch-semseg

2. 安装依赖

使用 pip 一键安装所需包。国内用户建议添加 -i 参数指定镜像源：

# 使用清华镜像源安装
pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 准备数据

下载您需要的数据集（如 CamVid, Pascal VOC, Cityscapes 等）。

• 数据集下载地址参考：官方数据集列表
• 解压数据包后，记下数据存放路径，后续需在配置文件中修改。

基本使用

1. 配置模型与训练参数

复制示例配置文件（如有）或直接新建 config.yaml，根据您的需求编辑关键参数。以下是一个最小化的配置示例：

# Model Configuration
model:
    arch: fcn8s  # 可选：fcn8s, unet, segnet, pspnet 等
    n_classes: 21 # 根据数据集类别数设置，Pascal VOC 为 21

# Data Configuration
data:
    dataset: pascal # 可选：pascal, camvid, cityscapes 等
    train_split: training
    val_split: validation
    path: /path/to/your/dataset # 修改为您的实际数据路径
    img_rows: 512
    img_cols: 1024

# Training Configuration
training:
    batch_size: 4
    train_iters: 1000
    val_interval: 100
    loss:
        name: cross_entropy
    optimizer:
        name: adam
        lr: 1.0e-3

2. 训练模型

使用配置好的 yaml 文件启动训练：

python train.py --config config.yaml

3. 验证模型

训练完成后，使用保存的模型权重进行验证：

python validate.py --config config.yaml --model_path /path/to/checkpoint.pkl

4. 测试单张图片

对自定义图片进行推理测试：

python test.py --model_path /path/to/checkpoint.pkl --dataset pascal --img_path ./test_image.jpg --out_path ./result.png

• --dcrf: 可选参数，启用 DenseCRF 后处理以优化边缘细节。