[Edição 3D] SPIn-NeRF: Segmentação multivisualização e reparo perceptivo (CVPR 2023)


Página inicial do projeto: spinnerf3d.github.io (incluindo artigos e códigos)

Resumo

编辑任务:从3D场景中移除不需要的对象

O NeRF tornou-se um novo método popular de síntese de visualizações . No entanto, editar e manipular cenas NeRF continua sendo um desafio. Uma das tarefas de edição importantes é remover objetos indesejados da cena 3D e corrigi-los para que a área de substituição seja visualmente plausível e consistente com seu contexto.

Nosso método: Dado um conjunto de anotações esparsas com imagens de pose e uma única imagem de entrada , obtenha primeiro a máscara de segmentação 3D do objeto de destino . Em seguida, introduzimos uma abordagem baseada em otimização perceptiva que utiliza o modelo de pintura de imagem 2D aprendido para traduzir suas informações em espaço 3D e garantir a consistência do ponto de vista.

O artigo apresenta um conjunto de dados de cenas reais para avaliar o método de pintura de cenas 3D.


提示:以下是本篇文章正文内容,下面案例可供参考

1. Introdução

Com o uso generalizado de nerfs, a necessidade de editar e manipular cenas nerf continuará a crescer, incluindo remoção de objetos e pintura de imagens de cenas 3D, semelhante à bem estudada tarefa de pintura de imagens 2D [23]. No entanto, o obstáculo não está apenas no próprio processo de pintura 3D, mas também na obtenção da máscara de segmentação de entrada . Primeiro, a cena NeRF é codificada implicitamente nos pesos do mapa neural, levando a uma representação emaranhada e não interpretável, que não é fácil de manipular. Além disso, a pintura de cenas 3D deve gerar uma aparência perceptivamente realista em uma determinada visualização e também deve preservar as propriedades 3D fundamentais, como consistência de aparência e plausibilidade geométrica entre visualizações . Este método precisa anotar várias imagens (e manter várias exibições consistentes), o que é um fardo para os usuários .

Uma alternativa atraente é esperar apenas um conjunto mínimo de anotações para uma única exibição . Isso motiva um método capaz de obter máscaras de segmentação 3D consistentes com a visão de objetos a partir de anotações esparsas de visão única.

Nosso método: primeiro a partir de um pequeno número de pontos de imagem definidos pelo usuário em um objeto alvo (e algumas amostras negativas externas), o algoritmo usa um modelo baseado em vídeo [4] para inicializar a máscara e ajustando a semântica NeRF [ 36,76,77] para levantá-lo em uma segmentação 3D coerente . Então, ao aplicar o pintor 2D pré-treinado [48] ao conjunto de imagens multi-view, um procedimento de ajuste NeRF personalizado é usado para reconstruir a cena pintada em 3D, usando uma perda de percepção [72] responsável pelas inconsistências na pintura interna 2D.

insira a descrição da imagem aqui

2. Trabalho relacionado

1. Pintura interna da imagem

As primeiras técnicas dependiam de patches [9, 16] e as abordagens neurais mais recentes otimizam a implementação e a reconstrução da percepção (por exemplo, [22, 27, 48]). Para melhorar a fidelidade visual, várias direções de pesquisa continuam a ser exploradas, incluindo treinamento adversário ([40,73]), melhorias de estrutura ([27,30,62,63]), saída multivariada ([73, 75]), multi- processamento de escala ([21, 58, 69]) e métricas perceptivas (como [22, 48, 72]).

Nosso trabalho utiliza o LaMa pré-treinado [48] , que aplica uma transformada no domínio da frequência [7] inspirada no Transformer [11, 12]. No entanto, nada disso levanta o problema para imagens 3D; portanto, pintar várias imagens de uma cena de maneira consistente continua sendo uma tarefa pouco explorada. Nosso método opera diretamente em 3D por meio de um modelo NeRF baseado em múltiplas visualizações.

2. Operação NeRF

