Ashkan Ganj 1 · Hang Su 2 · Tian Guo 1
1 Instituto Politécnico de Worcester 2 Investigación de Nvidia
Lanzamos una versión mejorada de HybridDepth, ¡ahora disponible con nuevas funciones y rendimiento optimizado!
Este trabajo presenta HybridDepth. HybridDepth es una solución práctica de estimación de profundidad basada en imágenes de pila focal capturadas desde una cámara. Este enfoque supera a los modelos de última generación en varios conjuntos de datos conocidos, incluidos NYU V2, DDFF12 y ARKitScenes.
30/10/2024 : Lanzamiento de la versión 2 de HybridDepth con rendimiento mejorado y pesos previamente entrenados.
2024-10-30 : Soporte integrado para TorchHub para facilitar la carga e inferencia del modelo.
25/07/2024 : Lanzamiento inicial de modelos previamente entrenados.
23/07/2024 : El repositorio de GitHub y el modelo HybridDepth se pusieron en marcha.
Comience rápidamente con HybridDepth utilizando el cuaderno Colab.
Puede seleccionar un modelo previamente entrenado directamente con TorchHub.
Modelos preentrenados disponibles:
HybridDepth_NYU5
: entrenado previamente en el conjunto de datos NYU Depth V2 utilizando una entrada de pila de 5 focos, con la rama DFF y la capa de refinamiento entrenadas.
HybridDepth_NYU10
: entrenado previamente en el conjunto de datos NYU Depth V2 utilizando una entrada de pila de 10 focos, con la rama DFF y la capa de refinamiento entrenadas.
HybridDepth_DDFF5
: previamente entrenado en el conjunto de datos DDFF utilizando una entrada de pila de 5 focos.
HybridDepth_NYU_PretrainedDFV5
: Preentrenado solo en la capa de refinamiento con el conjunto de datos NYU Depth V2 usando una pila de 5 focales, luego del entrenamiento previo con DFV.
model_name = 'HybridDepth_NYU_PretrainedDFV5' #cambia este modelo = torch.hub.load('cake-lab/HybridDepth', model_name, pretrained=True)model.eval()
Clona el repositorio e instala las dependencias:
clon de git https://github.com/cake-lab/HybridDepth.gitcd HybridDepth conda env crear -f entorno.yml conda activar profundidad híbrida
Descargue pesas previamente entrenadas:
Descargue los pesos para el modelo desde los enlaces a continuación y colóquelos en el directorio de checkpoints
:
HybridDepth_NYU_FocalStack5
HybridDepth_NYU_FocalStack10
HybridDepth_DDFF_FocalStack5
HybridDepth_NYU_PretrainedDFV_FocalStack5
Predicción
Para inferencia, puede ejecutar el siguiente código:
# Cargar el modelo checkpointmodel_path = 'checkpoints/NYUBest5.ckpt'model = DepthNetModule.load_from_checkpoint(model_path)model.eval()model = model.to('cuda')
Después de cargar el modelo, use el siguiente código para procesar las imágenes de entrada y obtener el mapa de profundidad:
Nota : actualmente, la función prepare_input_image
solo admite imágenes .jpg
. Modifique la función si necesita compatibilidad con otros formatos de imagen.
from utils.io import prepare_input_imagedata_dir = 'directorio de imágenes de pila focal' # Ruta a las imágenes de pila focal en una carpeta# Cargue las imágenes de pila focalfocal_stack, rgb_img, focus_dist = prepare_input_image(data_dir)# Ejecute inferencia con torch.no_grad(): out = model (rgb_img, focal_stack, focus_dist)metric_profundidad = out[0].squeeze().cpu().numpy() # La profundidad métrica
Primero descargue los pesos para el modelo desde los enlaces a continuación y colóquelos en el directorio de checkpoints
:
HybridDepth_NYU_FocalStack5
HybridDepth_NYU_FocalStack10
HybridDepth_DDFF_FocalStack5
HybridDepth_NYU_PretrainedDFV_FocalStack5
NYU Depth V2 : descargue el conjunto de datos siguiendo las instrucciones proporcionadas aquí.
DDFF12 : descargue el conjunto de datos siguiendo las instrucciones proporcionadas aquí.
ARKitScenes : descargue el conjunto de datos siguiendo las instrucciones proporcionadas aquí.
Configure el archivo de configuración config.yaml
en el directorio configs
. Los archivos preconfigurados para cada conjunto de datos están disponibles en el directorio configs
, donde puede especificar rutas, configuraciones de modelo y otros hiperparámetros. Aquí hay una configuración de ejemplo:
datos: class_path: dataloader.dataset.NYUDataModule # Ruta a su módulo de carga de datos en dataset.py init_args:nyuv2_data_root: "path/to/NYUv2" # Ruta a los datos específicos setimg_size: [480, 640] # Ajuste según los requisitos de su DataModuleremove_white_border: Truenum_workers: 0 # Establezca en 0 si usa datos sintéticosuse_labels: Truemodel: invert_ Depth: True # Establezca a True si el modelo genera Depthckpt_path invertido: puntos de control/checkpoint.ckpt
Especifique el archivo de configuración en el script test.sh
:
prueba python cli_run.py --config configs/config_file_name.yaml
Luego, ejecute la evaluación con:
guiones de cd sh evaluar.sh
Instale el paquete basado en CUDA necesario para la síntesis de imágenes:
instalación de python utils/synthetic/gauss_psf/setup.py
Esto instala el paquete necesario para sintetizar imágenes.
Configure el archivo de configuración config.yaml
en el directorio configs
, especificando la ruta del conjunto de datos, el tamaño del lote y otros parámetros de entrenamiento. A continuación se muestra una configuración de ejemplo para entrenar con el conjunto de datos NYUv2:
modelo: invert_profundidad: Verdadero # tasa de aprendizaje lr: 3e-4 # Ajustar según sea necesario # caída de peso wd: 0.001 # Ajuste según sea necesariodata: class_path: dataloader.dataset.NYUDataModule # Ruta a su módulo de carga de datos en dataset.py init_args:nyuv2_data_root: "path/to/NYUv2" # Conjunto de datos pathimg_size: [480, 640] # Ajustar para NYUDataModuleremove_white_border: Truebatch_size: 24 # Ajustar según la memoria disponiblenum_workers: 0 # Establecer en 0 si usa datos sintéticosuse_labels: Trueckpt_path: null
Especifique el archivo de configuración en el script train.sh
:
python cli_run.py tren --config configs/config_file_name.yaml
Ejecute el comando de entrenamiento:
guiones de cd sh tren.sh
Si nuestro trabajo le ayuda en su investigación, cítelo de la siguiente manera:
@misc{ganj2024hybrid Depthrobustmetric Depth, title={HybridDepth: fusión robusta de profundidad métrica aprovechando la profundidad del enfoque y los antecedentes de una sola imagen}, autor={Ashkan Ganj y Hang Su y Tian Guo}, año={2024}, eprint={2407.18443} , archivePrefix={arXiv}, clase primaria={cs.CV}, URL={https://arxiv.org/abs/2407.18443}, }