logo

渐进式语义分割 (MagNet)

MagNet是一个多尺度框架，通过在多个放大级别观察图像来解决局部歧义。它具有多个处理阶段，每个阶段对应一个放大级别，一个阶段的输出被输入到下一个阶段以进行从粗到细的信息传播。在城市景观、航拍场景和医学图像三个高分辨率数据集上的实验表明，MagNet在性能上始终显著优于现有最先进的方法。

MagNet模型架构和实验结果的详细信息可以在我们的以下论文中找到：

@inproceedings{m_Huynh-etal-CVPR21,
  author = {Chuong Huynh and Anh Tran and Khoa Luu and Minh Hoai},
  title = {Progressive Semantic Segmentation},
  year = {2021},
  booktitle = {Proceedings of the {IEEE} Conference on Computer Vision and Pattern Recognition (CVPR)},
}

当使用MagNet产生发表的结果或将其整合到其他软件中时，请引用我们的论文。

数据集

当前代码提供了在两个数据集上进行训练和评估的配置：Cityscapes和DeepGlobe。要准备数据集，请在./data目录中执行以下步骤：

Cityscapes数据集

在此页面注册账户并登录。
下载leftImg8bit_trainvaltest.zip和gtFine_trainvaltest.zip。
运行以下脚本将zip文件解压到正确位置：

sh ./prepare_cityscapes.sh

DeepGlobe数据集

在此页面注册账户并登录。
访问此页面并下载#1 Development阶段的Starting Kit。
运行以下脚本将zip文件解压到正确位置：

sh ./prepare_deepglobe.sh

如果您想使用自己的数据集进行训练/评估，请按照此文档中的步骤操作。

入门

要求

该框架在具有以下环境的机器上测试通过：

Python >= 3.6
CUDA >= 10.0

要安装依赖项，请运行以下命令：

pip install -r requirements.txt

预训练模型

预训练模型在数据集上的性能：

数据集	骨干网络	基线 IoU (%)	MagNet IoU (%)	MagNet-Fast IoU (%)	下载
Cityscapes	HRNetW18+OCR	63.24	68.20	67.37	骨干网络 refine_512x256 refine_1024x512 refine_2048x1024
DeepGlobe	Resnet50-FPN	67.22	72.10	68.22	骨干网络 refine

请手动下载预训练模型到./checkpoints目录，或运行以下脚本：

cd checkpoints
sh ./download_cityscapes.sh # 用于Cityscapes
# 或
sh ./download_deepglobe.sh # 用于DeepGlobe

使用方法

您可以运行这个Google Colab笔记本来测试我们的预训练模型，使用街景图像。请按照笔记本中的说明体验我们网络的性能。

如果您想在本地机器上测试我们的框架：

要测试Cityscapes图像，例如data/frankfurt_000001_003056_leftImg8bit.png：

使用MagNet细化：

python demo.py --dataset cityscapes \
               --image data/frankfurt_000001_003056_leftImg8bit.png \
               --scales 256-128,512-256,1024-512,2048-1024 \
               --crop_size 256 128 \
               --input_size 256 128 \
               --model hrnet18+ocr \
               --pretrained checkpoints/cityscapes_hrnet.pth \
               --pretrained_refinement checkpoints/cityscapes_refinement_512.pth checkpoints/cityscapes_refinement_1024.pth checkpoints/cityscapes_refinement_2048.pth \
               --num_classes 19 \
               --n_points 32768 \
               --n_patches -1 \
               --smooth_kernel 5 \
               --save_pred \
               --save_dir test_results/demo

# 或者简短地，您可以运行
sh scripts/cityscapes/demo_magnet.sh data/frankfurt_000001_003056_leftImg8bit.png

使用MagNet-Fast细化

python demo.py --dataset cityscapes \
               --image frankfurt_000001_003056_leftImg8bit.png \
               --scales 256-128,512-256,1024-512,2048-1024 \
               --crop_size 256 128 \
               --input_size 256 128 \
               --model hrnet18+ocr \
               --pretrained checkpoints/cityscapes_hrnet.pth \
               --pretrained_refinement checkpoints/cityscapes_refinement_512.pth checkpoints/cityscapes_refinement_1024.pth checkpoints/cityscapes_refinement_2048.pth \
               --num_classes 19 \
               --n_points 0.9 \
               --n_patches 4 \
               --smooth_kernel 5 \
               --save_pred \
               --save_dir test_results/demo