Com base na renderização de volume diferenciável [18, 52] e codificação posicional [13, 49, 53], os NeRFs [35] mostraram um desempenho notável na síntese de novas visualizações. Controlar a cena NeRF continua sendo um desafio. Clip-NeRF [55], Object-NeRF [64], LaTeRF [36] e outros [25, 32, 70] introduzem métodos para modificar e completar objetos representados por NeRFs; no entanto, seu desempenho é limitado a objetos simples, em vez do que cenas com confusão e textura significativas, ou se concentram em outras tarefas além da pintura interna geral (por exemplo, repintar ou deformar).

Mais intimamente relacionado com a nossa abordagem é o NeRF-In [31], um trabalho inédito contemporâneo que usa prioris de geometria e brilho para imagens 2D, mas não aborda inconsistências; em vez disso, ele usa uma simples perda de pixel, para reduzir o número de visualizações usadas para encaixe, reduzindo assim a qualidade da síntese da vista final . Semelhante ao uso de perda de pixel, o modelo contemporâneo Remove-NeRF [60] reduz a inconsistência excluindo visualizações com base em mecanismos incertos . Em contraste, nosso método é capaz de mapear representações NeRF de cenas desafiadoras do mundo real, incorporando informações 2D de maneira consistente com a visão por meio de um mecanismo de treinamento perceptivo [72]. Isso evita a pintura interna excessivamente restritiva (que geralmente leva ao desfoque).

3. Histórico: Conhecimento NeRF

Consulte o blog [reconstrução 3D] princípio NeRF + explicação do código


3. Método

Dado um conjunto de imagens RGB, I = {I i } n i=1 , e a pose 3D correspondente, g = {G i } n i=1 , e a matriz intrínseca da câmera, K, o modelo espera uma distribuição esparsa do usuário. vista "Fonte" anotada. A partir dessas entradas, geramos um modelo NeRF da cena, capaz de sintetizar uma imagem pintada a partir de qualquer novo ponto de vista.

Primeiro obtemos uma máscara 3D inicial de uma fonte de anotação de visão única (Seção 4.1.1) e, em seguida, ajustamos um NeRF semântico para melhorar a consistência e a qualidade da máscara (Seção 4.1.2). Em 4.2, descrevemos o método inpaint consistente com a exibição, que recebe como entrada uma exibição e uma máscara restaurada. Nosso método utiliza a saída do inpainter 2D [48] como prévias de aparência e geometria para supervisionar a adaptação de um novo NeRF.

3.1. Segmentação multivisualização

3.1.1 Inicialização da máscara

Denomine a visualização anotada do código-fonte como I 1 . Alimente as informações esparsas de objetos e exibições de origem para um modelo de segmentação interativo Edgeflow15] para estimar a máscara de objeto de origem inicial M 1 ^ \hat{M~1~}M1_ _  ^ . Em seguida, pegue a exibição de treinamento como uma sequência de vídeo, comM 1 ^ \hat{M~1~}M1_ _  ^ juntos fornecem um modelo de segmentação de instância de vídeo V [4, 57] para calcularinsira a descrição da imagem aqui, obter a máscara para cada exibição. A máscara inicial,{cMi} ni=1
geralmente é imprecisa perto da borda porque as exibições de treinamento não são quadros de vídeo adjacentes. Portanto, usamos o modelo semântico NeRF [36, 76, 77] para resolver a inconsistência e melhorar a máscara (4.1.2), obtendo assim a máscara de cada visualização de entrada para pintura interna (4.2).

3.1.2 Segmentação baseada em nerf

A entrada do módulo de segmentação de múltiplas visualizações é: imagem RGB, parâmetros intrínsecos e extrínsecos da câmera correspondente e máscara inicial, treine um NeRF semântico [76, consulte a Figura 2] para a rede; para o ponto x e a direção da visualização d, exceto a densidade σ (x) e cor c(x,d), que retorna um valor lógico s(x) de "objetividade" . Em seguida, obtenha a probabilidade alvo p(x) = Sigmoid(s(x)) . Usamos o Instant-NGP de rápida convergência [3, 37, 38] como a arquitetura NeRF.
insira a descrição da imagem aqui

