STEGO：通过提炼特征对应关系实现无监督语义分割

项目页面 | 论文 | 视频 | ICLR 2022

Mark Hamilton, Zhoutong Zhang, Bharath Hariharan, Noah Snavely, William T. Freeman

这是论文“通过提炼特征对应关系实现无监督语义分割”的官方实现。

目录

• 安装
• 评估
• 训练
  • 使用自定义数据
• 理解 STEGO
  • 无监督语义分割
  • 深度特征连接跨图像的物体
  • STEGO 架构
  • 结果
• 引用
• 联系

安装

克隆本仓库：

git clone https://github.com/mhamilton723/STEGO.git
cd STEGO

安装 Conda 环境

如果您尚未安装 conda，请访问 Anaconda 安装页面

conda env create -f environment.yml
conda activate stego

下载预训练模型

cd src
python download_models.py

下载数据集

首先，将 pytorch_data_dir 变量更改为您的 PyTorch 数据目录，即存储数据集的地方。

python download_datasets.py

下载完成后，请导航到您的 PyTorch 数据目录并解压文件：

cd /YOUR/PYTORCH/DATA/DIR
unzip cocostuff.zip
unzip cityscapes.zip
unzip potsdam.zip
unzip potsdamraw.zip

评估

要评估我们的预训练模型，请在 STEGO/src 中运行以下命令：

python eval_segmentation.py

可以通过编辑 STEGO/src/configs/eval_config.yml 来更改评估参数和模型。

训练

要从头开始训练 STEGO，请先为感兴趣的数据集生成 KNN 索引。请注意，这需要先生成裁剪后的数据集，并且您可能需要修改 crop_datasets.py 来指定要裁剪的数据集：

python crop_datasets.py
python precompute_knns.py

然后您可以在 STEGO/src 中运行以下命令：

python train_segmentation.py

超参数可以在 STEGO/src/configs/train_config.yml 中调整。

要使用 TensorBoard 监控训练过程，请从 STEGO 目录运行以下命令：

tensorboard --logdir logs

使用自定义数据

要使用您自己的数据集训练 STEGO，请在您的 PyTorch 数据根目录中创建一个具有以下结构的目录。请注意，如果您没有标签，请省略 labels 目录：

dataset_name
|── imgs
|   ├── train
|   |   |── unique_img_name_1.jpg
|   |   └── unique_img_name_2.jpg
|   └── val
|       |── unique_img_name_3.jpg
|       └── unique_img_name_4.jpg
└── labels
    ├── train
    |   |── unique_img_name_1.png
    |   └── unique_img_name_2.png
    └── val
        |── unique_img_name_3.png
        └── unique_img_name_4.png

接下来，在 STEGO/src/configs/train_config.yml 中设置以下参数：

dataset_name: "directory"
dir_dataset_name: "dataset_name"
dir_dataset_n_classes: 5 # 这是需要识别的对象类型数量

如果您希望通过裁剪来提高空间分辨率，可以使用我们的 裁剪工具。

最后，取消注释自定义数据集代码，并从 STEGO\src 运行 python precompute_knns.py，以生成自定义数据集所需的 KNN 信息。

现在您可以使用以下命令在自定义数据集上进行训练：

python train_segmentation.py

理解 STEGO

无监督语义分割

现实世界的图像往往杂乱无章，包含多种物体，这使得分类显得随意。此外，现实中的物体并不总是能被边界框完全包围。语义分割方法旨在通过为图像中的每个像素分配独立的类别标签来避免这些挑战。然而，传统的语义分割方法由于依赖于密集标注的图像而难以训练，而密集标注所需的时间可能是边界框或类别标注的 100 倍。这使得在人类事先不了解数据结构的领域中，收集大规模且多样化的数据集变得极为困难。我们通过仅利用自监督学习，以像素级语义分割的方式学习对象的本体论，从而绕过这些挑战。

深度特征连接跨图像的物体

自监督对比学习使算法能够在无需人工标注的情况下学习图像的智能表示。STEGO 在此基础上进一步发现，像 Caron 等人的 DINO 这样的自监督视觉 Transformer 所提取的特征已经能够捕捉到物体之间的关系。通过计算图像特征之间的余弦相似度，我们可以看到，诸如草地、摩托车和天空等语义区域会因特征相似性而“连接”在一起。

STEGO 架构

STEGO 无监督分割系统通过对比损失，将图像之间的对应关系提炼为一组类别标签来学习。具体而言，我们的目标是学习一种能够尊重对象之间所诱导对应关系的分割结果。为此，我们在 DINO ViT 主干网络之上训练了一个浅层分割网络，并引入了三项对比损失项，分别用于提炼图像与其自身、相似图像以及随机其他图像之间的关联。如果两个区域在深度特征上存在强耦合关系，我们就鼓励它们共享同一类别。

