head_detector
1.0.0
Orest Kupyn 13 · Eugene Khvedchenia 2 · Christian Rupprecht 1 ·
1 Université d'Oxford · 2 Université catholique ukrainienne · 3 piñatafarms AI
VGGheads est un ensemble de données entièrement synthétique à grande échelle pour la détection de la tête humaine et l'estimation du maillage 3D avec plus d'un million d'images générées avec des modèles de diffusion. Un modèle formé uniquement sur les données synthétiques se généralise bien au monde réel et est capable de détection simultanée des têtes et de la reconstruction des mailles de tête à partir d'une seule image en une seule étape.
Pour télécharger l'ensemble de données VGGheads, vous avez deux options:
pip install academictorrents
at-get 1ac36f16386061685ed303dea6f0d6179d2e2121ou utiliser Aria2c
aria2c --seed-time=0 --max-overall-download-limit=10M --file-allocation=none https://academictorrents.com/download/1ac36f16386061685ed303dea6f0d6179d2e2121.torrentLien de torrent complet
Nous vous recommandons d'utiliser la méthode Torrent car elle est généralement plus rapide et aide à réduire la charge sur nos serveurs.
wget https://thor.robots.ox.ac.uk/vgg-heads/VGGHeads.tar Cela téléchargera un fichier nommé VGGHeads.tar dans votre répertoire actuel.
Pour vérifier l'intégrité du fichier téléchargé, nous aurons besoin des sommes de contrôle MD5. Téléchargez-les en utilisant:
wget https://thor.robots.ox.ac.uk/vgg-heads/MD5SUMS Une fois les deux fichiers téléchargés, vérifiez l'intégrité du fichier VGGHeads.tar :
md5sum -c MD5SUMSSi le téléchargement a réussi et que le fichier est intact, vous devriez voir un message "OK".
Si la vérification a réussi, extraire le contenu du fichier TAR:
tar -xvf VGGHeads.tarCela extraitra le contenu de l'archive dans votre répertoire actuel.
Notes:
conda create --name vgg_heads python=3.10
conda activate vgg_headsgit clone https://github.com/KupynOrest/head_detector.git
cd head_detector
pip install -e ./Ou simplement installer
pip install git+https://github.com/KupynOrest/head_detector.gitPour tester le modèle VGGheads sur vos propres images, utilisez simplement ce code:
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. Vous pouvez exporter des maillages de tête sous forme de fichiers OBJ à l'aide de la méthode save_meshes :
# After getting predictions
save_folder = "path/to/save/folder"
predictions . save_meshes ( save_folder )Cela enregistrera les fichiers OBJ individuels pour chaque tête détectée dans le dossier spécifié.
Pour obtenir des cultures de tête alignées, utilisez la méthode 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 )Cela renvoie une liste de cultures de tête alignées que vous pouvez traiter ou enregistrer davantage.
Voici un exemple complet incorporant toutes les fonctionnalités:
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'" )Cet exemple étendu montre comment utiliser toutes les fonctionnalités du modèle VGGheads, y compris la détection de base de la tête, le dessin des résultats, l'exportation de mailles de tête et l'obtention de cultures de tête alignées.
De plus, les poids ONNX sont disponibles sur HuggingFace. L'exemple de l'inférence peut être trouvé à: Colab
Nous fournissons également une démo Gradio, que vous pouvez exécuter localement:
cd gradio
pip install -r requirements.txt
python app.py Vous pouvez spécifier les arguments --server_port , --share , --server_name pour satisfaire vos besoins!
Vérifiez yolo_head_training/Makefile pour des exemples de scripts de train.
Pour exécuter la formation sur toutes les données avec des données distribuées parallèles (DDP), utilisez la commande suivante:
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=DDPRemplacez les espaces réservés suivants:
NUM_GPUS : le nombre de GPU que vous souhaitez utiliser pour la formation.DATA_FOLDER : Le chemin du répertoire contenant votre ensemble de données extrait.Formation GPU unique: si vous utilisez un seul GPU, vous pouvez simplifier la commande:
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 Configuration personnalisée: vous pouvez modifier le paramètre --config-name pour utiliser différentes configurations de modèle. Vérifiez les fichiers de configuration dans le répertoire du projet pour les options disponibles.
Réglage des hyperparamètres: vous pouvez ajuster divers hyperparamètres en les ajoutant à la ligne de commande. Par exemple:
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 Reprise de la formation: si vous devez reprendre la formation à partir d'un point de contrôle, vous pouvez utiliser le drapeau 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=TrueVous pouvez surveiller les progrès de la formation à travers la sortie de la console. Envisagez d'utiliser des outils tels que Tensorboard pour une surveillance et une visualisation plus détaillées des mesures de formation.
Si vous trouvez VGGheads utile pour vos recherches et vos applications, veuillez nous citer en utilisant ce bibtex:
@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 } ,
}Ce travail est concédé sous licence Creative Commons Attribution-NonCommercial 4.0 International Licence.
