SegNext

إعادة النظر في تجزئة الصور التفاعلية مع زمن وصول منخفض وجودة عالية ومطالبات متنوعة

تنفيذ 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 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

قم بتشغيل واجهة المستخدم الرسومية التفاعلية مع الأوزان التي تم تنزيلها. تحتوي assets على صور للعرض التوضيحي.

 ./run_demo.sh

مجموعات البيانات

نقوم بتدريب واختبار طريقتنا على ثلاث مجموعات بيانات: DAVIS، وCOCO+LVIS، وHQSeg-44K.

مجموعة البيانات	وصف	رابط التحميل
ديفيس	345 صورة تحتوي كل منها على كائن واحد (اختبار)	DAVIS.zip (43 ميجابايت)
HQSeg-44K	44320 صورة (قطار)؛ 1537 صورة (فال)	الموقع الرسمي
كوكو+ليفيس*	99 ألف صورة مع 1.5 مليون نسخة (قطار)	صور 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 ميكرويوو	رقم C90	رقم C95	رقمF95	5 ميكرويوو	رقم C90	رقم C95	رقمF95
ج+ل	فيتب-ساكس1 (408 ميجابايت)	85.41	7.47	11.94	731	90.13	5.46	13.31	177
ج+ل	فيتب-ساكس2 (435 ميجابايت)	85.71	7.18	11.52	700	89.85	5.34	12.80	163
ج+ل+المقر الرئيسي	فيتب-ساكس2 (435 ميجابايت)	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 للتدريب.

 # 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 }
}