结果

我们在 CocoStuff、Cityscapes 和 Potsdam 语义分割数据集上评估了 STEGO 算法。由于这些方法未使用任何标注信息，我们采用匈牙利匹配算法来寻找聚类与数据集类别之间的最佳映射。实验结果表明，STEGO 能够以远高于现有方法 PiCIE 的空间分辨率和灵敏度，对复杂且杂乱的场景进行分割。这不仅带来了显著的定性提升，还使平均交并比（mIoU）提高了两倍以上。有关 Cityscapes 和 Potsdam 数据集上的结果，请参阅我们的论文：arXiv 链接。

引用

@inproceedings{hamilton2022unsupervised,
	title={通过提炼特征对应关系实现无监督语义分割},
	author={Mark Hamilton、Zhoutong Zhang、Bharath Hariharan、Noah Snavely、William T. Freeman},
	booktitle={国际表示学习大会},
	year={2022},
	url={https://openreview.net/forum?id=SaKO6z6Hl0c}
}

联系方式

如需反馈、提问或媒体咨询，请联系 Mark Hamilton。

STEGO 快速上手指南

STEGO 是一个基于无监督学习的语义分割工具，通过蒸馏特征对应关系，无需人工标注即可将图像中的物体划分为不同的语义类别。

环境准备

• 操作系统: Linux (推荐 Ubuntu)
• 依赖管理: Conda (必须)
• 硬件要求: 支持 CUDA 的 NVIDIA GPU (用于训练和推理)
• 前置知识: 熟悉 Python 基础及 PyTorch 深度学习框架

注意: 请确保已安装 Anaconda 或 Miniconda。如未安装，请访问 Anaconda 官网 下载。国内用户可使用清华或中科大镜像源加速 conda 包下载。

安装步骤

1. 克隆代码仓库

git clone https://github.com/mhamilton723/STEGO.git
cd STEGO

2. 创建并激活 Conda 环境

使用项目提供的配置文件创建环境：

conda env create -f environment.yml
conda activate stego

(国内用户若下载慢，可在 .condarc 中配置清华源：conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/)

3. 下载预训练模型

进入 src 目录并运行下载脚本：

cd src
python download_models.py

4. 下载数据集 (可选)

如需复现论文结果或进行标准评估，需下载数据集。 首先修改 download_datasets.py 中的 pytorch_data_dir 变量，指向你本地的 PyTorch 数据存放目录。

python download_datasets.py

下载完成后，进入数据目录解压文件：

cd /YOUR/PYTORCH/DATA/DIR
unzip cocostuff.zip
unzip cityscapes.zip
unzip potsdam.zip
unzip potsdamraw.zip

基本使用

场景一：评估预训练模型

最直接的使用方式是运行官方提供的预训练模型进行评估。

在 STEGO/src 目录下执行：

python eval_segmentation.py

• 自定义配置: 如需更改评估参数或切换模型，请编辑 STEGO/src/configs/eval_config.yml 文件。

场景二：在自己的数据上训练

如果你希望用自己的图片集训练 STEGO（无标签数据），请按以下步骤操作：

1. 准备数据目录结构 在你的 PyTorch 数据根目录下创建文件夹，结构如下（若无标签可省略 labels 文件夹）：

dataset_name
|── imgs
|   ├── train
|   |   |── unique_img_name_1.jpg
|   |   └── unique_img_name_2.jpg
|   └── val
|       |── unique_img_name_3.jpg
|       └── unique_img_name_4.jpg
└── labels (可选)
    ├── train
    |   |── unique_img_name_1.png
    |   └── unique_img_name_2.png
    └── val
        |── unique_img_name_3.png
        └── unique_img_name_4.png

2. 修改配置文件 编辑 STEGO/src/configs/train_config.yml，设置以下参数：

dataset_name: "directory"
dir_dataset_name: "dataset_name"  # 替换为你实际的文件夹名称
dir_dataset_n_classes: 5          # 设定你希望发现的物体类别数量

3. 预处理数据 生成裁剪后的数据集及 KNN 索引（这是训练的前置步骤）：

# 如需提高空间分辨率，可先运行裁剪工具
python crop_datasets.py 

# 生成 KNN 信息
python precompute_knns.py

注意：运行 precompute_knns.py 前，请确保在配置文件中取消注释自定义数据集相关的代码部分。

4. 开始训练

python train_segmentation.py

5. 监控训练过程 在新终端窗口中运行 TensorBoard 查看日志：

tensorboard --logdir logs