head_detector
1.0.0
Orest Kupyn 13 · Eugene Khvedchenia 2 · Christian Rupprecht 1 ·
1 Universität Oxford · 2 Ukrainische katholische Universität · 3 Piñatafarms AI
VGGHeads ist ein groß angelegter vollständig synthetischer Datensatz für die Erkennung menschlicher Kopf und 3D-Netzschätzung mit über 1 Million Bildern, die mit Diffusionsmodellen erzeugt werden. Ein Modell, das nur auf synthetischen Daten trainiert wurde, verallgemeinert sich gut auf die reale Welt und kann in einem einzigen Schritt die Rekonstruktion eines einzelnen Bildes gleichzeitig mit Kopfnetze und Kopfnetze gleichzeitig köpft.
Um den VGGHeads -Datensatz herunterzuladen, haben Sie zwei Optionen:
pip install academictorrents
at-get 1ac36f16386061685ed303dea6f0d6179d2e2121oder verwenden Sie ARIA2C
aria2c --seed-time=0 --max-overall-download-limit=10M --file-allocation=none https://academictorrents.com/download/1ac36f16386061685ed303dea6f0d6179d2e2121.torrentVoller Torrent -Link
Wir empfehlen, die Torrent -Methode zu verwenden, da sie normalerweise schneller ist und die Last auf unseren Servern reduziert.
wget https://thor.robots.ox.ac.uk/vgg-heads/VGGHeads.tar Dadurch wird eine Datei namens VGGHeads.tar in Ihr aktuelles Verzeichnis heruntergeladen.
Um die Integrität der heruntergeladenen Datei zu überprüfen, benötigen wir die MD5 -Prüfsummen. Laden Sie sie mit: herunterladen:
wget https://thor.robots.ox.ac.uk/vgg-heads/MD5SUMS Überprüfen Sie nach dem Herunterladen beider Dateien die Integrität der Datei VGGHeads.tar :
md5sum -c MD5SUMSWenn der Download erfolgreich war und die Datei intakt ist, sollten Sie eine "OK" -Meldung sehen.
Wenn die Überprüfung erfolgreich war, extrahieren Sie den Inhalt der TAR -Datei:
tar -xvf VGGHeads.tarDadurch wird der Inhalt des Archivs in Ihr aktuelles Verzeichnis ausgezeichnet.
Anmerkungen:
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 ./Oder einfach installieren
pip install git+https://github.com/KupynOrest/head_detector.gitVerwenden Sie diesen Code einfach, um das VGGHeads -Modell auf Ihren eigenen Bildern zu testen:
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. Sie können Kopfnetze als OBJ -Dateien mithilfe der Methode save_meshes exportieren:
# After getting predictions
save_folder = "path/to/save/folder"
predictions . save_meshes ( save_folder )Dadurch speichert einzelne OBJ -Dateien für jeden erkannten Kopf im angegebenen Ordner.
Um ausgerichtete Kopfpflanzen zu erhalten, verwenden Sie die Methode 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 )Dadurch wird eine Liste der ausgerichteten Kopfpflanzen zurückgegeben, die Sie weiter verarbeiten oder speichern können.
Hier ist ein vollständiges Beispiel, das alle Funktionen enthält:
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'" )Dieses erweiterte Beispiel zeigt, wie alle Merkmale des VGGHEADS -Modells verwendet werden, einschließlich grundlegender Kopferkennung, Zeichnenergebnisse, Exportieren von Kopfnetzen und Erlangung der ausgerichteten Kopfkulturen.
Zusätzlich sind die ONNX -Gewichte bei Huggingface erhältlich. Das Beispiel für die Schlussfolgerung findet sich bei: Colab: Colab
Wir bieten auch eine Gradio -Demo, die Sie lokal ausführen können:
cd gradio
pip install -r requirements.txt
python app.py Sie können die Argumente --server_port , --share , --server_name angeben, um Ihre Anforderungen zu erfüllen!
Überprüfen Sie yolo_head_training/Makefile für Beispiele für Zugskripte.
Verwenden Sie den folgenden Befehl, um das Training für alle Daten mit verteilten Daten parallel (DDP) auszuführen:
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=DDPErsetzen Sie die folgenden Platzhalter:
NUM_GPUS : Die Anzahl der GPUs, die Sie für das Training verwenden möchten.DATA_FOLDER : Der Pfad zum Verzeichnis, der Ihren extrahierten Datensatz enthält.Single GPU -Training: Wenn Sie eine einzelne GPU verwenden, können Sie den Befehl vereinfachen:
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 Benutzerdefinierte Konfiguration: Sie können den Parameter- --config-name so ändern, dass verschiedene Modellkonfigurationen verwendet werden. Überprüfen Sie die Konfigurationsdateien im Projektverzeichnis auf verfügbare Optionen.
Hyperparameter einstellen: Sie können verschiedene Hyperparameter einstellen, indem Sie sie zur Befehlszeile hinzufügen. Zum Beispiel:
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 Wiederaufnahme des Trainings: Wenn Sie das Training von einem Kontrollpunkt aus wieder aufnehmen müssen, können Sie das Flag training_hyperparams.resume verwenden:
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=TrueSie können den Trainingsfortschritt über die Konsolenausgabe überwachen. Erwägen Sie, Tools wie Tensorboard für eine detailliertere Überwachung und Visualisierung von Trainingsmetriken zu verwenden.
Wenn Sie VGGHeads für Ihre Forschung und Anwendungen nützlich finden, zitieren Sie uns bitte mit diesem 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 } ,
}Diese Arbeit ist im Rahmen einer Kreative Commons Attribution-Noncommercial 4.0 International Lizenz lizenziert.
