Pytorch-Implementierung für Papier: Neudenken der interaktiven Bildsegmentierung mit geringer Latenz, hoher Qualität und vielfältigen Eingabeaufforderungen, CVPR 2024.
Qin Liu, Jaemin Cho, Mohit Bansal, Marc Niethammer
UNC-Chapel Hill
Der Code wird mit python=3.10 , torch=2.2.0 , torchvision=0.17.0 getestet.
git clone https://github.com/uncbiag/SegNext
cd SegNext
Erstellen Sie nun eine neue Conda-Umgebung und installieren Sie die erforderlichen Pakete entsprechend.
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
Laden Sie zunächst drei Modellgewichte herunter: vitb_sax1 (408M), vitb_sax2 (435M) und vitb_sax2_ft (435M). Diese Gewichte werden automatisch im weights gespeichert.
python download.py
Führen Sie eine interaktive GUI mit den heruntergeladenen Gewichten aus. Die assets enthalten Bilder zur Demo.
./run_demo.sh
Wir trainieren und testen unsere Methode anhand von drei Datensätzen: DAVIS, COCO+LVIS und HQSeg-44K.
| Datensatz | Beschreibung | Download-Link |
|---|---|---|
| DAVIS | 345 Bilder mit je einem Objekt (Test) | DAVIS.zip (43 MB) |
| HQSeg-44K | 44320 Bilder (Zug); 1537 Bilder (Wert) | offizielle Seite |
| COCO+LVIS* | 99.000 Bilder mit 1,5 Millionen Instanzen (Zug) | Original-LVIS-Bilder + kombinierte Anmerkungen |
Vergessen Sie nicht, nach dem Herunterladen und Entpacken die Pfade zu den Datensätzen in config.yml zu ändern.
(*) Um COCO+LVIS vorzubereiten, müssen Sie das Original-LVIS v1.0 herunterladen und dann vorverarbeitete Anmerkungen herunterladen und entpacken, die durch die Kombination von COCO- und LVIS-Datensatz im Ordner mit LVIS v1.0 erhalten werden. (Die kombinierten Anmerkungen werden von RITM erstellt.)
Wir stellen ein Skript ( run_eval.sh ) zur Verfügung, um unsere vorgestellten Modelle auszuwerten. Der folgende Befehl führt die NoC-Auswertung für alle Testdatensätze aus.
python ./segnext/scripts/evaluate_model.py --gpus=0 --checkpoint=./weights/vitb_sa2_cocolvis_hq44k_epoch_0.pth --datasets=DAVIS,HQSeg44K
| Zug Datensatz | Modell | HQSeg-44K | DAVIS | ||||||
|---|---|---|---|---|---|---|---|---|---|
| 5 Mio. Einheiten | Nr. C90 | Nr. C95 | Nr.F95 | 5 Mio. Einheiten | Nr. C90 | Nr. C95 | Nr.F95 | ||
| 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+HQ | vitb-sax2 (435 MB) | 91,75 | 5.32 | 9.42 | 583 | 91,87 | 4.43 | 10.73 | 123 |
Informationen zur SAT-Latenzbewertung finden Sie unter eval_sat_latency.ipynb.
Wir stellen ein Skript ( run_train.sh ) zum Trainieren unserer Modelle auf dem HQSeg-44K-Datensatz bereit. Mit den folgenden Befehlen können Sie mit dem Training beginnen. Standardmäßig verwenden wir für das Training 4 A6000-GPUs.
# 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 }
}
