语言驱动的语义分割(LSeg)

这个仓库包含论文《语言驱动的语义分割》的官方 PyTorch 实现。

ICLR 2022

作者:

• Boyi Li
• Kilian Q. Weinberger
• Serge Belongie
• Vladlen Koltun
• Rene Ranftl

概述

我们提出了 LSeg,一种新颖的语言驱动语义图像分割模型。LSeg 使用文本编码器计算描述性输入标签(如"草"或"建筑")的嵌入,同时使用基于 transformer 的图像编码器计算输入图像的每像素密集嵌入。图像编码器通过对比目标进行训练,以将像素嵌入与相应语义类别的文本嵌入对齐。文本嵌入提供了一种灵活的标签表示,其中语义相似的标签映射到嵌入空间中的相似区域(例如,"猫"和"毛茸茸的")。这使得 LSeg 能够在测试时泛化到以前未见过的类别,而无需重新训练甚至不需要一个额外的训练样本。我们证明,与现有的零样本和少样本语义分割方法相比,我们的方法实现了极具竞争力的零样本性能,甚至在提供固定标签集时与传统分割算法的准确度相匹配。

请查看我们的视频演示(4k)以进一步展示 LSeg 的功能。

使用方法

安装

选项 1:

pip install -r requirements.txt

选项 2:

conda install ipython
pip install torch==1.7.1 torchvision==0.8.2 torchaudio==0.7.2
pip install git+https://github.com/zhanghang1989/PyTorch-Encoding/
pip install pytorch-lightning==1.3.5
pip install opencv-python
pip install imageio
pip install ftfy regex tqdm
pip install git+https://github.com/openai/CLIP.git
pip install altair
pip install streamlit
pip install --upgrade protobuf
pip install timm
pip install tensorboardX
pip install matplotlib
pip install test-tube
pip install wandb

数据准备

默认情况下,对于训练、测试和演示,我们使用 ADE20k。

python prepare_ade20k.py
unzip ../datasets/ADEChallengeData2016.zip

注意:对于演示,如果您想使用随机输入,可以忽略数据加载并注释此链接处的代码。

🌻 立即尝试演示

下载演示模型

👉 选项 1: 运行交互式应用

下载演示模型并将其放在 checkpoints 文件夹下,命名为 checkpoints/demo_e200.ckpt。

然后运行 streamlit run lseg_app.py

👉 选项 2: Jupyter Notebook

下载演示模型并将其放在 checkpoints 文件夹下,命名为 checkpoints/demo_e200.ckpt。

然后按照 lseg_demo.ipynb 来体验 LSeg。祝您使用愉快!

训练和测试示例

训练: 骨干网络 = ViT-L/16, 文本编码器来自 CLIP ViT-B/32

bash train.sh

测试: 骨干网络 = ViT-L/16, 文本编码器来自 CLIP ViT-B/32

bash test.sh

零样本实验

数据准备

请按照 HSNet 的说明,将所有数据集放在 data/Dataset_HSN 中

Pascal-5i

for fold in 0 1 2 3; do
python -u test_lseg_zs.py --backbone clip_resnet101 --module clipseg_DPT_test_v2 --dataset pascal \
--widehead --no-scaleinv --arch_option 0 --ignore_index 255 --fold ${fold} --nshot 0 \
--weights checkpoints/pascal_fold${fold}.ckpt 
done

COCO-20i

for fold in 0 1 2 3; do
python -u test_lseg_zs.py --backbone clip_resnet101 --module clipseg_DPT_test_v2 --dataset coco \
--widehead --no-scaleinv --arch_option 0 --ignore_index 255 --fold ${fold} --nshot 0 \
--weights checkpoints/pascal_fold${fold}.ckpt 
done

FSS

python -u test_lseg_zs.py --backbone clip_vitl16_384 --module clipseg_DPT_test_v2 --dataset fss \
--widehead --no-scaleinv --arch_option 0 --ignore_index 255 --fold 0 --nshot 0 \
--weights checkpoints/fss_l16.ckpt 

python -u test_lseg_zs.py --backbone clip_resnet101 --module clipseg_DPT_test_v2 --dataset fss \
--widehead --no-scaleinv --arch_option 0 --ignore_index 255 --fold 0 --nshot 0 \
--weights checkpoints/fss_rn101.ckpt 

模型库

如果您觉得这个仓库有用，请引用：

@inproceedings{
li2022languagedriven,
title={Language-driven Semantic Segmentation},
author={Boyi Li and Kilian Q Weinberger and Serge Belongie and Vladlen Koltun and Rene Ranftl},
booktitle={International Conference on Learning Representations},
year={2022},
url={https://openreview.net/forum?id=RriDjddCLN}
}

致谢

感谢以下项目的代码库：DPT、Pytorch_lightning、CLIP、Pytorch Encoding、Streamlit、Wandb