latent diffusion segmentation
1.0.0
Repo ini berisi implementasi LDMSeg Pytorch: pendekatan difusi laten sederhana untuk segmentasi panoptik dan inpainting Mask. Kode yang diberikan mencakup pelatihan dan evaluasi.
Pendekatan Difusi Laten Sederhana untuk Segmentasi Panoptik dan Pengecatan Masker
Wouter Van Gansbeke dan Bert De Brabandere
Makalah ini menyajikan pendekatan difusi laten bersyarat untuk mengatasi tugas segmentasi panoptik. Tujuannya adalah untuk menghilangkan kebutuhan akan arsitektur khusus (misalnya, jaringan proposal wilayah atau kueri objek), fungsi kerugian yang kompleks (misalnya, pencocokan Hongaria atau berdasarkan kotak pembatas), dan metode pasca-pemrosesan tambahan (misalnya, pengelompokan, NMS , atau menempelkan objek). Oleh karena itu, kami mengandalkan Difusi Stabil, yang merupakan kerangka kerja tanpa tugas. Pendekatan yang disajikan terdiri dari dua langkah: (1) memproyeksikan topeng segmentasi panoptik ke ruang laten dengan auto-encoder yang dangkal; (2) melatih model difusi dalam ruang laten, yang dikondisikan pada gambar RGB.
Kontribusi Utama : Kontribusi kami ada tiga:
Kode berjalan dengan versi Pytorch terbaru, misalnya 2.0. Selanjutnya, Anda dapat membuat lingkungan python dengan Anaconda:
conda create -n LDMSeg python=3.11
conda activate LDMSeg
Kami menyarankan untuk mengikuti instalasi otomatis (lihat tools/scripts/install_env.sh ). Jalankan perintah berikut untuk menginstal proyek dalam mode yang dapat diedit. Perhatikan bahwa semua dependensi akan diinstal secara otomatis. Karena ini mungkin tidak selalu berhasil (misalnya, karena masalah CUDA atau gcc), silakan lihat langkah-langkah instalasi manual.
python -m pip install -e .
pip install git+https://github.com/facebookresearch/detectron2.git
pip install git+https://github.com/cocodataset/panopticapi.gitPaket yang paling penting dapat diinstal dengan cepat dengan pip sebagai:
pip install torch torchvision einops # Main framework
pip install diffusers transformers xformers accelerate timm # For using pretrained models
pip install scipy opencv-python # For augmentations or loss
pip install pyyaml easydict hydra-core # For using config files
pip install termcolor wandb # For printing and logging Lihat data/environment.yml untuk salinan lingkungan saya. Kami juga mengandalkan beberapa dependensi dari detector2 dan panopticapi. Silakan ikuti dokumen mereka.
Saat ini kami mendukung kumpulan data COCO. Silakan ikuti dokumen untuk memasang gambar dan masker segmentasi panoptik yang sesuai. Juga, lihat direktori ldmseg/data/ untuk beberapa contoh pada dataset COCO. Sebagai catatan tambahan, struktur yang diadopsi harus cukup standar:
.
└── coco
├── annotations
├── panoptic_semseg_train2017
├── panoptic_semseg_val2017
├── panoptic_train2017 -> annotations/panoptic_train2017
├── panoptic_val2017 -> annotations/panoptic_val2017
├── test2017
├── train2017
└── val2017
Terakhir, ubah jalur di configs/env/root_paths.yml masing-masing ke root kumpulan data dan direktori keluaran yang Anda inginkan.
Pendekatan yang disajikan memiliki dua cabang: Pertama, kami melatih pembuat enkode otomatis untuk merepresentasikan peta segmentasi dalam ruang berdimensi lebih rendah (misalnya, 64x64). Selanjutnya, kita mulai dari Model Difusi Laten (LDM) yang telah dilatih sebelumnya, khususnya Difusi Stabil, untuk melatih model yang dapat menghasilkan topeng panoptik dari gambar RGB. Model dapat dilatih dengan menjalankan perintah berikut. Secara default kita akan melatih dataset COCO dengan file konfigurasi dasar yang ditentukan di tools/configs/base/base.yaml . Perhatikan bahwa file ini akan dimuat secara otomatis karena kami mengandalkan paket hydra .
python - W ignore tools / main_ae . py
datasets = coco
base . train_kwargs . fp16 = True
base . optimizer_name = adamw
base . optimizer_kwargs . lr = 1e-4
base . optimizer_kwargs . weight_decay = 0.05 Detail selengkapnya tentang penyampaian argumen dapat ditemukan di tools/scripts/train_ae.sh . Misalnya, saya menjalankan model ini untuk 50 ribu iterasi pada satu GPU sebesar 23 GB dengan total ukuran batch 16.
python - W ignore tools / main_ldm . py
datasets = coco
base . train_kwargs . gradient_checkpointing = True
base . train_kwargs . fp16 = True
base . train_kwargs . weight_dtype = float16
base . optimizer_zero_redundancy = True
base . optimizer_name = adamw
base . optimizer_kwargs . lr = 1e-4
base . optimizer_kwargs . weight_decay = 0.05
base . scheduler_kwargs . weight = 'max_clamp_snr'
base . vae_model_kwargs . pretrained_path = '$AE_MODEL' $AE_MODEL menunjukkan jalur ke model yang diperoleh dari langkah sebelumnya. Detail lebih lanjut tentang penyampaian argumen dapat ditemukan di tools/scripts/train_diffusion.sh . Misalnya, saya menjalankan model ini untuk 200 ribu iterasi pada 8 GPU berukuran 16 GB dengan total ukuran batch 256.
Kami berencana merilis beberapa model terlatih. Metrik PQ (kelas-agnostik) disediakan pada set validasi COCO.
| Model | #Param | Kumpulan data | Iter | PQ | meter persegi | RQ | Tautan unduhan |
|---|---|---|---|---|---|---|---|
| AE | ~2M | KELAPA | 66k | - | - | - | Unduh (23MB) |
| LDM | ~800 juta | KELAPA | 200k | 51.7 | 82.0 | 63.0 | Unduh (3,3 GB) |
Catatan: AE yang kurang kuat (yaitu, lapisan downsampling atau upsampling yang lebih sedikit) sering kali dapat menguntungkan proses inpainting, karena kami tidak melakukan penyesuaian tambahan.
Evaluasinya akan terlihat seperti:
python - W ignore tools / main_ldm . py
datasets = coco
base . sampling_kwargs . num_inference_steps = 50
base . eval_only = True
base . load_path = $ PRETRAINED_MODEL_PATH Anda dapat menambahkan parameter jika perlu. Ambang batas yang lebih tinggi seperti --base.eval_kwargs.count_th 700 atau --base.eval_kwargs.mask_th 0.9 dapat semakin meningkatkan jumlahnya. Namun, kami menggunakan nilai standar dengan menetapkan ambang batas sebesar 0,5 dan menghapus segmen dengan luas lebih kecil dari 512 untuk evaluasi.
Untuk mengevaluasi model terlatih dari atas, jalankan tools/scripts/eval.sh .
Di sini, kami memvisualisasikan hasilnya:
Jika Anda merasa repositori ini berguna untuk penelitian Anda, mohon pertimbangkan untuk mengutip makalah berikut:
@article { vangansbeke2024ldmseg ,
title = { a simple latent diffusion approach for panoptic segmentation and mask inpainting } ,
author = { Van Gansbeke, Wouter and De Brabandere, Bert } ,
journal = { arxiv preprint arxiv:2401.10227 } ,
year = { 2024 }
}Untuk pertanyaan apa pun, silakan hubungi penulis utama.
Perangkat lunak ini dirilis di bawah lisensi creative commons yang hanya mengizinkan penggunaan pribadi dan penelitian. Untuk lisensi komersial silakan hubungi penulis. Anda dapat melihat ringkasan lisensi di sini.
Saya berterima kasih atas semua repositori publik (lihat juga referensi dalam kode), dan khususnya untuk perpustakaan detector2 dan diffuser.
