latent diffusion segmentation
1.0.0
Этот репозиторий содержит реализацию LDMSeg на Pytorch: простой подход скрытой диффузии для паноптической сегментации и рисования по маске. Предоставленный код включает в себя как обучение, так и оценку.
Простой подход скрытой диффузии для паноптической сегментации и рисования маски
Воутер Ван Гансбеке и Берт Де Брабандер
В этой статье представлен подход условной скрытой диффузии для решения задачи паноптической сегментации. Цель состоит в том, чтобы исключить необходимость в специализированных архитектурах (например, сетях предложений регионов или объектных запросах), сложных функциях потерь (например, венгерском сопоставлении или на основе ограничивающих рамок) и дополнительных методах постобработки (например, кластеризации, NMS). или вставка объекта). В результате мы полагаемся на Stable Diffusion, который является платформой, не зависящей от задач. Представленный подход состоит из двух этапов: (1) проецирование паноптических масок сегментации в скрытое пространство с помощью неглубокого автокодировщика; (2) обучить модель диффузии в скрытом пространстве, обусловленную изображениями RGB.
Ключевой вклад : Наш вклад тройной:
Код работает с последними версиями Pytorch, например 2.0. Кроме того, вы можете создать среду Python с помощью Anaconda:
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 для получения копии моей среды. Мы также полагаемся на некоторые зависимости отDetectron2 и 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. Модели можно обучить, выполнив следующие команды. По умолчанию мы будем обучаться на наборе данных COCO с базовым файлом конфигурации, определенным tools/configs/base/base.yaml . Обратите внимание, что этот файл будет загружен автоматически, поскольку мы полагаемся на пакет hydra .
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 . Например, я запускаю эту модель на 50 тыс. итераций на одном графическом процессоре емкостью 23 ГБ с общим размером пакета 16.
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 . Например, я прогнал эту модель на 200 тысяч итераций на 8 графических процессорах по 16 ГБ с общим размером пакета 256.
Мы планируем выпустить несколько обученных моделей. Метрика PQ (независимая от класса) предоставляется в наборе проверки COCO.
| Модель | #Параметры | Набор данных | Итеры | ПК | КВ. | RQ | Ссылка для скачивания |
|---|---|---|---|---|---|---|---|
| АЕ | ~2 млн. | КОКО | 66 тыс. | - | - | - | Скачать (23 МБ) |
| ЛДМ | ~800М | КОКО | 200 тыс. | 51,7 | 82,0 | 63,0 | Скачать (3,3 ГБ) |
Примечание. Менее мощный 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 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 }
}По всем вопросам обращайтесь к основному автору.
Это программное обеспечение выпущено по лицензии Creative Commons, которая разрешает использование только в личных и исследовательских целях. Для получения коммерческой лицензии свяжитесь с авторами. Вы можете просмотреть сводную информацию о лицензии здесь.
Я благодарен за все публичные репозитории (см. также ссылки в коде), и в частности за библиотеки детектора2 и диффузоров.