# 或者简短地，您可以运行
sh scripts/cityscapes/demo_magnet_fast.sh data/frankfurt_000001_003056_leftImg8bit.png

所有结果将存储在test_results/demo/frankfurt_000001_003056_leftImg8bit中

要测试DeepGlobe图像，例如data/639004_sat.jpg：

使用MagNet细化：

python demo.py --dataset deepglobe \
               --image data/639004_sat.jpg \
               --scales 612-612,1224-1224,2448-2448 \
               --crop_size 612 612 \
               --input_size 508 508 \
               --model fpn \
               --pretrained checkpoints/deepglobe_fpn.pth \
               --pretrained_refinement checkpoints/deepglobe_refinement.pth \
               --num_classes 7 \
               --n_points 0.75 \
               --n_patches -1 \
               --smooth_kernel 11 \
               --save_pred \
               --save_dir test_results/demo
# 或者简单地运行
sh scripts/deepglobe/demo_magnet.sh data/639004_sat.jpg

使用MagNet-Fast细化

python demo.py --dataset deepglobe \
               --image data/639004_sat.jpg \
               --scales 612-612,1224-1224,2448-2448 \
               --crop_size 612 612 \
               --input_size 508 508 \
               --model fpn \
               --pretrained checkpoints/deepglobe_fpn.pth \
               --pretrained_refinement checkpoints/deepglobe_refinement.pth \
               --num_classes 7 \
               --n_points 0.9 \
               --n_patches 3 \
               --smooth_kernel 11 \
               --save_pred \
               --save_dir test_results/demo

# 或者简单地运行
sh scripts/deepglobe/demo_magnet_fast.sh data/639004_sat.jpg

所有结果将保存在 test_results/demo/639004_sat

训练

训练骨干网络

我们自定义了来自HRNet仓库的训练脚本来训练我们的骨干网络。请先进入 ./backbone 目录并运行以下脚本：

用于Cityscapes的HRNetW18V2+OCR

下载在ImageNet上预训练的权重：

# 在 ./backbone 目录下
cd pretrained_weights
wget https://public.vinai.io/chuonghm/hrnet_w18_v2_imagenet.pth

训练模型：

# 在 ./backbone 目录下
python train.py --cfg experiments/cityscapes/hrnet_ocr_w18_train_256x128_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml

训练日志保存在 ./log/cityscapes/HRNetW18_OCR。

训练后的骨干网络检查点保存在 ./output/cityscapes/hrnet_ocr_w18_train_256x128_sgd_lr1e-2_wd5e-4_bs_12_epoch484/best.pth。这个检查点用于进一步训练细化模块。

用于Deepglobe的Resnet50-FPN

训练模型：

# 在 ./backbone 目录下
python train.py --cfg experiments/deepglobe/resnet_fpn_train_612x612_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml

训练日志保存在 ./log/deepglobe/ResnetFPN。

训练后的骨干网络检查点保存在 ./output/deepglobe/resnet_fpn_train_612x612_sgd_lr1e-2_wd5e-4_bs_12_epoch484/best.pth。这个检查点用于进一步训练细化模块。

训练细化模块

训练可用的参数：

train.py [-h] --dataset DATASET [--root ROOT] [--datalist DATALIST]
                --scales SCALES --crop_size N [N ...] --input_size N [N ...]
                [--num_workers NUM_WORKERS] --model MODEL --num_classes
                NUM_CLASSES --pretrained PRETRAINED
                [--pretrained_refinement PRETRAINED_REFINEMENT [PRETRAINED_REFINEMENT ...]]
                --batch_size BATCH_SIZE [--log_dir LOG_DIR] --task_name
                TASK_NAME [--lr LR] [--momentum MOMENTUM] [--decay DECAY]
                [--gamma GAMMA] [--milestones N [N ...]] [--epochs EPOCHS]

