Ashkan Ganj 1 · Hang Su 2 · Tian Guo 1
1 Instituto Politécnico de Worcester 2 Pesquisa Nvidia
Lançamos uma versão melhorada do HybridDepth, agora disponível com novos recursos e desempenho otimizado!
Este trabalho apresenta o HybridDepth. HybridDepth é uma solução prática de estimativa de profundidade baseada em imagens de pilha focal capturadas de uma câmera. Essa abordagem supera os modelos de última geração em vários conjuntos de dados conhecidos, incluindo NYU V2, DDFF12 e ARKitScenes.
30/10/2024 : Lançada a versão 2 do HybridDepth com desempenho aprimorado e pesos pré-treinados.
30/10/2024 : Suporte integrado para TorchHub para fácil carregamento e inferência de modelos.
25/07/2024 : Lançamento inicial de modelos pré-treinados.
23/07/2024 : O repositório GitHub e o modelo HybridDepth foram lançados.
Comece rapidamente a usar o HybridDepth usando o notebook Colab.
Você pode selecionar um modelo pré-treinado diretamente com o TorchHub.
Modelos pré-treinados disponíveis:
HybridDepth_NYU5
: pré-treinado no conjunto de dados NYU Depth V2 usando uma entrada de pilha focal 5, com a ramificação DFF e a camada de refinamento treinadas.
HybridDepth_NYU10
: pré-treinado no conjunto de dados NYU Depth V2 usando uma entrada de pilha focal 10, com a ramificação DFF e a camada de refinamento treinadas.
HybridDepth_DDFF5
: pré-treinado no conjunto de dados DDFF usando uma entrada de pilha focal 5.
HybridDepth_NYU_PretrainedDFV5
: Pré-treinado apenas na camada de refinamento com conjunto de dados NYU Depth V2 usando uma pilha 5-focal, após pré-treinamento com DFV.
model_name = 'HybridDepth_NYU_PretrainedDFV5' #change thismodel = torch.hub.load('cake-lab/HybridDepth', model_name , pretrained=True)model.eval()
Clone o repositório e instale as dependências:
clone git https://github.com/cake-lab/HybridDepth.gitcd HybridDepth conda env create -f ambiente.yml conda ativar profundidade híbrida
Baixe pesos pré-treinados:
Baixe os pesos do modelo nos links abaixo e coloque-os no diretório de checkpoints
:
HybridDepth_NYU_FocalStack5
HybridDepth_NYU_FocalStack10
HybridDepth_DDFF_FocalStack5
HybridDepth_NYU_PretrainedDFV_FocalStack5
Previsão
Para inferência, você pode executar o seguinte código:
# Carregue o modelo checkpointmodel_path = 'checkpoints/NYUBest5.ckpt'model = DepthNetModule.load_from_checkpoint(model_path)model.eval()model = model.to('cuda')
Após carregar o modelo, use o seguinte código para processar as imagens de entrada e obter o mapa de profundidade:
Nota : Atualmente, a função prepare_input_image
suporta apenas imagens .jpg
. Modifique a função se precisar de suporte para outros formatos de imagem.
from utils.io import prepare_input_imagedata_dir = 'diretório de imagens da pilha focal' # Caminho para as imagens da pilha focal em uma pasta# Carregue as imagens da pilha focalfocal_stack, rgb_img, focus_dist = prepare_input_image(data_dir)# Execute a inferência com torch.no_grad(): out = model (rgb_img, focal_stack, focus_dist)metric_profundidade = out[0].squeeze().cpu().numpy() # A profundidade da métrica
Primeiro baixe os pesos para o modelo nos links abaixo e coloque-os no diretório de checkpoints
:
HybridDepth_NYU_FocalStack5
HybridDepth_NYU_FocalStack10
HybridDepth_DDFF_FocalStack5
HybridDepth_NYU_PretrainedDFV_FocalStack5
NYU Depth V2 : Baixe o conjunto de dados seguindo as instruções fornecidas aqui.
DDFF12 : Baixe o conjunto de dados seguindo as instruções fornecidas aqui.
ARKitScenes : Baixe o conjunto de dados seguindo as instruções fornecidas aqui.
Configure o arquivo de configuração config.yaml
no diretório configs
. Arquivos pré-configurados para cada conjunto de dados estão disponíveis no diretório configs
, onde você pode especificar caminhos, configurações de modelo e outros hiperparâmetros. Aqui está um exemplo de configuração:
data: class_path: dataloader.dataset.NYUDataModule # Caminho para o módulo dataloader em dataset.py init_args:nyuv2_data_root: "path/to/NYUv2" # Caminho para o datasetimg_size específico: [480, 640] # Ajuste com base nos requisitos do DataModuleremove_white_border: Truenum_workers: 0 # Defina como 0 se estiver usando datause_labels sintéticos: Truemodel: invert_profundidade: True # Definir como True se o modelo gerar profundidade invertida: pontos de verificação/checkpoint.ckpt
Especifique o arquivo de configuração no script test.sh
:
teste python cli_run.py --config configs/config_file_name.yaml
Em seguida, execute a avaliação com:
scripts de CD sh avaliar.sh
Instale o pacote necessário baseado em CUDA para síntese de imagem:
instalação de utilitários python/synthetic/gauss_psf/setup.py
Isto instala o pacote necessário para sintetizar imagens.
Configure o arquivo de configuração config.yaml
no diretório configs
, especificando o caminho do conjunto de dados, o tamanho do lote e outros parâmetros de treinamento. Abaixo está um exemplo de configuração para treinamento com o conjunto de dados NYUv2:
modelo: profundidade_invertida: Verdadeiro # taxa de aprendizagem lr: 3e-4 # Ajuste conforme necessário # perda de peso wd: 0.001 # Ajuste conforme necessáriodata: class_path: dataloader.dataset.NYUDataModule # Caminho para seu módulo dataloader em dataset.py init_args:nyuv2_data_root: "path/to/NYUv2" # Dataset pathimg_size: [480, 640] # Ajuste para NYUDataModuleremove_white_border: Truebatch_size: 24 # Ajuste com base no memorynum_workers disponível: 0 # Defina como 0 se estiver usando datause_labels sintéticos: Trueckpt_path: null
Especifique o arquivo de configuração no script train.sh
:
python cli_run.py train --config configs/config_file_name.yaml
Execute o comando de treinamento:
scripts de CD sh trem.sh
Se nosso trabalho o ajudar em sua pesquisa, cite-o da seguinte forma:
@misc{ganj2024hybridprofundobustmetricprofundidade, title={HybridDepth: Robust Metric Depth Fusion aproveitando a profundidade do foco e prioridades de imagem única}, autor={Ashkan Ganj e Hang Su e Tian Guo}, ano={2024}, eprint={2407.18443} , archivePrefix={arXiv}, classe primária={cs.CV}, url={https://arxiv.org/abs/2407.18443}, }