O valor lógico de objetividade desejado associado ao raio r , obtido pela renderização da densidade de pontos em r, em vez da cor:

insira a descrição da imagem aqui

Para simplificar, s(r(t i )) é denotado por si . Em seguida, use a perda de classificação para supervisionar um raio, a probabilidade de objetidade ,Pb(r)=Sigmoide
insira a descrição da imagem aqui

Ao treinar Lcls , a cor é separada para limitar a atualização supervisionada dos valores de logits . Isso impedirá a alteração da geometria existente, pois as atualizações de gradiente alterariam o valor de densidade σ.

insira a descrição da imagem aqui
Após a otimização, uma máscara de consistência 3D {Mi} n i=1 é obtida limitando a probabilidade alvo e mascarando pixels com probabilidade maior que 0,5 . Por fim, adotamos uma otimização em dois estágios para melhorar ainda mais a máscara: após obter a máscara 3D inicial, a máscara é renderizada a partir da visualização de treinamento e usada para supervisionar o modelo de segmentação multivisualização quadrática como estimativa inicial.

3.2. Pintura interna de múltiplas vistas

A Figura 3 mostra uma visão geral do Inpainting NeRF consistente com a exibição, treinado com a seguinte perda:
insira a descrição da imagem aqui

Entre eles, L' rec é a perda de reconstrução de pixels não mascarados, L LPIPS e L depth definem a percepção e perda de profundidade, e o peso é λ.

insira a descrição da imagem aqui

3.2.1 RGB anterior

O Inpaint NeRF usa a entrada RGB, os parâmetros intrínsecos e extrínsecos da câmera e a máscara de objeto correspondente para ajustar um NeRF com o objeto mascarado removido da cena. Primeiro, alimente cada par de imagem e máscara ( I i , M i ) para o INP do pintor de imagens:insira a descrição da imagem aqui

Como cada exibição é pintada de forma independente e usada diretamente para supervisionar o NeRF, isso levará a resultados borrados devido a inconsistências 3D (consulte a Figura 7). Neste artigo, em vez de usar o erro quadrático médio (MSE) para otimizar a área mascarada, usamos a perda perceptual L PIPS [72] para otimizar a parte mascarada da imagem enquanto ainda usamos o MSE para otimizar a parte não mascarada . A perda é calculado da seguinte forma:
insira a descrição da imagem aqui

Entre eles, B é o índice de lote entre 1 e n, I ^ \hat{I}EU^ ié a i-ésima vista renderizada usando NeRF. O Ipaint NeRF de múltiplas visualizações e o NeRF semântico do modelo de segmentação usam a mesma arquitetura (consulte a Figura 2)

3.2.2 Profundidade anterior

Mesmo com perda de percepção, as diferenças entre as visualizações do Inpaint podem guiar incorretamente o modelo para convergir para a geometria (por exemplo, a geometria "névoa" pode se formar perto da câmera para explicar informações diferentes de cada visualização). Portanto, usamos o mapa de profundidade do inpaint como um guia adicional para o modelo NeRF e separamos os pesos ao calcular a perda perceptiva e usamos a perda perceptiva para ajustar apenas a cor da cena . Para isso, usamos um NeRF otimizado para imagens contendo objetos indesejados e renderizamos os mapas de profundidade correspondentes às visualizações de treinamento. O mapa de profundidade é calculado substituindo a Equação 1 pela distância até a câmera em vez da cor do ponto:
insira a descrição da imagem aqui
Dê a profundidade renderizada a um inpainter para obter o mapa de profundidade desenhado:insira a descrição da imagem aqui

A pintura profunda usando LaMa [48], com resultados de qualidade suficientemente alta, é computada como uma etapa de pré-processamento. Esses mapas de profundidade são usados ​​para supervisionar a geometria NeRF desenhada, renderizando a distância L2 de profundidade:
insira a descrição da imagem aqui

3.2.3 Otimização baseada em patches