可选参数:
  -h, --help            显示帮助信息并退出
  --dataset DATASET     数据集名称：cityscapes, deepglobe (默认：None)
  --root ROOT           训练和测试图像的路径 (默认：)
  --datalist DATALIST   包含图像和标签路径的.txt文件路径 (默认：)
  --scales SCALES       尺度：w1-h1,w2-h2,... , 例如
                        512-512,1024-1024,2048-2048 (默认：None)
  --crop_size N [N ...]
                        裁剪大小，例如 256 128 (默认：None)
  --input_size N [N ...]
                        输入大小，例如 256 128 (默认：None)
  --num_workers NUM_WORKERS
                        数据加载器的工作进程数 (默认：1)
  --model MODEL         模型名称。可选：fpn, psp, hrnet18+ocr, hrnet48+ocr
                        (默认：None)
  --num_classes NUM_CLASSES
                        类别数 (默认：None)
  --pretrained PRETRAINED
                        预训练权重 (默认：None)
  --pretrained_refinement PRETRAINED_REFINEMENT [PRETRAINED_REFINEMENT ...]
                        细化模块的预训练权重 (默认：[''])
  --batch_size BATCH_SIZE
                        训练的批量大小 (默认：None)
  --log_dir LOG_DIR     存储日志文件的目录 (默认：runs)
  --task_name TASK_NAME
                        任务名称，实验名称。最终的日志路径为
                        <log_dir>/<task_name>/<timestamp> (默认：None)
  --lr LR               学习率 (默认：0.001)
  --momentum MOMENTUM   优化器的动量 (默认：0.9)
  --decay DECAY         优化器的权重衰减 (默认：0.0005)
  --gamma GAMMA         学习率调度器的gamma (默认：0.1)
  --milestones N [N ...]
                        降低学习率的里程碑 (默认：[10, 20, 30, 40, 45])
  --epochs EPOCHS       训练的轮数 (默认：50)

Cityscapes

要在Cityscapes数据集上训练MagNet，请运行以下示例脚本：

python train.py --dataset cityscapes \
                --root data/cityscapes \
                --datalist data/list/cityscapes/train.txt \
                --scales 256-128,512-256,1024-512,2048-1024 \
                --crop_size 256 128 \
                --input_size 256 128 \
                --num_workers 8 \
                --model hrnet18+ocr \
                --pretrained checkpoints/cityscapes_hrnet.pth \
                --num_classes 19 \
                --batch_size 8 \
                --task_name cityscapes_refinement \
                --lr 0.001

# 或者简单地运行以下脚本
sh scripts/cityscapes/train_magnet.sh

Deepglobe

要在Deepglobe数据集上训练MagNet，请运行以下示例脚本：

python train.py --dataset deepglobe \
                --root data/deepglobe \
                --datalist data/list/deepglobe/train.txt \
                --scales 612-612,1224-1224,2448-2448 \
                --crop_size 612 612 \
                --input_size 508 508 \
                --num_workers 8 \
                --model fpn \
                --pretrained checkpoints/deepglobe_fpn.pth \
                --num_classes 7 \
                --batch_size 8 \
                --task_name deepglobe_refinement \
                --lr 0.001

# 或者简单地运行以下脚本
sh scripts/deepglobe/train_magnet.sh

评估

测试可用的参数：

test.py [-h] --dataset DATASET [--root ROOT] [--datalist DATALIST]
               --scales SCALES --crop_size N [N ...] --input_size N [N ...]
               [--num_workers NUM_WORKERS] --model MODEL --num_classes
               NUM_CLASSES --pretrained PRETRAINED
               [--pretrained_refinement PRETRAINED_REFINEMENT [PRETRAINED_REFINEMENT ...]]
               [--image IMAGE] --sub_batch_size SUB_BATCH_SIZE
               [--n_patches N_PATCHES] --n_points N_POINTS
               [--smooth_kernel SMOOTH_KERNEL] [--save_pred]
               [--save_dir SAVE_DIR]
可选参数:
  -h, --help            显示此帮助信息并退出
  --dataset DATASET     数据集名称: cityscapes, deepglobe (默认: None)
  --root ROOT           训练和测试图像的路径 (默认: )
  --datalist DATALIST   包含图像和标签路径的.txt文件路径 (默认: )
  --scales SCALES       尺度: w1-h1,w2-h2,... , 例如 512-512,1024-1024,2048-2048 (默认: None)
  --crop_size N [N ...]
                        裁剪大小, 例如 256 128 (默认: None)
  --input_size N [N ...]
                        输入大小, 例如 256 128 (默认: None)
  --num_workers NUM_WORKERS
                        数据加载器的工作线程数 (默认: 1)
  --model MODEL         模型名称. 可选: fpn, psp, hrnet18+ocr, hrnet48+ocr (默认: None)
  --num_classes NUM_CLASSES
                        类别数量 (默认: None)
  --pretrained PRETRAINED
                        预训练权重 (默认: None)
  --pretrained_refinement PRETRAINED_REFINEMENT [PRETRAINED_REFINEMENT ...]
                        细化模块的预训练权重 (默认: [''])
  --image IMAGE         测试图像路径 (仅用于演示) (默认: None)
  --sub_batch_size SUB_BATCH_SIZE
                        patch处理的批量大小 (默认: None)
  --n_patches N_PATCHES
                        每个阶段要细化的patch数量. 如果n_patches=-1, 将细化所有patch (默认: -1)
  --n_points N_POINTS   每个阶段要细化的点数. 如果n_points < 1.0, 将作为总点数的比例 (默认: None)
  --smooth_kernel SMOOTH_KERNEL
                        应用于错误分数的模糊操作的核大小 (默认: 16)
  --save_pred           是否保存预测结果, 每张图像将包含: 原图, 真实标签, 粗糙预测, 精细预测 (默认: False)
  --save_dir SAVE_DIR   保存目录 (默认: test_results)

