语言驱动的语义分割(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

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

🌻 立即尝试演示

下载演示模型

<table> <thead> <tr style="text-align: right;"> <th>名称</th> <th>骨干网络</th> <th>文本编码器</th> <th>链接</th> </tr> </thead> <tbody> <tr> <td>演示模型</td> <th>ViT-L/16</th> <th>CLIP ViT-B/32</th> <td><a href="https://drive.google.com/file/d/1FTuHY1xPUkM-5gaDtMfgCl3D0gR89WV7/view?usp=sharing">下载</a></td> </tr> </tbody> </table>

👉 选项 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

模型库

<table> <thead> <tr style="text-align: right;"> <th>数据集</th> <th>折叠</th> <th>骨干网络</th> <th>文本编码器</th> <th>性能</th> <th>链接</th> </tr> </thead> <tbody> <tr> <th>pascal</th> <td>0</td> <th>ResNet101</th> <th>CLIP ViT-B/32</th> <th>52.8</th> <td><a href="https://drive.google.com/file/d/1y4z4_yNGlZtn6osaeN4ZjMs6c0vr0F3m/view?usp=sharing">下载</a></td> </tr> <tr> <th>pascal</th> <td>1</td> <th>ResNet101</th> <th>CLIP ViT-B/32</th> <th>53.8</th> <td><a href="https://drive.google.com/file/d/1UZzN8kWkH-G8v6P8xcBXEHRZlQKPRxrX/view?usp=sharing">下载</a></td> </tr> <tr> <th>pascal</th> <td>2</td> <th>ResNet101</th> <th>CLIP ViT-B/32</th> <th>44.4</th> <td><a href="https://drive.google.com/file/d/1KCq1JphSMvj8X78bkWbdNIFm5zYzLTMX/view?usp=sharing">下载</a></td> </tr> <tr> <th>pascal</th> <td>3</td> <th>ResNet101</th> <th>CLIP ViT-B/32</th> <th>38.5</th> <td><a href="https://drive.google.com/file/d/1A_fllOJqyBg0ZTJcm85Cn0NAcwQbXnhl/view?usp=sharing">下载</a></td> </tr> <tr> <th>coco</th> <td>0</td> <th>ResNet101</th> <th>CLIP ViT-B/32</th> <th>22.1</th> <td><a href="https://drive.google.com/file/d/1nSYO3XtAv4mzWi4x-MFUfk04cpBKJT38/view?usp=sharing">下载</a></td> </tr> <tr> <th>coco</th> <td>1</td> <th>ResNet101</th> <th>CLIP ViT-B/32</th> <th>25.1</th> <td><a href="https://drive.google.com/file/d/1w0vz3yjEi_ZLgECRrgtoLxEHugyJkrs5/view?usp=sharing">下载</a></td> </tr> <tr> <th>coco</th> <td>2</td> <th>ResNet101</th> <th>CLIP ViT-B/32</th> <th>24.9</th> <td><a href="https://drive.google.com/file/d/1wmHtmJLdta18XuWQv6oX9llidCll_HrD/view?usp=sharing">下载</a></td> </tr> <tr> <th>coco</th> <td>3</td> <th>ResNet101</th> <th>CLIP ViT-B/32</th> <th>21.5</th> <td><a href="https://drive.google.com/file/d/1dliBUSOog7taJxMmb9cdefKH4XOKChVJ/view?usp=sharing">下载</a></td> </tr> <tr> <th>fss</th> <td>-</td> <th>ResNet101</th> <th>CLIP ViT-B/32</th> <th>84.7</th> <td><a href="https://drive.google.com/file/d/1UIj49Wp1mAopPub5M6O4WW-Z79VB1bhw/view?usp=sharing">下载</a></td> </tr> <tr> <th>fss</th> <td>-</td> <th>ViT-L/16</th> <th>CLIP ViT-B/32</th> <th>87.8</th> <td><a href="https://drive.google.com/file/d/1Nplkc_JsHIS55d--K2vonOOC3HrppzYy/view?usp=sharing">下载</a></td> </tr> </tbody> </table>

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

@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}
}