ODISE：基于文本到图像扩散模型的开放词汇全景分割

ODISE：基于 Open-vocabulary DIffusion 的全景分割，利用预训练的文本-图像扩散模型和判别模型来实现开放词汇的全景分割。 它借助这两个模型的冻结表征，在真实场景中对任意类别进行全景分割。

本仓库是论文中介绍的 ODISE 的官方实现：

基于文本到图像扩散模型的开放词汇全景分割
Jiarui Xu、
Sifei Liu*、
Arash Vahdat*、
Wonmin Byeon、
Xiaolong Wang、
Shalini De Mello
CVPR 2023 精选论文。（共同第一作者）

如需商业合作，请访问我们的官网并提交表格：NVIDIA Research Licensing。

视觉结果

链接

• Jiarui Xu 的项目页面（包含更多视觉结果）
• HuggingFace 🤗 演示
• arXiv 页面

引用

如果您在研究中使用了我们的工作，请引用以下内容：

@article{xu2023odise,
  title={{Open-Vocabulary Panoptic Segmentation with Text-to-Image Diffusion Models}},
  author={Xu, Jiarui and Liu, Sifei and Vahdat, Arash and Byeon, Wonmin and Wang, Xiaolong and De Mello, Shalini},
  journal={arXiv preprint arXiv:2303.04803},
  year={2023}
}

环境配置

通过以下命令安装依赖：

conda create -n odise python=3.9
conda activate odise
conda install pytorch=1.13.1 torchvision=0.14.1 pytorch-cuda=11.6 -c pytorch -c nvidia
conda install -c "nvidia/label/cuda-11.6.1" libcusolver-dev
git clone git@github.com:NVlabs/ODISE.git 
cd ODISE
pip install -e .

（可选）安装 xformers 以高效实现 Transformer： 可以安装预编译版本：

pip install xformers==0.0.16

或者从最新源码构建：

# （可选）使构建速度更快
pip install ninja
# 如果在不同类型的 GPU 上运行和构建，请设置 TORCH_CUDA_ARCH_LIST
pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers
# （这可能需要数十分钟）

模型库

我们提供了两个经过预训练的 ODISE 模型，分别在 COCO 的整个训练集上使用标签或描述性文字进行监督训练。 ODISE 的预训练模型受 知识共享署名-非商业性使用-相同方式共享 4.0 国际许可协议 的约束。 每个模型包含 2810 万个可训练参数。这些模型的下载链接见下表。 首次运行 demo/demo.py 或推理脚本时，系统会自动将 ODISE 的预训练模型下载到您本地的 $HOME/.torch/iopath_cache/NVlabs/ODISE/releases/download/v1.0.0/ 目录中。

	ADE20K(A-150)			COCO			ADE20K-Full (A-847)	Pascal Context 59 (PC-59)	Pascal Context 459 (PC-459)	Pascal VOC 21 (PAS-21)	下载
	PQ	mAP	mIoU	PQ	mAP	mIoU	mIoU	mIoU	mIoU	mIoU
ODISE（标签）	22.6	14.4	29.9	55.4	46.0	65.2	11.1	57.3	14.5	84.6	checkpoint
ODISE（描述）	23.4	13.9	28.7	45.6	38.4	52.4	11.0	55.3	13.8	82.7	checkpoint

开始使用

请参阅 为 ODISE 准备数据集。

有关使用 ODISE 进行训练和推理的详细说明，请参阅 ODISE 使用指南。

演示

• 集成到 Huggingface Spaces 🤗 中，使用 Gradio 实现。体验网页版演示：

• 在 Google Colab 上运行演示：

重要提示：首次运行 demo/demo.py 脚本时，除了 ODISE 的预训练模型外，还会自动从其原始来源下载 Stable Diffusion v1.3 和 CLIP 的预训练模型，并分别保存到本地目录 $HOME/.torch/ 和 $HOME/.cache/clip 中。Stable Diffusion 和 CLIP 的预训练模型受各自原始许可协议的约束，分别来自 Stable Diffusion 和 CLIP。

• 通过命令行运行 ODISE 的演示：

  python demo/demo.py --input demo/examples/coco.jpg --output demo/coco_pred.jpg --vocab "black pickup truck, pickup truck; blue sky, sky"
  

  输出将保存为 demo/coco_pred.jpg。有关 demo/demo.py 的更多详细选项，请参阅 ODISE 使用指南。

• 在本地运行 Gradio 演示：

  python demo/app.py
  

致谢

代码主要基于 Detectron2、Stable Diffusion、Mask2Former、OpenCLIP 和 GLIDE。

感谢所有优秀的开源项目！

ODISE 快速上手指南

ODISE (Open-Vocabulary Panoptic Segmentation) 是一个利用预训练的文本 - 图像扩散模型和判别模型，实现开放词汇全景分割的 AI 工具。它无需针对特定类别重新训练，即可对任意类别的物体进行分割。

环境准备

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

• 操作系统: Linux (推荐 Ubuntu)
• GPU: 支持 CUDA 的 NVIDIA 显卡 (建议显存 8GB 以上)
• CUDA 版本: 11.6
• Python 版本: 3.9
• PyTorch 版本: 1.13.1

安装步骤

请依次执行以下命令来配置环境并安装 ODISE。

1. 创建并激活 Conda 环境

conda create -n odise python=3.9
conda activate odise

2. 安装 PyTorch 及 CUDA 依赖

使用官方源安装指定版本的 PyTorch 和 CUDA 工具包：

conda install pytorch=1.13.1 torchvision=0.14.1 pytorch-cuda=11.6 -c pytorch -c nvidia
conda install -c "nvidia/label/cuda-11.6.1" libcusolver-dev

3. 克隆代码库并安装 ODISE

git clone git@github.com:NVlabs/ODISE.git 
cd ODISE
pip install -e .

4. (可选) 安装 xformers 加速

为了获得更高效的 Transformer 推理速度，建议安装 xformers。您可以选择安装预编译版本（推荐）或从源码构建：

方式 A：安装预编译版本

pip install xformers==0.0.16

方式 B：从源码构建（耗时较长）

pip install ninja
pip install -v -U git+https://github.com/facebookresearch/xformers.git@main#egg=xformers

基本使用

首次运行脚本时，程序会自动下载 ODISE 预训练模型以及 Stable Diffusion v1.3 和 CLIP 的基础模型到本地缓存目录（$HOME/.torch/ 和 $HOME/.cache/clip），请确保网络连接畅通。

命令行演示

使用以下命令对单张图片进行开放词汇分割。您只需通过 --vocab 参数指定想要分割的物体类别描述即可。

python demo/demo.py --input demo/examples/coco.jpg --output demo/coco_pred.jpg --vocab "black pickup truck, pickup truck; blue sky, sky"

• --input: 输入图片路径。
• --output: 输出结果图片保存路径。
• --vocab: 需要分割的类别列表，多个类别可用分号分隔，同义词可用逗号分隔。

运行完成后，分割结果将保存在 demo/coco_pred.jpg。

本地 Web 界面演示

如果您希望使用图形化界面进行操作，可以启动本地的 Gradio Demo：

python demo/app.py

启动后，终端会显示一个本地访问地址（通常是 http://127.0.0.1:7860），在浏览器中打开该地址即可上传图片并输入文本提示词进行交互式分割。