论文 Rethinking Interactive Image Segmentation with Low Latency, High Quality, and Diverse Prompts（CVPR 2024）的 Pytorch 实现。
刘勤、赵在民、莫希特·班萨尔、马克·尼特哈默
北卡罗来纳大学教堂山分校
该代码使用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
使用下载的权重运行交互式 GUI。该assets包含演示图像。
./run_demo.sh
我们在三个数据集上训练和测试我们的方法：DAVIS、COCO+LVIS 和 HQSeg-44K。
| 数据集 | 描述 | 下载链接 |
|---|---|---|
| 戴维斯 | 345 张图像，每张图像一个对象（测试） | 戴维斯.zip (43 MB) |
| HQSeg-44K | 44320 张图像（火车）； 1537 张图片（验证值） | 官方网站 |
| 可可+LVIS* | 具有 150 万个实例的 99k 图像（训练） | 原始 LVIS 图像 + 组合注释 |
下载并解压后，不要忘记更改 config.yml 中数据集的路径。
(*) 要准备COCO+LVIS，您需要下载原始LVIS v1.0，然后将COCO和LVIS数据集组合得到的预处理注释下载并解压到LVIS v1.0的文件夹中。 （综合注释由 RITM 准备。）
我们提供了一个脚本（ 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-mIoU | 编号C90 | C95号 | F95号 | 5-mIoU | 编号C90 | C95号 | F95号 | ||
| 中+左 | vitb-sax1 (408 MB) | 85.41 | 7.47 | 11.94 | 第731章 | 90.13 | 5.46 | 13.31 | 177 |
| 中+左 | vitb-sax2 (435 MB) | 85.71 | 7.18 | 11.52 | 700 | 89.85 | 5.34 | 12.80 | 163 |
| C+L+总部 | vitb-sax2 (435 MB) | 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 GPU 进行训练。
# 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 }
}
