head_detector
1.0.0
Orest Kupyn 13 · Eugene Khvedchenia 2 · Christian Rupprecht 1 ·
1 Universidad de Oxford · 2 Universidad Católica ucraniana · 3 piñatafarms ai
VGGheads es un conjunto de datos totalmente sintético a gran escala para la detección de la cabeza humana y la estimación de malla 3D con más de 1 millón de imágenes generadas con modelos de difusión. Un modelo entrenado solo en datos sintéticos se generaliza bien al mundo real y es capaz de detección de cabezas simultáneas y reconstrucción de mallas de cabeza desde una sola imagen en un solo paso.
Para descargar el conjunto de datos VGGheads, tiene dos opciones:
pip install academictorrents
at-get 1ac36f16386061685ed303dea6f0d6179d2e2121o usar aria2c
aria2c --seed-time=0 --max-overall-download-limit=10M --file-allocation=none https://academictorrents.com/download/1ac36f16386061685ed303dea6f0d6179d2e2121.torrentEnlace de torrent completo
Recomendamos usar el método Torrent, ya que es típicamente más rápido y ayuda a reducir la carga en nuestros servidores.
wget https://thor.robots.ox.ac.uk/vgg-heads/VGGHeads.tar Esto descargará un archivo llamado VGGHeads.tar a su directorio actual.
Para verificar la integridad del archivo descargado, necesitaremos las suma de verificación MD5. Descarglos usando:
wget https://thor.robots.ox.ac.uk/vgg-heads/MD5SUMS Después de descargar ambos archivos, verifique la integridad del archivo VGGHeads.tar :
md5sum -c MD5SUMSSi la descarga fue exitosa y el archivo está intacto, debería ver un mensaje "OK".
Si la verificación fue exitosa, extraiga el contenido del archivo TAR:
tar -xvf VGGHeads.tarEsto extraerá el contenido del archivo en su directorio actual.
Notas:
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 ./O simplemente instalar
pip install git+https://github.com/KupynOrest/head_detector.gitPara probar el modelo VGGheads en sus propias imágenes, simplemente use este código:
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. Puede exportar mallas de cabeza como archivos OBJ utilizando el método save_meshes :
# After getting predictions
save_folder = "path/to/save/folder"
predictions . save_meshes ( save_folder )Esto guardará archivos OBJ individuales para cada cabezal detectado en la carpeta especificada.
Para obtener cultivos de cabeza alineados, use el método 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 )Esto devuelve una lista de cultivos de cabeza alineados que puede procesar o guardar más.
Aquí hay un ejemplo completo que incorpora todas las características:
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'" )Este ejemplo extendido demuestra cómo usar todas las características del modelo VGGheads, incluida la detección básica de la cabeza, los resultados del dibujo, la exportación de mallas de la cabeza y la obtención de cultivos de cabeza alineados.
Además, los pesos ONNX están disponibles en Huggingface. El ejemplo de la inferencia se puede encontrar en: Colab
También proporcionamos una demostración de Gradio, que puede ejecutar localmente:
cd gradio
pip install -r requirements.txt
python app.py ¡Puede especificar el --server_port , --share , --server_name argumentos para satisfacer sus necesidades!
Consulte yolo_head_training/Makefile para ver ejemplos de scripts de trenes.
Para ejecutar la capacitación en todos los datos con datos distribuidos paralelos (DDP), use el siguiente comando:
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=DDPReemplace los siguientes marcadores de posición:
NUM_GPUS : el número de GPU que desea usar para el entrenamiento.DATA_FOLDER : la ruta al directorio que contiene su conjunto de datos extraído.Entrenamiento único de GPU: si está utilizando una sola GPU, puede simplificar el comando:
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 Configuración personalizada: puede modificar el parámetro --config-name para usar diferentes configuraciones de modelo. Verifique los archivos de configuración en el directorio del proyecto para ver las opciones disponibles.
Ajuste de hiperparámetros: puede ajustar varios hiperparámetros agregándolos a la línea de comando. Por ejemplo:
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 Entrenamiento de reanudación: si necesita reanudar la capacitación desde un punto de control, puede usar la bandera de 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=TruePuede monitorear el progreso de la capacitación a través de la salida de la consola. Considere el uso de herramientas como TensorBoard para un monitoreo y visualización más detallados de las métricas de capacitación.
Si encuentra útiles para su investigación y aplicaciones, cíquanos usando este 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 } ,
}Este trabajo tiene licencia bajo una licencia internacional de Atribución de Commons Creative Commons.