否则,以下是使用我们预训练模型进行测试的示例脚本。

### Cityscapes

完整的MagNet细化:
```bash
python test.py --dataset cityscapes \
               --root data/cityscapes \
               --datalist data/list/cityscapes/val.txt \
               --scales 256-128,512-256,1024-512,2048-1024 \
               --crop_size 256 128 \
               --input_size 256 128 \
               --num_workers 8 \
               --model hrnet18+ocr \
               --pretrained checkpoints/cityscapes_hrnet.pth \
               --pretrained_refinement checkpoints/cityscapes_refinement_512.pth checkpoints/cityscapes_refinement_1024.pth checkpoints/cityscapes_refinement_2048.pth \
               --num_classes 19 \
               --sub_batch_size 1 \
               --n_points 32768 \
               --n_patches -1 \
               --smooth_kernel 5 \
               --save_pred \
               --save_dir test_results/cityscapes

# 或者简单地运行以下脚本
sh scripts/cityscapes/test_magnet.sh

MagNet-Fast细化:

python test.py --dataset cityscapes \
               --root data/cityscapes \
               --datalist data/list/cityscapes/val.txt \
               --scales 256-128,512-256,1024-512,2048-1024 \
               --crop_size 256 128 \
               --input_size 256 128 \
               --num_workers 8 \
               --model hrnet18+ocr \
               --pretrained checkpoints/cityscapes_hrnet.pth \
               --pretrained_refinement checkpoints/cityscapes_refinement_512.pth checkpoints/cityscapes_refinement_1024.pth checkpoints/cityscapes_refinement_2048.pth \
               --num_classes 19 \
               --sub_batch_size 1 \
               --n_points 0.9 \
               --n_patches 4 \
               --smooth_kernel 5 \
               --save_pred \
               --save_dir test_results/cityscapes_fast

# 或者简单地运行以下脚本
sh scripts/cityscapes/test_magnet_fast.sh

Deepglobe

完整的MagNet细化:

python test.py --dataset deepglobe \
               --root data/deepglobe \
               --datalist data/list/deepglobe/test.txt \
               --scales 612-612,1224-1224,2448-2448 \
               --crop_size 612 612 \
               --input_size 508 508 \
               --num_workers 8 \
               --model fpn \
               --pretrained checkpoints/deepglobe_fpn.pth \
               --pretrained_refinement checkpoints/deepglobe_refinement.pth \
               --num_classes 7 \
               --sub_batch_size 1 \
               --n_points 0.75 \
               --n_patches -1 \
               --smooth_kernel 11 \
               --save_pred \
               --save_dir test_results/deepglobe

# 或者简单地运行以下脚本
sh scripts/deepglobe/test_magnet.sh

MagNet-Fast细化:

python test.py --dataset deepglobe \
               --root data/deepglobe \
               --datalist data/list/deepglobe/test.txt \
               --scales 612-612,1224-1224,2448-2448 \
               --crop_size 612 612 \
               --input_size 508 508 \
               --num_workers 8 \
               --model fpn \
               --pretrained checkpoints/deepglobe_fpn.pth \
               --pretrained_refinement checkpoints/deepglobe_refinement.pth \
               --num_classes 7 \
               --sub_batch_size 1 \
               --n_points 0.9 \
               --n_patches 3 \
               --smooth_kernel 11 \
               --save_pred \
               --save_dir test_results/deepglobe_fast

# 或者简单地运行以下脚本
sh scripts/deepglobe/test_magnet_fast.sh

致谢

感谢High-resolution networks and Segmentation Transformer for Semantic Segmentation提供的主干网络训练脚本。

联系方式

如果您有任何问题,请发送电子邮件至minhchuong.itus@gmail.com或在此仓库创建一个问题。