낮은 대기 시간, 고품질 및 다양한 프롬프트로 대화형 이미지 분할에 대한 재고, CVPR 2024 논문을 위한 Pytorch 구현.
친 리우(Qin Liu), 조재민, 모히트 반살(Mohit Bansal), 마크 니담머(Marc Niethammer)
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
다운로드한 가중치를 사용하여 대화형 GUI를 실행합니다. assets 데모용 이미지가 포함되어 있습니다.
./run_demo.sh
우리는 DAVIS, COCO+LVIS 및 HQSeg-44K의 세 가지 데이터 세트에 대해 방법을 훈련하고 테스트합니다.
| 데이터세트 | 설명 | 다운로드 링크 |
|---|---|---|
| 데이비스 | 각각 하나의 개체가 포함된 345개 이미지(테스트) | DAVIS.zip(43MB) |
| HQSeg-44K | 44320개의 이미지(기차); 이미지 1,537개(발) | 공식 사이트 |
| 코코+LVIS* | 150만 개의 인스턴스가 포함된 99,000개의 이미지(열차) | 원본 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 | 데이비스 | ||||||
|---|---|---|---|---|---|---|---|---|---|
| 5mIoU | NoC90 | NoC95 | NoF95 | 5mIoU | NoC90 | NoC95 | NoF95 | ||
| C+L | vitb-sax1 (408MB) | 85.41 | 7.47 | 11.94 | 731 | 90.13 | 5.46 | 13.31 | 177 |
| C+L | vitb-sax2 (435MB) | 85.71 | 7.18 | 11.52 | 700 | 89.85 | 5.34 | 12.80 | 163 |
| C+L+HQ | vitb-sax2 (435MB) | 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 }
}
