การใช้ 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
ขั้นแรก ดาวน์โหลดตุ้มน้ำหนักโมเดลสามรายการ: 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 ภาพ (วาล) | เว็บไซต์อย่างเป็นทางการ |
| โคโค่+ลวิส* | 99,000 ภาพพร้อมอินสแตนซ์ 1.5M (ฝึก) | รูปภาพ 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 | NoC90 | NoC95 | NoF95 | 5-mIoU | NoC90 | NoC95 | NoF95 | ||
| ซี+แอล | 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 |
| ซี+แอล+สำนักงานใหญ่ | 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 คุณสามารถเริ่มการฝึกด้วยคำสั่งต่อไปนี้ ตามค่าเริ่มต้น เราใช้ A6000 GPU 4 ตัวในการฝึก
# 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 }
}