O cálculo da perda perceptiva (Equação 7) requer a apresentação da visão de entrada completa durante a otimização. Como a renderização de cada pixel requer várias passagens pelo MLP, esse é um processo caro para imagens de alta resolução, levando a problemas como (i) o tamanho do lote |B| deve ser pequeno, reduzindo o gráfico de computação na memória correspondente; (ii) Otimização lenta, mesmo com tamanho de lote |B| = 1.

Uma solução simples é renderizar uma imagem reduzida e compará-la com a imagem inpaint reduzida; no entanto, isso pode levar à perda de informações se o fator de redução de escala for grande. Com base no trabalho de imagem (por exemplo, SinGAN [46] e DPNN [14]) e trabalho 3D (por exemplo, ARF [71]), realizamos cálculos com base em patches; em vez de renderizar visualizações completas, renderizamos imagens menores e as comparamos com os patches correspondentes na imagem inpaint de acordo com a perda perceptiva.

Para ajustar regiões não mascaradas, L' rec (Eq 6) amostra raios apenas de pixels não mascarados . Ao separar as perdas perceptivas e de reconstrução, evitamos inconsistências dentro da máscara, ao mesmo tempo em que evitamos alterações desnecessárias no restante da cena.

3.2.4 Refinamento da máscara

Aqui, consideramos explorar ainda mais os dados de múltiplas visualizações para orientar a pintura interna da imagem. Em particular, partes das imagens de treinamento atualmente geradas pelo pintor interno 2D podem ser visíveis em outras visualizações; nesse caso, não há necessidade de alucinar esses detalhes, pois eles podem ser recuperados de outras visualizações. Para evitar essa pintura desnecessária, propomos um método de refinamento de máscara:

Para cada imagem de origem, profundidade e tupla de máscara (I s , D s , M s ), substituímos os pixels em I e d que são visíveis em pelo menos uma outra visualização para reduzir a máscara de origem , após esta etapa de refinamento , apenas partes I e d que estão obstruídos por objetos indesejados em todas as visualizações de treinamento permanecerão mascarados. Portanto, o inpainter tem que preencher uma área menor, melhorando assim o inpaint.


4. Experimente

insira a descrição da imagem aqui

Comparações qualitativas são feitas de nosso modelo de segmentação multi-view com Neural Volumetric Object Selection (NVOS) [44]), segmentação de vídeo [4] e máscaras anotadas por humanos (GT). Nossa abordagem de dois estágios refere-se à execução de nosso modelo de segmentação de múltiplas visualizações duas vezes, com a saída da primeira execução servindo como inicialização para a segunda execução (consulte 4.1.2). Nosso método é menos ruidoso que o NVOS, que também perde alguns segmentos do objeto (por exemplo, a linha mais baixa de flores na linha inferior), mas captura mais detalhes (por exemplo, bordas borradas de flores) do que apenas a segmentação de vídeo. Nossa abordagem de dois estágios ajuda a preencher alguns aspectos ausentes da saída de estágio único.

insira a descrição da imagem aqui

Observe que as visualizações compostas permanecem consistentes umas com as outras; no entanto, ainda existem efeitos dependentes da visualização (por exemplo, a iluminação das partes nuas do piano)

insira a descrição da imagem aqui

por NeRF (em imagens não mascaradas), NeRF-In, NeRF-In (com imagens de treinamento mascaradas únicas) e nosso método. O NeRF-In é visivelmente mais desfocado, enquanto o NeRF-In (único) tende a sofrer com detalhes mais próximos da borda da borda da máscara

5. Instalação e explicação do código

1. Instalação do projeto

Clone o projeto primeiro e depois instale as dependências

git clone
pip install -r requirements.txt
pip install -r lama/requirements.txt

Em seguida, prepare os dados, o seguinte é o link de dados oficial:

https://drive.google.com/drive/folders/1N7D4-6IutYD40v9lfXGSVbWrd47UdJEC?usp=share_link

Baixe o pacote compactado da estátua e descompacte-o com o seguinte comando

unzip statue.zip -d data

Seu formato de dados é o seguinte (imagens2 é uma resolução maior e o rótulo contém a máscara do alvo a ser removido):

