head_detector

Vggheads: Dataset sintetis skala besar untuk kepala manusia 3D

Orest Kupyn ¹³ · Eugene Khvedchenia ² · Christian Rupprecht ¹ ·

¹ Universitas Oxford · ² Universitas Katolik Ukraina · ³ Piñatafarms AI

Vggheads adalah dataset sintetis sepenuhnya skala besar untuk deteksi kepala manusia dan estimasi mesh 3D dengan lebih dari 1 juta gambar yang dihasilkan dengan model difusi. Model yang dilatih hanya pada data sintetis menggeneralisasi dengan baik ke dunia nyata dan mampu mendeteksi kepala simultan dan rekonstruksi kepala dari satu gambar dalam satu langkah.

Berita

• [2024/08/30] Rilis versi 0.1.0. Menambahkan contoh perataan kepala dan menghemat mesh sebagai .obj
• [2024/08/29] Kami merilis dataset, instruksi pelatihan, dan bobot ONNX !!
• [2024/08/09] Kami merilis pos pemeriksaan vggheads_l dan controlnet mesh
• [2024/07/26] Kami merilis versi awal basis kode, kertas, halaman web proyek, dan demo gambar !!

Instruksi Unduh Dataset Vggheads

1. Unduh dataset

Untuk mengunduh dataset Vggheads, Anda memiliki dua opsi:

1. Unduh Torrent (Metode Pilihan): Cara Mengunduh

pip install academictorrents
at-get 1ac36f16386061685ed303dea6f0d6179d2e2121

atau gunakan ARIA2C

aria2c --seed-time=0 --max-overall-download-limit=10M --file-allocation=none https://academictorrents.com/download/1ac36f16386061685ed303dea6f0d6179d2e2121.torrent

Tautan torrent penuh

Kami merekomendasikan menggunakan metode torrent karena biasanya lebih cepat dan membantu mengurangi beban pada server kami.

2. Unduh Langsung:

wget https://thor.robots.ox.ac.uk/vgg-heads/VGGHeads.tar

Ini akan mengunduh file bernama VGGHeads.tar ke direktori Anda saat ini.

2. Unduh checksum MD5

Untuk memverifikasi integritas file yang diunduh, kita akan memerlukan checksum MD5. Unduh mereka menggunakan:

wget https://thor.robots.ox.ac.uk/vgg-heads/MD5SUMS

3. Verifikasi unduhan

Setelah kedua file diunduh, verifikasi integritas file VGGHeads.tar :

md5sum -c MD5SUMS

Jika unduhan berhasil dan file utuh, Anda akan melihat pesan "OK".

4. Ekstrak dataset

Jika verifikasi berhasil, ekstrak isi file tar:

tar -xvf VGGHeads.tar

Ini akan mengekstraksi isi arsip ke dalam direktori Anda saat ini.

Catatan:

• Ukuran dataset sekitar 187 GB. Pastikan Anda memiliki ruang disk yang cukup sebelum mengunduh dan mengekstraksi.
• Proses unduhan dan ekstraksi mungkin membutuhkan waktu tergantung pada koneksi internet dan kecepatan komputer Anda.
• Jika Anda mengalami masalah selama proses unduhan atau ekstraksi, coba unduh lagi atau periksa utilitas tar sistem Anda.

Instalasi

Buat lingkungan virtual conda

conda create --name vgg_heads python=3.10
conda activate vgg_heads

Klon proyek dan instal paket

git clone https://github.com/KupynOrest/head_detector.git
cd head_detector

pip install -e ./

Atau cukup instal

pip install git+https://github.com/KupynOrest/head_detector.git

Penggunaan

Untuk menguji model vggheads pada gambar Anda sendiri cukup gunakan kode ini:

 from head_detector import HeadDetector
import cv2
detector = HeadDetector ()
image_path = "your_image.jpg"
predictions = detector ( image_path )
# predictions.heads contain a list of heads with .bbox, .vertices_3d, .head_pose params
result_image = predictions . draw () # draw heads on the image
cv2 . imwrite ( "result.png" , result_image ) # save result image to preview it.

Mengekspor kepala jerat

Anda dapat mengekspor jala kepala sebagai file OBJ menggunakan metode save_meshes :

 # After getting predictions
save_folder = "path/to/save/folder"
predictions . save_meshes ( save_folder )

Ini akan menyimpan file OBJ individual untuk setiap kepala yang terdeteksi di folder yang ditentukan.

Mendapatkan tanaman kepala yang selaras

Untuk mendapatkan tanaman kepala yang selaras, gunakan metode get_aligned_heads :

 # After getting predictions
aligned_heads = predictions . get_aligned_heads ()

# Process or save aligned head crops
for i , head in enumerate ( aligned_heads ):
    cv2 . imwrite ( f"aligned_head_ { i } .png" , head )

Ini mengembalikan daftar tanaman kepala yang selaras yang dapat Anda proses atau simpan lebih lanjut.

Contoh diperpanjang

