紙向けの 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
まず、3 つのモデル ウェイト、vitb_sax1 (408M)、vitb_sax2 (435M)、および vitb_sax2_ft (435M) をダウンロードします。これらの重みは自動的にweightsフォルダーに保存されます。
python download.py
ダウンロードした重みを使用して対話型 GUI を実行します。 assetsデモ用の画像が含まれています。
./run_demo.sh
DAVIS、COCO+LVIS、HQSeg-44K の 3 つのデータセットでメソッドをトレーニングし、テストします。
| データセット | 説明 | ダウンロードリンク |
|---|---|---|
| デイビス | それぞれ 1 つのオブジェクトを含む 345 枚の画像 (テスト) | デイビス.zip (43 MB) |
| HQセグ-44K | 画像44320枚（電車）。画像 1537 枚 (ヴァル) | 公式サイト |
| ココ+リヴィス* | 99k の画像と 150 万個のインスタンス (トレイン) | オリジナル 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
| 電車 データセット | モデル | HQセグ-44K | デイビス | ||||||
|---|---|---|---|---|---|---|---|---|---|
| 5ミリオU | NoC90 | NoC95 | NoF95 | 5ミリオU | NoC90 | NoC95 | NoF95 | ||
| C+L | vitb-sax1 (408 MB) | 85.41 | 7.47 | 11.94 | 731 | 90.13 | 5.46 | 13.31 | 177 |
| C+L | 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 を参照してください。
HQSeg-44K データセットでモデルをトレーニングするためのスクリプト ( run_train.sh ) を提供します。次のコマンドでトレーニングを開始できます。デフォルトでは、トレーニングに 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 }
}
