論文 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 }
}
