Реализация Pytorch для бумаги. Переосмысление интерактивной сегментации изображений с низкой задержкой, высоким качеством и разнообразными подсказками, CVPR 2024.
Цинь Лю, Джэмин Чо, Мохит Бансал, Марк Нитхаммер
UNC-Чапел-Хилл
Код протестирован с использованием python=3.10 , torch=2.2.0 , torchvision=0.17.0 .
git clone https://github.com/uncbiag/SegNext
cd SegNext
Теперь создайте новую среду conda и установите необходимые пакеты соответственно.
conda create -n segnext python=3.10
conda activate segnext
conda install pytorch==2.2.0 torchvision==0.17.0 torchaudio==2.2.0 pytorch-cuda=11.8 -c pytorch -c nvidia
pip install -r requirements.txt
Сначала загрузите три веса модели: vitb_sax1 (408M), vitb_sax2 (435M) и vitb_sax2_ft (435M). Эти веса будут автоматически сохранены в папке weights .
python download.py
Запустите интерактивный графический интерфейс с загруженными весами. assets содержат изображения для демонстрации.
./run_demo.sh
Мы обучаем и тестируем наш метод на трех наборах данных: DAVIS, COCO+LVIS и HQSeg-44K.
| Набор данных | Описание | Ссылка для скачивания |
|---|---|---|
| ДЭВИС | 345 изображений по одному объекту (тест) | ДЭВИС.zip (43 МБ) |
| HQSeg-44K | 44320 изображений (поезд); 1537 изображений (вал) | официальный сайт |
| КОКО+ЛВИС* | 99 тыс. изображений с 1,5 млн экземпляров (поезд) | оригинальные изображения LVIS + комбинированные аннотации |
Не забудьте после скачивания и распаковки изменить пути к датасетам в config.yml.
(*) Чтобы подготовить COCO+LVIS, вам необходимо загрузить исходный LVIS v1.0, затем загрузить и распаковать предварительно обработанные аннотации, полученные путем объединения набора данных COCO и LVIS, в папку с LVIS v1.0. (Комбинированные аннотации подготовлены РИТМ.)
Мы предоставляем скрипт ( run_eval.sh ) для оценки представленных нами моделей. Следующая команда запускает оценку NoC для всех наборов тестовых данных.
python ./segnext/scripts/evaluate_model.py --gpus=0 --checkpoint=./weights/vitb_sa2_cocolvis_hq44k_epoch_0.pth --datasets=DAVIS,HQSeg44K
| Тренироваться Набор данных | Модель | HQSeg-44K | ДЭВИС | ||||||
|---|---|---|---|---|---|---|---|---|---|
| 5-млн | №C90 | №C95 | №F95 | 5-млн | №C90 | №C95 | №F95 | ||
| С+Л | витб-сакс1 (408 МБ) | 85,41 | 7.47 | 11.94 | 731 | 90.13 | 5.46 | 13.31 | 177 |
| С+Л | витб-сакс2 (435 МБ) | 85,71 | 7.18 | 11.52 | 700 | 89,85 | 5.34 | 12.80 | 163 |
| С+Л+ШТАБ | витб-сакс2 (435 МБ) | 91,75 | 5.32 | 9.42 | 583 | 91,87 | 4.43 | 10.73 | 123 |
Для оценки задержки SAT см. eval_sat_latency.ipynb.
Мы предоставляем сценарий ( run_train.sh ) для обучения наших моделей на наборе данных HQSeg-44K. Вы можете начать тренировку с помощью следующих команд. По умолчанию для обучения мы используем 4 графических процессора A6000.
# train vitb-sax1 model on coco+lvis
MODEL_CONFIG=./segnext/models/default/plainvit_base1024_cocolvis_sax1.py
torchrun --nproc-per-node=4 --master-port 29504 ./segnext/train.py ${MODEL_CONFIG} --batch-size=16 --gpus=0,1,2,3
# train vitb-sax2 model on coco+lvis
MODEL_CONFIG=./segnext/models/default/plainvit_base1024_cocolvis_sax2.py
torchrun --nproc-per-node=4 --master-port 29505 ./segnext/train.py ${MODEL_CONFIG} --batch-size=16 --gpus=0,1,2,3
# finetune vitb-sax2 model on hqseg-44k
MODEL_CONFIG=./segnext/models/default/plainvit_base1024_hqseg44k_sax2.py
torchrun --nproc-per-node=4 --master-port 29506 ./segnext/train.py ${MODEL_CONFIG} --batch-size=12 --gpus=0,1,2,3 --weights ./weights/vitb_sa2_cocolvis_epoch_90.pth
@article { liu2024rethinking ,
title = { Rethinking Interactive Image Segmentation with Low Latency, High Quality, and Diverse Prompts } ,
author = { Liu, Qin and Cho, Jaemin and Bansal, Mohit and Niethammer, Marc } ,
journal = { arXiv preprint arXiv:2404.00741 } ,
year = { 2024 }
}