statue
├── images
│   ├── IMG_2707.jpg
│   ├── IMG_2708.jpg
│   ├── ...
│   └── IMG_2736.jpg
└── images_2
    ├── IMG_2707.png
    ├── IMG_2708.png
    ├── ...
    ├── IMG_2736.png
    └── label
        ├── IMG_2707.png
        ├── IMG_2708.png
        ├── ...
        └── IMG_2736.png

Se você deseja treinar seus próprios dados, pode instalar o colmap e usar o seguinte código para encontrar as informações de pose:

python imgs2poses.py <your_datadir>

2. Execute o NeRF original para obter as informações de profundidade:

rm -r LaMa_test_images/*  
rm -r output/label/*
python DS_NeRF/run_nerf.py --config DS_NeRF/configs/config.txt --render_factor 1 --prepare --i_weight 1000000000 --i_video 1000000000 --i_feat 4000 --N_iters 4001 --expname statue --datadir ./data/statue --factor 2 --N_gt 0

Posteriormente, as disparidades renderizadas (anti-profundidade) estarão na lama/LaMa_test_imagespasta e a tag de máscara de primeiro plano associada estará no arquivo lama/LaMa_test_images/label.

3. Execute o LaMa para gerar guias geométricas e aparentes

O LaMa combina principalmente a máscara, remove o alvo no mapa de profundidade e no mapa RGB e o preenche com o plano de fundo para obter uma imagem pintada muito realista.

cd lama

export TORCH_HOME=$(pwd) && export PYTHONPATH=$(pwd)
python bin/predict.py refine=True model.path=$(pwd)/big-lama indir=$(pwd)/LaMa_test_images outdir=$(pwd)/output

disparidades pintadas (mapa de profundidade modificado) em lama/output/label. Copie e coloque em data/statue/images_2/depth., o código é o seguinte:

dataset=statue
factor=2
rm -r ../data/$dataset/images_$factor/depth
mkdir ../data/$dataset/images_$factor/depth
cp ./output/label/*.png ../data/$dataset/images_$factor/depth

Em seguida, gere a imagem RGB pintada:

dataset=statue
factor=2

rm -r LaMa_test_images/*
rm -r output/label/*
cp ../data/$dataset/images_$factor/*.png LaMa_test_images
mkdir LaMa_test_images/label
cp ../data/$dataset/images_$factor/label/*.png LaMa_test_images/label
python bin/predict.py refine=True model.path=$(pwd)/big-lama indir=$(pwd)/LaMa_test_images outdir=$(pwd)/output
rm -r ../data/$dataset/images_$factor/lama_images
mkdir ../data/$dataset/images_$factor/lama_images
cp ../data/$dataset/images_$factor/*.png ../data/$dataset/images_$factor/lama_images
cp ./output/label/*.png ../data/$dataset/images_$factor/lama_images

A imagem RGB pintada está em lama/output/label, e copiada para data/statue/images_2/lama_images.

4. Execute o NeRF do pintor multivisualização

Usando o seguinte comando, a última otimização NeRF pintada será desdobrada. O vídeo do resultado da NeRF pintada será salvo todas as vezes i_video. Número total de vezes de treinamento N_iters.

python DS_NeRF/run_nerf.py --config DS_NeRF/configs/config.txt --i_feat 200 --lpips --i_weight 1000000000000 --i_video 1000 --N_iters 10001 --expname statue --datadir ./data/statue --N_gt 0 --factor $factor

5. Código-chave

  1. Leia os dados, pré-processe, obtenha algum GroundTruth

use_batching=True, significa sobrepor todos os dados da imagem juntos e amostrar os raios deles;

# 0.首先得到 colmap预处理得到的深度信息----------------------------------------------------------
if args.dataset_type == 'llff':
    if args.colmap_depth:
        depth_gts = load_colmap_depth( args.datadir, factor=args.factor, bd_factor=.75, prepare=args.prepare)  # n_img个list,每一个包含: depth(3879) coord(3879,2) weight(3879)归一化
    images, poses, bds, render_poses, i_test, masks, inpainted_depths, mask_indices = load_llff_data(args.datadir,  args.factor,  recenter=True,  bd_factor=.75,  spherify=args.spherify, prepare=args.prepare, args=args)     
    # mask代表前景掩码


H, W, focal = hwf 

# 1.创建 nerf model---------------------------------------------------------------------------
render_kwargs_train, render_kwargs_test, start, grad_vars, optimizer = create_nerf_tcnn(args)



# 2.得到 rays_rgb(射线o,d 加上img,mask)---------------------------------------------------------
labels = np.repeat(masks[:, None], 3, axis=1)  # [N, 3, H, W, 1]
rays = np.stack([get_rays_np(H, W, focal, p)
                   for p in poses[:, :3, :4]], 0)    # 维度[N, ro+rd, H, W, 3],比如 (29, 2, 378, 504, 3)
rays_rgb = np.concatenate([rays, images[:, None]], 1)    # (29, 2+1, 378, 504, 3)      ray融入了img,
rays_rgb = np.concatenate([rays_rgb, labels], -1)        # (29, 2+1, 378, 504, 3+1)    ray融入了img,mask

rays_rgb = np.reshape(rays_rgb, [-1, 3, 4])          # (29, 378, 504, 2+1, 3+1) --> (5524848, 3, 4)




# 3.得到 rays_inp (射线o,d 加上img,inpainted_depths)-----------------------------------------------
rays = np.stack([get_rays_np(H, W, focal, p)
                 for p in poses[:, :3, :4]], 0)  # [N, ro+rd, H, W, 3]  (29, 2, 378, 504, 3)
                 
labels = np.expand_dims(inpainted_depths, axis=-1)    # [N, H, W, 1]
labels = np.repeat(labels[:, None], 3, axis=1)        # [N, 3, H, W, 1]
      
rays_inp = np.concatenate([rays, images[:, None]], 1) # (29, 2+1, 378, 504, 3) ray融入了img
rays_inp = np.concatenate([rays_inp, labels], -1)     # (29, 2+1, 378, 504, 3+1) ray融入了img, inpainted_depths

rays_inp = np.reshape(rays_inp, [-1, 3, 4])           # (29, 378, 504, 2+1, 3+1) --> (5524848, 3, 4)




# 4.得到 rays_depth(colmap得到的4个深度值)-----------------------------------------------------------
if args.colmap_depth(True):
    for i in i_train(循环29张图):
        rays_depth = np.stack(get_rays_by_coord_np(H, W, focal, poses[i, :3, :4], depth_gts[i]['coord']),
                                      axis=0)  # 2 x N x 3
        # depth_gts[i]['coord']维度(27112),是图像上前景点的绝对坐标。这行代码,根据坐标得到对应的2711个射线
        # 具体函数定义如下:
        def get_rays_by_coord_np(H, W, focal, c2w, coords):
            i, j = (coords[:, 0] - W * 0.5) / focal, -(coords[:, 1] - H * 0.5) / focal  # (2711) (2711)
            dirs = np.stack([i, j, -np.ones_like(i)], -1)                               # (2711, 3) 
           rays_d = np.sum(dirs[..., np.newaxis, :] * c2w[:3, :3], -1)
           rays_o = np.broadcast_to(c2w[:3, -1], np.shape(rays_d))                      # c2w 最后一列是原点O
        return rays_o, rays_d                                                           # (2711, 3) (2711, 3) 


        depth_value = np.repeat(depth_gts[i]['depth'][:, None, None], 3, axis=2)   # (2711) 广播至(2711,1 ,3) 绝对深度
        weights = np.repeat( depth_gts[i]['weight'][:, None, None], 3, axis=2)     # (2711) 广播至(2711,1 ,3) 归一化权重
        rays_depth = np.concatenate([rays_depth, depth_value, weights], axis=1)    # (2711,4,3) 4分别代表 colmap估计的射线0+d,深度,权重
        rays_depth_list.append(rays_depth)
        
    rays_depth = np.concatenate(rays_depth_list, axis=0)      # (112377, 4 3) 29张图的前景射线,深度叠加在一起



# 5.得到 rays_rgb_clf 和 rays_inp--------------------------------------------------------------------
rays_rgb_clf = rays_rgb[rays_rgb[:, :, 3] == 0].reshape(-1, 3, 4)    # (5524848,3,4) -> (4827239,3,4) 只保留了mask==0的部分
rays_inp = rays_inp[rays_rgb[:, :, 3] != 0].reshape(-1, 3, 4)            # (5524848,3,4) -> (697609,3,4) 只保留了inpaint_mask!=0的部分

# 6.修改 rays_rgb
if not args.prepare:
   rays_rgb = rays_rgb[rays_rgb[:, :, 3] == 1].reshape(-1, 3, 4)        # (5524848,3,4) -> (23906,3,4)只保留了mask==0的部分



# 7.迭代产生 GroundTruth
raysRGB_iter = iter(DataLoader(RayDataset(rays_rgb), batch_size=N_rand, shuffle=True, num_workers=0,
                              generator=torch.Generator(device=device)))
raysINP_iter = iter(DataLoader(RayDataset(rays_inp), batch_size=N_rand, shuffle=True, num_workers=0,
                               generator=torch.Generator(device=device)))
raysRGBCLF_iter = iter(DataLoader(RayDataset(rays_rgb_clf), batch_size=N_rand, shuffle=True, num_workers=0,
                                  generator=torch.Generator(device=device)))

batch_inp = next(raysINP_iter).to(device)     # (1024,3,4)
batch = next(raysRGB_iter).to(device)         # (1024,3,4)
batch_clf = next(raysRGBCLF_iter).to(device)  # (1024,3,4)

# 7.1 target_s
batch_rays, target_s = batch[:2], batch[2]
target_s, label_s = target_s[:, :3], target_s[:, 3]           #(10243)(1024)射线的颜色和 mask(都是1)

# 7.2 target_inp
batch_inp, target_inp = batch_inp[:2], batch_inp[2]
target_inp, depth_inp = target_inp[:, :3], target_inp[:, 3]   #(10243)(1024)inpaint mask中的射线和深度

# 7.3 batch_clf
batch_rays_clf, target_clf = batch_clf[:2], batch_clf[2]      #(10243)(1024)只保留mask=0的前景 射线的颜色和 mask(都是0)

# 7.4 得到来自colmap的深度 gt
raysDepth_iter = iter(DataLoader(RayDataset(rays_depth), batch_size=N_rand, shuffle=True,  num_workers=0,   generator=torch.Generator(device=device))) if rays_depth is not None else None
batch_depth = next(raysDepth_iter).to(device)
batch_rays_depth = batch_depth[:2]        # 2 x B x 3
target_depth = batch_depth[2, :, 0]       # B=1024
ray_weights = batch_depth[3, :, 0]        # B=1024
  1. Leia a máscara do alvo a ser excluído e use cv2.dilate para expandir (principalmente considerando a máscara precisa, que pode não cobrir a sombra do alvo, etc.)
    for i, f in enumerate(mskfiles):    # mskfiles就是前景的mask,在label文件夹内
        try:
            msk = imread(f)
            msk = msk / msk.max()
            msk = cv2.resize( msk, (imgs.shape[1], imgs.shape[0]), interpolation=cv2.INTER_NEAREST)
                print(msk.shape)

            msk = cv2.dilate(msk, np.ones((5, 5), np.uint8), iterations=5)
            masks.append(msk)
            mask_indices.append(i)

2. A função NeRF_TCNN em DS_NeRF/run_nerf_helpers_tcnn.py mapeia xyz (posição) e d (direção do raio) para valores de cor e densidade:

Um total de 4 renderizações, cada entrada é um raio diferente:

rgb, disp, acc, depth, extras = render(H, W, focal, chunk=args.chunk, 
                                       rays=batch_rays_clf,   verbose=i < 10, retraw=True, **render_kwargs_train)  

rgb_complete, _, _, _, extras_complete = render(H, W, focal, chunk=args.chunk, 
                                               rays=batch_rays,  verbose=i < 10, retraw=True, detach_weights=True, **render_kwargs_train)

_, disp_inp, _, _, extras_inp = render(H, W, focal, chunk=args.chunk, 
                                      rays=batch_inp, verbose=i < 10, retraw=True,  **render_kwargs_train)                  # disp_inp (1024)预测每个点weight得到的深度图

_, _, _, depth_col, extras_col = render(H, W, focal, chunk=args.chunk, 
                                       rays=batch_rays_depth, verbose=i < 10, retraw=True, depths=target_depth,  **render_kwargs_train)             # depth_col (1024)      

O código de renderização é sempre o mesmo, o processo de encaminhamento específico é o seguinte

    def forward(self, input):
        x = input[:, :3]        # (655363)
        d = input[:, 3:]        # (655363)

        # sigma
        x = (x + self.bound) / (2 * self.bound)  # to [0, 1]  self.bound=100
        x = self.encoder(x)          # 16倍分辨率的hash表,得到(6553632)
        h = self.sigma_net(x)        # 得到(6553616)

        sigma = h[..., 0]
        geo_feat = h[..., 1:]

        # color
        d = (d + 1) / 2  # tcnn SH encoding requires inputs to be in [0, 1]
        d = self.encoder_dir(d)

        h = torch.cat([d, geo_feat], dim=-1)
        h = self.color_net(h)

        # sigmoid activation for rgb
        color = h                   # (65536,3)

        outputs = torch.cat([color, sigma[..., None]], -1)
        return outputs
  1. calcular perda
optimizer.zero_grad()
img_loss = img2mse(rgb, target_clf)    # 输入(1024,3) 输出:0.05 函数img2mse = lambda x, y: torch.mean((x - y) ** 2)
psnr = mse2psnr(img_loss)              # 函数 mse2psnr = lambda x: -10. * torch.log(x) / torch.log(torch.Tensor([10.]))


if not args.masked_NeRF and not args.object_removal:     # True
       img_loss += img2mse(rgb_complete, target_s)       # rgb_complete 是用batch_rays渲染出的颜色
       if 'rgb0' in extras_complete and not args.no_coarse:
                img_loss0 = img2mse(extras_complete['rgb0'], target_s)
                img_loss += img_loss0

# 3.计算深度损失
if args.depth_loss:
        if args.weighted_loss:
           depth_loss = img2mse(depth_col, target_depth)  # target_depth(1024),是colmap的深度

img_loss0 = img2mse(extras['rgb0'], target_clf)    # (10243)
loss = loss + img_loss0

A cada 300 iters, calcule a perda semântica LPIP


np.random.shuffle(idx)      # 一个batch中,随机抽取几张图片
random_poses = poses[idx]

# 按比例裁剪出patch。实验中缩放因子分别是28
patch_len = (hwf[0] / / lpips_render_factor / / patch_len_factor,
             hwf[1] / / lpips_render_factor / / patch_len_factor)


rgbs, disps, (Xs, Ys) = render_path(random_poses, hwf, args.chunk, render_kwargs_test,
                                    render_factor=lpips_render_factor,rgb_require_grad=True,need_alpha=False,detach_weights=True,patch_len=patch_len,masks=masks[idx])                    
# 得到 (4,23,31,3) (4,23,31) 以及随机选的batch个中心点坐标


prediction = ((rgbs[_] - 0.5) * 2).permute(2, 0, 1)[None, ...]            # (23,31,3)

target = ((images[idx[_]] - 0.5) * 2).permute(2, 0, 1)[None, ...]
target = transform(target)[:, :, Xs[_]:Xs[_] + patch_len[0], Ys[_]:Ys[_] + patch_len[1]]         #  (23,31,3)

# 计算损失(LPIPS就是预训练的vgg):
lpips_loss += LPIPS(prediction, target).mean()

Acho que você gosta

Origin blog.csdn.net/qq_45752541/article/details/132037526
Recomendado
Clasificación