用于全景分割和掩模修复的简单潜在扩散方法

该存储库包含 LDMSeg 的 Pytorch 实现：一种用于全景分割和掩模修复的简单潜在扩散方法。提供的代码包括培训和评估。

用于全景分割和掩模修复的简单潜在扩散方法
沃特·范·甘斯贝克和伯特·德·布拉班代尔

内容

介绍
安装
- 自动安装
- 手动安装
训练
- 第 1 步：训练自动编码器
- 第 2 步：训练 LDM
预训练模型
引文
执照
致谢

？介绍

本文提出了一种条件潜在扩散方法来解决全景分割的任务。目的是省略对专门架构（例如，区域提议网络或对象查询）、复杂损失函数（例如，匈牙利匹配或基于边界框）以及其他后处理方法（例如，聚类、NMS）的需求。，或对象粘贴）。因此，我们依赖稳定扩散，这是一个与任务无关的框架。所提出的方法包括两个步骤：（1）使用浅层自动编码器将全景分割掩模投影到潜在空间； (2) 以 RGB 图像为条件，在潜在空间中训练扩散模型。

主要贡献：我们的贡献有三重：

生成框架：我们提出了一种基于潜在扩散模型（LDM）的完全生成方法，用于全景分割。我们的方法建立在稳定扩散的基础上，力求简单并简化计算。我们首先研究与类无关的设置，以将全景分割从预定义的类中解放出来。
通用设计：我们的方法规避了大多数流行方法中存在的专用架构、复杂的损失函数和对象检测模块。这里，去噪目标忽略了对象查询、区域提议和匈牙利匹配的必要性。这种简单而通用的方法为未来扩展到各种密集预测任务（例如深度预测、显着性估计等）铺平了道路。
掩模修复：我们成功地将我们的方法应用于以场景为中心的数据集，并展示了其针对不同稀疏级别的掩模修复功能。该方法显示了全局掩模修复的有希望的结果。

？安装

该代码可在最新的 Pytorch 版本上运行，例如 2.0。此外，您可以使用 Anaconda 创建 python 环境：

 conda create -n LDMSeg python=3.11
conda activate LDMSeg

自动安装

我们建议遵循自动安装（请参阅tools/scripts/install_env.sh ）。运行以下命令以可编辑模式安装项目。请注意，所有依赖项都会自动安装。由于这可能并不总是有效（例如，由于 CUDA 或 gcc 问题），请查看手动安装步骤。

python -m pip install -e .
pip install git+https://github.com/facebookresearch/detectron2.git
pip install git+https://github.com/cocodataset/panopticapi.git

手动安装

最重要的软件包可以使用 pip 快速安装，如下所示：

pip install torch torchvision einops                            # Main framework
pip install diffusers transformers xformers accelerate timm     # For using pretrained models
pip install scipy opencv-python                                 # For augmentations or loss
pip install pyyaml easydict hydra-core                          # For using config files
pip install termcolor wandb                                     # For printing and logging

请参阅data/environment.yml以获取我的环境的副本。我们还依赖于 detectorron2 和 panopticapi 的一些依赖项。请遵循他们的文档。

?️ 数据集

我们目前支持 COCO 数据集。请按照文档安装图像及其相应的全景分割蒙版。另外，请查看ldmseg/data/目录，了解 COCO 数据集上的一些示例。顺便说一句，所采用的结构应该是相当标准的：

 .
└── coco
    ├── annotations
    ├── panoptic_semseg_train2017
    ├── panoptic_semseg_val2017
    ├── panoptic_train2017 -> annotations/panoptic_train2017
    ├── panoptic_val2017 -> annotations/panoptic_val2017
    ├── test2017
    ├── train2017
    └── val2017

最后但并非最不重要的一点是，将configs/env/root_paths.yml中的路径分别更改为数据集根目录和所需的输出目录。

⏳ 培训

所提出的方法是双管齐下的：首先，我们训练一个自动编码器来表示较低维空间（例如，64x64）中的分割图。接下来，我们从预训练的潜在扩散模型（LDM），特别是稳定扩散模型开始，训练一个可以从 RGB 图像生成全景掩模的模型。可以通过运行以下命令来训练模型。默认情况下，我们将使用tools/configs/base/base.yaml中定义的基本配置文件在 COCO 数据集上进行训练。请注意，由于我们依赖hydra包，因此该文件将自动加载。

步骤 1：在全景分割图上训练自动编码器

 python - W ignore tools / main_ae . py 
    datasets = coco 
    base . train_kwargs . fp16 = True 
    base . optimizer_name = adamw 
    base . optimizer_kwargs . lr = 1e-4 
    base . optimizer_kwargs . weight_decay = 0.05

有关传递参数的更多详细信息可以在tools/scripts/train_ae.sh中找到。例如，我在 23 GB 的单个 GPU 上运行此模型 50k 迭代，总批量大小为 16。

步骤 2：训练用于基于 RGB 图像的全景分割的 LDM

 python - W ignore tools / main_ldm . py 
    datasets = coco 
    base . train_kwargs . gradient_checkpointing = True 
    base . train_kwargs . fp16 = True 
    base . train_kwargs . weight_dtype = float16 
    base . optimizer_zero_redundancy = True 
    base . optimizer_name = adamw 
    base . optimizer_kwargs . lr = 1e-4 
    base . optimizer_kwargs . weight_decay = 0.05 
    base . scheduler_kwargs . weight = 'max_clamp_snr' 
    base . vae_model_kwargs . pretrained_path = '$AE_MODEL'

$AE_MODEL表示上一步获得的模型的路径。有关传递参数的更多详细信息可以在tools/scripts/train_diffusion.sh中找到。例如，我在 8 个 16 GB 的 GPU 上运行此模型 200k 迭代，总批量大小为 256。

预训练模型

我们计划发布几个经过训练的模型。（与类别无关的）PQ 指标在 COCO 验证集上提供。

模型	#参数	数据集	迭代器	质子Q	平方数	需求量	下载链接
AE	～2M	可可	66k	-	-	-	下载 (23 MB)
LDM	～800M	可可	20万	51.7	82.0	63.0	下载 (3.3 GB)

注意：功能较弱的 AE（即较少的下采样或上采样层）通常有利于修复，因为我们不执行额外的微调。

评估结果应该是这样的：

 python - W ignore tools / main_ldm . py 
    datasets = coco 
    base . sampling_kwargs . num_inference_steps = 50 
    base . eval_only = True 
    base . load_path = $ PRETRAINED_MODEL_PATH

如果需要，您可以添加参数。更高的阈值（例如--base.eval_kwargs.count_th 700或--base.eval_kwargs.mask_th 0.9可以进一步增加数字。然而，我们使用标准值，阈值设置为 0.5，并删除面积小于 512 的片段进行评估。

要评估上面的预训练模型，请运行tools/scripts/eval.sh 。

在这里，我们将结果可视化：

？引文

如果您发现此存储库对您的研究有用，请考虑引用以下论文：

 @article { vangansbeke2024ldmseg ,
  title = { a simple latent diffusion approach for panoptic segmentation and mask inpainting } ,
  author = { Van Gansbeke, Wouter and De Brabandere, Bert } ,
  journal = { arxiv preprint arxiv:2401.10227 } ,
  year = { 2024 }
}