Orest Kupyn 13 · Eugene Khvedchenia 2 ·基督徒Rupprecht 1 ·
1牛津大学· 2乌克兰天主教大学· 3PiñataFarmsAI
VGGHEADS是用于人头检测的大规模完全合成数据集和3D网格估计,具有超过100万张图像的扩散模型。仅在合成数据上训练的模型可以很好地推广到现实世界,并且能够同时从单个步骤中从单个图像中进行重建。

要下载vggheads数据集,您有两个选择:
pip install academictorrents
at-get 1ac36f16386061685ed303dea6f0d6179d2e2121或使用Aria2c
aria2c --seed-time=0 --max-overall-download-limit=10M --file-allocation=none https://academictorrents.com/download/1ac36f16386061685ed303dea6f0d6179d2e2121.torrent完整的洪流链接
我们建议使用洪流方法,因为它通常更快,并有助于减少服务器上的负载。
wget https://thor.robots.ox.ac.uk/vgg-heads/VGGHeads.tar这将下载一个名为VGGHeads.tar的文件到您当前的目录。
为了验证下载文件的完整性,我们需要MD5校验和。使用以下方式下载它们:
wget https://thor.robots.ox.ac.uk/vgg-heads/MD5SUMS下载两个文件后,验证VGGHeads.tar文件的完整性:
md5sum -c MD5SUMS如果下载成功并且文件完好无损,则应该看到一条“确定”消息。
如果验证成功,请提取焦油文件的内容:
tar -xvf VGGHeads.tar这将将存档的内容提取到您当前的目录中。
笔记:
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 ./或简单地安装
pip install git+https://github.com/KupynOrest/head_detector.git要在您自己的图像上测试VGGHEADS模型,只需使用以下代码:
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.您可以使用save_meshes方法将头部网格作为OBJ文件导出:
# After getting predictions
save_folder = "path/to/save/folder"
predictions . save_meshes ( save_folder )这将保存指定文件夹中每个检测到的头部的单个OBJ文件。
要获得对齐的头部作物,请使用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 )这返回了您可以进一步加工或保存的一排式头部作物列表。
这是一个结合所有功能的完整示例:
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'" )这个扩展的示例演示了如何使用VGGHEADS模型的所有功能,包括基本的头部检测,绘制结果,导出头部网格以及获得对齐的头部作物。
此外,在HuggingFace上还可以使用ONNX权重。可以在:COLAB上找到推理的例子
我们还提供了一个Gradio演示,您可以在本地运行:
cd gradio
pip install -r requirements.txt
python app.py您可以指定--server_port , --share , --server_name参数以满足您的需求!
检查yolo_head_training/Makefile以获取火车脚本的示例。
要使用并行分布式数据(DDP)对所有数据进行培训,请使用以下命令:
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替换以下占位符:
NUM_GPUS :您要用于培训的GPU数量。DATA_FOLDER :通往包含您提取的数据集的目录的路径。单个GPU培训:如果您使用的是单个GPU,则可以简化命令:
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自定义配置:您可以修改--config-name参数以使用不同的模型配置。检查项目目录中的配置文件以获取可用选项。
调整超参数:您可以通过将它们添加到命令行中调整各种超参数。例如:
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恢复培训:如果您需要从检查站恢复培训,则可以使用training_hyperparams.resume flag:
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您可以通过控制台输出监视训练进度。考虑使用张量板等工具,以更详细的监视和可视化培训指标。
如果您发现Vggheads对您的研究和应用程序有用,请使用此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 } ,
}这项工作是根据创意共享归因于非商业4.0国际许可证获得许可的。