head_detector
1.0.0
Орест Купин 13 · Евгений Хведхения 2 · Кристиан Рупрехт 1 ·
1 Оксфордский университет · 2 Украинский католический университет · 3 Piñatafarms AI
Vggheads является крупномасштабным полностью синтетическим набором данных для обнаружения головы человека и оценки 3D-сетки с более чем 1 миллионами изображений, сгенерированных с помощью диффузионных моделей. Модель, обученная только синтетическим данным, хорошо обобщается в реальном мире, и способна одновременно обнаружить головы, а реконструкция головки сетки с одного изображения за один шаг.
Чтобы загрузить набор данных 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:
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. Вы можете экспортировать сетки головки в виде файлов OBJ, используя метод save_meshes :
# 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, включая базовое обнаружение головки, результаты рисования, экспорт сетки головки и получение выровненных культур головы.
Кроме того, веса ONNX доступны на HuggingFace. Пример вывода можно найти на: 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 : количество графических процессоров, которые вы хотите использовать для обучения.DATA_FOLDER : путь к каталогу, содержащему ваш извлеченный набор данных.Обучение отдельного графического процессора: если вы используете один графический процессор, вы можете упростить команду:
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 для использования различных конфигураций модели. Проверьте файлы конфигурации в каталоге Project для доступных параметров.
Регулировка гиперпараметров: вы можете настроить различные гиперпараметры, добавив их в командную строку. Например:
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 :
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Вы можете отслеживать прогресс обучения через вывод консоли. Подумайте об использовании таких инструментов, как Tensorboard для более подробного мониторинга и визуализации метрик обучения.
Если вы найдете 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 } ,
}Эта работа лицензирована по международной лицензии Creative Commons Attribution-Noncommercial 4.0.