Berikut contoh lengkap yang menggabungkan semua fitur:

 from head_detector import HeadDetector
import cv2
import os

# Initialize the detector
detector = HeadDetector ()

# Specify the path to your image
image_path = "your_image.jpg"

# Get predictions
predictions = detector ( image_path )

# Draw heads on the image
result_image = predictions . draw ()
cv2 . imwrite ( "result.png" , result_image )

# Save head meshes
save_folder = "head_meshes"
os . makedirs ( save_folder , exist_ok = True )
predictions . save_meshes ( save_folder )

# Get and save aligned head crops
aligned_heads = predictions . get_aligned_heads ()
for i , head in enumerate ( aligned_heads ):
    cv2 . imwrite ( f"aligned_head_ { i } .png" , head )

print ( f"Detected { len ( predictions . heads ) } heads." )
print ( f"Result image saved as 'result.png'" )
print ( f"Head meshes saved in ' { save_folder } ' folder" )
print ( f"Aligned head crops saved as 'aligned_head_*.png'" )

Contoh yang diperluas ini menunjukkan cara menggunakan semua fitur model VGGHEADS, termasuk deteksi kepala dasar, hasil menggambar, mengekspor jerat kepala, dan mendapatkan tanaman kepala yang selaras.

Selain itu, bobot ONNX tersedia di Huggingface. Contoh inferensi dapat ditemukan di: colab

Demo Gradio

Kami juga menyediakan demo gradio, yang dapat Anda jalankan secara lokal:

 cd gradio
pip install -r requirements.txt
python app.py

Anda dapat menentukan argumen --server_port , --share , --server_name untuk memenuhi kebutuhan Anda!

Pelatihan

Periksa yolo_head_training/Makefile untuk contoh skrip kereta api.

Untuk menjalankan pelatihan pada semua data dengan data paralel terdistribusi (DDP), gunakan perintah berikut:

torchrun --standalone --nnodes=1 --nproc_per_node=NUM_GPUS train.py --config-name=yolo_heads_l 
    dataset_params.train_dataset_params.data_dir=DATA_FOLDER/large 
    dataset_params.val_dataset_params.data_dir=DATA_FOLDER/large 
    num_gpus=NUM_GPUS multi_gpu=DDP

Ganti placeholder berikut:

• NUM_GPUS : Jumlah GPU yang ingin Anda gunakan untuk pelatihan.
• DATA_FOLDER : Jalur ke direktori yang berisi dataset yang diekstraksi.

Opsi pelatihan tambahan

1. Pelatihan GPU tunggal: Jika Anda menggunakan satu GPU, Anda dapat menyederhanakan perintah:

   python train.py --config-name=yolo_heads_l 
       dataset_params.train_dataset_params.data_dir=DATA_FOLDER/large 
       dataset_params.val_dataset_params.data_dir=DATA_FOLDER/large

2. Konfigurasi Kustom: Anda dapat memodifikasi parameter --config-name untuk menggunakan konfigurasi model yang berbeda. Periksa file konfigurasi di direktori proyek untuk opsi yang tersedia.

3. Menyesuaikan hyperparameters: Anda dapat menyesuaikan berbagai hiperparameter dengan menambahkannya ke baris perintah. Misalnya:

   python train.py --config-name=yolo_heads_l 
       dataset_params.train_dataset_params.data_dir=DATA_FOLDER/large 
       dataset_params.val_dataset_params.data_dir=DATA_FOLDER/large 
       training_hyperparams.initial_lr=0.001 
       training_hyperparams.max_epochs=100

4. Pelatihan Lanjutkan: Jika Anda perlu melanjutkan pelatihan dari pos pemeriksaan, Anda dapat menggunakan bendera training_hyperparams.resume :

   python train.py --config-name=yolo_heads_l 
       dataset_params.train_dataset_params.data_dir=DATA_FOLDER/large 
       dataset_params.val_dataset_params.data_dir=DATA_FOLDER/large 
       training_hyperparams.resume=True

Pelatihan Pemantauan

Anda dapat memantau kemajuan pelatihan melalui output konsol. Pertimbangkan untuk menggunakan alat seperti Tensorboard untuk pemantauan yang lebih rinci dan visualisasi metrik pelatihan.

Mengutip

Jika Anda menemukan vgghead berguna untuk penelitian dan aplikasi Anda, silakan mengutip kami menggunakan Bibtex ini:

 @article { vggheads ,
      title = { VGGHeads: A Large-Scale Synthetic Dataset for 3D Human Heads } ,
      author = { Orest Kupyn and Eugene Khvedchenia and Christian Rupprecht } ,
      year = { 2024 } ,
      eprint = { 2407.18245 } ,
      archivePrefix = { arXiv } ,
      primaryClass = { cs.CV } ,
      url = { https://arxiv.org/abs/2407.18245 } ,
}

Karya ini dilisensikan di bawah lisensi internasional Atribut-Nonkomersial 4.0 Creative Commons.