Ashkan Ganj 1 · Hang Su 2 · Tian Guo 1
1 Institut polytechnique de Worcester 2 Recherche Nvidia
Nous avons publié une version améliorée d'HybridDepth, désormais disponible avec de nouvelles fonctionnalités et des performances optimisées !
Ce travail présente HybridDepth. HybridDepth est une solution pratique d'estimation de profondeur basée sur des images de pile focale capturées à partir d'une caméra. Cette approche surpasse les modèles de pointe sur plusieurs ensembles de données bien connus, notamment NYU V2, DDFF12 et ARKitScenes.
30/10/2024 : Sortie de la version 2 d'HybridDepth avec des performances améliorées et des poids pré-entraînés.
30/10/2024 : Prise en charge intégrée de TorchHub pour un chargement et une inférence faciles du modèle.
2024-07-25 : Sortie initiale des modèles pré-entraînés.
2024-07-23 : Le référentiel GitHub et le modèle HybridDepth ont été mis en ligne.
Démarrez rapidement avec HybridDepth à l'aide du notebook Colab.
Vous pouvez sélectionner un modèle pré-entraîné directement avec TorchHub.
Modèles pré-entraînés disponibles :
HybridDepth_NYU5
: pré-entraîné sur l'ensemble de données NYU Depth V2 à l'aide d'une entrée de pile à 5 focales, avec la branche DFF et la couche de raffinement entraînées.
HybridDepth_NYU10
: pré-entraîné sur l'ensemble de données NYU Depth V2 à l'aide d'une entrée de pile à 10 focales, avec la branche DFF et la couche de raffinement entraînées.
HybridDepth_DDFF5
: Pré-entraîné sur l'ensemble de données DDFF à l'aide d'une entrée de pile à 5 focales.
HybridDepth_NYU_PretrainedDFV5
: Pré-entraînement uniquement sur la couche de raffinement avec l'ensemble de données NYU Depth V2 à l'aide d'une pile à 5 focales, après un pré-entraînement avec DFV.
model_name = 'HybridDepth_NYU_PretrainedDFV5' #change thismodel = torch.hub.load('cake-lab/HybridDepth', model_name, pretrained=True)model.eval()
Clonez le dépôt et installez les dépendances :
clone git https://github.com/cake-lab/HybridDepth.gitcd HybridDepth conda env créer -f environnement.yml conda active la profondeur hybride
Téléchargez les poids pré-entraînés :
Téléchargez les poids du modèle à partir des liens ci-dessous et placez-les dans le répertoire checkpoints
:
HybridDepth_NYU_FocalStack5
HybridDepth_NYU_FocalStack10
HybridDepth_DDFF_FocalStack5
HybridDepth_NYU_PretrainedDFV_FocalStack5
Prédiction
Pour l'inférence, vous pouvez exécuter le code suivant :
# Charger le modèle checkpointmodel_path = 'checkpoints/NYUBest5.ckpt'model = DepthNetModule.load_from_checkpoint(model_path)model.eval()model = model.to('cuda')
Après avoir chargé le modèle, utilisez le code suivant pour traiter les images d'entrée et obtenir la carte de profondeur :
Remarque : Actuellement, la fonction prepare_input_image
ne prend en charge que les images .jpg
. Modifiez la fonction si vous avez besoin de la prise en charge d'autres formats d'image.
from utils.io import prepare_input_imagedata_dir = 'focal stack images directory' # Chemin d'accès aux images de la pile focale dans un dossier # Charger les images de la pile focalefocal_stack, rgb_img, focus_dist = prepare_input_image(data_dir)# Exécuter l'inférence avec torch.no_grad() : out = model (rgb_img, focal_stack, focus_dist)metric_degree = out[0].squeeze().cpu().numpy() # La profondeur métrique
Veuillez d'abord télécharger les poids du modèle à partir des liens ci-dessous et les placer dans le répertoire checkpoints
:
HybridDepth_NYU_FocalStack5
HybridDepth_NYU_FocalStack10
HybridDepth_DDFF_FocalStack5
HybridDepth_NYU_PretrainedDFV_FocalStack5
NYU Depth V2 : Téléchargez l'ensemble de données en suivant les instructions fournies ici.
DDFF12 : Téléchargez l'ensemble de données en suivant les instructions fournies ici.
ARKitScenes : Téléchargez l'ensemble de données en suivant les instructions fournies ici.
Configurez le fichier de configuration config.yaml
dans le répertoire configs
. Les fichiers préconfigurés pour chaque ensemble de données sont disponibles dans le répertoire configs
, où vous pouvez spécifier des chemins, des paramètres de modèle et d'autres hyperparamètres. Voici un exemple de configuration :
data: class_path: dataloader.dataset.NYUDataModule # Chemin d'accès à votre module dataloader dans dataset.py init_args:nyuv2_data_root : "path/to/NYUv2" # Chemin d'accès au datasetimg_size spécifique : [480, 640] # Ajuster en fonction des exigences de votre DataModuleremove_white_border : Truenum_workers : 0 # Définir sur 0 si vous utilisez des données synthétiquesuse_labels : Truemodel : invert_degree : True # Définir à True si le modèle génère un deepckpt_path inversé : points de contrôle/checkpoint.ckpt
Spécifiez le fichier de configuration dans le script test.sh
:
python cli_run.py test --config configs/config_file_name.yaml
Ensuite, exécutez l'évaluation avec :
scripts de CD sh évaluer.sh
Installez le package basé sur CUDA requis pour la synthèse d'images :
installation de python utils/synthetic/gauss_psf/setup.py
Cela installe le package nécessaire à la synthèse des images.
Configurez le fichier de configuration config.yaml
dans le répertoire configs
, en spécifiant le chemin de l'ensemble de données, la taille du lot et d'autres paramètres de formation. Vous trouverez ci-dessous un exemple de configuration pour l'entraînement avec l'ensemble de données NYUv2 :
modèle : invert_profondeur : Vrai # taux d'apprentissage gd: 3e-4 # Ajuster si nécessaire # perte de poids wd: 0.001 # Ajustez selon les besoins data: class_path: dataloader.dataset.NYUDataModule # Chemin d'accès à votre module dataloader dans dataset.py init_args:nyuv2_data_root : "path/to/NYUv2" # Ensemble de données pathimg_size : [480, 640] # Ajuster pour NYUDataModuleremove_white_border : Truebatch_size : 24 # Ajuster en fonction de la mémoire disponiblenum_workers : 0 # Définir sur 0 si vous utilisez des données synthétiquesuse_labels : Trueckpt_path : null
Spécifiez le fichier de configuration dans le script train.sh
:
python cli_run.py train --config configs/config_file_name.yaml
Exécutez la commande d'entraînement :
scripts de CD sh train.sh
Si notre travail vous aide dans vos recherches, veuillez le citer comme suit :
@misc{ganj2024hybriddegreerobustmetricdegree, title={HybridDepth : Fusion de profondeur métrique robuste en tirant parti de la profondeur de la mise au point et des priorités à image unique}, author={Ashkan Ganj et Hang Su et Tian Guo}, année={2024}, eprint={2407.18443} , archivePrefix={arXiv}, PrimaryClass={cs.CV}, url={https://arxiv.org/abs/2407.18443}, }