head_detector
1.0.0
Orest Kupyn 13 · Eugene Khvedchenia 2 · Christian Rupprecht 1 ·
1 Universidade de Oxford · 2 Universidade Católica Ucraniana · 3 PiñataFarms AI
O VGGHEADS é um conjunto de dados totalmente sintético em larga escala para detecção de cabeça humana e estimativa de malha 3D, com mais de 1 milhão de imagens geradas com modelos de difusão. Um modelo treinado apenas em dados sintéticos generaliza bem para o mundo real e é capaz de detecção de cabeças simultâneas e reconstrução de malhas da cabeça de uma única imagem em uma única etapa.
Para baixar o conjunto de dados VGGHEADS, você tem duas opções:
pip install academictorrents
at-get 1ac36f16386061685ed303dea6f0d6179d2e2121ou use ARIA2C
aria2c --seed-time=0 --max-overall-download-limit=10M --file-allocation=none https://academictorrents.com/download/1ac36f16386061685ed303dea6f0d6179d2e2121.torrentLink completo de torrent
Recomendamos o uso do método torrent, pois normalmente é mais rápido e ajuda a reduzir a carga em nossos servidores.
wget https://thor.robots.ox.ac.uk/vgg-heads/VGGHeads.tar Isso baixará um arquivo chamado VGGHeads.tar para o seu diretório atual.
Para verificar a integridade do arquivo baixado, precisaremos das somas de verificação do MD5. Baixe -os usando:
wget https://thor.robots.ox.ac.uk/vgg-heads/MD5SUMS Depois que os dois arquivos forem baixados, verifique a integridade do arquivo VGGHeads.tar :
md5sum -c MD5SUMSSe o download foi bem -sucedido e o arquivo estiver intacto, você deverá ver uma mensagem "OK".
Se a verificação foi bem -sucedida, extraia o conteúdo do arquivo TAR:
tar -xvf VGGHeads.tarIsso extrairá o conteúdo do arquivo para o seu diretório atual.
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 ./Ou simplesmente instalar
pip install git+https://github.com/KupynOrest/head_detector.gitPara testar o modelo VGGHEADS em suas próprias imagens, basta usar 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. Você pode exportar malhas de cabeça como arquivos OBJ usando o método save_meshes :
# After getting predictions
save_folder = "path/to/save/folder"
predictions . save_meshes ( save_folder )Isso salvará arquivos OBJ individuais para cada cabeça detectada na pasta especificada.
Para obter culturas de cabeça alinhadas, use o 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 )Isso retorna uma lista de culturas de cabeça alinhadas que você pode processar ou salvar ainda mais.
Aqui está um exemplo completo incorporando todos os recursos:
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 exemplo estendido demonstra como usar todos os recursos do modelo VGGHEADS, incluindo detecção básica de cabeça, resultados de desenho, exportação de malhas da cabeça e obtenção de culturas de cabeça alinhadas.
Além disso, os pesos ONNX estão disponíveis no HuggingFace. O exemplo da inferência pode ser encontrado em: colab
Também fornecemos uma demonstração de graduação, que você pode executar localmente:
cd gradio
pip install -r requirements.txt
python app.py Você pode especificar os argumentos --server_port , --share , --server_name para satisfazer suas necessidades!
Verifique yolo_head_training/Makefile para obter exemplos de scripts de trem.
Para executar o treinamento em todos os dados com dados distribuídos paralelos (DDP), use o seguinte 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=DDPSubstitua os seguintes espaços reservados:
NUM_GPUS : o número de GPUs que você deseja usar para treinamento.DATA_FOLDER : o caminho para o diretório que contém seu conjunto de dados extraído.Treinamento de GPU único: se você estiver usando uma única GPU, poderá simplificar o 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 Configuração personalizada: você pode modificar o parâmetro --config-name para usar diferentes configurações de modelo. Verifique os arquivos de configuração no diretório do projeto para obter opções disponíveis.
Ajustando os hiperparâmetros: você pode ajustar vários hiperparâmetros adicionando -os à linha de comando. Por exemplo:
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 Retomando o treinamento: se você precisar retomar o treinamento do ponto de verificação, poderá usar o sinalizador training_hyperparams.resume : 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=TrueVocê pode monitorar o progresso do treinamento através da saída do console. Considere o uso de ferramentas como o Tensorboard para um monitoramento e visualização mais detalhados das métricas de treinamento.
Se você achar vggheads úteis para sua pesquisa e aplicativos, cite -nos 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 trabalho é licenciado sob uma licença Creative Commons Attribution-NonCommercial 4.0 International.
