+ We introduce a camera relocalization pipeline that combines absolute pose regression (APR) and direct feature matching. By incorporating exposure-adaptive novel view synthesis, our method successfully addresses photometric distortions in outdoor environments that existing photometric-based methods fail to handle. With domain-invariant feature matching, our solution improves pose regression accuracy using semi-supervised learning on unlabeled data. In particular, the pipeline consists of two components: Novel View Synthesizer and DFNet. The former synthesizes novel views compensating for changes in exposure and the latter regresses camera poses and extracts robust features that close the domain gap between real images and synthetic ones. Furthermore, we introduce an online synthetic data generation scheme. We show that these approaches effectively enhance camera pose estimation both in indoor and outdoor scenes. Hence, our method achieves a state-of-the-art accuracy by outperforming existing single-image APR methods by as much as 56%, comparable to 3D structure-based methods.
+
+
+
Supplementary Video
+
+
+
+
+
+
+
+
+
+
6-DOF Camera Pose Regression
+
+
+
+
+
+
+
+
+
+
+
Histogram-assisted NeRF
+
+
+
+
+
+
+
+
+
Citation
+
+ @inproceedings{chen2022dfnet,
+ title={{DFN}et: {E}nhance Absolute Pose Regression with Direct Feature Matching},
+ author={Chen, Shuai and Li, Xinghui and Wang, Zirui and Prisacariu, Victor},
+ booktitle={ECCV},
+ year={2022}
+ }
+
+
+
+
+
+
+
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index fd08bd8..0000000
--- a/requirements.txt
+++ /dev/null
@@ -1,107 +0,0 @@
-Package Version
------------------------ -------------------
-absl-py 0.14.0
-antlr4-python3-runtime 4.8
-backcall 0.2.0
-beautifulsoup4 4.9.3
-brotlipy 0.7.0
-cached-property 1.5.2
-cachetools 4.2.4
-certifi 2021.5.30
-cffi 1.14.5
-chardet 3.0.4
-conda 4.10.1
-conda-build 3.21.4
-conda-package-handling 1.7.3
-ConfigArgParse 1.5.2
-cryptography 3.4.7
-cycler 0.10.0
-decorator 5.0.9
-dnspython 2.1.0
-efficientnet-pytorch 0.7.1
-einops 0.3.2
-filelock 3.0.12
-fvcore 0.1.5.post20210924
-glob2 0.7
-google-auth 1.35.0
-google-auth-oauthlib 0.4.6
-grpcio 1.41.0
-h5py 3.5.0
-idna 2.10
-imageio 2.9.0
-imageio-ffmpeg 0.4.5
-importlib-metadata 4.8.1
-iopath 0.1.9
-ipython 7.22.0
-ipython-genutils 0.2.0
-jedi 0.17.0
-Jinja2 3.0.0
-kiwisolver 1.3.2
-kornia 0.6.2
-libarchive-c 2.9
-Markdown 3.3.4
-MarkupSafe 2.0.1
-matplotlib 3.3.2
-mkl-fft 1.3.0
-mkl-random 1.2.1
-mkl-service 2.3.0
-numpy 1.20.2
-oauthlib 3.1.1
-olefile 0.46
-omegaconf 2.1.1
-opencv-python 4.4.0.46
-packaging 21.3
-parso 0.8.2
-pexpect 4.8.0
-pickleshare 0.7.5
-Pillow 8.2.0
-pip 21.1.2
-pkginfo 1.7.0
-portalocker 2.3.2
-prompt-toolkit 3.0.17
-protobuf 3.18.0
-psutil 5.8.0
-ptyprocess 0.7.0
-pyasn1 0.4.8
-pyasn1-modules 0.2.8
-pycosat 0.6.3
-pycparser 2.20
-Pygments 2.9.0
-pykalman 0.9.5
-pyOpenSSL 19.1.0
-pyparsing 2.4.7
-PySocks 1.7.1
-python-dateutil 2.8.2
-python-etcd 0.4.5
-pytorch3d 0.3.0
-pytz 2021.1
-PyYAML 5.4.1
-requests 2.24.0
-requests-oauthlib 1.3.0
-rsa 4.7.2
-ruamel-yaml-conda 0.15.100
-scipy 1.7.3
-setuptools 52.0.0.post20210125
-six 1.16.0
-soupsieve 2.2.1
-tabulate 0.8.9
-tensorboard 2.6.0
-tensorboard-data-server 0.6.1
-tensorboard-plugin-wit 1.8.0
-termcolor 1.1.0
-torch 1.11.0+cu113
-torchelastic 0.2.0
-torchsummary 1.5.1
-torchtext 0.10.0
-torchvision 0.10.0
-tqdm 4.51.0
-traitlets 5.0.5
-transforms3d 0.3.1
-trimesh 3.9.32
-typing-extensions 3.7.4.3
-urllib3 1.25.11
-wcwidth 0.2.5
-Werkzeug 2.0.1
-wheel 0.35.1
-yacs 0.1.8
-zipp 3.6.0
diff --git a/script/config_dfnet.txt b/script/config_dfnet.txt
deleted file mode 100755
index a46c9f1..0000000
--- a/script/config_dfnet.txt
+++ /dev/null
@@ -1,51 +0,0 @@
-############################################### NeRF-Hist training example Cambridge ###############################################
-model_name=dfnet
-basedir=../logs/kings
-expname=nerfh
-datadir=../data/Cambridge/KingsCollege
-dataset_type=Cambridge
-trainskip=2 # train
-testskip=1 # train
-df=2
-load_pose_avg_stats=True
-NeRFH=True
-epochs=2000
-encode_hist=True
-tinyimg=True
-DFNet=True
-tripletloss=True
-featurenet_batch_size=4 # batch size, 4 or 8
-random_view_synthesis=True
-rvs_refresh_rate=20
-rvs_trans=3
-rvs_rotation=7.5
-d_max=1
-# pretrain_model_path = ../logs/kings/dfnet/checkpoint-0604-0.2688.pt # add your trained model for eval
-# eval=True # add this for eval
-
-############################################### NeRF-Hist training example 7-Scenes ###############################################
-# model_name=dfnet
-# basedir=../logs/heads
-# expname=nerfh
-# datadir=../data/7Scenes/heads
-# dataset_type=7Scenes
-# trainskip=5 # train
-# testskip=1 #train
-# df=2 # train
-# load_pose_avg_stats=True
-# NeRFH=True
-# epochs=2000
-# encode_hist=True
-# batch_size=1 # NeRF loader batch size
-# tinyimg=True
-# DFNet=True
-# tripletloss=True
-# featurenet_batch_size=4
-# val_batch_size=8 # new
-# random_view_synthesis=True
-# rvs_refresh_rate=20
-# rvs_trans=0.2
-# rvs_rotation=10
-# d_max=0.2
-# # pretrain_model_path = ../logs/heads/dfnet/checkpoint-0888-0.0025.pt # add your trained model for eval
-# # eval=True # add this for eval
\ No newline at end of file
diff --git a/script/config_dfnetdm.txt b/script/config_dfnetdm.txt
deleted file mode 100755
index 2582240..0000000
--- a/script/config_dfnetdm.txt
+++ /dev/null
@@ -1,56 +0,0 @@
-############################################### NeRF-Hist training example Cambridge ###############################################
-model_name=dfnetdm
-expname=nerfh
-basedir=../logs/kings # change this if change scenes
-datadir=../data/Cambridge/KingsCollege # change this if change scenes
-dataset_type=Cambridge
-pretrain_model_path=../logs/kings/dfnet/checkpoint-0604-0.2688.pt # this is your trained dfnet model for pose regression
-pretrain_featurenet_path=../logs/kings/dfnet/checkpoint-0604-0.2688.pt # this is your trained dfnet model for feature extraction
-trainskip=2 # train
-testskip=1 # train
-df=2
-load_pose_avg_stats=True
-NeRFH=True
-encode_hist=True
-freezeBN=True
-featuremetric=True
-pose_only=3
-svd_reg=True
-combine_loss = True
-combine_loss_w = [0., 0., 1.]
-finetune_unlabel=True
-i_eval=20
-DFNet=True
-val_on_psnr=True
-feature_matching_lvl = [0]
-# eval=True # add this for eval
-# pretrain_model_path=../logs/kings/dfnetdm/checkpoint-0267-17.1446.pt # add the trained model for eval
-
-
-############################################### NeRF-Hist training example 7-Scenes ###############################################
-# model_name=dfnetdm
-# expname=nerfh
-# basedir=../logs/heads
-# datadir=../data/7Scenes/heads
-# dataset_type=7Scenes
-# pretrain_model_path=../logs/heads/dfnet/checkpoint-0888-0.0025.pt # this is your trained dfnet model for pose regression
-# pretrain_featurenet_path=../logs/heads/dfnet/checkpoint-0888-0.0025.pt # this is your trained dfnet model for feature extraction
-# trainskip=5 # train
-# testskip=1 #train
-# df=2
-# load_pose_avg_stats=True
-# NeRFH=True
-# encode_hist=True
-# freezeBN=True
-# featuremetric=True
-# pose_only=3
-# svd_reg=True
-# combine_loss = True
-# combine_loss_w = [0., 0., 1.]
-# finetune_unlabel=True
-# i_eval=20
-# DFNet=True
-# val_on_psnr=True
-# feature_matching_lvl = [0]
-# # eval=True # add this for eval
-# # pretrain_model_path=../logs/heads/dfnetdm/checkpoint-0317-17.5881.pt # add the trained model for eval
\ No newline at end of file
diff --git a/script/config_nerfh.txt b/script/config_nerfh.txt
deleted file mode 100755
index 262ca23..0000000
--- a/script/config_nerfh.txt
+++ /dev/null
@@ -1,28 +0,0 @@
-############################################### NeRF-Hist training example Cambridge ###############################################
-expname=nerfh
-basedir=../logs/kings
-datadir=../data/Cambridge/KingsCollege
-dataset_type=Cambridge
-lrate_decay=5
-trainskip=2
-testskip=1
-df=4
-load_pose_avg_stats=True
-NeRFH=True
-encode_hist=True
-# render_test=True # add this for eval
-
-############################################### NeRF-Hist training example 7-Scenes ###############################################
-# expname=nerfh
-# basedir=../logs/heads
-# datadir=../data/7Scenes/heads
-# dataset_type=7Scenes
-# lrate_decay=0.754
-# trainskip=5
-# testskip=50
-# df=4
-# load_pose_avg_stats=True
-# NeRFH=True
-# encode_hist=True
-# # testskip=1 # add this for eval
-# # render_test=True # add this for eval
\ No newline at end of file
diff --git a/script/dm/__init__.py b/script/dm/__init__.py
deleted file mode 100755
index e69de29..0000000
diff --git a/script/dm/callbacks.py b/script/dm/callbacks.py
deleted file mode 100755
index 7e37a2e..0000000
--- a/script/dm/callbacks.py
+++ /dev/null
@@ -1,106 +0,0 @@
-import numpy as np
-import torch
-from tqdm import tqdm
-import os, pdb
-
-def Callback():
- #TODO: Callback func. https://dzlab.github.io/dl/2019/03/16/pytorch-training-loop/
- def __init__(self): pass
- def on_train_begin(self): pass
- def on_train_end(self): pass
- def on_epoch_begin(self): pass
- def on_epoch_end(self): pass
- def on_batch_begin(self): pass
- def on_batch_end(self): pass
- def on_loss_begin(self): pass
- def on_loss_end(self): pass
- def on_step_begin(self): pass
- def on_step_end(self): pass
-
-class EarlyStopping:
- """Early stops the training if validation loss doesn't improve after a given patience."""
- # source https://blog.csdn.net/qq_37430422/article/details/103638681
- def __init__(self, args, patience=50, verbose=False, delta=0):
- """
- Args:
- patience (int): How long to wait after last time validation loss improved.
- Default: 50
- verbose (bool): If True, prints a message for each validation loss improvement.
- Default: False
- delta (float): Minimum change in the monitored quantity to qualify as an improvement.
- Default: 0
- """
- self.val_on_psnr = args.val_on_psnr
- self.patience = patience
- self.verbose = verbose
- self.counter = 0
- self.best_score = None
- self.early_stop = False
- self.val_loss_min = np.Inf
- self.delta = delta
-
- self.basedir = args.basedir
- self.model_name = args.model_name
-
- self.out_folder = os.path.join(self.basedir, self.model_name)
- self.ckpt_save_path = os.path.join(self.out_folder, 'checkpoint.pt')
- if not os.path.isdir(self.out_folder):
- os.mkdir(self.out_folder)
-
- def __call__(self, val_loss, model, epoch=-1, save_multiple=False, save_all=False, val_psnr=None):
-
- # find maximum psnr
- if self.val_on_psnr:
- score = val_psnr
- if self.best_score is None:
- self.best_score = score
- self.save_checkpoint(val_psnr, model, epoch=epoch, save_multiple=save_multiple)
- elif score < self.best_score + self.delta:
- self.counter += 1
-
- if self.counter >= self.patience:
- self.early_stop = True
-
- if save_all: # save all ckpt
- self.save_checkpoint(val_psnr, model, epoch=epoch, save_multiple=True, update_best=False)
- else: # save best ckpt only
- self.best_score = score
- self.save_checkpoint(val_psnr, model, epoch=epoch, save_multiple=save_multiple)
- self.counter = 0
-
- # find minimum loss
- else:
- score = -val_loss
- if self.best_score is None:
- self.best_score = score
- self.save_checkpoint(val_loss, model, epoch=epoch, save_multiple=save_multiple)
- elif score < self.best_score + self.delta:
- self.counter += 1
-
- if self.counter >= self.patience:
- self.early_stop = True
-
- if save_all: # save all ckpt
- self.save_checkpoint(val_loss, model, epoch=epoch, save_multiple=True, update_best=False)
- else: # save best ckpt only
- self.best_score = score
- self.save_checkpoint(val_loss, model, epoch=epoch, save_multiple=save_multiple)
- self.counter = 0
-
- def save_checkpoint(self, val_loss, model, epoch=-1, save_multiple=False, update_best=True):
- '''Saves model when validation loss decrease.'''
- if self.verbose:
- tqdm.write(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}). Saving model ...')
- ckpt_save_path = self.ckpt_save_path
- if save_multiple:
- ckpt_save_path = ckpt_save_path[:-3]+f'-{epoch:04d}-{val_loss:.4f}.pt'
-
- torch.save(model.state_dict(), ckpt_save_path)
- if update_best:
- self.val_loss_min = val_loss
-
- def isBestModel(self):
- ''' Check if current model the best one.
- get early stop counter, if counter==0: it means current model has the best validation loss
- '''
- return self.counter==0
\ No newline at end of file
diff --git a/script/dm/direct_pose_model.py b/script/dm/direct_pose_model.py
deleted file mode 100755
index 3ae69bc..0000000
--- a/script/dm/direct_pose_model.py
+++ /dev/null
@@ -1,423 +0,0 @@
-import math
-import os
-import time
-import pprint
-import random
-import pdb
-import cv2
-from copy import deepcopy
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.nn.init
-import numpy as np
-from torchvision import models
-# from torch.utils.tensorboard import SummaryWriter
-from tqdm import tqdm
-
-from dm.pose_model import preprocess_data, get_error_in_q
-# from dm.prepare_data import prepare_data
-from models.nerfw import create_nerf, to8b, img2mse, mse2psnr
-from models.ray_utils import get_rays
-from models.rendering import render
-from utils.utils import freeze_bn_layer_train
-''' FeatureNet models '''
-# from feature.model import PoseNetV2 as FeatureNet
-# from feature.model import EfficientNetB3 as FeatureNet
-from feature.dfnet import DFNet as FeatureNet
-from feature.dfnet import DFNet_s as FeatureNet3
-# from feature.efficientnet import EfficientNetB3 as FeatureNet
-# from feature.efficientnet import EfficientNetB0 as FeatureNet
-
-''' PoseNet Models '''
-# from feature.efficientnet import EfficientNetB0 as PoseNet
-from feature.dfnet import DFNet as PoseNet
-from feature.dfnet import DFNet_s as PoseNet3
-from torchvision.utils import save_image
-
-def disable_model_grad(model):
- ''' set whole model to requires_grad=False, this is for nerf model '''
- print("disable_model_grad...")
- for module in model.modules():
- # print("this is a layer:", module)
- if hasattr(module, 'weight'):
- module.weight.requires_grad_(False)
- if hasattr(module, 'bias'):
- module.bias.requires_grad_(False)
- return model
-
-def inference_pose_regression(args, data, device, model):
- """
- Inference the Pose Regression Network
- Inputs:
- args: parsed argument
- data: Input image in shape (batchsize, channels, H, W)
- device: gpu device
- model: PoseNet model
- Outputs:
- pose: Predicted Pose in shape (batchsize, 3, 4)
- """
- inputs = data.to(device)
- if args.preprocess_ImgNet:
- inputs = preprocess_data(inputs, device)
- predict_pose = model(inputs)
- pose = predict_pose.reshape(args.batch_size, 3, 4)
-
- if args.svd_reg:
- # pdb.set_trace()
- R_torch = pose[:,:3,:3].clone() # debug
- u,s,v=torch.svd(R_torch)
- Rs = torch.matmul(u, v.transpose(-2,-1))
- pose[:,:3,:3] = Rs
- return pose
-
-def rgb_loss(rgb, target, extras):
- ''' Compute RGB MSE Loss, original from NeRF Paper '''
- # Compute MSE loss between predicted and true RGB.
- img_loss = img2mse(rgb, target)
- loss = img_loss
-
- # Add MSE loss for coarse-grained model
- if 'rgb0' in extras:
- img_loss0 = img2mse(extras['rgb0'], target)
- loss += img_loss0
- return loss
-
-def PoseLoss(args, pose_, pose, device):
- loss_func = nn.MSELoss()
- predict_pose = pose_.reshape(args.batch_size, 12).to(device) # maynot need reshape
- pose_loss = loss_func(predict_pose, pose)
- return pose_loss
-
-def load_exisiting_model(args, isFeatureNet=False):
- ''' Load a pretrained PoseNet model '''
- if isFeatureNet==False: # load posenet
- if args.DFNet_s:
- model = PoseNet3()
- model.load_state_dict(torch.load(args.pretrain_model_path))
- else:
- model = PoseNet()
- model.load_state_dict(torch.load(args.pretrain_model_path))
- return model
- else:
- if args.DFNet_s: # load featurenet
- model = FeatureNet3()
- model.load_state_dict(torch.load(args.pretrain_featurenet_path))
- else:
- model=FeatureNet()
- model.load_state_dict(torch.load(args.pretrain_featurenet_path))
- return model
-
-def prepare_batch_render(args, pose, batch_size, target_, H, W, focal, half_res=True, rand=True):
- ''' Break batch of images into rays '''
- target_ = target_.permute(0, 2, 3, 1).numpy()#.squeeze(0) # convert to numpy image
- if half_res:
- N_rand = batch_size * (H//2) * (W//2)
- target_half = np.stack([cv2.resize(target_[i], (H//2, W//2), interpolation=cv2.INTER_AREA) for i in range(batch_size)], 0)
- target_half = torch.Tensor(target_half)
-
- rays = torch.stack([torch.stack(get_rays(H//2, W//2, focal/2, pose[i]), 0) for i in range(batch_size)], 0) # [N, ro+rd, H, W, 3] (130, 2, 100, 100, 3)
- rays_rgb = torch.cat((rays, target_half[:, None, ...]), 1)
-
- else:
- # N_rand = batch_size * H * W
- N_rand = args.N_rand
- target_ = torch.Tensor(target_)
- rays = torch.stack([torch.stack(get_rays(H, W, focal, pose[i]), 0) for i in range(batch_size)], 0) # [N, ro+rd, H, W, 3] (130, 2, 200, 200, 3)
- # [N, ro+rd+rgb, H, W, 3]
- rays_rgb = torch.cat([rays, target_[:, None, ...]], 1)
-
- # [N, H, W, ro+rd+rgb, 3]
- rays_rgb = rays_rgb.permute(0, 2, 3, 1, 4)
-
- # [(N-1)*H*W, ro+rd+rgb, 3]
- rays_rgb = torch.reshape(rays_rgb, (-1, 3, 3))
-
- if 1:
- #print('shuffle rays')
- rays_rgb = rays_rgb[torch.randperm(rays_rgb.shape[0])]
-
- # Random over all images
- batch = rays_rgb[:N_rand].permute(1, 0 , 2) # [B, 2+1, 3*?] # (4096, 3, 3)
- batch_rays, target_s = batch[:2], batch[2] # [2, 4096, 3], [4096, 3]
-
- return batch_rays, target_s
-
-def fix_coord_supp(args, pose, world_setup_dict, device=None):
- # this function needs to be fixed because it is taking args.pose_scale
- '''supplementary fix_coord() for direct matching
- Inputs:
- args: parsed argument
- pose: pose [N, 3, 4]
- device: cpu or gpu
- Outputs:
- pose: converted Pose in shape [N, 3, 4]
- '''
- sc=world_setup_dict['pose_scale'] # manual tuned factor, align with colmap scale
- if device is None:
- move_all_cam_vec = torch.Tensor(world_setup_dict['move_all_cam_vec'])
- else:
- move_all_cam_vec = torch.Tensor(world_setup_dict['move_all_cam_vec']).to(device)
- sc2 = world_setup_dict['pose_scale2']
- pose[:,:3,3] *= sc
- # move center of camera pose
- pose[:, :3, 3] += move_all_cam_vec
- pose[:,:3,3] *= sc2
- return pose
-
-def eval_on_batch(args, data, model, pose, img_idx, hwf, half_res, device, **render_kwargs_test):
- ''' Perform 1 step of eval'''
- with torch.no_grad():
- H, W, focal = hwf
- target_ = deepcopy(data)
- pose_ = inference_pose_regression(args, data, device, model)
- device_cpu = torch.device('cpu')
- pose_ = pose_.to(device_cpu) # put predict pose back to cpu
- pose_nerf = pose_.clone()
-
- if args.NeRFH:
- # rescale the predicted pose to nerf scales
- pose_nerf = fix_coord_supp(args, pose_nerf)
-
- batch_rays, target = prepare_batch_render(args, pose_nerf, args.batch_size, target_, H, W, focal, half_res)
- batch_rays = batch_rays.to(device)
- target = target.to(device)
- pose = pose.to(device)
- img_idx = img_idx.to(device)
-
- # every new tensor from onward is in GPU
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- if half_res:
- rgb, disp, acc, extras = render(H//2, W//2, focal/2, chunk=args.chunk, rays=batch_rays, img_idx=img_idx, **render_kwargs_test)
- else:
- rgb, disp, acc, extras = render(H, W, focal, chunk=args.chunk, rays=batch_rays, img_idx=img_idx, **render_kwargs_test)
-
- loss = PoseLoss(args, pose_, pose, device)
- psnr = mse2psnr(img2mse(rgb, target))
-
- # end of every new tensor from onward is in GPU
- torch.set_default_tensor_type('torch.FloatTensor')
-
- iter_loss = loss.to(device_cpu).detach().numpy()
- iter_loss = np.array([iter_loss])
-
- iter_psnr = psnr.to(device_cpu).detach().numpy()
- return iter_loss, iter_psnr
-
-def eval_on_epoch(args, data_loaders, model, hwf, half_res, device, **render_kwargs_test):
- ''' Perform 1 epoch of training with batch '''
- model.eval()
- batch_size = 1
-
- train_dl, val_dl, test_dl = data_loaders
-
- total_loss = []
- total_psnr = []
-
- #### Core optimization loop #####
- for data, pose, img_idx in val_dl:
- # training one step with batch_size = args.batch_size
- loss, psnr = eval_on_batch(args, data, model, pose, img_idx, hwf, half_res, device, **render_kwargs_test)
- total_loss.append(loss.item())
- total_psnr.append(psnr.item())
- total_loss_mean = np.mean(total_loss)
- total_psnr_mean = np.mean(total_psnr)
- return total_loss_mean, total_psnr_mean
-
-def train_on_batch(args, data, model, pose, img_idx, hwf, optimizer, half_res, device, **render_kwargs_test):
- ''' Perform 1 step of training'''
- H, W, focal = hwf
- target_ = deepcopy(data)
- pose_ = inference_pose_regression(args, data, device, model)
- device_cpu = torch.device('cpu')
- pose_ = pose_.to(device_cpu) # put predict pose back to cpu
- pose_nerf = pose_.clone()
- if args.NeRFH:
- # rescale the predicted pose to nerf scales
- pose_nerf = fix_coord_supp(args, pose_nerf)
-
- batch_rays, target = prepare_batch_render(args, pose_nerf, args.batch_size, target_, H, W, focal, half_res)
- batch_rays = batch_rays.to(device)
- target = target.to(device)
- pose = pose.to(device)
- img_idx = img_idx.to(device)
-
- # # every new tensor from onward is in GPU
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- if half_res:
- rgb, disp, acc, extras = render(H//2, W//2, focal/2, chunk=args.chunk, rays=batch_rays, img_idx=img_idx, **render_kwargs_test)
- else:
- rgb, disp, acc, extras = render(H, W, focal, chunk=args.chunk, rays=batch_rays, img_idx=img_idx, **render_kwargs_test)
-
- ### Loss Design Here ###
- # Compute RGB MSE Loss
- photo_loss = rgb_loss(rgb, target, extras)
-
- # Compute Combine Loss if needed
- if args.combine_loss:
- pose_loss = PoseLoss(args, pose_, pose, device)
- loss = args.combine_loss_w[0] * pose_loss + args.combine_loss_w[1] * photo_loss
- # print("img_idx: {}, pose_loss: {}, rgb_loss {}, total_loss {}".format(img_idx, args.combine_loss_w[0] * pose_loss, args.combine_loss_w[1] * photo_loss, loss))
-
- ### Loss Design End
- loss.backward()
- optimizer.step()
- optimizer.zero_grad()
- psnr = mse2psnr(img2mse(rgb, target))
-
- # end of every new tensor from onward is in GPU
- torch.set_default_tensor_type('torch.FloatTensor')
-
- iter_loss = loss.to(device_cpu).detach().numpy()
- iter_loss = np.array([iter_loss])
-
- iter_psnr = psnr.to(device_cpu).detach().numpy()
- return iter_loss, iter_psnr
-
-
-def train_on_epoch(args, data_loaders, model, hwf, optimizer, half_res, device, **render_kwargs_test):
- ''' Perform 1 epoch of training with batch '''
- model.train()
- model = freeze_bn_layer_train(model)
-
- # Prepare dataloaders for PoseNet, each batch contains (image, pose)
- train_dl, val_dl, test_dl = data_loaders
- total_loss = []
- total_psnr = []
-
- #### Core optimization loop #####
- for data, pose, img_idx in train_dl:
- # print("img_idx: {}, pose: {}".format(img_idx, pose) )
- # training one step with batch_size = args.batch_size
- loss, psnr = train_on_batch(args, data, model, pose, img_idx, hwf, optimizer, half_res, device, **render_kwargs_test)
- total_loss.append(loss.item())
- total_psnr.append(psnr.item())
- total_loss_mean = np.mean(total_loss)
- total_psnr_mean = np.mean(total_psnr)
- return total_loss_mean, total_psnr_mean
-
-def save_val_result_7Scenes(args, epoch, val_dl, model, hwf, half_res, device, num_samples=1, **render_kwargs_test):
- ''' Perform inference on a random val image and save the result '''
- half_res=False # save half res image to reduce computational speed
- model.eval()
- i = 0
- for batch in val_dl:
- if args.NeRFH:
- data, pose, img_idx = batch
- else:
- data, pose = batch
- if i >= num_samples:
- return
- H, W, focal = hwf
- target_ = deepcopy(data)
- img_idx = img_idx.to(device)
-
- pose_ = inference_pose_regression(args, data, device, model)
- device_cpu = torch.device('cpu')
- pose_ = pose_.to(device_cpu) # put predict_pose back to cpu
-
- pose_nerf = pose_.clone()
- if args.NeRFH:
- # rescale the predicted pose to nerf scales
- pose_nerf = fix_coord_supp(args, pose_nerf)
-
- # every new tensor from onward is in GPU
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- with torch.no_grad():
- if half_res:
- rgb, disp, acc, extras = render(H//2, W//2, focal/2, chunk=args.chunk, c2w=pose_nerf[0].to(device), img_idx=img_idx, **render_kwargs_test)
- else:
- rgb, disp, acc, extras = render(H, W, focal, chunk=args.chunk, c2w=pose_nerf[0].to(device), img_idx=img_idx, **render_kwargs_test)
-
- ### Set Save Dir ###
- if num_samples <=1:
- out_folder = os.path.join(args.basedir, args.model_name, 'val_imgs')
- else:
- out_folder = os.path.join(args.basedir, args.model_name, 'val_imgs_batches')
- if not os.path.isdir(out_folder):
- os.mkdir(out_folder)
-
- if half_res:
- target_img = F.interpolate(target_, scale_factor=0.5, mode='area').permute(0,2,3,1).reshape(H//2,W//2,3)
- rgb_img = rgb.reshape(H//2,W//2,3)
- else:
- target_img = target_.permute(0,2,3,1).reshape(H,W,3)
- rgb_img = rgb.reshape(H,W,3)
- target_img_to_save = to8b(target_img.to(device_cpu).detach().numpy())
- rgb_img_to_save = to8b(rgb_img.to(device_cpu).detach().numpy())
- # save NeRF Rendered RGB Img
- import imageio
- if num_samples <= 1:
- imageio.imwrite(os.path.join(out_folder, '{0:04d}_gt.png'.format(epoch)), target_img_to_save)
- imageio.imwrite(os.path.join(out_folder, '{0:04d}.png'.format(epoch)), rgb_img_to_save)
- else:
- imageio.imwrite(os.path.join(out_folder, '{0:04d}_gt.png'.format(i)), target_img_to_save)
- imageio.imwrite(os.path.join(out_folder, '{0:04d}.png'.format(i)), rgb_img_to_save)
-
- # end of every new tensor from onward is in GPU
- torch.set_default_tensor_type('torch.FloatTensor')
- i = i+1
-
-
-def train_nerf_tracking(args, model, optimizer, i_split, hwf, near, far, device, early_stopping, images=None, poses_train=None, train_dl=None, val_dl=None, test_dl=None):
- ''' finetune pretrained PoseNet using NeRF '''
- half_res = False # direct-pn paper settings
- # half_res = True # This half_res is to further downsample the output image of nerf to 100x100
- # Prepare dataloaders for PoseNet, each batch contains (image, pose)
- if args.dataset_type != '7Scenes' and args.dataset_type != 'Cambridge': # blender dataset
- train_dl, val_dl, test_dl = prepare_data(args, images, poses_train, i_split)
- _, render_kwargs_test, _, grad_vars, _ = create_nerf(args) # llff nerf
- else:
- # load NeRF model
- _, render_kwargs_test, start, grad_vars, _ = create_nerf(args)
- global_step = start
- if args.reduce_embedding==2:
- render_kwargs_test['i_epoch'] = global_step
-
- # # only freeze 150MB? effectiveness of this towards training result is yet unknown
- render_kwargs_test['embedding_a'] = disable_model_grad(render_kwargs_test['embedding_a'])
- render_kwargs_test['embedding_t'] = disable_model_grad(render_kwargs_test['embedding_t'])
- render_kwargs_test['network_fn'] = disable_model_grad(render_kwargs_test['network_fn'])
- render_kwargs_test['network_fine'] = disable_model_grad(render_kwargs_test['network_fine'])
-
- data_loaders = [train_dl, val_dl, test_dl]
- bds_dict = {
- 'near' : near,
- 'far' : far,
- }
- # render_kwargs_train.update(bds_dict)
- render_kwargs_test.update(bds_dict)
- i_train, i_val, i_test = i_split
-
- N_epoch = 2001
- print('Begin')
- print('TRAIN views are', i_train)
- print('TEST views are', i_test)
- print('VAL views are', i_val)
- time0 = time.time()
-
- model_log = tqdm(total=0, position = 1, bar_format='{desc}')
- for epoch in tqdm(range(N_epoch), desc='epochs'):
- #train 1 epoch with batch_size = 1
- loss, psnr = train_on_epoch(args, data_loaders, model, hwf, optimizer, half_res, device, **render_kwargs_test)
-
- val_loss, val_psnr = eval_on_epoch(args, data_loaders, model, hwf, half_res, device, **render_kwargs_test)
-
- tqdm.write('At epoch {0:4d} : train loss: {1:.4f}, train psnr: {2:.4f}, val loss: {3:.4f}, val psnr: {4:.4f}'.format(epoch, loss, psnr, val_loss, val_psnr))
-
- # check wether to early stop
- early_stopping(val_loss, model, epoch=epoch, save_multiple=(not args.no_save_multiple), save_all=args.save_all_ckpt)
- if early_stopping.early_stop:
- print("Early stopping")
- break
- model_log.set_description_str(f'Best val loss: {early_stopping.val_loss_min:.4f}')
-
- if (epoch % 1 == 0) and (args.pose_only != 1):
- ### run one single image save the result
- if args.dataset_type == '7Scenes' or args.dataset_type == 'llff' or args.dataset_type == 'Cambridge':
- save_val_result_7Scenes(args, epoch, val_dl, model, hwf, half_res, device, **render_kwargs_test)
-
- if epoch % args.i_eval == 0:
- # calculate position and angular error
- get_error_in_q(args, val_dl, model, len(val_dl.dataset), device, batch_size=1)
diff --git a/script/dm/options.py b/script/dm/options.py
deleted file mode 100755
index 6a1c8ae..0000000
--- a/script/dm/options.py
+++ /dev/null
@@ -1,138 +0,0 @@
-import configargparse
-def config_parser():
- parser = configargparse.ArgumentParser()
- parser.add_argument("-f", "--fff", help="a dummy argument to fool ipython", default="1")
- parser.add_argument("--multi_gpu", action='store_true', help='use multiple gpu on the server')
-
- # 7Scenes
- parser.add_argument("--trainskip", type=int, default=1, help='will load 1/N images from train sets, useful for large datasets like 7 Scenes')
- parser.add_argument("--df", type=float, default=1., help='image downscale factor')
- parser.add_argument("--reduce_embedding", type=int, default=-1, help='fourier embedding mode: -1: paper default, \
- 0: reduce by half, 1: remove embedding, 2: DNeRF embedding')
- parser.add_argument("--epochToMaxFreq", type=int, default=-1, help='DNeRF embedding mode: (based on DNeRF paper): \
- hyper-parameter for when α should reach the maximum number of frequencies m')
- parser.add_argument("--render_pose_only", action='store_true', help='render a spiral video for 7 Scene')
- parser.add_argument("--save_pose_avg_stats", action='store_true', help='save a pose avg stats to unify NeRF, posenet, nerf tracking training')
- parser.add_argument("--load_pose_avg_stats", action='store_true', help='load precomputed pose avg stats to unify NeRF, posenet, nerf tracking training')
- parser.add_argument("--finetune_unlabel", action='store_true', help='finetune unlabeled sequence like MapNet')
- parser.add_argument("--i_eval", type=int, default=50, help='frequency of eval posenet result')
- parser.add_argument("--save_all_ckpt", action='store_true', help='save all ckpts for each epoch')
- parser.add_argument("--train_local_nerf", type=int, default=-1, help='train local NeRF with ith training sequence only, ie. Stairs can pick 0~3')
- parser.add_argument("--render_video_train", action='store_true', help='render train set NeRF and save as video, make sure i_eval is True')
- parser.add_argument("--render_video_test", action='store_true', help='render val set NeRF and save as video, make sure i_eval is True')
- parser.add_argument("--no_DNeRF_viewdir", action='store_true', default=False, help='will not use DNeRF in viewdir encoding')
- parser.add_argument("--val_on_psnr", action='store_true', default=False, help='EarlyStopping with max validation psnr')
- parser.add_argument("--feature_matching_lvl", nargs='+', type=int, default=[0,1,2],
- help='lvl of features used for feature matching, default use lvl 0, 1, 2')
-
- ##################### PoseNet Settings ########################
- parser.add_argument("--pose_only", type=int, default=0, help='posenet type to train, \
- 1: train baseline posenet, 2: posenet+nerf manual optimize, \
- 3: featurenet,')
- parser.add_argument("--learning_rate", type=float, default=0.00001, help='learning rate')
- parser.add_argument("--batch_size", type=int, default=1, help='train posenet only')
- parser.add_argument("--pretrain_model_path", type=str, default='', help='model path of pretrained pose regrssion model')
- parser.add_argument("--pretrain_featurenet_path", type=str, default='', help='model path of pretrained featurenet model')
- parser.add_argument("--model_name", type=str, help='pose model output folder name')
- parser.add_argument("--combine_loss", action='store_true',
- help='combined l2 pose loss + rgb mse loss')
- parser.add_argument("--combine_loss_w", nargs='+', type=float, default=[0.5, 0.5],
- help='weights of combined loss ex, [0.5 0.5], \
- default None, only use when combine_loss is True')
- parser.add_argument("--patience", nargs='+', type=int, default=[200, 50], help='set training schedule for patience [EarlyStopping, reduceLR]')
- parser.add_argument("--resize_factor", type=int, default=2, help='image resize downsample ratio')
- parser.add_argument("--freezeBN", action='store_true', help='Freeze the Batch Norm layer at training PoseNet')
- parser.add_argument("--preprocess_ImgNet", action='store_true', help='Normalize input data for PoseNet')
- parser.add_argument("--eval", action='store_true', help='eval model')
- parser.add_argument("--no_save_multiple", action='store_true', help='default, save multiple posenet model, if true, save only one posenet model')
- parser.add_argument("--resnet34", action='store_true', default=False, help='use resnet34 backbone instead of mobilenetV2')
- parser.add_argument("--efficientnet", action='store_true', default=False, help='use efficientnet-b3 backbone instead of mobilenetV2')
- parser.add_argument("--efficientnet_block", type=int, default=6, help='choose which features from feature block (0-6) of efficientnet to use')
- parser.add_argument("--dropout", type=float, default=0.5, help='dropout rate for resnet34 backbone')
- parser.add_argument("--DFNet", action='store_true', default=False, help='use DFNet')
- parser.add_argument("--DFNet_s", action='store_true', default=False, help='use accelerated DFNet, performance is similar to DFNet but slightly faster')
- parser.add_argument("--val_batch_size", type=int, default=1, help='batch_size for validation, higher number leads to faster speed')
-
- ##################### NeRF Settings ########################
- parser.add_argument('--config', is_config_file=True, help='config file path')
- parser.add_argument("--expname", type=str, help='experiment name')
- parser.add_argument("--basedir", type=str, default='../logs/', help='where to store ckpts and logs')
- parser.add_argument("--datadir", type=str, default='./data/llff/fern', help='input data directory')
-
- # training options
- parser.add_argument("--netdepth", type=int, default=8, help='layers in network')
- parser.add_argument("--netwidth", type=int, default=128, help='channels per layer')
- parser.add_argument("--netdepth_fine", type=int, default=8, help='layers in fine network')
- parser.add_argument("--netwidth_fine", type=int, default=128, help='channels per layer in fine network')
- parser.add_argument("--N_rand", type=int, default=1536, help='batch size (number of random rays per gradient step)')
- parser.add_argument("--lrate", type=float, default=5e-4, help='learning rate')
- parser.add_argument("--lrate_decay", type=int, default=250, help='exponential learning rate decay (in 1000 steps)')
- parser.add_argument("--chunk", type=int, default=1024*32, help='number of rays processed in parallel, decrease if running out of memory')
- parser.add_argument("--netchunk", type=int, default=1024*64, help='number of pts sent through network in parallel, decrease if running out of memory')
- parser.add_argument("--no_batching", action='store_true', help='only take random rays from 1 image at a time')
- parser.add_argument("--no_reload", action='store_true', help='do not reload weights from saved ckpt')
- parser.add_argument("--ft_path", type=str, default=None, help='specific weights npy file to reload for coarse network')
- parser.add_argument("--no_grad_update", default=True, help='do not update nerf in training')
- parser.add_argument("--per_channel", default=False, action='store_true', help='using per channel cosine similarity loss instead of per pixel, defualt False')
-
- # NeRF-Hist training options
- parser.add_argument("--NeRFH", action='store_true', help='new implementation for NeRFH')
- parser.add_argument("--N_vocab", type=int, default=1000,
- help='''number of vocabulary (number of images)
- in the dataset for nn.Embedding''')
- parser.add_argument("--fix_index", action='store_true', help='fix training frame index as 0')
- parser.add_argument("--encode_hist", default=False, action='store_true', help='encode histogram instead of frame index')
- parser.add_argument("--hist_bin", type=int, default=10, help='image histogram bin size')
- parser.add_argument("--in_channels_a", type=int, default=50, help='appearance embedding dimension, hist_bin*N_a when embedding histogram')
- parser.add_argument("--in_channels_t", type=int, default=20, help='transient embedding dimension, hist_bin*N_tau when embedding histogram')
- parser.add_argument("--svd_reg", default=False, action='store_true', help='use svd regularize output at training')
-
- # rendering options
- parser.add_argument("--N_samples", type=int, default=64, help='number of coarse samples per ray')
- parser.add_argument("--N_importance", type=int, default=64,help='number of additional fine samples per ray')
- parser.add_argument("--perturb", type=float, default=1.,help='set to 0. for no jitter, 1. for jitter')
- parser.add_argument("--use_viewdirs", action='store_true', default=True, help='use full 5D input instead of 3D')
- parser.add_argument("--i_embed", type=int, default=0, help='set 0 for default positional encoding, -1 for none')
- parser.add_argument("--multires", type=int, default=10, help='log2 of max freq for positional encoding (3D location)')
- parser.add_argument("--multires_views", type=int, default=4, help='log2 of max freq for positional encoding (2D direction)')
- parser.add_argument("--raw_noise_std", type=float, default=0., help='std dev of noise added to regularize sigma_a output, 1e0 recommended')
-
- parser.add_argument("--render_only", action='store_true', help='do not optimize, reload weights and render out render_poses path')
- parser.add_argument("--render_test", action='store_true', help='render the test set instead of render_poses path')
- parser.add_argument("--render_factor", type=int, default=0, help='downsampling factor to speed up rendering, set 4 or 8 for fast preview')
-
- # legacy mesh options
- parser.add_argument("--mesh_only", action='store_true', help='do not optimize, reload weights and save mesh to a file')
- parser.add_argument("--mesh_grid_size", type=int, default=80,help='number of grid points to sample in each dimension for marching cubes')
-
- # training options
- parser.add_argument("--precrop_iters", type=int, default=0,help='number of steps to train on central crops')
- parser.add_argument("--precrop_frac", type=float,default=.5, help='fraction of img taken for central crops')
-
- # dataset options
- parser.add_argument("--dataset_type", type=str, default='llff', help='options: llff / blender / deepvoxels')
- parser.add_argument("--testskip", type=int, default=1, help='will load 1/N images from test/val sets, useful for large datasets like deepvoxels')
-
- ## legacy blender flags
- parser.add_argument("--white_bkgd", action='store_true', help='set to render synthetic data on a white bkgd (always use for dvoxels)')
- parser.add_argument("--half_res", action='store_true', help='load blender synthetic data at 400x400 instead of 800x800')
-
- ## llff flags
- parser.add_argument("--factor", type=int, default=8, help='downsample factor for LLFF images')
- parser.add_argument("--no_ndc", action='store_true', help='do not use normalized device coordinates (set for non-forward facing scenes)')
- parser.add_argument("--lindisp", action='store_true', help='sampling linearly in disparity rather than depth')
- parser.add_argument("--spherify", action='store_true', help='set for spherical 360 scenes')
- parser.add_argument("--llffhold", type=int, default=8, help='will take every 1/N images as LLFF test set, paper uses 8')
- parser.add_argument("--no_bd_factor", action='store_true', default=False, help='do not use bd factor')
-
- # featruremetric supervision
- parser.add_argument("--featuremetric", action='store_true', help='use featuremetric supervision if true')
-
- # logging/saving options
- parser.add_argument("--i_print", type=int, default=1, help='frequency of console printout and metric loggin')
- parser.add_argument("--i_img", type=int, default=500, help='frequency of tensorboard image logging')
- parser.add_argument("--i_weights", type=int, default=200, help='frequency of weight ckpt saving')
- parser.add_argument("--i_testset", type=int, default=200, help='frequency of testset saving')
- parser.add_argument("--i_video", type=int, default=50000, help='frequency of render_poses video saving')
-
- return parser
\ No newline at end of file
diff --git a/script/dm/pose_model.py b/script/dm/pose_model.py
deleted file mode 100755
index 4365e28..0000000
--- a/script/dm/pose_model.py
+++ /dev/null
@@ -1,391 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.nn.init
-import numpy as np
-from torchvision import models
-from efficientnet_pytorch import EfficientNet
-from torch.utils.tensorboard import SummaryWriter
-from tqdm import tqdm
-import pdb
-import matplotlib.pyplot as plt
-
-import math
-import time
-
-import pytorch3d.transforms as transforms
-
-def preprocess_data(inputs, device):
- # normalize inputs according to https://pytorch.org/hub/pytorch_vision_mobilenet_v2/
- mean = torch.Tensor([0.485, 0.456, 0.406]).to(device) # per channel subtraction
- std = torch.Tensor([0.229, 0.224, 0.225]).to(device) # per channel division
- inputs = (inputs - mean[None,:,None,None])/std[None,:,None,None]
- return inputs
-
-def filter_hook(m, g_in, g_out):
- g_filtered = []
- for g in g_in:
- g = g.clone()
- g[g != g] = 0
- g_filtered.append(g)
- return tuple(g_filtered)
-
-def vis_pose(vis_info):
- '''
- visualize predicted pose result vs. gt pose
- '''
- pdb.set_trace()
- pose = vis_info['pose']
- pose_gt = vis_info['pose_gt']
- theta = vis_info['theta']
- ang_threshold=10
- seq_num = theta.shape[0]
- # # create figure object
- # plot translation traj.
- fig = plt.figure(figsize = (8,6))
- plt.subplots_adjust(left=0, bottom=0, right=1, top=1)
- ax1 = fig.add_axes([0, 0.2, 0.9, 0.85], projection='3d')
- ax1.scatter(pose[10:,0],pose[10:,1],zs=pose[10:,2], c='r', s=3**2,depthshade=0) # predict
- ax1.scatter(pose_gt[:,0], pose_gt[:,1], zs=pose_gt[:,2], c='g', s=3**2,depthshade=0) # GT
- ax1.scatter(pose[0:10,0],pose[0:10,1],zs=pose[0:10,2], c='k', s=3**2,depthshade=0) # predict
- # ax1.plot(pose[:,0],pose[:,1],zs=pose[:,2], c='r') # predict
- # ax1.plot(pose_gt[:,0], pose_gt[:,1], zs=pose_gt[:,2], c='g') # GT
- ax1.view_init(30, 120)
- ax1.set_xlabel('x (m)')
- ax1.set_ylabel('y (m)')
- ax1.set_zlabel('z (m)')
- # ax1.set_xlim(-10, 10)
- # ax1.set_ylim(-10, 10)
- # ax1.set_zlim(-10, 10)
-
- ax1.set_xlim(-1, 1)
- ax1.set_ylim(-1, 1)
- ax1.set_zlim(-1, 1)
-
- # ax1.set_xlim(-3, 3)
- # ax1.set_ylim(-3, 3)
- # ax1.set_zlim(-3, 3)
-
- # plot angular error
- ax2 = fig.add_axes([0.1, 0.05, 0.75, 0.2])
- err = theta.reshape(1, seq_num)
- err = np.tile(err, (20, 1))
- ax2.imshow(err, vmin=0,vmax=ang_threshold, aspect=3)
- ax2.set_yticks([])
- ax2.set_xticks([0, seq_num*1/5, seq_num*2/5, seq_num*3/5, seq_num*4/5, seq_num])
- fname = './vis_pose.png'
- plt.savefig(fname, dpi=50)
-
-def compute_error_in_q(args, dl, model, device, results, batch_size=1):
- use_SVD=True # Turn on for Direct-PN and Direct-PN+U reported result, despite it makes minuscule differences
- time_spent = []
- predict_pose_list = []
- gt_pose_list = []
- ang_error_list = []
- pose_result_raw = []
- pose_GT = []
- i = 0
-
- for batch in dl:
- if args.NeRFH:
- data, pose, img_idx = batch
- else:
- data, pose = batch
- data = data.to(device) # input
- pose = pose.reshape((batch_size,3,4)).numpy() # label
-
- if args.preprocess_ImgNet:
- data = preprocess_data(data, device)
-
- if use_SVD:
- # using SVD to make sure predict rotation is normalized rotation matrix
- with torch.no_grad():
- if args.featuremetric:
- _, predict_pose = model(data)
- else:
- predict_pose = model(data)
-
- R_torch = predict_pose.reshape((batch_size, 3, 4))[:,:3,:3] # debug
- predict_pose = predict_pose.reshape((batch_size, 3, 4)).cpu().numpy()
-
- R = predict_pose[:,:3,:3]
- res = R@np.linalg.inv(R)
- # print('R@np.linalg.inv(R):', res)
-
- u,s,v=torch.svd(R_torch)
- Rs = torch.matmul(u, v.transpose(-2,-1))
- predict_pose[:,:3,:3] = Rs[:,:3,:3].cpu().numpy()
- else:
- start_time = time.time()
- # inference NN
- with torch.no_grad():
- predict_pose = model(data)
- predict_pose = predict_pose.reshape((batch_size, 3, 4)).cpu().numpy()
- time_spent.append(time.time() - start_time)
-
- pose_q = transforms.matrix_to_quaternion(torch.Tensor(pose[:,:3,:3]))#.cpu().numpy() # gnd truth in quaternion
- pose_x = pose[:, :3, 3] # gnd truth position
- predicted_q = transforms.matrix_to_quaternion(torch.Tensor(predict_pose[:,:3,:3]))#.cpu().numpy() # predict in quaternion
- predicted_x = predict_pose[:, :3, 3] # predict position
- pose_q = pose_q.squeeze()
- pose_x = pose_x.squeeze()
- predicted_q = predicted_q.squeeze()
- predicted_x = predicted_x.squeeze()
-
- #Compute Individual Sample Error
- q1 = pose_q / torch.linalg.norm(pose_q)
- q2 = predicted_q / torch.linalg.norm(predicted_q)
- d = torch.abs(torch.sum(torch.matmul(q1,q2)))
- d = torch.clamp(d, -1., 1.) # acos can only input [-1~1]
- theta = (2 * torch.acos(d) * 180/math.pi).numpy()
- error_x = torch.linalg.norm(torch.Tensor(pose_x-predicted_x)).numpy()
- results[i,:] = [error_x, theta]
- #print ('Iteration: {} Error XYZ (m): {} Error Q (degrees): {}'.format(i, error_x, theta))
-
- # save results for visualization
- predict_pose_list.append(predicted_x)
- gt_pose_list.append(pose_x)
- ang_error_list.append(theta)
- pose_result_raw.append(predict_pose)
- pose_GT.append(pose)
- i += 1
- # pdb.set_trace()
- predict_pose_list = np.array(predict_pose_list)
- gt_pose_list = np.array(gt_pose_list)
- ang_error_list = np.array(ang_error_list)
- pose_result_raw = np.asarray(pose_result_raw)[:,0,:,:]
- pose_GT = np.asarray(pose_GT)[:,0,:,:]
- vis_info_ret = {"pose": predict_pose_list, "pose_gt": gt_pose_list, "theta": ang_error_list, "pose_result_raw": pose_result_raw, "pose_GT": pose_GT}
- return results, vis_info_ret
-
-# # pytorch
-def get_error_in_q(args, dl, model, sample_size, device, batch_size=1):
- ''' Convert Rotation matrix to quaternion, then calculate the location errors. original from PoseNet Paper '''
- model.eval()
-
- results = np.zeros((sample_size, 2))
- results, vis_info = compute_error_in_q(args, dl, model, device, results, batch_size)
- median_result = np.median(results,axis=0)
- mean_result = np.mean(results,axis=0)
-
- # standard log
- print ('Median error {}m and {} degrees.'.format(median_result[0], median_result[1]))
- print ('Mean error {}m and {} degrees.'.format(mean_result[0], mean_result[1]))
-
- # timing log
- #print ('Avg execution time (sec): {:.3f}'.format(np.mean(time_spent)))
-
- # standard log2
- # num_translation_less_5cm = np.asarray(np.where(results[:,0]<0.05))[0]
- # num_rotation_less_5 = np.asarray(np.where(results[:,1]<5))[0]
- # print ('translation error less than 5cm {}/{}.'.format(num_translation_less_5cm.shape[0], results.shape[0]))
- # print ('rotation error less than 5 degree {}/{}.'.format(num_rotation_less_5.shape[0], results.shape[0]))
- # print ('results:', results)
-
- # save for direct-pn paper log
- # if 0:
- # filename='Direct-PN+U_' + args.datadir.split('/')[-1] + '_result.txt'
- # np.savetxt(filename, predict_pose)
-
- # visualize results
- # vis_pose(vis_info)
-
-class EfficientNetB3(nn.Module):
- ''' EfficientNet-B3 backbone,
- model ref: https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/efficientnet_pytorch/model.py
- '''
- def __init__(self, feat_dim=12):
- super(EfficientNetB3, self).__init__()
- self.backbone_net = EfficientNet.from_pretrained('efficientnet-b3')
- self.feature_extractor = self.backbone_net.extract_features
- self.avgpool = nn.AdaptiveAvgPool2d(1)
- self.fc_pose = nn.Linear(1536, feat_dim) # 1280 for efficientnet-b0, 1536 for efficientnet-b3
-
- def forward(self, Input):
- x = self.feature_extractor(Input)
- x = self.avgpool(x)
- x = x.reshape(x.size(0), -1)
- predict = self.fc_pose(x)
- return predict
-
-# PoseNet (SE(3)) w/ mobilev2 backbone
-class PoseNetV2(nn.Module):
- def __init__(self, feat_dim=12):
- super(PoseNetV2, self).__init__()
- self.backbone_net = models.mobilenet_v2(pretrained=True)
- self.feature_extractor = self.backbone_net.features
- self.avgpool = nn.AdaptiveAvgPool2d(1)
- self.fc_pose = nn.Linear(1280, feat_dim)
-
- def forward(self, Input):
- x = self.feature_extractor(Input)
- x = self.avgpool(x)
- x = x.reshape(x.size(0), -1)
- predict = self.fc_pose(x)
- # pdb.set_trace()
- return predict
-
-# PoseNet (SE(3)) w/ resnet34 backnone. We found dropout layer is unnecessary, so we set droprate as 0 in reported results.
-class PoseNet_res34(nn.Module):
- def __init__(self, droprate=0.5, pretrained=True,
- feat_dim=2048):
- super(PoseNet_res34, self).__init__()
- self.droprate = droprate
-
- # replace the last FC layer in feature extractor
- self.feature_extractor = models.resnet34(pretrained=True)
- self.feature_extractor.avgpool = nn.AdaptiveAvgPool2d(1)
- fe_out_planes = self.feature_extractor.fc.in_features
- self.feature_extractor.fc = nn.Linear(fe_out_planes, feat_dim)
- self.fc_pose = nn.Linear(feat_dim, 12)
-
- # initialize
- if pretrained:
- init_modules = [self.feature_extractor.fc]
- else:
- init_modules = self.modules()
-
- for m in init_modules:
- if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
- nn.init.kaiming_normal_(m.weight.data)
- if m.bias is not None:
- nn.init.constant_(m.bias.data, 0)
-
- def forward(self, x):
- x = self.feature_extractor(x)
- x = F.relu(x)
- if self.droprate > 0:
- x = F.dropout(x, p=self.droprate)
- predict = self.fc_pose(x)
- return predict
-
-
-# from MapNet paper CVPR 2018
-class PoseNet(nn.Module):
- def __init__(self, feature_extractor, droprate=0.5, pretrained=True,
- feat_dim=2048, filter_nans=False):
- super(PoseNet, self).__init__()
- self.droprate = droprate
-
- # replace the last FC layer in feature extractor
- self.feature_extractor = models.resnet34(pretrained=True)
- self.feature_extractor.avgpool = nn.AdaptiveAvgPool2d(1)
- fe_out_planes = self.feature_extractor.fc.in_features
- self.feature_extractor.fc = nn.Linear(fe_out_planes, feat_dim)
-
- self.fc_xyz = nn.Linear(feat_dim, 3)
- self.fc_wpqr = nn.Linear(feat_dim, 3)
- if filter_nans:
- self.fc_wpqr.register_backward_hook(hook=filter_hook)
- # initialize
- if pretrained:
- init_modules = [self.feature_extractor.fc, self.fc_xyz, self.fc_wpqr]
- else:
- init_modules = self.modules()
-
- for m in init_modules:
- if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
- nn.init.kaiming_normal_(m.weight.data)
- if m.bias is not None:
- nn.init.constant_(m.bias.data, 0)
-
- def forward(self, x):
- x = self.feature_extractor(x)
- x = F.relu(x)
- if self.droprate > 0:
- x = F.dropout(x, p=self.droprate)
-
- xyz = self.fc_xyz(x)
- wpqr = self.fc_wpqr(x)
- return torch.cat((xyz, wpqr), 1)
-
-class MapNet(nn.Module):
- """
- Implements the MapNet model (green block in Fig. 2 of paper)
- """
- def __init__(self, mapnet):
- """
- :param mapnet: the MapNet (two CNN blocks inside the green block in Fig. 2
- of paper). Not to be confused with MapNet, the model!
- """
- super(MapNet, self).__init__()
- self.mapnet = mapnet
-
- def forward(self, x):
- """
- :param x: image blob (N x T x C x H x W)
- :return: pose outputs
- (N x T x 6)
- """
- s = x.size()
- x = x.view(-1, *s[2:])
- poses = self.mapnet(x)
- poses = poses.view(s[0], s[1], -1)
- return poses
-
-def eval_on_epoch(args, dl, model, optimizer, loss_func, device):
- model.eval()
- val_loss_epoch = []
- for data, pose in dl:
- inputs = data.to(device)
- labels = pose.to(device)
- if args.preprocess_ImgNet:
- inputs = preprocess_data(inputs, device)
- predict = model(inputs)
- loss = loss_func(predict, labels)
- val_loss_epoch.append(loss.item())
- val_loss_epoch_mean = np.mean(val_loss_epoch)
- return val_loss_epoch_mean
-
-
-def train_on_epoch(args, dl, model, optimizer, loss_func, device):
- model.train()
- train_loss_epoch = []
- for data, pose in dl:
- inputs = data.to(device) # (N, Ch, H, W) ~ (4,3,200,200), 7scenes [4, 3, 256, 341] wierd shape...
- labels = pose.to(device)
- if args.preprocess_ImgNet:
- inputs = preprocess_data(inputs, device)
-
- predict = model(inputs)
- loss = loss_func(predict, labels)
- loss.backward()
- optimizer.step()
- optimizer.zero_grad()
- train_loss_epoch.append(loss.item())
- train_loss_epoch_mean = np.mean(train_loss_epoch)
- return train_loss_epoch_mean
-
-def train_posenet(args, train_dl, val_dl, model, epochs, optimizer, loss_func, scheduler, device, early_stopping):
- writer = SummaryWriter()
- model_log = tqdm(total=0, position=1, bar_format='{desc}')
- for epoch in tqdm(range(epochs), desc='epochs'):
-
- # train 1 epoch
- train_loss = train_on_epoch(args, train_dl, model, optimizer, loss_func, device)
- writer.add_scalar("Loss/train", train_loss, epoch)
-
- # validate every epoch
- val_loss = eval_on_epoch(args, val_dl, model, optimizer, loss_func, device)
- writer.add_scalar("Loss/val", val_loss, epoch)
-
- # reduce LR on plateau
- scheduler.step(val_loss)
- writer.add_scalar("lr", optimizer.param_groups[0]['lr'], epoch)
-
- # logging
- tqdm.write('At epoch {0:6d} : train loss: {1:.4f}, val loss: {2:.4f}'.format(epoch, train_loss, val_loss))
-
- # check wether to early stop
- early_stopping(val_loss, model, epoch=epoch, save_multiple=(not args.no_save_multiple), save_all=args.save_all_ckpt)
- if early_stopping.early_stop:
- print("Early stopping")
- break
-
- model_log.set_description_str(f'Best val loss: {early_stopping.val_loss_min:.4f}')
-
- if epoch % args.i_eval == 0:
- get_error_in_q(args, val_dl, model, len(val_dl.dataset), device, batch_size=1)
-
-
- writer.flush()
diff --git a/script/dm/prepare_data.py b/script/dm/prepare_data.py
deleted file mode 100755
index ff75ef5..0000000
--- a/script/dm/prepare_data.py
+++ /dev/null
@@ -1,84 +0,0 @@
-import utils.set_sys_path
-import torch
-from torch.utils.data import TensorDataset, DataLoader
-from torchvision import transforms
-import numpy as np
-
-from dataset_loaders.load_llff import load_llff_data
-to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8)
-
-def prepare_data(args, images, poses_train, i_split):
- ''' prepare data for ready to train posenet, return dataloaders '''
- #TODO: Convert GPU friendly data generator later: https://stanford.edu/~shervine/blog/pytorch-how-to-generate-data-parallel
- #TODO: Probably a better implementation style here: https://github.com/PyTorchLightning/pytorch-lightning
-
- i_train, i_val, i_test = i_split
-
- img_train = torch.Tensor(images[i_train]).permute(0, 3, 1, 2) # now shape is [N, CH, H, W]
- pose_train = torch.Tensor(poses_train[i_train])
-
- trainset = TensorDataset(img_train, pose_train)
- train_dl = DataLoader(trainset, batch_size=args.batch_size, shuffle=True)
-
- img_val = torch.Tensor(images[i_val]).permute(0, 3, 1, 2) # now shape is [N, CH, H, W]
- pose_val = torch.Tensor(poses_train[i_val])
-
- valset = TensorDataset(img_val, pose_val)
- val_dl = DataLoader(valset)
-
- img_test = torch.Tensor(images[i_test]).permute(0, 3, 1, 2) # now shape is [N, CH, H, W]
- pose_test = torch.Tensor(poses_train[i_test])
-
- testset = TensorDataset(img_test, pose_test)
- test_dl = DataLoader(testset)
-
- return train_dl, val_dl, test_dl
-
-def load_dataset(args):
- ''' load posenet training data '''
- if args.dataset_type == 'llff':
- if args.no_bd_factor:
- bd_factor = None
- else:
- bd_factor = 0.75
- images, poses, bds, render_poses, i_test = load_llff_data(args.datadir, args.factor,
- recenter=True, bd_factor=bd_factor,
- spherify=args.spherify)
-
- hwf = poses[0,:3,-1]
- poses = poses[:,:3,:4]
- print('Loaded llff', images.shape, render_poses.shape, hwf, args.datadir)
- if not isinstance(i_test, list):
- i_test = [i_test]
-
- if args.llffhold > 0:
- print('Auto LLFF holdout,', args.llffhold)
- i_test = np.arange(images.shape[0])[::args.llffhold]
-
- i_val = i_test
- i_train = np.array([i for i in np.arange(int(images.shape[0])) if
- (i not in i_test and i not in i_val)])
-
- print('DEFINING BOUNDS')
- if args.no_ndc:
- near = np.ndarray.min(bds) * .9
- far = np.ndarray.max(bds) * 1.
-
- else:
- near = 0.
- far = 1.
-
- if args.finetune_unlabel:
- i_train = i_test
- i_split = [i_train, i_val, i_test]
- else:
- print('Unknown dataset type', args.dataset_type, 'exiting')
- return
-
- poses_train = poses[:,:3,:].reshape((poses.shape[0],12)) # get rid of last row [0,0,0,1]
- print("images.shape {}, poses_train.shape {}".format(images.shape, poses_train.shape))
-
- INPUT_SHAPE = images[0].shape
- print("=====================================================================")
- print("INPUT_SHAPE:", INPUT_SHAPE)
- return images, poses_train, render_poses, hwf, i_split, near, far
\ No newline at end of file
diff --git a/script/feature/dfnet.py b/script/feature/dfnet.py
deleted file mode 100755
index 8b86831..0000000
--- a/script/feature/dfnet.py
+++ /dev/null
@@ -1,273 +0,0 @@
-import torch
-from torch import nn
-import torch.nn.functional as F
-from torchvision import models
-from typing import List
-
-# VGG-16 Layer Names and Channels
-vgg16_layers = {
- "conv1_1": 64,
- "relu1_1": 64,
- "conv1_2": 64,
- "relu1_2": 64,
- "pool1": 64,
- "conv2_1": 128,
- "relu2_1": 128,
- "conv2_2": 128,
- "relu2_2": 128,
- "pool2": 128,
- "conv3_1": 256,
- "relu3_1": 256,
- "conv3_2": 256,
- "relu3_2": 256,
- "conv3_3": 256,
- "relu3_3": 256,
- "pool3": 256,
- "conv4_1": 512,
- "relu4_1": 512,
- "conv4_2": 512,
- "relu4_2": 512,
- "conv4_3": 512,
- "relu4_3": 512,
- "pool4": 512,
- "conv5_1": 512,
- "relu5_1": 512,
- "conv5_2": 512,
- "relu5_2": 512,
- "conv5_3": 512,
- "relu5_3": 512,
- "pool5": 512,
-}
-
-class AdaptLayers(nn.Module):
- """Small adaptation layers.
- """
-
- def __init__(self, hypercolumn_layers: List[str], output_dim: int = 128):
- """Initialize one adaptation layer for every extraction point.
-
- Args:
- hypercolumn_layers: The list of the hypercolumn layer names.
- output_dim: The output channel dimension.
- """
- super(AdaptLayers, self).__init__()
- self.layers = []
- channel_sizes = [vgg16_layers[name] for name in hypercolumn_layers]
- for i, l in enumerate(channel_sizes):
- layer = nn.Sequential(
- nn.Conv2d(l, 64, kernel_size=1, stride=1, padding=0),
- nn.ReLU(),
- nn.Conv2d(64, output_dim, kernel_size=5, stride=1, padding=2),
- nn.BatchNorm2d(output_dim),
- )
- self.layers.append(layer)
- self.add_module("adapt_layer_{}".format(i), layer) # ex: adapt_layer_0
-
- def forward(self, features: List[torch.tensor]):
- """Apply adaptation layers. # here is list of three levels of features
- """
-
- for i, _ in enumerate(features):
- features[i] = getattr(self, "adapt_layer_{}".format(i))(features[i])
- return features
-
-class DFNet(nn.Module):
- ''' DFNet implementation '''
- default_conf = {
- 'hypercolumn_layers': ["conv1_2", "conv3_3", "conv5_3"],
- 'output_dim': 128,
- }
- mean = [0.485, 0.456, 0.406]
- std = [0.229, 0.224, 0.225]
-
- def __init__(self, feat_dim=12, places365_model_path=''):
- super(DFNet, self).__init__()
-
- self.layer_to_index = {k: v for v, k in enumerate(vgg16_layers.keys())}
- self.hypercolumn_indices = [self.layer_to_index[n] for n in self.default_conf['hypercolumn_layers']] # [2, 14, 28]
-
- # Initialize architecture
- vgg16 = models.vgg16(pretrained=True)
-
- self.encoder = nn.Sequential(*list(vgg16.features.children()))
-
- self.scales = []
- current_scale = 0
- for i, layer in enumerate(self.encoder):
- if isinstance(layer, torch.nn.MaxPool2d):
- current_scale += 1
- if i in self.hypercolumn_indices:
- self.scales.append(2**current_scale)
-
- ## adaptation layers, see off branches from fig.3 in S2DNet paper
- self.adaptation_layers = AdaptLayers(self.default_conf['hypercolumn_layers'], self.default_conf['output_dim'])
-
- # pose regression layers
- self.avgpool = nn.AdaptiveAvgPool2d(1)
- self.fc_pose = nn.Linear(512, feat_dim)
-
- def forward(self, x, return_feature=False, isSingleStream=False, return_pose=True, upsampleH=240, upsampleW=427):
- '''
- inference DFNet. It can regress camera pose as well as extract intermediate layer features.
- :param x: image blob (2B x C x H x W) two stream or (B x C x H x W) single stream
- :param return_feature: whether to return features as output
- :param isSingleStream: whether it's an single stream inference or siamese network inference
- :param upsampleH: feature upsample size H
- :param upsampleW: feature upsample size W
- :return feature_maps: (2, [B, C, H, W]) or (1, [B, C, H, W]) or None
- :return predict: [2B, 12] or [B, 12]
- '''
- # normalize input data
- mean, std = x.new_tensor(self.mean), x.new_tensor(self.std)
- x = (x - mean[:, None, None]) / std[:, None, None]
-
- ### encoder ###
- feature_maps = []
- for i in range(len(self.encoder)):
- x = self.encoder[i](x)
-
- if i in self.hypercolumn_indices:
- feature = x.clone()
- feature_maps.append(feature)
-
- if i==self.hypercolumn_indices[-1]:
- if return_pose==False:
- predict = None
- break
-
- ### extract and process intermediate features ###
- if return_feature:
- feature_maps = self.adaptation_layers(feature_maps) # (3, [B, C, H', W']), H', W' are different in each layer
-
- if isSingleStream: # not siamese network style inference
- feature_stacks = []
- for f in feature_maps:
- feature_stacks.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(f))
- feature_maps = [torch.stack(feature_stacks)] # (1, [3, B, C, H, W])
- else: # siamese network style inference
- feature_stacks_t = []
- feature_stacks_r = []
- for f in feature_maps:
- # split real and nerf batches
- batch = f.shape[0] # should be target batch_size + rgb batch_size
- feature_t = f[:batch//2]
- feature_r = f[batch//2:]
-
- feature_stacks_t.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(feature_t)) # GT img
- feature_stacks_r.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(feature_r)) # render img
- feature_stacks_t = torch.stack(feature_stacks_t) # [3, B, C, H, W]
- feature_stacks_r = torch.stack(feature_stacks_r) # [3, B, C, H, W]
- feature_maps = [feature_stacks_t, feature_stacks_r] # (2, [3, B, C, H, W])
- else:
- feature_maps = None
-
- if return_pose==False:
- return feature_maps, predict
-
- ### pose regression head ###
- x = self.avgpool(x)
- x = x.reshape(x.size(0), -1)
- predict = self.fc_pose(x)
-
- return feature_maps, predict
-
-class DFNet_s(nn.Module):
- ''' A slight accelerated version of DFNet, we experimentally found this version's performance is similar to original DFNet but inferences faster '''
- default_conf = {
- 'hypercolumn_layers': ["conv1_2"],
- 'output_dim': 128,
- }
- mean = [0.485, 0.456, 0.406]
- std = [0.229, 0.224, 0.225]
-
- def __init__(self, feat_dim=12, places365_model_path=''):
- super(DFNet_s, self).__init__()
-
- self.layer_to_index = {k: v for v, k in enumerate(vgg16_layers.keys())}
- self.hypercolumn_indices = [self.layer_to_index[n] for n in self.default_conf['hypercolumn_layers']] # [2, 14, 28]
-
- # Initialize architecture
- vgg16 = models.vgg16(pretrained=True)
-
- self.encoder = nn.Sequential(*list(vgg16.features.children()))
-
- self.scales = []
- current_scale = 0
- for i, layer in enumerate(self.encoder):
- if isinstance(layer, torch.nn.MaxPool2d):
- current_scale += 1
- if i in self.hypercolumn_indices:
- self.scales.append(2**current_scale)
-
- ## adaptation layers, see off branches from fig.3 in S2DNet paper
- self.adaptation_layers = AdaptLayers(self.default_conf['hypercolumn_layers'], self.default_conf['output_dim'])
-
- # pose regression layers
- self.avgpool = nn.AdaptiveAvgPool2d(1)
- self.fc_pose = nn.Linear(512, feat_dim)
-
- def forward(self, x, return_feature=False, isSingleStream=False, return_pose=True, upsampleH=240, upsampleW=427):
- '''
- inference DFNet_s. It can regress camera pose as well as extract intermediate layer features.
- :param x: image blob (2B x C x H x W) two stream or (B x C x H x W) single stream
- :param return_feature: whether to return features as output
- :param isSingleStream: whether it's an single stream inference or siamese network inference
- :param upsampleH: feature upsample size H
- :param upsampleW: feature upsample size W
- :return feature_maps: (2, [B, C, H, W]) or (1, [B, C, H, W]) or None
- :return predict: [2B, 12] or [B, 12]
- '''
-
- # normalize input data
- mean, std = x.new_tensor(self.mean), x.new_tensor(self.std)
- x = (x - mean[:, None, None]) / std[:, None, None]
-
- ### encoder ###
- feature_maps = []
- for i in range(len(self.encoder)):
- x = self.encoder[i](x)
-
- if i in self.hypercolumn_indices:
- feature = x.clone()
- feature_maps.append(feature)
-
- if i==self.hypercolumn_indices[-1]:
- if return_pose==False:
- predict = None
- break
-
- ### extract and process intermediate features ###
- if return_feature:
- feature_maps = self.adaptation_layers(feature_maps) # (3, [B, C, H', W']), H', W' are different in each layer
-
- if isSingleStream: # not siamese network style inference
- feature_stacks = []
- for f in feature_maps:
- feature_stacks.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(f))
- feature_maps = [torch.stack(feature_stacks)] # (1, [3, B, C, H, W])
- else: # siamese network style inference
- feature_stacks_t = []
- feature_stacks_r = []
- for f in feature_maps:
- # split real and nerf batches
- batch = f.shape[0] # should be target batch_size + rgb batch_size
- feature_t = f[:batch//2]
- feature_r = f[batch//2:]
-
- feature_stacks_t.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(feature_t)) # GT img
- feature_stacks_r.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(feature_r)) # render img
- feature_stacks_t = torch.stack(feature_stacks_t) # [3, B, C, H, W]
- feature_stacks_r = torch.stack(feature_stacks_r) # [3, B, C, H, W]
- feature_maps = [feature_stacks_t, feature_stacks_r] # (2, [3, B, C, H, W])
- else:
- feature_maps = None
-
- if return_pose==False:
- return feature_maps, predict
-
- ### pose regression head ###
- x = self.avgpool(x)
- x = x.reshape(x.size(0), -1)
- predict = self.fc_pose(x)
-
- return feature_maps, predict
\ No newline at end of file
diff --git a/script/feature/direct_feature_matching.py b/script/feature/direct_feature_matching.py
deleted file mode 100755
index 1ca029a..0000000
--- a/script/feature/direct_feature_matching.py
+++ /dev/null
@@ -1,471 +0,0 @@
-import time
-import pdb
-import cv2
-from copy import deepcopy
-import os
-import torch
-import torch.nn as nn
-import torch.nn.init
-import numpy as np
-from tqdm import tqdm
-
-from dm.pose_model import preprocess_data, get_error_in_q
-# from dm.prepare_data import prepare_data
-from dm.direct_pose_model import fix_coord_supp
-from models.nerfw import create_nerf, to8b, img2mse, mse2psnr
-from models.ray_utils import get_rays
-from models.rendering import render
-from utils.utils import freeze_bn_layer_train
-from feature.model import PoseNetV2 as FeatureNet
-from torchvision.utils import save_image
-from utils.utils import plot_features, save_image_saliancy
-
-def tmp_plot2(target_in, rgb_in, features_target, features_rgb, i=0):
- '''
- print 1 pair of batch of salient feature map
- :param: target_in [B, 3, H, W]
- :param: rgb_in [B, 3, H, W]
- :param: features_target [B, C, H, W]
- :param: features_rgb [B, C, H, W]
- :param: frame index i of batch
- '''
- print("for debug only...")
- pdb.set_trace()
- save_image(target_in[i], './tmp/target_in.png')
- save_image(rgb_in[i], './tmp/rgb_in.png')
- features_t = features_target[i].clone()[:, None, :, :]
- features_r = features_rgb[i].clone()[:, None, :, :]
- save_image_saliancy(features_t, './tmp/target', True)
- save_image_saliancy(features_r, './tmp/rgb', True)
-
-def preprocess_features_for_loss(feature):
- '''
- transform output features from the network to required shape for computing loss
- :param: feature [L, B, C, H, W] # L stands for level of features (we currently use 3)
- return feature' [B,L*C,H,W]
- '''
- feature = feature.permute(1,0,2,3,4)
- B, L, C, H, W = feature.size()
- feature = feature.reshape((B,L*C,H,W))
- return feature
-
-def disable_model_grad(model):
- ''' set whole model to requires_grad=False, this is for nerf model '''
- # print("disable_model_grad...")
- for module in model.modules():
- # print("this is a layer:", module)
- if hasattr(module, 'weight'):
- module.weight.requires_grad_(False)
- if hasattr(module, 'bias'):
- module.bias.requires_grad_(False)
- return model
-
-def inference_pose_regression(args, data, device, model, retFeature=False, isSingleStream=True, return_pose=True):
- """
- Inference the Pose Regression Network
- Inputs:
- args: parsed argument
- data: Input image in shape (batchsize, channels, H, W)
- device: gpu device
- model: PoseNet model
- Outputs:
- pose: Predicted Pose in shape (batchsize, 3, 4)
- """
- inputs = data.to(device)
- _,_,H,W = data.size()
- if args.preprocess_ImgNet:
- inputs = preprocess_data(inputs, device)
- if args.DFNet:
- features, predict_pose = model(inputs, return_feature=retFeature, isSingleStream=isSingleStream, return_pose=return_pose, upsampleH=H, upsampleW=W) # features: , predict_pose: [1, 12]
- else:
- features, predict_pose = model(inputs, isTrain=retFeature, isSingleStream=isSingleStream) # features: (1, [1, 1, 320, 8, 14]), predict_pose: [1, 12]
-
- if return_pose==False:
- return features, predict_pose
-
- pose = predict_pose.reshape(inputs.shape[0], 3, 4)
-
- if args.svd_reg:
- R_torch = pose[:,:3,:3].clone()
- u,s,v=torch.svd(R_torch)
- Rs = torch.matmul(u, v.transpose(-2,-1))
- pose[:,:3,:3] = Rs
- return features, pose
-
-def rgb_loss(rgb, target, extras):
- ''' Compute RGB MSE Loss, original from NeRF Paper '''
- # Compute MSE loss between predicted and true RGB.
- img_loss = img2mse(rgb, target)
- loss = img_loss
- return loss
-
-def normalize_features(tensor, value_range=None, scale_each: bool = False):
- ''' Find unit norm of channel wise feature
- :param: tensor, img tensor (C,H,W)
- '''
- tensor = tensor.clone() # avoid modifying tensor in-place
- C,H,W = tensor.size()
-
- # normlaize the features with l2 norm
- tensor = tensor.reshape(C, H*W)
- tensor = torch.nn.functional.normalize(tensor)
- return tensor
-
-def feature_loss(feature_rgb, feature_target, img_in=True, per_channel=False):
- ''' Compute Feature MSE Loss
- :param: feature_rgb, [C,H,W] or [C, N_rand]
- :param: feature_target, [C,H,W] or [C, N_rand]
- :param: img_in, True: input is feature maps, False: input is rays
- :param: random, True: randomly using per pixel or per channel cossimilarity loss
- '''
- if img_in:
- C,H,W = feature_rgb.size()
- fr = feature_rgb.reshape(C, H*W)
- ft = feature_target.reshape(C, H*W)
- else:
- fr = feature_rgb
- ft = feature_target
-
- # cosine loss
- if per_channel:
- cos = torch.nn.CosineSimilarity(dim=0, eps=1e-6)
- else:
- cos = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
- loss = 1 - cos(fr, ft).mean()
-
- return loss
-
-def PoseLoss(args, pose_, pose, device):
- loss_func = nn.MSELoss()
- predict_pose = pose_.reshape(args.batch_size, 12).to(device) # maynot need reshape
- pose_loss = loss_func(predict_pose, pose)
- return pose_loss
-
-def prepare_batch_render(args, pose, batch_size, target_, H, W, focal, half_res=True, rand=True):
- ''' Break batch of images into rays '''
- target_ = target_.permute(0, 2, 3, 1).numpy() # convert to numpy image
- if half_res:
- N_rand = batch_size * (H//2) * (W//2)
- target_half = np.stack([cv2.resize(target_[i], (H//2, W//2), interpolation=cv2.INTER_AREA) for i in range(batch_size)], 0)
- target_half = torch.Tensor(target_half)
-
- rays = torch.stack([torch.stack(get_rays(H//2, W//2, focal/2, pose[i]), 0) for i in range(batch_size)], 0) # [N, ro+rd, H, W, 3] (130, 2, 100, 100, 3)
- rays_rgb = torch.cat((rays, target_half[:, None, ...]), 1)
-
- else:
- # N_rand = batch_size * H * W
- N_rand = args.N_rand
- target_ = torch.Tensor(target_)
- rays = torch.stack([torch.stack(get_rays(H, W, focal, pose[i]), 0) for i in range(batch_size)], 0) # [N, ro+rd, H, W, 3] (130, 2, 200, 200, 3)
- # [N, ro+rd+rgb, H, W, 3]
- rays_rgb = torch.cat([rays, target_[:, None, ...]], 1)
-
- # [N, H, W, ro+rd+rgb, 3]
- rays_rgb = rays_rgb.permute(0, 2, 3, 1, 4)
-
- # [(N-1)*H*W, ro+rd+rgb, 3]
- rays_rgb = torch.reshape(rays_rgb, (-1, 3, 3))
-
- if 1:
- rays_rgb = rays_rgb[torch.randperm(rays_rgb.shape[0])]
-
- # Random over all images
- batch = rays_rgb[:N_rand].permute(1, 0 , 2) # [B, 2+1, 3*?] # (4096, 3, 3)
- batch_rays, target_s = batch[:2], batch[2] # [2, 4096, 3], [4096, 3]
-
- return batch_rays, target_s
-
-def eval_on_batch(args, data, model, feat_model, pose, img_idx, hwf, half_res, device, world_setup_dict, **render_kwargs_test):
- ''' Perform 1 step of eval'''
- with torch.no_grad():
- H, W, focal = hwf
- _, pose_ = inference_pose_regression(args, data, device, model)
- device_cpu = torch.device('cpu')
- pose_ = pose_.to(device_cpu) # put predict pose back to cpu
- pose_nerf = pose_.clone()
-
- if args.NeRFH:
- # rescale the predicted pose to nerf scales
- pose_nerf = fix_coord_supp(args, pose_nerf, world_setup_dict, device=device_cpu)
-
- half_res=False # no need to use half_res for inference
- batch_rays, target = prepare_batch_render(args, pose_nerf, args.batch_size, data, H, W, focal, half_res)
- batch_rays = batch_rays.to(device)
- target = target.to(device)
- pose = pose.to(device)
- img_idx = img_idx.to(device)
- pose_nerf = pose_nerf.to(device)
-
- # every new tensor from onward is in GPU
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- rgb, disp, acc, extras = render(H, W, focal, chunk=args.chunk, rays=batch_rays, img_idx=img_idx, **render_kwargs_test)
-
- loss = PoseLoss(args, pose_, pose, device)
- psnr = mse2psnr(img2mse(rgb, target))
-
- # end of every new tensor from onward is in GPU
- torch.set_default_tensor_type('torch.FloatTensor')
-
- iter_loss = loss.to(device_cpu).detach().numpy()
- iter_loss = np.array([iter_loss])
-
- iter_psnr = psnr.to(device_cpu).detach().numpy()
- return iter_loss, iter_psnr
-
-def eval_on_epoch(args, data_loaders, model, feat_model, hwf, half_res, device, world_setup_dict, **render_kwargs_test):
- ''' Perform 1 epoch of training with batch '''
- model.eval()
- batch_size = 1
-
- train_dl, val_dl, test_dl = data_loaders
-
- total_loss = []
- total_psnr = []
-
- #### Core optimization loop #####
- for data, pose, img_idx in val_dl:
- # training one step with batch_size = args.batch_size
- loss, psnr = eval_on_batch(args, data, model, feat_model, pose, img_idx, hwf, half_res, device, world_setup_dict, **render_kwargs_test)
- total_loss.append(loss.item())
- total_psnr.append(psnr.item())
- total_loss_mean = np.mean(total_loss)
- total_psnr_mean = np.mean(total_psnr)
- return total_loss_mean, total_psnr_mean
-
-def train_on_feature_batch(args, data, model, feat_model, pose, img_idx, hwf, optimizer, device, world_setup_dict, **render_kwargs_test):
- ''' Perform 1 step of training using scheme1 '''
- batch_size_iter = data.shape[0]
-
- H, W, focal = hwf
- data = data.to(device) # [1, 3, 240, 427]
-
- # pose regression module
- _, pose_ = inference_pose_regression(args, data, device, model, retFeature=False) # here returns predicted pose [1, 3, 4] # real img features and predicted pose # features: (1, [3, 1, 128, 240, 427]), predict_pose: [1, 3, 4]
- pose_nerf = pose_.clone()
-
- # rescale the predicted pose to nerf scales
- pose_nerf = fix_coord_supp(args, pose_nerf, world_setup_dict, device=device)
-
- pose = pose.to(device)
- img_idx = img_idx.to(device)
- # every new tensor from onward is in GPU, here memory cost is a bottleneck
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
-
- # here is single frame
- target = data.permute(0,2,3,1) # [B,H,W,C]
- rays_o_list=[]
- rays_d_list=[]
- img_idx_list=[]
- N_rand = args.N_rand
- for i in range(pose_nerf.shape[0]):
- rays_o, rays_d = get_rays(H, W, focal, pose_nerf[i]) # (H, W, 3), (H, W, 3)
- rays_o_list.append(rays_o)
- rays_d_list.append(rays_d)
- img_idx_list.append(img_idx[i].repeat(N_rand,1))
- rays_o_batch = torch.stack(rays_o_list)
- rays_d_batch = torch.stack(rays_d_list)
- img_idx_batch = torch.cat(img_idx_list)
-
- # randomly select coords
- coords = torch.stack(torch.meshgrid(torch.linspace(0, H-1, H), torch.linspace(0, W-1, W), indexing='ij'), -1) # (H, W, 2)
- coords = torch.reshape(coords, [-1,2]) # (H * W, 2)
- select_inds = np.random.choice(coords.shape[0], size=[N_rand], replace=False) # (N_rand,)
- select_coords = coords[select_inds].long() # (N_rand, 2)
-
- # fetch from coords
- rays_o = rays_o_batch[:, select_coords[:, 0], select_coords[:, 1]]
- rays_d = rays_d_batch[:, select_coords[:, 0], select_coords[:, 1]]
- rays_o = rays_o.reshape(rays_o.shape[0]*rays_o.shape[1], 3) # (B*N_rand, 3)
- rays_d = rays_d.reshape(rays_d.shape[0]*rays_d.shape[1], 3) # (B*N_rand, 3)
- batch_rays = torch.stack([rays_o, rays_d], 0)
- target_s = target[:,select_coords[:, 0], select_coords[:, 1]].reshape(batch_size_iter*N_rand,3) # (B*N_rand, 3)
-
- rgb_feature, disp, acc, extras = render(H, W, focal, chunk=args.chunk, rays=batch_rays, img_idx=img_idx_batch, **render_kwargs_test)
- # rgb_feature is rgb 3 + features 128
- rgb = rgb_feature[...,:3] # [B*N_rand, 3]
- feature = rgb_feature[...,3:].reshape(batch_size_iter, N_rand, args.out_channel_size-3)[None, ...].permute(0,1,3,2) # [lvl, B, C, N_rand] assuming lvl size = 1
-
- # inference featurenet
- target_in = target.permute(0,3,1,2)
- features, _ = feat_model(target_in, True, True, H, W) # features: (1, [3,B,C,H,W])
-
- # get features_target, # now choose 1st level feature only
- feature_target = features[0][0] # [B,C,H,W]
- feature_target = feature_target[None, 0:, :, select_coords[:, 0], select_coords[:, 1]] # # [lvl, B, C, N_rand] assuming lvl size = 1
-
- ### Loss Design Here ###
- # Compute RGB MSE Loss
- photo_loss = rgb_loss(rgb, target_s, extras)
-
- feat_loss = feature_loss(feature[0,0], feature_target[0,0], img_in=False, per_channel=args.per_channel) # TODO: questionable implementation. Here we assume lvl size=1, batch_size=1
-
- # Compute Combine Loss if needed
- if args.combine_loss:
- pose_loss = PoseLoss(args, pose_, pose, device)
- loss = args.combine_loss_w[0] * pose_loss + args.combine_loss_w[1] * photo_loss + args.combine_loss_w[2] * feat_loss
-
- ### Loss Design End
- loss.backward()
- optimizer.step()
- optimizer.zero_grad()
- psnr = mse2psnr(img2mse(rgb, target_s))
-
- # end of every new tensor from onward is in GPU
- torch.set_default_tensor_type('torch.FloatTensor')
- device_cpu = torch.device('cpu')
- iter_loss = loss.to(device_cpu).detach().numpy()
- iter_loss = np.array([iter_loss])
-
- iter_psnr = psnr.to(device_cpu).detach().numpy()
- return iter_loss, iter_psnr
-
-def train_on_batch(args, data, model, feat_model, pose, img_idx, hwf, optimizer, half_res, device, world_setup_dict, **render_kwargs_test):
- ''' Perform 1 step of training '''
-
- H, W, focal = hwf
- data = data.to(device) # [1, 3, 240, 427] non_blocking=True
-
- # pose regression module
- _, pose_ = inference_pose_regression(args, data, device, model, retFeature=False)
- pose_nerf = pose_.clone()
-
- # direct matching module
- # rescale the predicted pose to nerf scales
- pose_nerf = fix_coord_supp(args, pose_nerf, world_setup_dict, device=device)
-
- pose = pose.to(device)
- img_idx = img_idx.to(device)
- # every new tensor from onward is in GPU, here memory cost is a bottleneck
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
-
- if half_res:
- rgb, disp, acc, extras = render(H//4, W//4, focal/4, chunk=args.chunk, c2w=pose_nerf[0,:3,:4], img_idx=img_idx, **render_kwargs_test)
- # convert rgb to B,C,H,W
- rgb = rgb[None,...].permute(0,3,1,2)
- # upsample rgb to hwf size
- rgb = torch.nn.Upsample(size=(H, W), mode='bicubic')(rgb)
- # # convert rgb back to H,W,C format
- # rgb = rgb[0].permute(1,2,0)
- else:
- rgb, disp, acc, extras = render(H, W, focal, chunk=args.chunk, c2w=pose_nerf[0,:3,:4], img_idx=img_idx, **render_kwargs_test)
- rgb = rgb[None,...].permute(0,3,1,2)
-
- # feature metric module
- feature_list, _ = inference_pose_regression(args, torch.cat([data, rgb]), device, feat_model, retFeature=True, isSingleStream=False, return_pose=False)
- feature_target = feature_list[0]
- feature_rgb = feature_list[1]
-
- ### Loss Design Here ###
- # Compute RGB MSE Loss
- photo_loss = rgb_loss(rgb, data, extras)
-
- # Compute Feature MSE Loss
- indices = torch.tensor(args.feature_matching_lvl)
- feature_rgb = torch.index_select(feature_rgb, 0, indices)
- feature_target = torch.index_select(feature_target, 0, indices)
-
- feature_rgb = preprocess_features_for_loss(feature_rgb)
- feature_target = preprocess_features_for_loss(feature_target)
-
- feat_loss = feature_loss(feature_rgb[0], feature_target[0], per_channel=args.per_channel)
-
- # Compute Combine Loss if needed
- if args.combine_loss:
- pose_loss = PoseLoss(args, pose_, pose, device)
- loss = args.combine_loss_w[0] * pose_loss + args.combine_loss_w[1] * photo_loss + args.combine_loss_w[2] * feat_loss
-
- ### Loss Design End
- loss.backward()
- optimizer.step()
- optimizer.zero_grad()
- psnr = mse2psnr(img2mse(rgb, data))
-
- # end of every new tensor from onward is in GPU
- torch.set_default_tensor_type('torch.FloatTensor')
- device_cpu = torch.device('cpu')
- iter_loss = loss.to(device_cpu).detach().numpy()
- iter_loss = np.array([iter_loss])
-
- iter_psnr = psnr.to(device_cpu).detach().numpy()
- return iter_loss, iter_psnr
-
-def train_on_epoch(args, data_loaders, model, feat_model, hwf, optimizer, half_res, device, world_setup_dict, **render_kwargs_test):
- ''' Perform 1 epoch of training with batch '''
- model.train()
- model = freeze_bn_layer_train(model)
-
- # Prepare dataloaders for PoseNet, each batch contains (image, pose)
- train_dl, val_dl, test_dl = data_loaders
- total_loss = []
- total_psnr = []
-
- #### Core optimization loop #####
- for data, pose, img_idx in train_dl:
- loss, psnr = train_on_batch(args, data, model, feat_model, pose, img_idx, hwf, optimizer, half_res, device, world_setup_dict, **render_kwargs_test)
-
- total_loss.append(loss.item())
- total_psnr.append(psnr.item())
- total_loss_mean = np.mean(total_loss)
- total_psnr_mean = np.mean(total_psnr)
- return total_loss_mean, total_psnr_mean
-
-def train_feature_matching(args, model, feat_model, optimizer, i_split, hwf, near, far, device, early_stopping, images=None, poses_train=None, train_dl=None, val_dl=None, test_dl=None):
- ''' finetune pretrained PoseNet using NeRF '''
- # half_res = False # direct-pn paper settings
- half_res = True # debug
-
- # load NeRF model
- _, render_kwargs_test, start, grad_vars, _ = create_nerf(args)
- global_step = start
- if args.reduce_embedding==2:
- render_kwargs_test['i_epoch'] = global_step
-
- data_loaders = [train_dl, val_dl, test_dl]
- bds_dict = {
- 'near' : near,
- 'far' : far,
- }
- # render_kwargs_train.update(bds_dict)
- render_kwargs_test.update(bds_dict)
- i_train, i_val, i_test = i_split
-
- render_kwargs_test['embedding_a'] = disable_model_grad(render_kwargs_test['embedding_a'])
- render_kwargs_test['embedding_t'] = disable_model_grad(render_kwargs_test['embedding_t'])
- render_kwargs_test['network_fn'] = disable_model_grad(render_kwargs_test['network_fn'])
- render_kwargs_test['network_fine'] = disable_model_grad(render_kwargs_test['network_fine'])
-
- N_epoch = 2001
- print('Begin')
- print('TRAIN views are', i_train)
- print('TEST views are', i_test)
- print('VAL views are', i_val)
-
- world_setup_dict = {
- 'pose_scale' : train_dl.dataset.pose_scale,
- 'pose_scale2' : train_dl.dataset.pose_scale2,
- 'move_all_cam_vec' : train_dl.dataset.move_all_cam_vec,
- }
-
- time0 = time.time()
-
- model_log = tqdm(total=0, position = 1, bar_format='{desc}')
- for epoch in tqdm(range(N_epoch), desc='epochs'):
- #train 1 epoch with batch_size = 1, 15% speed up for DFNet_s
- loss, psnr = train_on_epoch(args, data_loaders, model, feat_model, hwf, optimizer, half_res, device, world_setup_dict, **render_kwargs_test)
-
- # 26% speed up for DFNet_s
- val_loss, val_psnr = eval_on_epoch(args, data_loaders, model, feat_model, hwf, half_res, device, world_setup_dict, **render_kwargs_test)
-
-
- tqdm.write('At epoch {0:4d} : train loss: {1:.4f}, train psnr: {2:.4f}, val loss: {3:.4f}, val psnr: {4:.4f}'.format(epoch, loss, psnr, val_loss, val_psnr))
-
- # check wether to early stop
- early_stopping(val_loss, model, epoch=epoch, save_multiple=(not args.no_save_multiple), save_all=args.save_all_ckpt, val_psnr=val_psnr)
- if early_stopping.early_stop:
- print("Early stopping")
- break
- model_log.set_description_str(f'Best val loss: {early_stopping.val_loss_min:.4f}')
-
- if epoch % args.i_eval == 0:
- # calculate position and angular error
- get_error_in_q(args, val_dl, model, len(val_dl.dataset), device, batch_size=1)
diff --git a/script/feature/efficientnet.py b/script/feature/efficientnet.py
deleted file mode 100755
index a9edd92..0000000
--- a/script/feature/efficientnet.py
+++ /dev/null
@@ -1,285 +0,0 @@
-import torch
-from torch import nn
-import torch.nn.functional as F
-import copy, pdb
-from typing import List
-from efficientnet_pytorch import EfficientNet
-
-# efficientnet-B3 Layer Names and Channels
-EB3_layers = {
- "reduction_1": 24, # torch.Size([2, 24, 120, 213])
- "reduction_2": 32, # torch.Size([2, 32, 60, 106])
- "reduction_3": 48, # torch.Size([2, 48, 30, 53])
- "reduction_4": 136, # torch.Size([2, 136, 15, 26])
- "reduction_5": 384, # torch.Size([2, 384, 8, 13])
- "reduction_6": 1536, # torch.Size([2, 1536, 8, 13])
-}
-
-# efficientnet-B0 Layer Names and Channels
-EB0_layers = {
- "reduction_1": 16, # torch.Size([2, 16, 120, 213])
- "reduction_2": 24, # torch.Size([2, 24, 60, 106])
- "reduction_3": 40, # torch.Size([2, 40, 30, 53])
- "reduction_4": 112, # torch.Size([2, 112, 15, 26])
- "reduction_5": 320, # torch.Size([2, 320, 8, 13])
- "reduction_6": 1280, # torch.Size([2, 1280, 8, 13])
-}
-
-class AdaptLayers(nn.Module):
- """Small adaptation layers.
- """
-
- def __init__(self, hypercolumn_layers: List[str], output_dim: int = 128):
- """Initialize one adaptation layer for every extraction point.
-
- Args:
- hypercolumn_layers: The list of the hypercolumn layer names.
- output_dim: The output channel dimension.
- """
- super(AdaptLayers, self).__init__()
- self.layers = []
- channel_sizes = [EB3_layers[name] for name in hypercolumn_layers]
- for i, l in enumerate(channel_sizes):
- layer = nn.Sequential(
- nn.Conv2d(l, 64, kernel_size=1, stride=1, padding=0),
- nn.ReLU(),
- nn.Conv2d(64, output_dim, kernel_size=5, stride=1, padding=2),
- nn.BatchNorm2d(output_dim),
- )
- self.layers.append(layer)
- self.add_module("adapt_layer_{}".format(i), layer) # ex: adapt_layer_0
-
- def forward(self, features: List[torch.tensor]):
- """Apply adaptation layers. # here is list of three levels of features
- """
-
- for i, _ in enumerate(features):
- features[i] = getattr(self, "adapt_layer_{}".format(i))(features[i])
- return features
-
-class EfficientNetB3(nn.Module):
- ''' DFNet with EB3 backbone '''
- default_conf = {
- # 'hypercolumn_layers': ["reduction_1", "reduction_3", "reduction_6"],
- 'hypercolumn_layers': ["reduction_1", "reduction_3", "reduction_5"],
- # 'hypercolumn_layers': ["reduction_2", "reduction_4", "reduction_6"],
- 'output_dim': 128,
- }
- mean = [0.485, 0.456, 0.406]
- std = [0.229, 0.224, 0.225]
-
- def __init__(self, feat_dim=12, places365_model_path=''):
- super(EfficientNetB3, self).__init__()
- # Initialize architecture
- self.backbone_net = EfficientNet.from_pretrained('efficientnet-b3')
- self.feature_extractor = self.backbone_net.extract_endpoints
-
- # self.feature_block_index = [1, 3, 6] # same as the 'hypercolumn_layers'
- self.feature_block_index = [1, 3, 5] # same as the 'hypercolumn_layers'
- # self.feature_block_index = [2, 4, 6] # same as the 'hypercolumn_layers'
-
- ## adaptation layers, see off branches from fig.3 in S2DNet paper
- self.adaptation_layers = AdaptLayers(self.default_conf['hypercolumn_layers'], self.default_conf['output_dim'])
-
- # pose regression layers
- self.avgpool = nn.AdaptiveAvgPool2d(1)
- self.fc_pose = nn.Linear(1536, feat_dim)
-
- def forward(self, x, return_feature=False, isSingleStream=False, upsampleH=120, upsampleW=213):
- '''
- inference DFNet. It can regress camera pose as well as extract intermediate layer features.
- :param x: image blob (2B x C x H x W) two stream or (B x C x H x W) single stream
- :param return_feature: whether to return features as output
- :param isSingleStream: whether it's an single stream inference or siamese network inference
- :param upsampleH: feature upsample size H
- :param upsampleW: feature upsample size W
- :return feature_maps: (2, [B, C, H, W]) or (1, [B, C, H, W]) or None
- :return predict: [2B, 12] or [B, 12]
- '''
- # normalize input data
- mean, std = x.new_tensor(self.mean), x.new_tensor(self.std)
- x = (x - mean[:, None, None]) / std[:, None, None]
-
- ### encoder ###
- feature_maps = []
- list_x = self.feature_extractor(x)
-
- x = list_x['reduction_6'] # features to save
- for i in self.feature_block_index:
- fe = list_x['reduction_'+str(i)].clone()
- feature_maps.append(fe)
-
- ### extract and process intermediate features ###
- if return_feature:
- feature_maps = self.adaptation_layers(feature_maps) # (3, [B, C, H', W']), H', W' are different in each layer
-
- pdb.set_trace()
- if isSingleStream: # not siamese network style inference
- feature_stacks = []
-
- for f in feature_maps:
- feature_stacks.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(f))
- feature_maps = [torch.stack(feature_stacks)] # (1, [3, B, C, H, W])
- else: # siamese network style inference
- feature_stacks_t = []
- feature_stacks_r = []
-
- for f in feature_maps:
- # split real and nerf batches
- batch = f.shape[0] # should be target batch_size + rgb batch_size
- feature_t = f[:batch//2]
- feature_r = f[batch//2:]
-
- feature_stacks_t.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(feature_t)) # GT img
- feature_stacks_r.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(feature_r)) # render img
- feature_stacks_t = torch.stack(feature_stacks_t) # [3, B, C, H, W]
- feature_stacks_r = torch.stack(feature_stacks_r) # [3, B, C, H, W]
- feature_maps = [feature_stacks_t, feature_stacks_r] # (2, [3, B, C, H, W])
-
- else:
- feature_maps = None
-
- ### pose regression head ###
- x = self.avgpool(x)
- x = x.reshape(x.size(0), -1)
- predict = self.fc_pose(x)
-
- return feature_maps, predict
-
-class AdaptLayers2(nn.Module):
- """Small adaptation layers.
- """
-
- def __init__(self, hypercolumn_layers: List[str], output_dim: int = 128):
- """Initialize one adaptation layer for every extraction point.
-
- Args:
- hypercolumn_layers: The list of the hypercolumn layer names.
- output_dim: The output channel dimension.
- """
- super(AdaptLayers2, self).__init__()
- self.layers = []
- channel_sizes = [EB0_layers[name] for name in hypercolumn_layers]
- for i, l in enumerate(channel_sizes):
- layer = nn.Sequential(
- nn.Conv2d(l, 64, kernel_size=1, stride=1, padding=0),
- nn.ReLU(),
- nn.Conv2d(64, output_dim, kernel_size=5, stride=1, padding=2),
- nn.BatchNorm2d(output_dim),
- )
- self.layers.append(layer)
- self.add_module("adapt_layer_{}".format(i), layer) # ex: adapt_layer_0
-
- def forward(self, features: List[torch.tensor]):
- """Apply adaptation layers. # here is list of three levels of features
- """
-
- for i, _ in enumerate(features):
- features[i] = getattr(self, "adapt_layer_{}".format(i))(features[i])
- return features
-
-class EfficientNetB0(nn.Module):
- ''' DFNet with EB0 backbone, feature levels can be customized '''
- default_conf = {
- # 'hypercolumn_layers': ["reduction_1", "reduction_3", "reduction_6"],
- 'hypercolumn_layers': ["reduction_1", "reduction_3", "reduction_5"],
- # 'hypercolumn_layers': ["reduction_2", "reduction_4", "reduction_6"],
- # 'hypercolumn_layers': ["reduction_1"],
- 'output_dim': 128,
- }
- mean = [0.485, 0.456, 0.406]
- std = [0.229, 0.224, 0.225]
-
- def __init__(self, feat_dim=12, places365_model_path=''):
- super(EfficientNetB0, self).__init__()
- # Initialize architecture
- self.backbone_net = EfficientNet.from_pretrained('efficientnet-b0')
- self.feature_extractor = self.backbone_net.extract_endpoints
-
- # self.feature_block_index = [1, 3, 6] # same as the 'hypercolumn_layers'
- self.feature_block_index = [1, 3, 5] # same as the 'hypercolumn_layers'
- # self.feature_block_index = [2, 4, 6] # same as the 'hypercolumn_layers'
- # self.feature_block_index = [1]
-
- ## adaptation layers, see off branches from fig.3 in S2DNet paper
- self.adaptation_layers = AdaptLayers2(self.default_conf['hypercolumn_layers'], self.default_conf['output_dim'])
-
- # pose regression layers
- self.avgpool = nn.AdaptiveAvgPool2d(1)
- self.fc_pose = nn.Linear(1280, feat_dim)
-
- def forward(self, x, return_feature=False, isSingleStream=False, return_pose=False, upsampleH=120, upsampleW=213):
- '''
- inference DFNet. It can regress camera pose as well as extract intermediate layer features.
- :param x: image blob (2B x C x H x W) two stream or (B x C x H x W) single stream
- :param return_feature: whether to return features as output
- :param isSingleStream: whether it's an single stream inference or siamese network inference
- :param return_pose: TODO: if only return_pose, we don't need to compute return_feature part
- :param upsampleH: feature upsample size H
- :param upsampleW: feature upsample size W
- :return feature_maps: (2, [B, C, H, W]) or (1, [B, C, H, W]) or None
- :return predict: [2B, 12] or [B, 12]
- '''
- # normalize input data
- mean, std = x.new_tensor(self.mean), x.new_tensor(self.std)
- x = (x - mean[:, None, None]) / std[:, None, None]
-
- ### encoder ###
- feature_maps = []
- list_x = self.feature_extractor(x)
-
- x = list_x['reduction_6'] # features to save
- for i in self.feature_block_index:
- fe = list_x['reduction_'+str(i)].clone()
- feature_maps.append(fe)
-
- ### extract and process intermediate features ###
- if return_feature:
- feature_maps = self.adaptation_layers(feature_maps) # (3, [B, C, H', W']), H', W' are different in each layer
-
- if isSingleStream: # not siamese network style inference
- feature_stacks = []
-
- for f in feature_maps:
- feature_stacks.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(f))
- feature_maps = [torch.stack(feature_stacks)] # (1, [3, B, C, H, W])
- else: # siamese network style inference
- feature_stacks_t = []
- feature_stacks_r = []
-
- for f in feature_maps:
-
- # split real and nerf batches
- batch = f.shape[0] # should be target batch_size + rgb batch_size
- feature_t = f[:batch//2]
- feature_r = f[batch//2:]
-
- feature_stacks_t.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(feature_t)) # GT img
- feature_stacks_r.append(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(feature_r)) # render img
- feature_stacks_t = torch.stack(feature_stacks_t) # [3, B, C, H, W]
- feature_stacks_r = torch.stack(feature_stacks_r) # [3, B, C, H, W]
- feature_maps = [feature_stacks_t, feature_stacks_r] # (2, [3, B, C, H, W])
-
- else:
- feature_maps = None
-
- ### pose regression head ###
- x = self.avgpool(x)
- x = x.reshape(x.size(0), -1)
- predict = self.fc_pose(x)
-
- return feature_maps, predict
-
-def main():
- """
- test model
- """
- from torchsummary import summary
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
- feat_model = EfficientNetB3()
- # feat_model.load_state_dict(torch.load(''))
- feat_model.to(device)
- summary(feat_model, (3, 240, 427))
-
-if __name__ == '__main__':
- main()
diff --git a/script/feature/misc.py b/script/feature/misc.py
deleted file mode 100755
index f476936..0000000
--- a/script/feature/misc.py
+++ /dev/null
@@ -1,519 +0,0 @@
-import torch
-from torch import nn
-import pdb
-# from torchvision import transforms, datasets
-import pytorch3d.transforms as transforms
-import numpy as np
-import math
-from utils.utils import plot_features, save_image_saliancy
-from dm.direct_pose_model import fix_coord_supp
-from dm.pose_model import vis_pose
-from models.rendering import render, render_path
-import time
-import os
-import imageio
-from copy import deepcopy
-
-img2mse = lambda x, y : torch.mean((x - y) ** 2)
-mse2psnr = lambda x : -10. * torch.log(x) / torch.log(torch.Tensor([10.]))
-to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8)
-
-# translation z axis
-trans_t = lambda t : np.array([
- [1,0,0,0],
- [0,1,0,0],
- [0,0,1,t],
- [0,0,0,1]], dtype=np.float)
-
-# x rotation
-rot_phi = lambda phi : np.array([
- [1,0,0,0],
- [0,np.cos(phi),-np.sin(phi),0],
- [0,np.sin(phi), np.cos(phi),0],
- [0,0,0,1]], dtype=np.float)
-
-# y rotation
-rot_theta = lambda th : np.array([
- [np.cos(th),0,-np.sin(th),0],
- [0,1,0,0],
- [np.sin(th),0, np.cos(th),0],
- [0,0,0,1]], dtype=np.float)
-
-# z rotation
-rot_psi = lambda psi : np.array([
- [np.cos(psi),-np.sin(psi),0,0],
- [np.sin(psi),np.cos(psi),0,0],
- [0,0,1,0],
- [0,0,0,1]], dtype=np.float)
-
-def compute_error_in_q(args, dl, model, device, results, batch_size=1):
- use_SVD=True # Turn on for Direct-PN and Direct-PN+U reported result, despite it makes minuscule differences
- time_spent = []
- predict_pose_list = []
- gt_pose_list = []
- ang_error_list = []
- i = 0
- for batch in dl:
- if args.NeRFH:
- data, pose, img_idx = batch
- else:
- data, pose = batch
- data = data.to(device) # input
- pose = pose.reshape((batch_size,3,4)).numpy() # label
-
- if use_SVD:
- # using SVD to make sure predict rotation is normalized rotation matrix
- with torch.no_grad():
- _, predict_pose = model(data)
- R_torch = predict_pose.reshape((batch_size, 3, 4))[:,:3,:3] # debug
- predict_pose = predict_pose.reshape((batch_size, 3, 4)).cpu().numpy()
-
- R = predict_pose[:,:3,:3]
- res = R@np.linalg.inv(R)
- # print('R@np.linalg.inv(R):', res)
-
- u,s,v=torch.svd(R_torch)
- Rs = torch.matmul(u, v.transpose(-2,-1))
- predict_pose[:,:3,:3] = Rs[:,:3,:3].cpu().numpy()
- else:
- start_time = time.time()
- # inference NN
- with torch.no_grad():
- predict_pose = model(data)
- predict_pose = predict_pose.reshape((batch_size, 3, 4)).cpu().numpy()
- time_spent.append(time.time() - start_time)
-
- pose_q = transforms.matrix_to_quaternion(torch.Tensor(pose[:,:3,:3]))#.cpu().numpy() # gnd truth in quaternion
- pose_x = pose[:, :3, 3] # gnd truth position
- predicted_q = transforms.matrix_to_quaternion(torch.Tensor(predict_pose[:,:3,:3]))#.cpu().numpy() # predict in quaternion
- predicted_x = predict_pose[:, :3, 3] # predict position
- pose_q = pose_q.squeeze()
- pose_x = pose_x.squeeze()
- predicted_q = predicted_q.squeeze()
- predicted_x = predicted_x.squeeze()
-
- #Compute Individual Sample Error
- q1 = pose_q / torch.linalg.norm(pose_q)
- q2 = predicted_q / torch.linalg.norm(predicted_q)
- d = torch.abs(torch.sum(torch.matmul(q1,q2)))
- d = torch.clamp(d, -1., 1.) # acos can only input [-1~1]
- theta = (2 * torch.acos(d) * 180/math.pi).numpy()
- error_x = torch.linalg.norm(torch.Tensor(pose_x-predicted_x)).numpy()
- results[i,:] = [error_x, theta]
- #print ('Iteration: {} Error XYZ (m): {} Error Q (degrees): {}'.format(i, error_x, theta))
-
- # save results for visualization
- predict_pose_list.append(predicted_x)
- gt_pose_list.append(pose_x)
- ang_error_list.append(theta)
- i += 1
- predict_pose_list = np.array(predict_pose_list)
- gt_pose_list = np.array(gt_pose_list)
- ang_error_list = np.array(ang_error_list)
- vis_info_ret = {"pose": predict_pose_list, "pose_gt": gt_pose_list, "theta": ang_error_list}
- return results, vis_info_ret
-
-# # pytorch
-def get_error_in_q(args, dl, model, sample_size, device, batch_size=1):
- ''' Convert Rotation matrix to quaternion, then calculate the location errors. original from PoseNet Paper '''
- model.eval()
-
- results = np.zeros((sample_size, 2))
- results, vis_info = compute_error_in_q(args, dl, model, device, results, batch_size)
- median_result = np.median(results,axis=0)
- mean_result = np.mean(results,axis=0)
-
- # standard log
- print ('Median error {}m and {} degrees.'.format(median_result[0], median_result[1]))
- print ('Mean error {}m and {} degrees.'.format(mean_result[0], mean_result[1]))
-
- # visualize results
- # vis_pose(vis_info)
-
-def get_render_error_in_q(args, model, sample_size, device, targets, rgbs, poses, batch_size=1):
- ''' use nerf render imgs instead of use real imgs '''
- model.eval()
-
- results = np.zeros((sample_size, 2))
- print("to be implement...")
-
- predict_pose_list = []
- gt_pose_list = []
- ang_error_list = []
-
- for i in range(sample_size):
- data = rgbs[i:i+1].permute(0,3,1,2)
- pose = poses[i:i+1].reshape(batch_size, 12)
-
- data = data.to(device) # input
- pose = pose.reshape((batch_size,3,4)).numpy() # label
-
- # using SVD to make sure predict rotation is normalized rotation matrix
- with torch.no_grad():
- _, predict_pose = model(data)
- R_torch = predict_pose.reshape((batch_size, 3, 4))[:,:3,:3] # debug
- predict_pose = predict_pose.reshape((batch_size, 3, 4)).cpu().numpy()
-
- R = predict_pose[:,:3,:3]
- res = R@np.linalg.inv(R)
- # print('R@np.linalg.inv(R):', res)
-
- u,s,v=torch.svd(R_torch)
- Rs = torch.matmul(u, v.transpose(-2,-1))
- predict_pose[:,:3,:3] = Rs[:,:3,:3].cpu().numpy()
-
- pose_q = transforms.matrix_to_quaternion(torch.Tensor(pose[:,:3,:3]))#.cpu().numpy() # gnd truth in quaternion
- pose_x = pose[:, :3, 3] # gnd truth position
- predicted_q = transforms.matrix_to_quaternion(torch.Tensor(predict_pose[:,:3,:3]))#.cpu().numpy() # predict in quaternion
- predicted_x = predict_pose[:, :3, 3] # predict position
- pose_q = pose_q.squeeze()
- pose_x = pose_x.squeeze()
- predicted_q = predicted_q.squeeze()
- predicted_x = predicted_x.squeeze()
-
- #Compute Individual Sample Error
- q1 = pose_q / torch.linalg.norm(pose_q)
- q2 = predicted_q / torch.linalg.norm(predicted_q)
- d = torch.abs(torch.sum(torch.matmul(q1,q2)))
- d = torch.clamp(d, -1., 1.) # acos can only input [-1~1]
- theta = (2 * torch.acos(d) * 180/math.pi).numpy()
- error_x = torch.linalg.norm(torch.Tensor(pose_x-predicted_x)).numpy()
- results[i,:] = [error_x, theta]
- #print ('Iteration: {} Error XYZ (m): {} Error Q (degrees): {}'.format(i, error_x, theta))
-
- # save results for visualization
- predict_pose_list.append(predicted_x)
- gt_pose_list.append(pose_x)
- ang_error_list.append(theta)
-
- predict_pose_list = np.array(predict_pose_list)
- gt_pose_list = np.array(gt_pose_list)
- ang_error_list = np.array(ang_error_list)
- vis_info_ret = {"pose": predict_pose_list, "pose_gt": gt_pose_list, "theta": ang_error_list}
-
- median_result = np.median(results,axis=0)
- mean_result = np.mean(results,axis=0)
-
- # standard log
- print ('Median error {}m and {} degrees.'.format(median_result[0], median_result[1]))
- print ('Mean error {}m and {} degrees.'.format(mean_result[0], mean_result[1]))
-
- return
-
-def render_nerfw_imgs(args, dl, hwf, device, render_kwargs_test, world_setup_dict):
- ''' render nerfw imgs, save unscaled pose and results'''
- H, W, focal = hwf
- target_list = []
- rgb_list = []
- pose_list = []
- img_idx_list = []
- # profiling code
- # time0 = time.time()
-
- # inference nerfw and save rgb, target, pose
- for batch_idx, (target, pose, img_idx) in enumerate(dl):
- if batch_idx % 10 == 0:
- print("renders {}/total {}".format(batch_idx, len(dl.dataset)))
-
- target = target[0].permute(1,2,0).to(device) # (240,360,3)
- pose = pose.reshape(3,4) # reshape to 3x4 rot matrix
-
- img_idx = img_idx.to(device)
- pose_nerf = pose.clone()
-
- # rescale the predicted pose to nerf scales
- pose_nerf = fix_coord_supp(args, pose_nerf[None,...], world_setup_dict)
-
- # generate nerf image
- with torch.no_grad():
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- if args.tinyimg:
- rgb, _, _, _ = render(int(H//args.tinyscale), int(W//args.tinyscale), focal/args.tinyscale, chunk=args.chunk, c2w=pose_nerf[0,:3,:4].to(device), retraw=True, img_idx=img_idx, **render_kwargs_test)
- # convert rgb to B,C,H,W
- rgb = rgb[None,...].permute(0,3,1,2)
- # upsample rgb to hwf size
- rgb = torch.nn.Upsample(size=(H, W), mode='bicubic')(rgb)
- # convert rgb back to H,W,C format
- rgb = rgb[0].permute(1,2,0)
-
- else:
- rgb, _, _, _ = render(H, W, focal, chunk=args.chunk, c2w=pose_nerf[0,:3,:4].to(device), retraw=True, img_idx=img_idx, **render_kwargs_test)
- torch.set_default_tensor_type('torch.FloatTensor')
-
- target_list.append(target.cpu())
- rgb_list.append(rgb.cpu())
- pose_list.append(pose.cpu())
- img_idx_list.append(img_idx.cpu())
-
- targets = torch.stack(target_list).detach()
- rgbs = torch.stack(rgb_list).detach()
- poses = torch.stack(pose_list).detach()
- img_idxs = torch.stack(img_idx_list).detach()
- return targets, rgbs, poses, img_idxs
-
-def render_virtual_imgs(args, pose_perturb, img_idxs, hwf, device, render_kwargs_test, world_setup_dict):
- ''' render nerfw imgs, save unscaled pose and results'''
- H, W, focal = hwf
- rgb_list = []
-
- # inference nerfw and save rgb, target, pose
- for batch_idx in range(pose_perturb.shape[0]):
- # if batch_idx % 10 == 0:
- # print("renders RVS {}/total {}".format(batch_idx, pose_perturb.shape[0]))
-
- pose = pose_perturb[batch_idx]
- img_idx = img_idxs[batch_idx].to(device)
- pose_nerf = pose.clone()
-
- # rescale the predicted pose to nerf scales
- pose_nerf = fix_coord_supp(args, pose_nerf[None,...].cpu(), world_setup_dict)
-
- # generate nerf image
- with torch.no_grad():
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- if args.tinyimg:
- rgb, _, _, _ = render(int(H//args.tinyscale), int(W//args.tinyscale), focal/args.tinyscale, chunk=args.chunk, c2w=pose_nerf[0,:3,:4].to(device), retraw=False, img_idx=img_idx, **render_kwargs_test)
- # convert rgb to B,C,H,W
- rgb = rgb[None,...].permute(0,3,1,2)
- # upsample rgb to hwf size
- rgb = torch.nn.Upsample(size=(H, W), mode='bicubic')(rgb)
- # convert rgb back to H,W,C format
- rgb = rgb[0].permute(1,2,0)
-
- else:
- rgb, _, _, _ = render(H, W, focal, chunk=args.chunk, c2w=pose_nerf[0,:3,:4].to(device), retraw=False, img_idx=img_idx, **render_kwargs_test)
- torch.set_default_tensor_type('torch.FloatTensor')
- rgb_list.append(rgb.cpu())
-
- rgbs = torch.stack(rgb_list).detach()
- return rgbs
-
-
-class UnNormalize(object):
- def __init__(self, mean, std):
- self.mean = mean
- self.std = std
-
- def __call__(self, tensor):
- """
- Args:
- tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
- Returns:
- Tensor: Normalized image.
- """
- batch_size = tensor.size(0)
- for i in range(batch_size):
- for t, m, s in zip(tensor[i], self.mean, self.std):
- t.mul_(s).add_(m)
- # The normalize code -> t.sub_(m).div_(s)
- return tensor
-
-class L2Loss(nn.Module):
- """
- simple L_2-Loss
- """
- def __init__(self):
- super(L2Loss, self).__init__()
-
- def forward(self, input, target):
- return torch.sqrt(torch.mean(torch.abs(input - target).pow(2)))
-
-def PoseLoss(args, pose_, pose, device):
- loss_func = nn.MSELoss()
- predict_pose = pose_.to(device) # maynot need reshape
- pose_loss = loss_func(predict_pose, pose)
- return pose_loss
-
-def to_img(x):
- x = (x + 1.) * 0.5
- x = x.clamp(0, 1)
- return x
-
-def masked_loss(criterion, f1, f2, valid_mask):
- '''
- compute loss only in masked region
- :param criterion: loss function
- :param f1: [3, batch_size, H, W]
- :param f2: [3, batch_size, H, W]
-
- :param valid_mask: [batch_size, H, W]
- :return:
- loss
- '''
- # apply mask on f1
- f1 = f1 * (valid_mask)
- f2 = f2
-
- # # print masked img features f1, and unwarped features f2
- # plot_features(f1, 'f1.png', False)
- # plot_features(f2, 'f2.png', False)
- loss = criterion(f1, f2)
- loss = (loss * valid_mask).sum() # take sum of loss
- loss = loss / (valid_mask.sum()) # take mean of loss
- return loss
-
-def triplet_loss(f1, f2, margin=1.):
- '''
- naive implementation of triplet loss
- :param criterion: loss function
- :param f1: [lvl, B, C, H, W]
- :param f2: [lvl, B, C, H, W]
- :return:
- loss
- '''
- criterion = nn.TripletMarginLoss(margin=margin, reduction='mean')
- anchor = f1
- positive = f2
- negative = torch.roll(f2, shifts=1, dims=1)
- loss = criterion(anchor, positive, negative)
- return loss
-
-def triplet_loss_hard_negative_mining(f1, f2, margin=1.):
- '''
- triplet loss with hard negative mining, inspired by http://www.bmva.org/bmvc/2016/papers/paper119/paper119.pdf section3.3
- :param criterion: loss function
- :param f1: [lvl, B, C, H, W]
- :param f2: [lvl, B, C, H, W]
- :return:
- loss
- '''
- criterion = nn.TripletMarginLoss(margin=margin, reduction='mean')
- anchor = f1
- anchor_negative = torch.roll(f1, shifts=1, dims=1)
- positive = f2
- negative = torch.roll(f2, shifts=1, dims=1)
-
- # select in-triplet hard negative, reference: section3.3
- mse = nn.MSELoss(reduction='mean')
- with torch.no_grad():
- case1 = mse(anchor, negative)
- case2 = mse(positive, anchor_negative)
-
- # perform anchor swap if necessary
- if case1 < case2:
- loss = criterion(anchor, positive, negative)
- else:
- loss = criterion(positive, anchor, anchor_negative)
- return loss
-
-def triplet_loss_hard_negative_mining_plus(f1, f2, margin=1.):
- '''
- triplet loss with hard negative mining, four cases. inspired by http://www.bmva.org/bmvc/2016/papers/paper119/paper119.pdf section3.3
- :param criterion: loss function
- :param f1: [lvl, B, C, H, W]
- :param f2: [lvl, B, C, H, W]
- :return:
- loss
- '''
- criterion = nn.TripletMarginLoss(margin=margin, reduction='mean')
- anchor = f1
- anchor_negative = torch.roll(f1, shifts=1, dims=1)
- positive = f2
- negative = torch.roll(f2, shifts=1, dims=1)
-
- # select in-triplet hard negative, reference: section3.3
- mse = nn.MSELoss(reduction='mean')
- with torch.no_grad():
- case1 = mse(anchor, negative)
- case2 = mse(positive, anchor_negative)
- case3 = mse(anchor, anchor_negative)
- case4 = mse(positive, negative)
- distance_list = torch.stack([case1,case2,case3,case4])
- loss_case = torch.argmin(distance_list)
-
- # perform anchor swap if necessary
- if loss_case == 0:
- loss = criterion(anchor, positive, negative)
- elif loss_case == 1:
- loss = criterion(positive, anchor, anchor_negative)
- elif loss_case == 2:
- loss = criterion(anchor, positive, anchor_negative)
- elif loss_case == 3:
- loss = criterion(positive, anchor, negative)
- else:
- raise NotImplementedError
- return loss
-
-def perturb_rotation(c2w, theta, phi, psi=0):
- last_row = np.array([[0, 0, 0, 1]])# np.tile(np.array([0, 0, 0, 1]), (1, 1)) # (N_images, 1, 4)
- c2w = np.concatenate([c2w, last_row], 0) # (N_images, 4, 4) homogeneous coordinate
- # print("c2w", c2w)
- c2w = rot_phi(phi/180.*np.pi) @ c2w
- c2w = rot_theta(theta/180.*np.pi) @ c2w
- c2w = rot_psi(psi/180.*np.pi) @ c2w
- c2w = c2w[:3,:4]
- # print("updated c2w", c2w)
- return c2w
-
-def perturb_single_render_pose(poses, x, angle):
- """
- Inputs:
- poses: (3, 4)
- x: translational perturb range
- angle: rotation angle perturb range in degrees
- Outputs:
- new_c2w: (N_views, 3, 4) new poses
- """
- c2w=poses
-
- N_views = 1 # number of views
- new_c2w = np.zeros((N_views, 3, 4))
-
- for i in range(N_views):
- new_c2w[i] = c2w
- loc = deepcopy(new_c2w[i,:,3]) # this is a must
- # print("new_c2w[i,:,3]", new_c2w[i,:,3])
-
- # perturb rotation pose
- rot_rand=np.random.uniform(-angle,angle,3) # in degrees
- theta, phi, psi = rot_rand
-
- new_c2w[i] = perturb_rotation(new_c2w[i], theta, phi, psi)
- # print("rot new_c2w[i,:,3]", new_c2w[i,:,3])
-
- # perturb translational pose
- trans_rand = np.random.uniform(-x,x,3) # random number of 3 axis pose perturbation
- # print("trans_rand", trans_rand)
-
- # # normalize 3 axis perturbation to x, this will make sure sum of 3-axis perturbation to be x constant
- # trans_rand = trans_rand / abs(trans_rand).sum() * x
- new_c2w[i,:,3] = loc + trans_rand # perturb pos between -1 to 1
- # print("new new_c2w[i,:,3]", new_c2w[i,:,3])
-
- return new_c2w
-
-def perturb_single_render_pose_norm(poses, x, angle):
- """
- Inputs:
- poses: (3, 4)
- x: translational perturb range
- angle: rotation angle perturb range in degrees
- Outputs:
- new_c2w: (N_views, 3, 4) new poses
- """
- c2w=poses
-
- N_views = 1 # number of views
- new_c2w = np.zeros((N_views, 3, 4))
-
- # perturb translational pose
- for i in range(N_views):
- new_c2w[i] = c2w
- trans_rand = np.random.uniform(-x,x,3) # random number of 3 axis pose perturbation
-
- # normalize 3 axis perturbation to x, this will make sure sum of 3-axis perturbation to be x constant
- trans_rand = trans_rand / abs(trans_rand).sum() * x
- new_c2w[i,:,3] = new_c2w[i,:,3] + trans_rand # perturb pos between -1 to 1
-
- # perturb rotation pose
- theta=np.random.uniform(-angle,angle,1) # in degrees
- phi=np.random.uniform(-angle,angle,1) # in degrees
- psi=np.random.uniform(-angle,angle,1) # in degrees
-
- rot_rand = np.array([theta, phi, psi])
- rot_rand = rot_rand / abs(rot_rand).sum() * angle
- theta, phi, psi = rot_rand
-
- # normalize 3 axis perturbation to x, this will make sure sum of 3-axis perturbation to be x constant
- new_c2w[i] = perturb_rotation(new_c2w[i], theta, phi, psi)
- return new_c2w
diff --git a/script/feature/model.py b/script/feature/model.py
deleted file mode 100755
index 52de312..0000000
--- a/script/feature/model.py
+++ /dev/null
@@ -1,509 +0,0 @@
-import torch
-from torch import nn
-import torch.nn.functional as F
-from torchvision import models
-import copy, pdb
-from efficientnet_pytorch import EfficientNet
-from typing import List
-
-class ConvAutoencoder(nn.Module):
- def __init__(self):
- super(ConvAutoencoder, self).__init__()
- ## encoder layers ##
- # conv layer (depth from 3 --> 16), 3x3 kernels
- self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
- # conv layer (depth from 16 --> 4), 3x3 kernels
- self.conv2 = nn.Conv2d(16, 4, 3, padding=1)
- # pooling layer to reduce x-y dims by two; kernel and stride of 2
- self.pool = nn.MaxPool2d(2, 2)
-
- ## decoder layers ##
- ## a kernel of 2 and a stride of 2 will increase the spatial dims by 2
- self.t_conv1 = nn.ConvTranspose2d(4, 16, 2, stride=2)
- self.t_conv2 = nn.ConvTranspose2d(16, 3, 2, stride=2)
-
- def forward(self, x):
- ## encode ##
- # add hidden layers with relu activation function
- # and maxpooling after
- x = F.relu(self.conv1(x))
- x = self.pool(x)
- # add second hidden layer
- x = F.relu(self.conv2(x))
- x = self.pool(x) # compressed representation
-
- ## decode ##
- # add transpose conv layers, with relu activation function
- x = F.relu(self.t_conv1(x))
- # # output layer (with tanh for scaling from -1 to 1)
- x = F.tanh(self.t_conv2(x))
- # output layer (with tanh for scaling from 0 to 1)
- # x = F.sigmoid(self.t_conv2(x))
-
- return x
-
-class autoencoder_vgg1(nn.Module): # psnr 20.84
- def __init__(self):
- super(autoencoder_vgg1, self).__init__()
- self.encoder = nn.Sequential( # like the Composition layer you built
- nn.Conv2d(3, 16, 3, stride=2, padding=1),
- nn.ReLU(),
- nn.Conv2d(16, 32, 3, stride=2, padding=1),
- nn.ReLU(),
- nn.Conv2d(32, 64, 7)
- )
- self.decoder = nn.Sequential(
- nn.ConvTranspose2d(64, 32, 7),
- nn.ReLU(),
- nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1, output_padding=1),
- nn.ReLU(),
- nn.ConvTranspose2d(16, 3, 3, stride=2, padding=1, output_padding=1),
- # nn.Sigmoid() # if img is 0 to 1
- nn.Tanh() # if img is -1 to 1
- )
- def forward(self, x):
- encode = self.encoder(x)
- decode = self.decoder(encode)
- return encode, decode
-
-class autoencoder_vgg2(nn.Module): # psnr 25.41
- ''' add skip connections '''
- def __init__(self):
- super(autoencoder_vgg2, self).__init__()
- conv1 = nn.Sequential(
- nn.Conv2d(3, 16, 3, stride=2, padding=1),
- nn.ReLU(inplace=True),
- )
- conv2 = nn.Sequential(
- nn.Conv2d(16, 32, 3, stride=2, padding=1),
- nn.ReLU(inplace=True),
- )
- conv3 = nn.Sequential(
- nn.Conv2d(32, 64, 7)
- )
- self.encoder = nn.Sequential(
- conv1, conv2, conv3
- )
-
- deconv1 = nn.Sequential(
- nn.ConvTranspose2d(64, 32, 7),
- nn.ReLU(inplace=True),
- )
- deconv2 = nn.Sequential(
- nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1, output_padding=1),
- nn.ReLU(inplace=True),
- )
- deconv3 = nn.Sequential(
- nn.ConvTranspose2d(16, 3, 3, stride=2, padding=1, output_padding=1),
- nn.Tanh()
- )
- self.decoder = nn.Sequential(deconv1, deconv2, deconv3)
- def forward(self, x):
- feat1 = self.encoder[0](x)
- feat2 = self.encoder[1](feat1)
- x = self.encoder[2](feat2)
-
- x = self.decoder[0](x)
- x = x + feat2
- x = self.decoder[1](x)
- x = x + feat1
- x = self.decoder[2](x)
- return None, x
-
-class autoencoder_vgg3(nn.Module): # psnr: 37.77
- ''' vgg encoder '''
- def __init__(self):
- super(autoencoder_vgg3, self).__init__()
- self.encoder = models.vgg19(pretrained=True).features
- # receptive field not equal? so maybe it does not work very well
- self.decoder = nn.Sequential(
- nn.ConvTranspose2d(512, 512, 2, stride=2), # (b, 512, 14, 14)
- nn.ReLU(True),
- nn.ConvTranspose2d(512, 256, 4, stride=4), # (b, 256, 56, 56)
- nn.ReLU(True),
- nn.ConvTranspose2d(256, 64, 4, stride=4), # (b, 64, 224, 224)
- nn.ReLU(True),
- nn.Conv2d(64, 3, 3, stride=1, padding=1),
- nn.Tanh() # MSELoss
- # nn.Sigmoid() # BCELoss
- )
-
- def forward(self, x):
- feat = []
- feat_out = []
- for i in range(len(self.encoder)):
- # print("layer {} encoder layer: {}".format(i, self.encoder[i]))
- x = self.encoder[i](x)
- if i == 35: # ReLU-36
- feat.append(x)
- elif i == 17: # ReLU-17
- feat.append(x)
- elif i == 3: # ReLU-4
- feat.append(x)
- for i in range(len(self.decoder)):
- # print("layer {} decoder layer: {}".format(i, self.decoder[i]))
- x = self.decoder[i](x)
- if i == 1:
- x = x + feat[2]
- feat_out.append(x)
- elif i == 3:
- x = x + feat[1]
- feat_out.append(x)
- elif i == 5:
- x = x + feat[0]
- feat_out.append(x)
- return feat_out, x
-
-class autoencoder_vgg4(nn.Module): # 35.54 PSNR 36.05 BCELoss (120x120)
- ''' vgg encoder with bilinear upsampling'''
- def __init__(self):
- super(autoencoder_vgg4, self).__init__()
- self.encoder = models.vgg19(pretrained=True).features
- # receptive field not equal? so maybe it does not work very well
- self.decoder = nn.Sequential(
- # (b, 512, 14, 14)
- # nn.UpsamplingBilinear2d(scale_factor=2), # upsample to feature map's size
- nn.Conv2d(512, 512, 3, stride=1, padding=1),
- nn.ReLU(True),
- # (b, 256, 56, 56)
- # nn.UpsamplingBilinear2d(scale_factor=4),
- nn.Conv2d(512, 256, 3, stride=1, padding=1),
- nn.ReLU(True),
- # (b, 64, 224, 224)
- # nn.UpsamplingBilinear2d(scale_factor=4),
- nn.Conv2d(256, 64, 3, stride=1, padding=1),
- nn.ReLU(True),
- nn.Conv2d(64, 3, 3, stride=1, padding=1),
- # nn.Tanh() # MSELoss
- nn.Sigmoid() # BCELoss
- )
- def forward(self, x):
- # pdb.set_trace()
- feat = []
- feat_out = []
- for i in range(len(self.encoder)):
- # print("layer {} encoder layer: {}".format(i, self.encoder[i]))
- x = self.encoder[i](x)
- if i == 35: # ReLU-36
- feat.append(x)
- elif i == 17: # ReLU-17
- feat.append(x)
- elif i == 3: # ReLU-4
- feat.append(x)
-
- for i in range(len(self.decoder)):
- # print("layer {} decoder layer: {}".format(i, self.decoder[i]))
- x = self.decoder[i](x)
- if i == 1:
- _, _, h, w = feat[2].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[2]
- feat_out.append(x)
- elif i == 3:
- _, _, h, w = feat[1].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[1]
- feat_out.append(x)
- elif i == 5:
- _, _, h, w = feat[0].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[0]
- feat_out.append(x)
- return feat_out, x
-
-class autoencoder_vgg5(nn.Module): # 36.78 PSNR
- ''' vgg encoder with bilinear upsampling'''
- def __init__(self):
- super(autoencoder_vgg5, self).__init__()
- self.encoder = models.vgg19(pretrained=True).features
- self.decoder = nn.Sequential(
- # (b, 512, 14, 14)
- nn.Conv2d(512, 512, 3, stride=1, padding=1),
- nn.ReLU(True),
- # (b, 512, 28, 28)
- nn.Conv2d(512, 512, 3, stride=1, padding=1),
- nn.ReLU(True),
- # (b, 256, 56, 56)
- nn.Conv2d(512, 256, 3, stride=1, padding=1),
- nn.ReLU(True),
- # (b, 128, 112, 112)
- nn.Conv2d(256, 128, 3, stride=1, padding=1),
- nn.ReLU(True),
- # (b, 64, 224, 224)
- nn.Conv2d(128, 64, 3, stride=1, padding=1),
- nn.ReLU(True),
- nn.Conv2d(64, 3, 3, stride=1, padding=1),
- # nn.Tanh() # MSELoss
- nn.Sigmoid() # BCELoss
- )
- def forward(self, x):
- # pdb.set_trace()
- feat = []
- feat_out = [] # we only use high level features
- for i in range(len(self.encoder)):
- # print("layer {} encoder layer: {}".format(i, self.encoder[i]))
- x = self.encoder[i](x)
- if i == 3: # ReLU-4
- # pdb.set_trace()
- feat.append(x)
- elif i == 8: # ReLU-9
- # pdb.set_trace()
- feat.append(x)
- elif i == 17: # ReLU-18
- # pdb.set_trace()
- feat.append(x)
- elif i == 26: # ReLU-27
- # pdb.set_trace()
- feat.append(x)
- elif i == 35: # ReLU-36
- # pdb.set_trace()
- feat.append(x)
- # pdb.set_trace()
- for i in range(len(self.decoder)):
- # print("layer {} decoder layer: {}".format(i, self.decoder[i]))
- x = self.decoder[i](x)
- if i == 1:
- # pdb.set_trace()
- _, _, h, w = feat[4].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[4]
- elif i == 3:
- # pdb.set_trace()
- _, _, h, w = feat[3].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[3]
- elif i == 5:
- # pdb.set_trace()
- _, _, h, w = feat[2].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[2]
- feat_out.append(x)
- elif i == 7:
- # pdb.set_trace()
- _, _, h, w = feat[1].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[1]
- feat_out.append(x)
- elif i == 9:
- # pdb.set_trace()
- _, _, h, w = feat[0].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[0]
- feat_out.append(x)
- return feat_out, x
-
-class autoencoder_vgg6(nn.Module): # robust feature extractors
- ''' vgg encoder with bilinear upsampling'''
- def __init__(self):
- super(autoencoder_vgg6, self).__init__()
- self.encoder = models.vgg19(pretrained=True).features
- self.decoder = nn.Sequential(
- # (b, 512, 14, 14)
- nn.Conv2d(512, 512, 3, stride=1, padding=1),
- nn.ReLU(True),
- # (b, 512, 28, 28)
- nn.Conv2d(512, 512, 3, stride=1, padding=1),
- nn.ReLU(True),
- # (b, 256, 56, 56)
- nn.Conv2d(512, 256, 3, stride=1, padding=1),
- nn.ReLU(True),
- # (b, 128, 112, 112)
- nn.Conv2d(256, 128, 3, stride=1, padding=1),
- nn.ReLU(True),
- # (b, 64, 224, 224)
- nn.Conv2d(128, 64, 3, stride=1, padding=1),
- nn.ReLU(True),
- # nn.Conv2d(64, 3, 3, stride=1, padding=1),
- # nn.Tanh() # MSELoss
- # nn.Sigmoid() # BCELoss
- )
-
- def forward(self, x, upsampleH, upsampleW): #
- feat = []
- feat_out = [] # we only use high level features
- for i in range(len(self.encoder)):
- # print("layer {} encoder layer: {}".format(i, self.encoder[i]))
- x = self.encoder[i](x)
- if i == 3: # ReLU-4
- feat.append(x)
- elif i == 8: # ReLU-9
- feat.append(x)
- elif i == 17: # ReLU-18
- feat.append(x)
- elif i == 26: # ReLU-27
- feat.append(x)
- elif i == 35: # ReLU-36
- feat.append(x)
-
- for i in range(len(self.decoder)):
- # print("layer {} decoder layer: {}".format(i, self.decoder[i]))
- x = self.decoder[i](x)
- if i == 1:
- _, _, h, w = feat[4].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[4]
- elif i == 3:
- _, _, h, w = feat[3].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[3]
- elif i == 5:
- _, _, h, w = feat[2].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[2]
- feature = torch.mean(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(x), dim=1)
- feat_out.append(feature)
- elif i == 7:
- _, _, h, w = feat[1].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[1]
- feature = torch.mean(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(x), dim=1)
- feat_out.append(feature)
- elif i == 9:
- _, _, h, w = feat[0].shape
- x = nn.UpsamplingBilinear2d(size=(h,w))(x)
- x = x + feat[0]
- feature = torch.mean(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(x), dim=1)
- feat_out.append(feature)
- return feat_out, x
-
-class autoencoder_vgg7(nn.Module): # no decoder
- ''' vgg encoder with bilinear upsampling'''
- def __init__(self):
- super(autoencoder_vgg7, self).__init__()
- self.encoder = models.vgg19(pretrained=True).features
-
- def forward(self, x, upsampleH=224, upsampleW=224):
- feat_out = [] # we only use high level features
- for i in range(len(self.encoder)):
- # print("layer {} encoder layer: {}".format(i, self.encoder[i]))
- x = self.encoder[i](x)
- if i == 3: # ReLU-4
- feature = torch.mean(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(x), dim=1)
- feat_out.append(feature)
- elif i == 8: # ReLU-9
- feature = torch.mean(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(x), dim=1)
- feat_out.append(feature)
- elif i == 17: # ReLU-18
- feature = torch.mean(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(x), dim=1)
- feat_out.append(feature)
- return feat_out, x
-
-# PoseNet (SE(3)) w/ mobilev2 backbone
-class PoseNetV2(nn.Module):
- def __init__(self, feat_dim=12):
- super(PoseNetV2, self).__init__()
- self.backbone_net = models.mobilenet_v2(pretrained=True)
- self.feature_extractor = self.backbone_net.features
- self.avgpool = nn.AdaptiveAvgPool2d(1)
- self.fc_pose = nn.Linear(1280, feat_dim)
-
- def _aggregate_feature(self, x, upsampleH, upsampleW):
- '''
- assume target and nerf rgb are inferenced at the same time,
- slice target batch and nerf batch and aggregate features
- :param x: image blob (2B x C x H x W)
- :param upsampleH: New H
- :param upsampleW: New W
- :return feature: (2 x B x H x W)
- '''
- batch = x.shape[0] # should be target batch_size + rgb batch_size
- feature_t = torch.mean(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(x[:batch//2]), dim=1)
- feature_r = torch.mean(torch.nn.UpsamplingBilinear2d(size=(upsampleH, upsampleW))(x[batch//2:]), dim=1)
- feature = torch.stack([feature_t, feature_r])
- return feature
-
- def _aggregate_feature2(self, x):
- '''
- assume target and nerf rgb are inferenced at the same time,
- slice target batch and nerf batch and output stacked features
- :param x: image blob (2B x C x H x W)
- :return feature: (2 x B x C x H x W)
- '''
- batch = x.shape[0] # should be target batch_size + rgb batch_size
- feature_t = x[:batch//2]
- feature_r = x[batch//2:]
- feature = torch.stack([feature_t, feature_r])
- return feature
-
- def forward(self, x, upsampleH=224, upsampleW=224, isTrain=False, isSingleStream=False):
- '''
- Currently under dev.
- :param x: image blob ()
- :param upsampleH: New H obsolete
- :param upsampleW: New W obsolete
- :param isTrain: True to extract features, False only return pose prediction. Really should be isExtractFeature
- :param isSingleStrea: True to inference single img, False to inference two imgs in siemese network fashion
- '''
- feat_out = [] # we only use high level features
- for i in range(len(self.feature_extractor)):
- # print("layer {} encoder layer: {}".format(i, self.feature_extractor[i]))
- x = self.feature_extractor[i](x)
-
- if isTrain: # collect aggregate features
- if i >= 17 and i <= 17: # 17th block
- if isSingleStream:
- feature = torch.stack([x])
- else:
- feature = self._aggregate_feature2(x)
- feat_out.append(feature)
- x = self.avgpool(x)
- x = x.reshape(x.size(0), -1)
- predict = self.fc_pose(x)
- return feat_out, predict
-
-class EfficientNetB3(nn.Module):
- ''' EfficientNet-B3 backbone,
- model ref: https://github.com/lukemelas/EfficientNet-PyTorch/blob/master/efficientnet_pytorch/model.py
- '''
- def __init__(self, feat_dim=12, feature_block=6):
- super(EfficientNetB3, self).__init__()
- self.backbone_net = EfficientNet.from_pretrained('efficientnet-b3')
- self.feature_block = feature_block # determine which block's feature to use, max=6
- if self.feature_block == 6:
- self.feature_extractor = self.backbone_net.extract_features
- else:
- self.feature_extractor = self.backbone_net.extract_endpoints
-
- # self.feature_extractor = self.backbone_net.extract_endpoints # it can restore middle layer
- self.avgpool = nn.AdaptiveAvgPool2d(1)
- self.fc_pose = nn.Linear(1536, feat_dim) # 1280 for efficientnet-b0, 1536 for efficientnet-b3
-
- def _aggregate_feature2(self, x):
- '''
- assume target and nerf rgb are inferenced at the same time,
- slice target batch and nerf batch and output stacked features
- :param x: image blob (2B x C x H x W)
- :return feature: (2 x B x C x H x W)
- '''
- batch = x.shape[0] # should be target batch_size + rgb batch_size
- feature_t = x[:batch//2]
- feature_r = x[batch//2:]
- feature = torch.stack([feature_t, feature_r])
- return feature
-
- def forward(self, x, return_feature=False, isSingleStream=False):
- '''
- Currently under dev.
- :param x: image blob ()
- :param return_feature: True to extract features, False only return pose prediction. Really should be isExtractFeature
- :param isSingleStream: True to inference single img, False to inference two imgs in siemese network fashion
- '''
- # pdb.set_trace()
- feat_out = [] # we only use high level features
- if self.feature_block == 6:
- x = self.feature_extractor(x)
- fe = x.clone() # features to save
- else:
- list_x = self.feature_extractor(x)
- fe = list_x['reduction_'+str(self.feature_block)]
- x = list_x['reduction_6'] # features to save
- if return_feature:
- if isSingleStream:
- feature = torch.stack([fe])
- else:
- feature = self._aggregate_feature2(fe)
- feat_out.append(feature)
- x = self.avgpool(x)
- x = x.reshape(x.size(0), -1)
- predict = self.fc_pose(x)
- return feat_out, predict
\ No newline at end of file
diff --git a/script/feature/options.py b/script/feature/options.py
deleted file mode 100755
index 0724df7..0000000
--- a/script/feature/options.py
+++ /dev/null
@@ -1,143 +0,0 @@
-import configargparse
-def config_parser():
- parser = configargparse.ArgumentParser()
- parser.add_argument("-f", "--fff", help="a dummy argument to fool ipython", default="1")
- parser.add_argument("--multi_gpu", action='store_true', help='use multiple gpu on the server')
- parser.add_argument('--config', is_config_file=True, help='config file path')
- parser.add_argument("--expname", type=str, help='experiment name')
- parser.add_argument("--basedir", type=str, default='../logs', help='where to store ckpts and logs')
- parser.add_argument("--datadir", type=str, default='./data/llff/fern', help='input data directory')
- parser.add_argument("--places365_model_path", type=str, default='', help='ckpt path of places365 pretrained model')
-
- # 7Scenes
- parser.add_argument("--trainskip", type=int, default=1, help='will load 1/N images from train sets, useful for large datasets like 7 Scenes')
- parser.add_argument("--df", type=float, default=1., help='image downscale factor')
- parser.add_argument("--reduce_embedding", type=int, default=-1, help='fourier embedding mode: -1: paper default, \
- 0: reduce by half, 1: remove embedding, 2: DNeRF embedding')
- parser.add_argument("--epochToMaxFreq", type=int, default=-1, help='DNeRF embedding mode: (based on Nerfie paper): \
- hyper-parameter for when α should reach the maximum number of frequencies m')
- parser.add_argument("--render_pose_only", action='store_true', help='render a spiral video for 7 Scene')
- parser.add_argument("--save_pose_avg_stats", action='store_true', help='save a pose avg stats to unify NeRF, posenet, direct-pn training')
- parser.add_argument("--load_pose_avg_stats", action='store_true', help='load precomputed pose avg stats to unify NeRF, posenet, nerf tracking training')
- parser.add_argument("--train_local_nerf", type=int, default=-1, help='train local NeRF with ith training sequence only, ie. Stairs can pick 0~3')
- parser.add_argument("--render_video_train", action='store_true', help='render train set NeRF and save as video, make sure render_test is True')
- parser.add_argument("--render_video_test", action='store_true', help='render val set NeRF and save as video, make sure render_test is True')
- parser.add_argument("--frustum_overlap_th", type=float, help='frustsum overlap threshold')
- parser.add_argument("--no_DNeRF_viewdir", action='store_true', default=False, help='will not use DNeRF in viewdir encoding')
- parser.add_argument("--load_unique_view_stats", action='store_true', help='load unique views frame index')
- parser.add_argument("--finetune_unlabel", action='store_true', help='finetune unlabeled sequence like MapNet')
- parser.add_argument("--i_eval", type=int, default=20, help='frequency of eval posenet result')
- parser.add_argument("--save_all_ckpt", action='store_true', help='save all ckpts for each epoch')
- parser.add_argument("--val_on_psnr", action='store_true', default=False, help='EarlyStopping with max validation psnr')
-
- # NeRF training options
- parser.add_argument("--netdepth", type=int, default=8, help='layers in network')
- parser.add_argument("--netwidth", type=int, default=128, help='channels per layer')
- parser.add_argument("--netdepth_fine", type=int, default=8, help='layers in fine network')
- parser.add_argument("--netwidth_fine", type=int, default=128, help='channels per layer in fine network')
- parser.add_argument("--N_rand", type=int, default=1536, help='batch size (number of random rays per gradient step)')
- parser.add_argument("--lrate", type=float, default=5e-4, help='learning rate')
- parser.add_argument("--lrate_decay", type=float, default=250, help='exponential learning rate decay (in 1000 steps)')
- parser.add_argument("--chunk", type=int, default=1024*32, help='number of rays processed in parallel, decrease if running out of memory')
- parser.add_argument("--netchunk", type=int, default=1024*64, help='number of pts sent through network in parallel, decrease if running out of memory')
- parser.add_argument("--no_batching", action='store_true', default=True, help='only take random rays from 1 image at a time')
- parser.add_argument("--no_reload", action='store_true', help='do not reload weights from saved ckpt')
- parser.add_argument("--ft_path", type=str, default=None, help='specific weights npy file to reload for coarse network')
-
- # NeRF-Hist training options
- parser.add_argument("--NeRFH", action='store_true', default=True, help='new implementation for NeRFH, please add --encode_hist')
- parser.add_argument("--N_vocab", type=int, default=1000,
- help='''number of vocabulary (number of images)
- in the dataset for nn.Embedding''')
- parser.add_argument("--fix_index", action='store_true', help='fix training frame index as 0')
- parser.add_argument("--encode_hist", default=False, action='store_true', help='encode histogram instead of frame index')
- parser.add_argument("--hist_bin", type=int, default=10, help='image histogram bin size')
- parser.add_argument("--in_channels_a", type=int, default=50, help='appearance embedding dimension, hist_bin*N_a when embedding histogram')
- parser.add_argument("--in_channels_t", type=int, default=20, help='transient embedding dimension, hist_bin*N_tau when embedding histogram')
- parser.add_argument("--svd_reg", default=False, action='store_true', help='use svd regularize output at training')
-
- # NeRF rendering options
- parser.add_argument("--N_samples", type=int, default=64, help='number of coarse samples per ray')
- parser.add_argument("--N_importance", type=int, default=64,help='number of additional fine samples per ray')
- parser.add_argument("--perturb", type=float, default=1.,help='set to 0. for no jitter, 1. for jitter')
- parser.add_argument("--use_viewdirs", action='store_true', default=True, help='use full 5D input instead of 3D')
- parser.add_argument("--i_embed", type=int, default=0, help='set 0 for default positional encoding, -1 for none')
- parser.add_argument("--multires", type=int, default=10, help='log2 of max freq for positional encoding (3D location)')
- parser.add_argument("--multires_views", type=int, default=4, help='log2 of max freq for positional encoding (2D direction)')
- parser.add_argument("--raw_noise_std", type=float, default=0., help='std dev of noise added to regularize sigma_a output, 1e0 recommended')
- parser.add_argument("--render_only", action='store_true', help='do not optimize, reload weights and render out render_poses path')
- parser.add_argument("--render_test", action='store_true', help='render the test set instead of render_poses path')
- parser.add_argument("--render_factor", type=int, default=0, help='downsampling factor to speed up rendering, set 4 or 8 for fast preview')
- parser.add_argument("--no_grad_update", action='store_true', default=False, help='do not update nerf in training')
- parser.add_argument("--tinyimg", action='store_true', default=False, help='render nerf img in a tiny scale image, this is a temporal compromise for direct feature matching, must FIX later')
- parser.add_argument("--tinyscale", type=float, default=4., help='determine the scale of downsizing nerf rendring, must FIX later')
-
- ##################### PoseNet Settings ########################
- parser.add_argument("--pose_only", type=int, default=1, help='posenet type to train, \
- 1: train baseline posenet, 2: posenet+nerf manual optimize, \
- 3: VLocNet, 4: DGRNet')
- parser.add_argument("--learning_rate", type=float, default=0.0001, help='learning rate')
- parser.add_argument("--batch_size", type=int, default=1, help='dataloader batch size, Attention: this is NOT the actual training batch size, \
- please use --featurenet_batch_size for training')
- parser.add_argument("--featurenet_batch_size", type=int, default=8, help='featurenet training batch size, choose smaller batch size')
- parser.add_argument("--pretrain_model_path", type=str, default='', help='model path of pretrained model')
- parser.add_argument("--model_name", type=str, help='pose model output folder name')
- parser.add_argument("--combine_loss_w", nargs='+', type=float, default=[1, 1, 1],
- help='weights of combined loss ex, [1, 1, 1], \
- default None, only use when combine_loss is True')
- parser.add_argument("--patience", nargs='+', type=int, default=[200, 50], help='set training schedule for patience [EarlyStopping, reduceLR]')
- parser.add_argument("--resize_factor", type=int, default=2, help='image resize downsample ratio')
- parser.add_argument("--freezeBN", action='store_true', help='Freeze the Batch Norm layer at training PoseNet')
- parser.add_argument("--preprocess_ImgNet", action='store_true', help='Normalize input data for PoseNet')
- parser.add_argument("--eval", action='store_true', help='eval model')
- parser.add_argument("--no_save_multiple", action='store_true', help='default, save multiple posenet model, if true, save only one posenet model')
- parser.add_argument("--resnet34", action='store_true', default=False, help='use resnet34 backbone instead of mobilenetV2')
- parser.add_argument("--efficientnet", action='store_true', default=False, help='use efficientnet-b3 backbone instead of mobilenetV2')
- parser.add_argument("--dropout", type=float, default=0.5, help='dropout rate for resnet34 backbone')
- parser.add_argument("--DFNet", action='store_true', default=False, help='use DFNet')
- parser.add_argument("--DFNet_s", action='store_true', default=False, help='use accelerated DFNet, performance is similar to DFNet but slightly faster')
- parser.add_argument("--featurelossonly", action='store_true', default=False, help='only use feature loss to train feature extraction')
- parser.add_argument("--random_view_synthesis", action='store_true', default=False, help='add random view synthesis')
- parser.add_argument("--rvs_refresh_rate", type=int, default=2, help='re-synthesis new views per X epochs')
- parser.add_argument("--rvs_trans", type=float, default=5, help='jitter range for rvs on translation')
- parser.add_argument("--rvs_rotation", type=float, default=1.2, help='jitter range for rvs on rotation, this is in log_10 uniform range, log(15) = 1.2')
- parser.add_argument("--d_max", type=float, default=1, help='rvs bounds d_max')
- parser.add_argument("--val_batch_size", type=int, default=1, help='batch_size for validation, higher number leads to faster speed')
-
- # legacy mesh options
- parser.add_argument("--mesh_only", action='store_true', help='do not optimize, reload weights and save mesh to a file')
- parser.add_argument("--mesh_grid_size", type=int, default=80,help='number of grid points to sample in each dimension for marching cubes')
-
- # training options
- parser.add_argument("--precrop_iters", type=int, default=0,help='number of steps to train on central crops')
- parser.add_argument("--precrop_frac", type=float,default=.5, help='fraction of img taken for central crops')
- parser.add_argument("--epochs", type=int, default=600,help='number of epochs to train')
- parser.add_argument("--poselossonly", action='store_true', help='eval model')
- parser.add_argument("--tripletloss", action='store_true', help='use triplet loss at training featurenet, this is to prevent catastophic failing')
- parser.add_argument("--triplet_margin", type=float,default=1., help='triplet loss margin hyperparameter')
- parser.add_argument("--render_feature_only", action='store_true', default=False, help='render features and save to a path')
-
- # dataset options
- parser.add_argument("--dataset_type", type=str, default='llff', help='options: llff / 7Scenes')
- parser.add_argument("--testskip", type=int, default=1, help='will load 1/N images from test/val sets, useful for large datasets like deepvoxels')
-
- ## legacy blender flags
- parser.add_argument("--white_bkgd", action='store_true', help='set to render synthetic data on a white bkgd (always use for dvoxels)')
- parser.add_argument("--half_res", action='store_true', help='load blender synthetic data at 400x400 instead of 800x800')
-
- ## llff flags
- parser.add_argument("--factor", type=int, default=8, help='downsample factor for LLFF images')
- parser.add_argument("--no_ndc", action='store_true', help='do not use normalized device coordinates (set for non-forward facing scenes)')
- parser.add_argument("--lindisp", action='store_true', help='sampling linearly in disparity rather than depth')
- parser.add_argument("--spherify", action='store_true', help='set for spherical 360 scenes')
- parser.add_argument("--llffhold", type=int, default=8, help='will take every 1/N images as LLFF test set, paper uses 8')
- parser.add_argument("--no_bd_factor", action='store_true', default=False, help='do not use bd factor')
-
- # logging/saving options
- parser.add_argument("--i_print", type=int, default=1, help='frequency of console printout and metric loggin')
- parser.add_argument("--i_img", type=int, default=500, help='frequency of tensorboard image logging')
- parser.add_argument("--i_weights", type=int, default=200, help='frequency of weight ckpt saving')
- parser.add_argument("--i_testset", type=int, default=200, help='frequency of testset saving')
- parser.add_argument("--i_video", type=int, default=50000, help='frequency of render_poses video saving')
-
- return parser
\ No newline at end of file
diff --git a/script/models/__init__.py b/script/models/__init__.py
deleted file mode 100755
index e69de29..0000000
diff --git a/script/models/losses.py b/script/models/losses.py
deleted file mode 100755
index cd518d1..0000000
--- a/script/models/losses.py
+++ /dev/null
@@ -1,60 +0,0 @@
-import torch
-from torch import nn
-import pdb
-
-class ColorLoss(nn.Module):
- def __init__(self, coef=1):
- super().__init__()
- self.coef = coef
- self.loss = nn.MSELoss(reduction='mean')
-
- def forward(self, inputs, targets):
- loss = self.loss(inputs['rgb_coarse'], targets)
- if 'rgb_fine' in inputs:
- loss += self.loss(inputs['rgb_fine'], targets)
-
- return self.coef * loss
-
-
-class NerfWLoss(nn.Module):
- """
- Equation 13 in the NeRF-W paper.
- Name abbreviations:
- c_l: coarse color loss
- f_l: fine color loss (1st term in equation 13)
- b_l: beta loss (2nd term in equation 13)
- s_l: sigma loss (3rd term in equation 13)
- targets # [N, 3]
- inputs['rgb_coarse'] # [N, 3]
- inputs['rgb_fine'] # [N, 3]
- inputs['beta'] # [N]
- inputs['transient_sigmas'] # [N, 2*N_Samples]
- :return:
- """
- def __init__(self, coef=1, lambda_u=0.01):
- """
- lambda_u: in equation 13
- """
- super().__init__()
- self.coef = coef
- self.lambda_u = lambda_u
-
- def forward(self, inputs, targets, use_hier_rgbs=False, rgb_h=None, rgb_w=None):
-
- ret = {}
- ret['c_l'] = 0.5 * ((inputs['rgb_coarse']-targets)**2).mean()
- if 'rgb_fine' in inputs:
- if 'beta' not in inputs: # no transient head, normal MSE loss
- ret['f_l'] = 0.5 * ((inputs['rgb_fine']-targets)**2).mean()
- else:
- ret['f_l'] = ((inputs['rgb_fine']-targets)**2/(2*inputs['beta'].unsqueeze(1)**2)).mean()
- ret['b_l'] = 3 + torch.log(inputs['beta']).mean() # +3 to make it positive
- ret['s_l'] = self.lambda_u * inputs['transient_sigmas'].mean()
-
- for k, v in ret.items():
- ret[k] = self.coef * v
-
- return ret
-
-loss_dict = {'color': ColorLoss,
- 'nerfw': NerfWLoss}
\ No newline at end of file
diff --git a/script/models/metrics.py b/script/models/metrics.py
deleted file mode 100755
index fb10988..0000000
--- a/script/models/metrics.py
+++ /dev/null
@@ -1,20 +0,0 @@
-import torch
-from kornia.losses import ssim as dssim
-
-def mse(image_pred, image_gt, valid_mask=None, reduction='mean'):
- value = (image_pred-image_gt)**2
- if valid_mask is not None:
- value = value[valid_mask]
- if reduction == 'mean':
- return torch.mean(value)
- return value
-
-def psnr(image_pred, image_gt, valid_mask=None, reduction='mean'):
- return -10*torch.log10(mse(image_pred, image_gt, valid_mask, reduction))
-
-def ssim(image_pred, image_gt, reduction='mean'):
- """
- image_pred and image_gt: (1, 3, H, W)
- """
- dssim_ = dssim(image_pred, image_gt, 3, reduction) # dissimilarity in [0, 1]
- return 1-2*dssim_ # in [-1, 1]
\ No newline at end of file
diff --git a/script/models/nerf.py b/script/models/nerf.py
deleted file mode 100755
index 412b0ef..0000000
--- a/script/models/nerf.py
+++ /dev/null
@@ -1,372 +0,0 @@
-import os
-import torch
-torch.autograd.set_detect_anomaly(True)
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-
-# Misc
-img2mse = lambda x, y : torch.mean((x - y) ** 2)
-mse2psnr = lambda x : -10. * torch.log(x) / torch.log(torch.Tensor([10.]))
-to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8)
-
-def batchify(fn, chunk):
- """Constructs a version of 'fn' that applies to smaller batches.
- """
- if chunk is None:
- return fn
- def ret(inputs):
- return torch.cat([fn(inputs[i:i+chunk]) for i in range(0, inputs.shape[0], chunk)], 0)
- return ret
-
-
-def run_network(inputs, viewdirs, fn, embed_fn, embeddirs_fn, netchunk=1024*64):
- """Prepares inputs and applies network 'fn'.
- """
- inputs_flat = torch.reshape(inputs, [-1, inputs.shape[-1]])
- embedded = embed_fn(inputs_flat)
-
- if viewdirs is not None:
- input_dirs = viewdirs[:,None].expand(inputs.shape)
- input_dirs_flat = torch.reshape(input_dirs, [-1, input_dirs.shape[-1]])
- embedded_dirs = embeddirs_fn(input_dirs_flat)
- embedded = torch.cat([embedded, embedded_dirs], -1)
-
- outputs_flat = batchify(fn, netchunk)(embedded)
- outputs = torch.reshape(outputs_flat, list(inputs.shape[:-1]) + [outputs_flat.shape[-1]])
- return outputs
-
-def run_network_DNeRF(inputs, viewdirs, fn, embed_fn, embeddirs_fn, netchunk=1024*64, epoch=None, no_DNeRF_viewdir=False):
- """Prepares inputs and applies network 'fn'.
- """
- if epoch<0 or epoch==None:
- print("Error: run_network_DNeRF(): Invalid epoch")
- sys.exit()
- inputs_flat = torch.reshape(inputs, [-1, inputs.shape[-1]])
- embedded = embed_fn(inputs_flat, epoch)
- # add weighted function here
- if viewdirs is not None:
- input_dirs = viewdirs[:,None].expand(inputs.shape)
- input_dirs_flat = torch.reshape(input_dirs, [-1, input_dirs.shape[-1]])
-
- if no_DNeRF_viewdir:
- embedded_dirs = embeddirs_fn(input_dirs_flat)
- else:
- embedded_dirs = embeddirs_fn(input_dirs_flat, epoch)
- embedded = torch.cat([embedded, embedded_dirs], -1)
-
- outputs_flat = batchify(fn, netchunk)(embedded)
- outputs = torch.reshape(outputs_flat, list(inputs.shape[:-1]) + [outputs_flat.shape[-1]])
- return outputs
-
-
-# Positional encoding (section 5.1)
-class Embedder:
- def __init__(self, **kwargs):
- self.kwargs = kwargs
- self.N_freqs = 0
- self.N = -1 # epoch to max frequency, for Nerfie embedding only
- self.create_embedding_fn()
-
- def create_embedding_fn(self):
- embed_fns = []
- d = self.kwargs['input_dims']
- out_dim = 0
- if self.kwargs['include_input']:
- embed_fns.append(lambda x : x)
- out_dim += d
-
- max_freq = self.kwargs['max_freq_log2']
- self.N_freqs = self.kwargs['num_freqs']
-
- if self.kwargs['log_sampling']:
- freq_bands = 2.**torch.linspace(0., max_freq, steps=self.N_freqs) # tensor([ 1., 2., 4., 8., 16., 32., 64., 128., 256., 512.])
- else:
- freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=self.N_freqs)
-
- for freq in freq_bands: # 10 iters for 3D location, 4 iters for 2D direction
- for p_fn in self.kwargs['periodic_fns']:
- embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq))
- out_dim += d
- self.embed_fns = embed_fns
- self.out_dim = out_dim
-
- def embed(self, inputs):
- # inputs [65536, 3]
- if self.kwargs['max_freq_log2'] != 0:
- ret = torch.cat([fn(inputs) for fn in self.embed_fns], -1) # cos, sin embedding # ret.shape [65536, 63]
- else:
- ret = inputs
- return ret
-
- def get_embed_weight(self, epoch, num_freqs, N):
- ''' Nerfie Paper Eq.(8) '''
- alpha = num_freqs * epoch / N
- W_j = []
- for i in range(num_freqs):
- tmp = torch.clamp(torch.Tensor([alpha - i]), 0, 1)
- tmp2 = (1 - torch.cos(torch.Tensor([np.pi]) * tmp)) / 2
- W_j.append(tmp2)
- return W_j
-
- def embed_DNeRF(self, inputs, epoch):
- ''' Nerfie paper section 3.5 Coarse-to-Fine Deformation Regularization '''
- # get weight for each frequency band j
- W_j = self.get_embed_weight(epoch, self.N_freqs, self.N) # W_j: [W_0, W_1, W_2, ..., W_{m-1}]
-
- # Fourier embedding
- out = []
- for fn in self.embed_fns: # 17, embed_fns:[input, cos, sin, cos, sin, ..., cos, sin]
- out.append(fn(inputs))
-
- # apply weighted positional encoding, only to cos&sins
- for i in range(len(W_j)):
- out[2*i+1] = W_j[i] * out[2*i+1]
- out[2*i+2] = W_j[i] * out[2*i+2]
- ret = torch.cat(out, -1)
- return ret
-
- def update_N(self, N):
- self.N=N
-
-
-def get_embedder(multires, i=0, reduce_mode=-1, epochToMaxFreq=-1):
- if i == -1:
- return nn.Identity(), 3
-
- if reduce_mode == 0:
- # reduce embedding
- embed_kwargs = {
- 'include_input' : True,
- 'input_dims' : 3,
- 'max_freq_log2' : (multires-1)//2,
- 'num_freqs' : multires//2,
- 'log_sampling' : True,
- 'periodic_fns' : [torch.sin, torch.cos],
- }
- elif reduce_mode == 1:
- # remove embedding
- embed_kwargs = {
- 'include_input' : True,
- 'input_dims' : 3,
- 'max_freq_log2' : 0,
- 'num_freqs' : 0,
- 'log_sampling' : True,
- 'periodic_fns' : [torch.sin, torch.cos],
- }
- elif reduce_mode == 2:
- # DNeRF embedding
- embed_kwargs = {
- 'include_input' : True,
- 'input_dims' : 3,
- 'max_freq_log2' : multires-1,
- 'num_freqs' : multires,
- 'log_sampling' : True,
- 'periodic_fns' : [torch.sin, torch.cos],
- }
- else:
- # paper default
- embed_kwargs = {
- 'include_input' : True,
- 'input_dims' : 3,
- 'max_freq_log2' : multires-1,
- 'num_freqs' : multires,
- 'log_sampling' : True,
- 'periodic_fns' : [torch.sin, torch.cos],
- }
-
- embedder_obj = Embedder(**embed_kwargs)
- if reduce_mode == 2:
- embedder_obj.update_N(epochToMaxFreq)
- embed = lambda x, epoch, eo=embedder_obj: eo.embed_DNeRF(x, epoch)
- else:
- embed = lambda x, eo=embedder_obj : eo.embed(x)
- return embed, embedder_obj.out_dim, embedder_obj# 63 for pos, 27 for view dir
-
-# Model
-class NeRF(nn.Module):
- def __init__(self, D=8, W=256, input_ch=3, input_ch_views=3, output_ch=4, skips=[4], use_viewdirs=False):
- """
- """
- super(NeRF, self).__init__()
- self.D = D
- self.W = W
- self.input_ch = input_ch
- self.input_ch_views = input_ch_views
- self.skips = skips
- self.use_viewdirs = use_viewdirs
-
- self.pts_linears = nn.ModuleList(
- [nn.Linear(input_ch, W)] + [nn.Linear(W, W) if i not in self.skips else nn.Linear(W + input_ch, W) for i in range(D-1)])
-
- ### Implementation according to the official code release (https://github.com/bmild/nerf/blob/master/run_nerf_helpers.py#L104-L105)
- self.views_linears = nn.ModuleList([nn.Linear(input_ch_views + W, W//2)])
-
- ### Implementation according to the NeRF paper
- # self.views_linears = nn.ModuleList(
- # [nn.Linear(input_ch_views + W, W//2)] + [nn.Linear(W//2, W//2) for i in range(D//2)])
-
- if use_viewdirs:
- self.feature_linear = nn.Linear(W, W)
- self.alpha_linear = nn.Linear(W, 1)
- self.rgb_linear = nn.Linear(W//2, 3)
- else:
- self.output_linear = nn.Linear(W, output_ch)
-
- def forward(self, x):
- input_pts, input_views = torch.split(x, [self.input_ch, self.input_ch_views], dim=-1)
- h = input_pts
- for i, l in enumerate(self.pts_linears):
- h = self.pts_linears[i](h)
- h = F.relu(h)
- if i in self.skips:
- h = torch.cat([input_pts, h], -1)
-
- if self.use_viewdirs:
- alpha = self.alpha_linear(h)
- feature = self.feature_linear(h)
- h = torch.cat([feature, input_views], -1)
-
- for i, l in enumerate(self.views_linears):
- h = self.views_linears[i](h)
- h = F.relu(h)
-
- rgb = self.rgb_linear(h)
- outputs = torch.cat([rgb, alpha], -1)
- else:
- outputs = self.output_linear(h)
-
- return outputs
-
- def load_weights_from_keras(self, weights):
- assert self.use_viewdirs, "Not implemented if use_viewdirs=False"
-
- # Load pts_linears
- for i in range(self.D):
- idx_pts_linears = 2 * i
- self.pts_linears[i].weight.data = torch.from_numpy(np.transpose(weights[idx_pts_linears]))
- self.pts_linears[i].bias.data = torch.from_numpy(np.transpose(weights[idx_pts_linears+1]))
-
- # Load feature_linear
- idx_feature_linear = 2 * self.D
- self.feature_linear.weight.data = torch.from_numpy(np.transpose(weights[idx_feature_linear]))
- self.feature_linear.bias.data = torch.from_numpy(np.transpose(weights[idx_feature_linear+1]))
-
- # Load views_linears
- idx_views_linears = 2 * self.D + 2
- self.views_linears[0].weight.data = torch.from_numpy(np.transpose(weights[idx_views_linears]))
- self.views_linears[0].bias.data = torch.from_numpy(np.transpose(weights[idx_views_linears+1]))
-
- # Load rgb_linear
- idx_rbg_linear = 2 * self.D + 4
- self.rgb_linear.weight.data = torch.from_numpy(np.transpose(weights[idx_rbg_linear]))
- self.rgb_linear.bias.data = torch.from_numpy(np.transpose(weights[idx_rbg_linear+1]))
-
- # Load alpha_linear
- idx_alpha_linear = 2 * self.D + 6
- self.alpha_linear.weight.data = torch.from_numpy(np.transpose(weights[idx_alpha_linear]))
- self.alpha_linear.bias.data = torch.from_numpy(np.transpose(weights[idx_alpha_linear+1]))
-
-def create_nerf(args):
- """Instantiate NeRF's MLP model.
- """
- if args.reduce_embedding==2: # use DNeRF embedding
- embed_fn, input_ch, embedder_obj = get_embedder(args.multires, args.i_embed, args.reduce_embedding, args.epochToMaxFreq) # input_ch.shape=63
- else:
- embed_fn, input_ch, _ = get_embedder(args.multires, args.i_embed, args.reduce_embedding) # input_ch.shape=63
-
- input_ch_views = 0
- embeddirs_fn = None
- if args.use_viewdirs:
- if args.reduce_embedding==2: # use DNeRF embedding
- if args.no_DNeRF_viewdir: # no DNeRF embedding for viewdir
- embeddirs_fn, input_ch_views, _ = get_embedder(args.multires_views, args.i_embed)
- else:
- embeddirs_fn, input_ch_views, embedddirs_obj = get_embedder(args.multires_views, args.i_embed, args.reduce_embedding, args.epochToMaxFreq)
- else:
- embeddirs_fn, input_ch_views, _ = get_embedder(args.multires_views, args.i_embed, args.reduce_embedding) # input_ch_views.shape=27
- output_ch = 5 if args.N_importance > 0 else 4
- skips = [4]
- model = NeRF(D=args.netdepth, W=args.netwidth, input_ch=input_ch, output_ch=output_ch, skips=skips,
- input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs)
- device = torch.device("cuda")
- if args.multi_gpu:
- model = torch.nn.DataParallel(model).to(device)
- else:
- model = model.to(device)
- grad_vars = list(model.parameters())
-
- model_fine = None
- if args.N_importance > 0:
- model_fine = NeRF(D=args.netdepth_fine, W=args.netwidth_fine, input_ch=input_ch, output_ch=output_ch, skips=skips,
- input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs)
- if args.multi_gpu:
- model_fine = torch.nn.DataParallel(model_fine).to(device)
- else:
- model_fine = model_fine.to(device)
- grad_vars += list(model_fine.parameters())
-
- if args.reduce_embedding==2: # use DNeRF embedding
- network_query_fn = lambda inputs, viewdirs, network_fn, epoch: run_network_DNeRF(inputs, viewdirs, network_fn,
- embed_fn=embed_fn,
- embeddirs_fn=embeddirs_fn,
- netchunk=args.netchunk,
- epoch=epoch, no_DNeRF_viewdir=args.no_DNeRF_viewdir)
- else:
- network_query_fn = lambda inputs, viewdirs, network_fn : run_network(inputs, viewdirs, network_fn,
- embed_fn=embed_fn,
- embeddirs_fn=embeddirs_fn,
- netchunk=args.netchunk)
-
- # Create optimizer
- optimizer = torch.optim.Adam(params=grad_vars, lr=args.lrate, betas=(0.9, 0.999))
-
- start = 0
- basedir = args.basedir
- expname = args.expname
-
- # Load checkpoints
- if args.ft_path is not None and args.ft_path!='None':
- ckpts = [args.ft_path]
- else:
- ckpts = [os.path.join(basedir, expname, f) for f in sorted(os.listdir(os.path.join(basedir, expname))) if 'tar' in f]
-
- print('Found ckpts', ckpts)
- if len(ckpts) > 0 and not args.no_reload:
- ckpt_path = ckpts[-1]
- print('Reloading from', ckpt_path)
- ckpt = torch.load(ckpt_path)
-
- start = ckpt['global_step']
- optimizer.load_state_dict(ckpt['optimizer_state_dict'])
- # Load model
- model.load_state_dict(ckpt['network_fn_state_dict'])
-
- if model_fine is not None:
- model_fine.load_state_dict(ckpt['network_fine_state_dict'])
-
- ##########################
-
- render_kwargs_train = {
- 'network_query_fn' : network_query_fn,
- 'perturb' : args.perturb,
- 'N_importance' : args.N_importance,
- 'network_fine' : model_fine,
- 'N_samples' : args.N_samples,
- 'network_fn' : model,
- 'use_viewdirs' : args.use_viewdirs,
- 'white_bkgd' : args.white_bkgd,
- 'raw_noise_std' : args.raw_noise_std,
- }
-
- # NDC only good for LLFF-style forward facing data
- if args.dataset_type != 'llff' or args.no_ndc:
- print('Not ndc!')
- render_kwargs_train['ndc'] = False
- render_kwargs_train['lindisp'] = args.lindisp
-
- render_kwargs_test = {k : render_kwargs_train[k] for k in render_kwargs_train}
- render_kwargs_test['perturb'] = False
- render_kwargs_test['raw_noise_std'] = 0.
-
- return render_kwargs_train, render_kwargs_test, start, grad_vars, optimizer
\ No newline at end of file
diff --git a/script/models/nerfw.py b/script/models/nerfw.py
deleted file mode 100755
index 63eee23..0000000
--- a/script/models/nerfw.py
+++ /dev/null
@@ -1,502 +0,0 @@
-import os
-import torch
-import torch.nn as nn
-import numpy as np
-from einops import repeat
-from models.nerf import run_network, NeRF
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Misc
-img2mse = lambda x, y : torch.mean((x - y) ** 2)
-mse2psnr = lambda x : -10. * torch.log(x) / torch.log(torch.Tensor([10.]))
-to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8)
-
-def run_network_NeRFW(inputs, viewdirs, ts, fn, embed_fn, embeddirs_fn,
- typ, embedding_a, embedding_t, output_transient,
- netchunk=1024*64, test_time=False):
- ''' We need a new query function, Coarse = NeRF, Fine = NeRF-W
- Inputs:
- inputs: torch.Tensor() [N_rays,N_samples,3]
- viewdirs: torch.Tensor() [N_rays, 3]
- ts: latent code from img_idxs [N_rays]
- fn: NeRFW object
- embed_fn: embedder for position
- embeddirs_fn: embedder for view directions
- typ: 'coarse' or 'fine'
- embedding_a: NeRFW appearance embedding layer
- embedding_t: NeRFW transient embedding layer
- output_transient: True/False
- netchunk: chunk size to inference
- test_time: True/False
- '''
- out_chunks = []
- N_rays, N_samples = inputs.shape[0], inputs.shape[1]
-
- # embed inputs like NeRF
- if typ == 'coarse' and test_time:
- inputs_flat = torch.reshape(inputs, [-1, inputs.shape[-1]])
-
- # Feed NeRF coarse train
- for i in range(0, inputs_flat.shape[0], netchunk):
- embedded_inputs = [embed_fn(inputs_flat[i: i+netchunk])]
- out_chunks += [fn(torch.cat(embedded_inputs, 1), sigma_only=True)]
- out = torch.cat(out_chunks, 0) # [N_rays*N_samples, 4]
- out = torch.reshape(out, list(inputs.shape[:-1]) + [out.shape[-1]]) # [N_rays, N_samples, 4]
- return out
- if typ == 'coarse': # case: coarse + train
- inputs_flat = torch.reshape(inputs, [-1, inputs.shape[-1]])
-
- input_dirs = viewdirs[:,None].expand(inputs.shape)
- input_dirs_flat = torch.reshape(input_dirs, [-1, input_dirs.shape[-1]])
-
- # Feed NeRF coarse train
- for i in range(0, inputs_flat.shape[0], netchunk):
- embedded_inputs = [embed_fn(inputs_flat[i: i+netchunk]), embeddirs_fn(input_dirs_flat[i:i+netchunk])]
- out_chunks += [fn(torch.cat(embedded_inputs, 1), output_transient=output_transient)]
-
- out = torch.cat(out_chunks, 0) # [N_rays*N_samples, 4]
- out = torch.reshape(out, list(inputs.shape[:-1]) + [out.shape[-1]]) # [N_rays, N_samples, 4]
- return out
-
- elif typ == 'fine':
- inputs_flat = torch.reshape(inputs, [-1, inputs.shape[-1]])
-
- input_dirs = viewdirs[:,None].expand(inputs.shape)
- input_dirs_flat = torch.reshape(input_dirs, [-1, input_dirs.shape[-1]])
-
- # embed encode_a and encode_t
- a_embedded = embedding_a(ts.long())
-
- # if output_transient:
- t_embedded = embedding_t(ts.long())
-
- # for embedding histogram experiement
- if len(a_embedded.size()) == 3:
- a_embedded = a_embedded.reshape(N_rays, -1) # [N_rays, hist_bin*N_a]
- if len(t_embedded.size()) == 3:
- t_embedded = t_embedded.reshape(N_rays, -1) # [N_rays, hist_bin*N_tau]
-
- a_embedded_ = repeat(a_embedded, 'n1 c -> (n1 n2) c', n2=N_samples)
- t_embedded_ = repeat(t_embedded, 'n1 c -> (n1 n2) c', n2=N_samples)
-
- # Feed NeRF-W fine train
- for i in range(0, inputs_flat.shape[0], netchunk):
- # inputs for original NeRF
- embedded_inputs = [embed_fn(inputs_flat[i: i+netchunk]), embeddirs_fn(input_dirs_flat[i:i+netchunk])]
-
- # additional inputs for NeRF-W
- embedded_inputs += [a_embedded_[i:i+netchunk]]
- embedded_inputs += [t_embedded_[i:i+netchunk]]
-
- out_chunks += [fn(torch.cat(embedded_inputs, 1), output_transient=output_transient)]
- out = torch.cat(out_chunks, 0) # [N_rays*N_samples, 9]
- out = torch.reshape(out, list(inputs.shape[:-1]) + [out.shape[-1]]) # [N_rays, N_samples, 9]
- return out
-
-# Positional encoding (section 5.1 of NERF)
-class Embedder:
- def __init__(self, **kwargs):
- self.kwargs = kwargs
- self.N_freqs = 0
- self.N = -1 # epoch to max frequency, for Nerfie embedding only
- self.create_embedding_fn()
-
- def create_embedding_fn(self):
- embed_fns = []
- d = self.kwargs['input_dims']
- out_dim = 0
- if self.kwargs['include_input']:
- embed_fns.append(lambda x : x)
- out_dim += d
-
- max_freq = self.kwargs['max_freq_log2']
- self.N_freqs = self.kwargs['num_freqs']
-
- if self.kwargs['log_sampling']:
- freq_bands = 2.**torch.linspace(0., max_freq, steps=self.N_freqs) # tensor([ 1., 2., 4., 8., 16., 32., 64., 128., 256., 512.])
- else:
- freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=self.N_freqs)
-
- for freq in freq_bands: # 10 iters for 3D location, 4 iters for 2D direction
- for p_fn in self.kwargs['periodic_fns']:
- embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq))
- out_dim += d
- self.embed_fns = embed_fns
- self.out_dim = out_dim
-
- def embed(self, inputs):
- if self.kwargs['max_freq_log2'] != 0:
- ret = torch.cat([fn(inputs) for fn in self.embed_fns], -1) # cos, sin embedding
- else:
- ret = inputs
- return ret
-
- def get_embed_weight(self, epoch, num_freqs, N):
- ''' Nerfie Paper Eq.(8) '''
- alpha = num_freqs * epoch / N
- W_j = []
- for i in range(num_freqs):
- tmp = torch.clamp(torch.Tensor([alpha - i]), 0, 1)
- tmp2 = (1 - torch.cos(torch.Tensor([np.pi]) * tmp)) / 2
- W_j.append(tmp2)
- return W_j
-
- def embed_DNeRF(self, inputs, epoch):
- ''' Nerfie paper section 3.5 Coarse-to-Fine Deformation Regularization '''
- # get weight for each frequency band j
- W_j = self.get_embed_weight(epoch, self.N_freqs, self.N) # W_j: [W_0, W_1, W_2, ..., W_{m-1}]
-
- # Fourier embedding
- out = []
- for fn in self.embed_fns: # 17, embed_fns:[input, cos, sin, cos, sin, ..., cos, sin]
- out.append(fn(inputs))
-
- # apply weighted positional encoding, only to cos&sins
- for i in range(len(W_j)):
- out[2*i+1] = W_j[i] * out[2*i+1]
- out[2*i+2] = W_j[i] * out[2*i+2]
- ret = torch.cat(out, -1)
- return ret
-
- def update_N(self, N):
- self.N=N
-
-
-def get_embedder(multires, i=0, reduce_mode=-1, epochToMaxFreq=-1):
- if i == -1:
- return nn.Identity(), 3
-
- if reduce_mode == 0:
- # reduce embedding
- embed_kwargs = {
- 'include_input' : True,
- 'input_dims' : 3,
- 'max_freq_log2' : (multires-1)//2,
- 'num_freqs' : multires//2,
- 'log_sampling' : True,
- 'periodic_fns' : [torch.sin, torch.cos],
- }
- elif reduce_mode == 1:
- # remove embedding
- embed_kwargs = {
- 'include_input' : True,
- 'input_dims' : 3,
- 'max_freq_log2' : 0,
- 'num_freqs' : 0,
- 'log_sampling' : True,
- 'periodic_fns' : [torch.sin, torch.cos],
- }
- elif reduce_mode == 2:
- # DNeRF embedding
- embed_kwargs = {
- 'include_input' : True,
- 'input_dims' : 3,
- 'max_freq_log2' : multires-1,
- 'num_freqs' : multires,
- 'log_sampling' : True,
- 'periodic_fns' : [torch.sin, torch.cos],
- }
- else:
- # paper default
- embed_kwargs = {
- 'include_input' : True,
- 'input_dims' : 3,
- 'max_freq_log2' : multires-1,
- 'num_freqs' : multires,
- 'log_sampling' : True,
- 'periodic_fns' : [torch.sin, torch.cos],
- }
-
- embedder_obj = Embedder(**embed_kwargs)
- if reduce_mode == 2:
- embedder_obj.update_N(epochToMaxFreq)
- embed = lambda x, epoch, eo=embedder_obj: eo.embed_DNeRF(x, epoch)
- else:
- embed = lambda x, eo=embedder_obj : eo.embed(x)
- return embed, embedder_obj.out_dim, embedder_obj# 63 for pos, 27 for view dir
-
-# Model
-class NeRFW(nn.Module):
- def __init__(self, typ,
- D=8, W=256, skips=[4],
- in_channels_xyz=63, in_channels_dir=27,
- encode_appearance=False, in_channels_a=48,
- encode_transient=False, in_channels_t=16,
- beta_min=0.1, out_ch_size=3):
- """
- ---Parameters for the original NeRF---
- D: number of layers for density (sigma) encoder
- W: number of hidden units in each layer
- skips: add skip connection in the Dth layer
- in_channels_xyz: number of input channels for xyz (3+3*10*2=63 by default)
- in_channels_dir: number of input channels for direction (3+3*4*2=27 by default)
- in_channels_t: number of input channels for t
-
- ---Parameters for NeRF-W (used in fine model only as per section 4.3)---
- ---cf. Figure 3 of the paper---
- encode_appearance: whether to add appearance encoding as input (NeRF-A)
- in_channels_a: appearance embedding dimension. n^(a) in the paper
- encode_transient: whether to add transient encoding as input (NeRF-U)
- in_channels_t: transient embedding dimension. n^(tau) in the paper
- beta_min: minimum pixel color variance
- """
- super().__init__()
- torch.manual_seed(0)
- self.typ = typ
- self.D = D
- self.W = W
- self.skips = skips
- self.in_channels_xyz = in_channels_xyz
- self.in_channels_dir = in_channels_dir
- self.encode_appearance = False if typ=='coarse' else encode_appearance
- self.in_channels_a = in_channels_a if encode_appearance else 0
- self.encode_transient = False if typ=='coarse' else encode_transient
- self.in_channels_t = in_channels_t
- self.beta_min = beta_min
-
- # xyz encoding layers
- for i in range(D):
- if i == 0:
- layer = nn.Linear(in_channels_xyz, W)
- elif i in skips:
- layer = nn.Linear(W+in_channels_xyz, W)
- else:
- layer = nn.Linear(W, W)
- layer = nn.Sequential(layer, nn.ReLU(True))
- setattr(self, f"xyz_encoding_{i+1}", layer)
- self.xyz_encoding_final = nn.Linear(W, W)
-
- # direction encoding layers
- self.dir_encoding = nn.Sequential(
- nn.Linear(W+in_channels_dir+self.in_channels_a, W//2), nn.ReLU(True))
-
- # static output layers
- self.static_sigma = nn.Sequential(nn.Linear(W, 1), nn.Softplus())
-
- if out_ch_size == 3: #NeRF, NeRFW output rgb
- self.static_rgb = nn.Sequential(nn.Linear(W//2, 3), nn.Sigmoid())
- else: # output features like giraffe
- self.static_rgb = nn.Sequential(nn.Linear(W//2, out_ch_size))
-
- if self.encode_transient:
- # transient encoding layers
- self.transient_encoding = nn.Sequential(
- nn.Linear(W+in_channels_t, W//2), nn.ReLU(True),
- nn.Linear(W//2, W//2), nn.ReLU(True),
- nn.Linear(W//2, W//2), nn.ReLU(True),
- nn.Linear(W//2, W//2), nn.ReLU(True))
- # transient output layers
- self.transient_sigma = nn.Sequential(nn.Linear(W//2, 1), nn.Softplus())
- if out_ch_size == 3: #NeRF, NeRFW output rgb
- self.transient_rgb = nn.Sequential(nn.Linear(W//2, 3), nn.Sigmoid())
- else: # output features like giraffe
- self.transient_rgb = nn.Sequential(nn.Linear(W//2, out_ch_size))
- self.transient_beta = nn.Sequential(nn.Linear(W//2, 1), nn.Softplus())
-
- def forward(self, x, sigma_only=False, output_transient=True):
- """
- Encodes input (xyz+dir) to rgb+sigma (not ready to render yet).
- For rendering this ray, please see rendering.py
-
- Inputs:
- x: the embedded vector of position (+ direction + appearance + transient)
- sigma_only: whether to infer sigma only.
- has_transient: whether to infer the transient component.
-
- Outputs (concatenated):
- if sigma_ony:
- static_sigma
- elif output_transient:
- static_rgb, static_sigma, transient_rgb, transient_sigma, transient_beta
- else:
- static_rgb, static_sigma
- """
- if sigma_only:
- input_xyz = x
- elif output_transient:
- input_xyz, input_dir_a, input_t = \
- torch.split(x, [self.in_channels_xyz,
- self.in_channels_dir+self.in_channels_a,
- self.in_channels_t], dim=-1)
- else:
- input_xyz, input_dir_a = \
- torch.split(x, [self.in_channels_xyz,
- self.in_channels_dir+self.in_channels_a], dim=-1)
- xyz_ = input_xyz
- for i in range(self.D):
- if i in self.skips:
- xyz_ = torch.cat([input_xyz, xyz_], 1)
- xyz_ = getattr(self, f"xyz_encoding_{i+1}")(xyz_)
-
- static_sigma = self.static_sigma(xyz_) # (B, 1)
- if sigma_only:
- return static_sigma
-
- xyz_encoding_final = self.xyz_encoding_final(xyz_)
- dir_encoding_input = torch.cat([xyz_encoding_final, input_dir_a], 1)
- dir_encoding = self.dir_encoding(dir_encoding_input)
- static_rgb = self.static_rgb(dir_encoding) # (B, 3)
- static = torch.cat([static_rgb, static_sigma], 1) # (B, 4)
-
- if not output_transient:
- return static
-
- transient_encoding_input = torch.cat([xyz_encoding_final, input_t], 1)
- transient_encoding = self.transient_encoding(transient_encoding_input)
- transient_sigma = self.transient_sigma(transient_encoding) # (B, 1)
- transient_rgb = self.transient_rgb(transient_encoding) # (B, 3)
- transient_beta = self.transient_beta(transient_encoding) # (B, 1)
-
- transient = torch.cat([transient_rgb, transient_sigma,
- transient_beta], 1) # (B, 5)
-
- return torch.cat([static, transient], 1) # (B, 9)
-
-def create_nerf(args):
- """Instantiate NeRF's MLP model.
- """
-
- # initialize embedding functions
- if args.reduce_embedding==2: # use DNeRF embedding
- embed_fn, input_ch, embedder_obj = get_embedder(args.multires, args.i_embed, args.reduce_embedding, args.epochToMaxFreq) # input_ch.shape=63
- else:
- embed_fn, input_ch, _ = get_embedder(args.multires, args.i_embed, args.reduce_embedding) # input_ch.shape=63
-
- input_ch_views = 0
- embeddirs_fn = None
- if args.use_viewdirs:
- if args.reduce_embedding==2: # use DNeRF embedding
- if args.no_DNeRF_viewdir: # no DNeRF embedding for viewdir
- raise NotImplementedError
- embeddirs_fn, input_ch_views, _ = get_embedder(args.multires_views, args.i_embed) # currently not used
- else:
- embeddirs_fn, input_ch_views, embedddirs_obj = get_embedder(args.multires_views, args.i_embed, args.reduce_embedding, args.epochToMaxFreq)
- else:
- embeddirs_fn, input_ch_views, _ = get_embedder(args.multires_views, args.i_embed, args.reduce_embedding) # input_ch_views.shape=27
- output_ch = 5 if args.N_importance > 0 else 4
- skips = [4]
-
- device = torch.device("cuda")
-
- encode_a = True # static appearance
- encode_t = True # transient
- if encode_a:
- if args.encode_hist: # experiemental embedding histogram
- embedding_a = torch.nn.Embedding(args.N_vocab, 5)
- embedding_a = embedding_a.to(device)
- if encode_t:
- if args.encode_hist: # experiemental embedding histogram
- embedding_t = torch.nn.Embedding(args.N_vocab, 2)
- embedding_t = embedding_t.to(device)
-
- # initialize NeRF model
- if args.NeRFH:
- model = NeRFW('coarse', D=args.netdepth, W=args.netwidth, skips=skips, in_channels_xyz=input_ch, in_channels_dir=input_ch_views)
- else:
- model = NeRF(D=args.netdepth, W=args.netwidth, input_ch=input_ch, output_ch=output_ch, skips=skips, input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs)
-
- if args.multi_gpu:
- model = torch.nn.DataParallel(model).to(device)
- else:
- model = model.to(device)
- grad_vars = list(model.parameters())
-
- model_fine = None
-
- if args.N_importance > 0:
- if args.NeRFH:
- model_fine = NeRFW('fine', D=args.netdepth, W=args.netwidth, skips=skips,
- in_channels_xyz=input_ch, in_channels_dir=input_ch_views,
- encode_appearance=True, encode_transient=True,
- in_channels_a=args.in_channels_a, in_channels_t=args.in_channels_t)
- else:
- model_fine = NeRF(D=args.netdepth_fine, W=args.netwidth_fine, input_ch=input_ch, output_ch=output_ch, skips=skips,
- input_ch_views=input_ch_views, use_viewdirs=args.use_viewdirs)
- if args.multi_gpu:
- model_fine = torch.nn.DataParallel(model_fine).to(device)
- else:
- model_fine = model_fine.to(device)
- grad_vars += list(model_fine.parameters())
- grad_vars += list(embedding_a.parameters())
- grad_vars += list(embedding_t.parameters())
-
- if args.NeRFH:
- network_query_fn = lambda inputs, viewdirs, ts, network_fn, \
- typ, embedding_a, embedding_t, output_transient, test_time : run_network_NeRFW(inputs, viewdirs, ts, network_fn,
- embed_fn=embed_fn,
- embeddirs_fn=embeddirs_fn,
- typ=typ,
- embedding_a=embedding_a,
- embedding_t=embedding_t,
- output_transient=output_transient,
- netchunk=args.netchunk,
- test_time=test_time)
- else:
- network_query_fn = lambda inputs, viewdirs, network_fn : run_network(inputs, viewdirs, network_fn,
- embed_fn=embed_fn,
- embeddirs_fn=embeddirs_fn,
- netchunk=args.netchunk)
-
- # Create optimizer if NeRF is need to be trained, otherwise returns None
- if args.no_grad_update:
- grad_vars = None
- optimizer = None
- else:
- optimizer = torch.optim.Adam(params=grad_vars, lr=args.lrate, betas=(0.9, 0.999))
-
- start = 0
- basedir = args.basedir
- expname = args.expname
-
- # Load checkpoints
- if args.ft_path is not None and args.ft_path!='None':
- ckpts = [args.ft_path]
- else:
- ckpts = [os.path.join(basedir, expname, f) for f in sorted(os.listdir(os.path.join(basedir, expname))) if 'tar' in f]
-
- print('Found ckpts', ckpts)
- if len(ckpts) > 0 and not args.no_reload:
- ckpt_path = ckpts[-1]
- print('Reloading from', ckpt_path)
- ckpt = torch.load(ckpt_path)
-
- start = ckpt['global_step']
- # optimizer.load_state_dict(ckpt['optimizer_state_dict'])
- # Load model
- model.load_state_dict(ckpt['network_fn_state_dict'])
-
- if model_fine is not None:
- model_fine.load_state_dict(ckpt['network_fine_state_dict'])
- embedding_a.load_state_dict(ckpt['embedding_a_state_dict'])
- embedding_t.load_state_dict(ckpt['embedding_t_state_dict'])
-
- ##########################
-
- render_kwargs_train = {
- 'network_query_fn' : network_query_fn,
- 'perturb' : args.perturb,
- 'N_importance' : args.N_importance,
- 'network_fine' : model_fine,
- 'N_samples' : args.N_samples,
- 'network_fn' : model,
- 'use_viewdirs' : args.use_viewdirs,
- 'white_bkgd' : args.white_bkgd,
- 'raw_noise_std' : args.raw_noise_std,
- 'embedding_a' : embedding_a,
- 'embedding_t' : embedding_t,
- 'test_time' : False
- }
-
- # NDC only good for LLFF-style forward facing data
- if args.dataset_type != 'llff' or args.no_ndc:
- print('Not ndc!')
- render_kwargs_train['ndc'] = False
- render_kwargs_train['lindisp'] = args.lindisp
-
- render_kwargs_test = {k : render_kwargs_train[k] for k in render_kwargs_train}
- render_kwargs_test['perturb'] = False
- render_kwargs_test['raw_noise_std'] = 0.
- render_kwargs_test['test_time'] = True
-
- return render_kwargs_train, render_kwargs_test, start, grad_vars, optimizer
\ No newline at end of file
diff --git a/script/models/options.py b/script/models/options.py
deleted file mode 100755
index 3b2284c..0000000
--- a/script/models/options.py
+++ /dev/null
@@ -1,100 +0,0 @@
-import configargparse
-def config_parser():
- parser = configargparse.ArgumentParser()
- parser.add_argument("-f", "--fff", help="a dummy argument to fool ipython", default="1")
- parser.add_argument("--device", type=int, default=-1, help='CUDA_VISIBLE_DEVICES')
- parser.add_argument("--multi_gpu", action='store_true', help='use multiple gpu on the server')
- parser.add_argument('--config', is_config_file=True, help='config file path')
- parser.add_argument("--expname", type=str, help='experiment name')
- parser.add_argument("--basedir", type=str, default='../logs', help='where to store ckpts and logs')
- parser.add_argument("--datadir", type=str, default='./data/llff/fern', help='input data directory')
-
- # 7Scenes
- parser.add_argument("--trainskip", type=int, default=1, help='will load 1/N images from train sets, useful for large datasets like 7 Scenes')
- parser.add_argument("--df", type=float, default=1., help='image downscale factor')
- parser.add_argument("--reduce_embedding", type=int, default=-1, help='fourier embedding mode: -1: paper default, \
- 0: reduce by half, 1: remove embedding, 2: DNeRF embedding')
- parser.add_argument("--epochToMaxFreq", type=int, default=-1, help='DNeRF embedding mode: (based on Nerfie paper): \
- hyper-parameter for when α should reach the maximum number of frequencies m')
- parser.add_argument("--render_pose_only", action='store_true', help='render a spiral video for 7 Scene')
- parser.add_argument("--save_pose_avg_stats", action='store_true', help='save a pose avg stats to unify NeRF, posenet, direct-pn training')
- parser.add_argument("--load_pose_avg_stats", action='store_true', help='load precomputed pose avg stats to unify NeRF, posenet, nerf tracking training')
- parser.add_argument("--train_local_nerf", type=int, default=-1, help='train local NeRF with ith training sequence only, ie. Stairs can pick 0~3')
- parser.add_argument("--render_video_train", action='store_true', help='render train set NeRF and save as video, make sure render_test is True')
- parser.add_argument("--render_video_test", action='store_true', help='render val set NeRF and save as video, make sure render_test is True')
- parser.add_argument("--frustum_overlap_th", type=float, help='frustsum overlap threshold')
- parser.add_argument("--no_DNeRF_viewdir", action='store_true', default=False, help='will not use DNeRF in viewdir encoding')
- parser.add_argument("--load_unique_view_stats", action='store_true', help='load unique views frame index')
-
- # NeRF training options
- parser.add_argument("--netdepth", type=int, default=8, help='layers in network')
- parser.add_argument("--netwidth", type=int, default=128, help='channels per layer')
- parser.add_argument("--netdepth_fine", type=int, default=8, help='layers in fine network')
- parser.add_argument("--netwidth_fine", type=int, default=128, help='channels per layer in fine network')
- parser.add_argument("--N_rand", type=int, default=1536, help='batch size (number of random rays per gradient step)')
- parser.add_argument("--lrate", type=float, default=5e-4, help='learning rate')
- parser.add_argument("--lrate_decay", type=float, default=250, help='exponential learning rate decay (in 1000 steps)')
- parser.add_argument("--chunk", type=int, default=1024*32, help='number of rays processed in parallel, decrease if running out of memory')
- parser.add_argument("--netchunk", type=int, default=1024*64, help='number of pts sent through network in parallel, decrease if running out of memory')
- parser.add_argument("--no_batching", action='store_true', help='only take random rays from 1 image at a time')
- parser.add_argument("--no_reload", action='store_true', help='do not reload weights from saved ckpt')
- parser.add_argument("--ft_path", type=str, default=None, help='specific weights npy file to reload for coarse network')
- parser.add_argument("--no_grad_update", action='store_true', default=False, help='do not update nerf in training')
-
- # NeRF-Hist training options
- parser.add_argument("--NeRFH", action='store_true', help='my implementation for NeRFH, to enable NeRF-Hist training, please make sure to add --encode_hist, otherwise it is similar to NeRFW')
- parser.add_argument("--N_vocab", type=int, default=1000,
- help='''number of vocabulary (number of images)
- in the dataset for nn.Embedding''')
- parser.add_argument("--fix_index", action='store_true', help='fix training frame index as 0')
- parser.add_argument("--encode_hist", default=False, action='store_true', help='encode histogram instead of frame index')
- parser.add_argument("--hist_bin", type=int, default=10, help='image histogram bin size')
- parser.add_argument("--in_channels_a", type=int, default=50, help='appearance embedding dimension, hist_bin*N_a when embedding histogram')
- parser.add_argument("--in_channels_t", type=int, default=20, help='transient embedding dimension, hist_bin*N_tau when embedding histogram')
-
- # NeRF rendering options
- parser.add_argument("--N_samples", type=int, default=64, help='number of coarse samples per ray')
- parser.add_argument("--N_importance", type=int, default=64,help='number of additional fine samples per ray')
- parser.add_argument("--perturb", type=float, default=1.,help='set to 0. for no jitter, 1. for jitter')
- parser.add_argument("--use_viewdirs", default=True, action='store_true', help='use full 5D input instead of 3D')
- parser.add_argument("--i_embed", type=int, default=0, help='set 0 for default positional encoding, -1 for none')
- parser.add_argument("--multires", type=int, default=10, help='log2 of max freq for positional encoding (3D location)')
- parser.add_argument("--multires_views", type=int, default=4, help='log2 of max freq for positional encoding (2D direction)')
- parser.add_argument("--raw_noise_std", type=float, default=0., help='std dev of noise added to regularize sigma_a output, 1e0 recommended')
- parser.add_argument("--render_only", action='store_true', help='do not optimize, reload weights and render out render_poses path')
- parser.add_argument("--render_test", action='store_true', help='render the test set instead of render_poses path')
- parser.add_argument("--render_factor", type=int, default=0, help='downsampling factor to speed up rendering, set 4 or 8 for fast preview')
-
- # legacy mesh options
- parser.add_argument("--mesh_only", action='store_true', help='do not optimize, reload weights and save mesh to a file')
- parser.add_argument("--mesh_grid_size", type=int, default=80,help='number of grid points to sample in each dimension for marching cubes')
-
- # training options
- parser.add_argument("--precrop_iters", type=int, default=0,help='number of steps to train on central crops')
- parser.add_argument("--precrop_frac", type=float,default=.5, help='fraction of img taken for central crops')
- parser.add_argument("--epochs", type=int, default=600,help='number of epochs to train')
-
- # dataset options
- parser.add_argument("--dataset_type", type=str, default='llff', help='options: llff / 7Scenes')
- parser.add_argument("--testskip", type=int, default=1, help='will load 1/N images from test/val sets, useful for large datasets like deepvoxels')
-
- ## legacy blender flags
- parser.add_argument("--white_bkgd", action='store_true', help='set to render synthetic data on a white bkgd (always use for dvoxels)')
- # parser.add_argument("--half_res", action='store_true', help='load blender synthetic data at 400x400 instead of 800x800')
-
- ## llff flags
- parser.add_argument("--factor", type=int, default=8, help='downsample factor for LLFF images')
- parser.add_argument("--no_ndc", action='store_true', help='do not use normalized device coordinates (set for non-forward facing scenes)')
- parser.add_argument("--lindisp", action='store_true', help='sampling linearly in disparity rather than depth')
- parser.add_argument("--spherify", action='store_true', help='set for spherical 360 scenes')
- parser.add_argument("--llffhold", type=int, default=8, help='will take every 1/N images as LLFF test set, paper uses 8')
- parser.add_argument("--no_bd_factor", action='store_true', default=False, help='do not use bd factor')
-
- # logging/saving options
- parser.add_argument("--i_print", type=int, default=1, help='frequency of console printout and metric loggin')
- parser.add_argument("--i_img", type=int, default=500, help='frequency of tensorboard image logging')
- parser.add_argument("--i_weights", type=int, default=200, help='frequency of weight ckpt saving')
- parser.add_argument("--i_testset", type=int, default=200, help='frequency of testset saving')
- parser.add_argument("--i_video", type=int, default=50000, help='frequency of render_poses video saving')
-
- return parser
\ No newline at end of file
diff --git a/script/models/ray_utils.py b/script/models/ray_utils.py
deleted file mode 100755
index c5de215..0000000
--- a/script/models/ray_utils.py
+++ /dev/null
@@ -1,44 +0,0 @@
-import torch
-import numpy as np
-
-# Ray helpers
-def get_rays(H, W, focal, c2w):
- i, j = torch.meshgrid(torch.linspace(0, W-1, W), torch.linspace(0, H-1, H), indexing='ij') # pytorch's meshgrid has indexing='ij'
- i = i.t()
- j = j.t()
- dirs = torch.stack([(i-W*.5)/focal, -(j-H*.5)/focal, -torch.ones_like(i)], -1)
-
- # Rotate ray directions from camera frame to the world frame
- rays_d = torch.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1) # dot product, equals to: [c2w.dot(dir) for dir in dirs]
- # Translate camera frame's origin to the world frame. It is the origin of all rays.
- rays_o = c2w[:3,-1].expand(rays_d.shape)
- return rays_o, rays_d # rays_o (100,100,3), rays_d (100,100,3)
-
-
-def get_rays_np(H, W, focal, c2w):
- i, j = np.meshgrid(np.arange(W, dtype=np.float32), np.arange(H, dtype=np.float32), indexing='xy')
- dirs = np.stack([(i-W*.5)/focal, -(j-H*.5)/focal, -np.ones_like(i)], -1)
- # Rotate ray directions from camera frame to the world frame
- rays_d = np.sum(dirs[..., np.newaxis, :] * c2w[:3,:3], -1) # dot product, equals to: [c2w.dot(dir) for dir in dirs]
- # Translate camera frame's origin to the world frame. It is the origin of all rays.
- rays_o = np.broadcast_to(c2w[:3,-1], np.shape(rays_d))
- return rays_o, rays_d
-
-def ndc_rays(H, W, focal, near, rays_o, rays_d):
- # Shift ray origins to near plane
- t = -(near + rays_o[...,2]) / rays_d[...,2] # t_n = −(n + o_z)/d_z move o to the ray's intersection with near plane
- rays_o = rays_o + t[...,None] * rays_d
-
- # Projection Formular (20)
- o0 = -1./(W/(2.*focal)) * rays_o[...,0] / rays_o[...,2]
- o1 = -1./(H/(2.*focal)) * rays_o[...,1] / rays_o[...,2]
- o2 = 1. + 2. * near / rays_o[...,2]
- # Formular (21)
- d0 = -1./(W/(2.*focal)) * (rays_d[...,0]/rays_d[...,2] - rays_o[...,0]/rays_o[...,2])
- d1 = -1./(H/(2.*focal)) * (rays_d[...,1]/rays_d[...,2] - rays_o[...,1]/rays_o[...,2])
- d2 = -2. * near / rays_o[...,2]
-
- rays_o = torch.stack([o0,o1,o2], -1) # o'
- rays_d = torch.stack([d0,d1,d2], -1) # d'
-
- return rays_o, rays_d
\ No newline at end of file
diff --git a/script/models/rendering.py b/script/models/rendering.py
deleted file mode 100755
index 5d9f876..0000000
--- a/script/models/rendering.py
+++ /dev/null
@@ -1,530 +0,0 @@
-import os
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-import imageio
-import time
-import math
-from tqdm import tqdm, trange
-from models.ray_utils import *
-from models.nerf import img2mse, mse2psnr,to8b
-from einops import rearrange, reduce, repeat
-import pdb
-
-'''
-Render Procedure:
-render()->batchify_rays()->render_rays()->raw2outputs()->sample_pdf()
-'''
-
-DEBUG = False
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Hierarchical sampling (section 5.2)
-def sample_pdf(bins, weights, N_samples, det=False, pytest=False):
- # Get pdf
- weights = weights + 1e-5 # prevent nans
- pdf = weights / torch.sum(weights, -1, keepdim=True)
- cdf = torch.cumsum(pdf, -1)
- cdf = torch.cat([torch.zeros_like(cdf[...,:1]), cdf], -1) # (batch, len(bins))
-
- # Take uniform samples
- if det:
- u = torch.linspace(0., 1., steps=N_samples)
- u = u.expand(list(cdf.shape[:-1]) + [N_samples])
- else:
- u = torch.rand(list(cdf.shape[:-1]) + [N_samples])
-
- # Pytest, overwrite u with numpy's fixed random numbers
- if pytest:
- np.random.seed(0)
- new_shape = list(cdf.shape[:-1]) + [N_samples]
- if det:
- u = np.linspace(0., 1., N_samples)
- u = np.broadcast_to(u, new_shape)
- else:
- u = np.random.rand(*new_shape)
- u = torch.Tensor(u)
-
- # Invert CDF
- u = u.contiguous()
- inds = torch.searchsorted(cdf.detach(), u, right=True)
- below = torch.max(torch.zeros_like(inds-1), inds-1)
- above = torch.min((cdf.shape[-1]-1) * torch.ones_like(inds), inds)
- inds_g = torch.stack([below, above], -1) # (batch, N_samples, 2)
-
- matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
- cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
- bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
-
- denom = (cdf_g[...,1]-cdf_g[...,0])
- denom = torch.where(denom<1e-5, torch.ones_like(denom), denom)
- t = (u-cdf_g[...,0])/denom
- samples = bins_g[...,0] + t * (bins_g[...,1]-bins_g[...,0])
-
- return samples
-
-def raw2outputs(raw, z_vals, rays_d, raw_noise_std=0, white_bkgd=False, pytest=False):
- '''
- Implementation of original NeRF
- seems like it is different with rendering implementation from nerf-w
- https://github.com/kwea123/nerf_pl/blob/nerfw/models/rendering.py
- Inputs:
- raw: torch.Tensor() [N_rays, N_samples, 4]
- '''
-
- # Function for computing density from model prediction. This value is
- # strictly between [0, 1].
- raw2alpha = lambda raw, dists, act_fn=F.relu: 1.-torch.exp(-act_fn(raw)*dists)
-
- # Compute 'distance' (in time) between each integration time along a ray.
- dists = z_vals[...,1:] - z_vals[...,:-1]
-
- # The 'distance' from the last integration time is infinity.
- dists = torch.cat([dists, torch.Tensor([1e10]).expand(dists[...,:1].shape)], -1) # [N_rays, N_samples] the last delta is infinity (1e10)
-
- # Multiply each distance by the norm of its corresponding direction ray
- # to convert to real world distance (accounts for non-unit directions).
- dists = dists * torch.norm(rays_d[...,None,:], dim=-1) # why raw2outputs_NeRFW doesn't have this step?
-
- # Extract RGB of each sample position along each ray. (NeRFW sigmoided in network)
- rgb = torch.sigmoid(raw[...,:3]) # [N_rays, N_samples, 3]
-
- # Add noise to model's predictions for density. Can be used to
- # regularize network during training (prevents floater artifacts).
- noise = 0.
- if raw_noise_std > 0.:
- noise = torch.randn(raw[...,3].shape) * raw_noise_std
-
- # Overwrite randomly sampled data if pytest
- if pytest:
- np.random.seed(0)
- noise = np.random.rand(*list(raw[...,3].shape)) * raw_noise_std
- noise = torch.Tensor(noise)
-
- # Predict density of each sample along each ray. Higher values imply
- # higher likelihood of being absorbed at this point.
- alpha = raw2alpha(raw[...,3] + noise, dists) # [N_rays, N_samples]
-
- # Compute weight for RGB of each sample along each ray. A cumprod() is
- # used to express the idea of the ray not having reflected up to this
- # sample yet.
- # weights = alpha * tf.math.cumprod(1.-alpha + 1e-10, -1, exclusive=True)
- weights = alpha * torch.cumprod(torch.cat([torch.ones((alpha.shape[0], 1)), 1.-alpha + 1e-10], -1), -1)[:, :-1]
-
- # Computed weighted color of each sample along each ray.
- rgb_map = torch.sum(weights[...,None] * rgb, -2) # [N_rays, 3]
-
- # Estimated depth map is expected distance.
- depth_map = torch.sum(weights * z_vals, -1)
-
- # Disparity map is inverse depth.
- disp_map = 1./torch.max(1e-10 * torch.ones_like(depth_map), depth_map / torch.sum(weights, -1))
-
- # Sum of weights along each ray. This value is in [0, 1] up to numerical error.
- acc_map = torch.sum(weights, -1)
-
- if white_bkgd:
- rgb_map = rgb_map + (1.-acc_map[...,None])
-
- return rgb_map, disp_map, acc_map, weights, depth_map
-
-def raw2outputs_NeRFW(raw, z_vals, rays_d, raw_noise_std=0, output_transient=False, beta_min=0.1, white_bkgd=False, test_time=False, static_only=True, typ="coarse"):
- ''' Convert NeRFW fine network output to rendered colors
- This version is implemented in nerf_pl https://github.com/kwea123/nerf_pl/tree/nerfw
- Inputs:
- raw: torch.Tensor() [N_rays, N_samples, 9]
-
- '''
-
- if typ=="coarse" and test_time:
- static_sigmas = raw[..., 0]
- transient_sigmas = None
- else:
- if output_transient==False:
- N_rays, N_samples, ch = raw.size()
- ch_rgbs = ch - 1 # rgb sigma - sigma channel. for rgb: ch_rgbs = 3, for feature: ch_rgbs = args.out_channel_size
- else:
- N_rays, N_samples, ch = raw.size()
- ch_rgbs = (ch - 3) // 2 # rgb sigma - static_sigma - transient_sigmas - transient_betas channels
-
- static_rgbs = raw[..., :ch_rgbs] # (N_rays, N_samples, 3), [..., 0:3]
- static_sigmas = raw[..., ch_rgbs] # (N_rays, N_samples) [..., 3]
- if output_transient:
- transient_rgbs = raw[..., ch_rgbs + 1: 2 * ch_rgbs + 1] # [..., 4:7]
- transient_sigmas = raw[..., 2 * ch_rgbs + 1] # [..., 7]
- transient_betas = raw[..., 2 * ch_rgbs + 2] # [..., 8]
- else:
- transient_sigmas = None
-
- # Convert these values using volume rendering
- deltas = z_vals[:, 1:] - z_vals[:, :-1] # (N_rays, N_samples_-1)
- delta_inf = 1e2 * torch.ones_like(deltas[:, :1]) # (N_rays, 1) the last delta is infinity, nerf used 1e10
-
- # In original NeRF, Multiply each distance by the norm of its corresponding direction ray
- # but not in this implementation
- deltas = torch.cat([deltas, delta_inf], -1) # (N_rays, N_samples_) [32768, 128]
-
- if output_transient:
- static_alphas = 1-torch.exp(-deltas*static_sigmas)
- transient_alphas = 1-torch.exp(-deltas*transient_sigmas)
- alphas = 1-torch.exp(-deltas*(static_sigmas+transient_sigmas))
- else:
- noise = torch.randn_like(static_sigmas) * raw_noise_std
- alphas = 1-torch.exp(-deltas*torch.relu(static_sigmas+noise))
-
- alphas_shifted = \
- torch.cat([torch.ones_like(alphas[:, :1]), 1-alphas], -1) # [1, 1-a1, 1-a2, ...]
- transmittance = torch.cumprod(alphas_shifted[:, :-1], -1) # [1, 1-a1, (1-a1)(1-a2), ...]
-
- if output_transient:
- static_weights = static_alphas * transmittance
- transient_weights = transient_alphas * transmittance
-
- weights = alphas * transmittance
- weights_sum = reduce(weights, 'n1 n2 -> n1', 'sum')
- acc_map = weights_sum
-
- if typ=="coarse" and test_time:
- rgb_map=None
- disp_map=None
- depth_map=None
- beta=None
- return rgb_map, disp_map, acc_map, weights, depth_map, transient_sigmas, beta
-
- if output_transient:
- static_rgb_map = reduce(rearrange(static_weights, 'n1 n2 -> n1 n2 1')*static_rgbs,
- 'n1 n2 c -> n1 c', 'sum') # _rgb_fine_static
- if white_bkgd:
- static_rgb_map += 1-rearrange(weights_sum, 'n -> n 1')
-
- transient_rgb_map = \
- reduce(rearrange(transient_weights, 'n1 n2 -> n1 n2 1')*transient_rgbs,
- 'n1 n2 c -> n1 c', 'sum') # _rgb_fine_transient
- beta = reduce(transient_weights*transient_betas, 'n1 n2 -> n1', 'sum')
-
- # Add beta_min AFTER the beta composition. Different from eq 10~12 in the paper.
- # See "Notes on differences with the paper" in README.
- beta += beta_min
-
- # the rgb maps here are when both fields exist
- rgb_fine = static_rgb_map + transient_rgb_map
- rgb_map = rgb_fine
-
- if test_time and static_only: # output static rgbs only, need testing
- # Compute static rgbs when only one field exists.
- # The result is different from when both fields exist, since the transimttance
- # will change.
- static_alphas_shifted = \
- torch.cat([torch.ones_like(static_alphas[:, :1]), 1-static_alphas], -1)
- static_transmittance = torch.cumprod(static_alphas_shifted[:, :-1], -1)
- static_weights_ = static_alphas * static_transmittance
- static_rgb_map_ = \
- reduce(rearrange(static_weights_, 'n1 n2 -> n1 n2 1')*static_rgbs,
- 'n1 n2 c -> n1 c', 'sum')
- if white_bkgd:
- static_rgb_map_ += 1-rearrange(weights_sum, 'n -> n 1')
- rgb_fine = static_rgb_map_
- depth_map = reduce(static_weights_*z_vals, 'n1 n2 -> n1', 'sum')
- disp_map = 1./torch.max(1e-10 * torch.ones_like(depth_map), depth_map / torch.sum(weights, -1))
- return rgb_map, disp_map, acc_map, weights, depth_map, transient_sigmas, beta
-
- else:
- rgb_map = reduce(rearrange(weights, 'n1 n2 -> n1 n2 1')*static_rgbs,
- 'n1 n2 c -> n1 c', 'sum')
- if white_bkgd:
- rgb_map += 1-rearrange(weights_sum, 'n -> n 1')
-
- beta=torch.Tensor([0]).repeat(weights_sum.shape[0]) # this is useless, just to fake a return
-
- # compute depth_map and disp_map
- depth_map = torch.sum(weights * z_vals, -1)
- disp_map = 1./torch.max(1e-10 * torch.ones_like(depth_map), depth_map / torch.sum(weights, -1))
- return rgb_map, disp_map, acc_map, weights, depth_map, transient_sigmas, beta
-
-def render_rays(ray_batch,
- network_fn,
- network_query_fn,
- N_samples,
- retraw=False,
- lindisp=False,
- perturb=0.,
- N_importance=0,
- network_fine=None,
- white_bkgd=False,
- raw_noise_std=0.,
- verbose=False,
- pytest=False,
- i_epoch=-1,
- embedding_a=None,
- embedding_t=None,
- test_time=False):
- N_rays = ray_batch.shape[0]
- rays_o, rays_d = ray_batch[:,0:3], ray_batch[:,3:6] # [N_rays, 3] each
- viewdirs = ray_batch[:,8:11] if ray_batch.shape[-1] > 8 else None
- bounds = torch.reshape(ray_batch[...,6:8], [-1,1,2])
- near, far = bounds[...,0], bounds[...,1] # [-1,1]
- img_idxs = ray_batch[...,11:] # same as ts from nerf_pl-nerfw code [N_rays]
-
- t_vals = torch.linspace(0., 1., steps=N_samples)
- if not lindisp:
- z_vals = near * (1.-t_vals) + far * (t_vals) # sample in depth
- else:
- z_vals = 1./(1./near * (1.-t_vals) + 1./far * (t_vals)) # sample in diparity disparity = 1/depth
-
- z_vals = z_vals.expand([N_rays, N_samples])
-
- if perturb > 0.:
- # get intervals between samples
- mids = .5 * (z_vals[...,1:] + z_vals[...,:-1]) # find mid points of each interval
- upper = torch.cat([mids, z_vals[...,-1:]], -1)
- lower = torch.cat([z_vals[...,:1], mids], -1)
- # stratified samples in those intervals
- t_rand = torch.rand(z_vals.shape) # randomly choose in intervals
-
- z_vals = lower + (upper - lower) * t_rand
-
- pts = rays_o[...,None,:] + rays_d[...,None,:] * z_vals[...,:,None] # [N_rays, N_samples, 3] # Sample 3D point print it out, [batchsize, 64 samples, 3 axis]
- # inference coarse MLPs
- if i_epoch>=0: # Coarse-to-fine
- raw = network_query_fn(pts, viewdirs, None, network_fn, 'coarse', None, None, False, test_time=test_time, epoch=i_epoch)
- else:
- raw = network_query_fn(pts, viewdirs, None, network_fn, 'coarse', None, None, False, test_time=test_time)
-
-
- rgb_map, disp_map, acc_map, weights, depth_map, _, _ = raw2outputs_NeRFW(raw, z_vals, rays_d, raw_noise_std, white_bkgd, test_time=test_time, typ="coarse")
- if N_importance > 0:
-
- rgb_map_0, disp_map_0, acc_map_0 = rgb_map, disp_map, acc_map
-
- z_vals_mid = .5 * (z_vals[...,1:] + z_vals[...,:-1])
- z_samples = sample_pdf(z_vals_mid, weights[...,1:-1], N_importance, det=(perturb==0.), pytest=pytest)
- z_samples = z_samples.detach()
-
- z_vals, _ = torch.sort(torch.cat([z_vals, z_samples], -1), -1)
- pts = rays_o[...,None,:] + rays_d[...,None,:] * z_vals[...,:,None] # [N_rays, N_samples + N_importance, 3]
-
- output_transient = True
-
- # inference fine MLPs
- if i_epoch>=0: # Coarse-to-fine
- raw = network_query_fn(pts, viewdirs, img_idxs, network_fine, 'fine', embedding_a, embedding_t, output_transient, test_time=test_time, epoch=i_epoch)
- else:
- raw = network_query_fn(pts, viewdirs, img_idxs, network_fine, 'fine', embedding_a, embedding_t, output_transient, test_time=test_time)
-
-
- rgb_map, disp_map, acc_map, weights, depth_map, transient_sigmas, beta = raw2outputs_NeRFW(raw, z_vals, rays_d, raw_noise_std, output_transient, network_fine.beta_min, white_bkgd, test_time, typ="fine")
-
- ret = {'rgb_map' : rgb_map, 'disp_map' : disp_map, 'acc_map' : acc_map}
- if retraw:
- ret['raw'] = raw
-
- if N_importance > 0 and test_time:
- pass
-
- elif N_importance > 0:
- ret['rgb0'] = rgb_map_0
- ret['disp0'] = disp_map_0
- ret['acc0'] = acc_map_0
- ret['z_std'] = torch.std(z_samples, dim=-1, unbiased=False) # [N_rays]
- ret['transient_sigmas'] = transient_sigmas
- ret['beta'] = beta
-
- for k in ret:
- if (torch.isnan(ret[k]).any() or torch.isinf(ret[k]).any()) and DEBUG:
- print(f"! [Numerical Error] {k} contains nan or inf.")
-
- return ret
-
-def batchify_rays(rays_flat, chunk=1024*32, **kwargs):
- """Render rays in smaller minibatches to avoid OOM.
- """
- all_ret = {}
- for i in range(0, rays_flat.shape[0], chunk):
- ret = render_rays(rays_flat[i:i+chunk], **kwargs)
- for k in ret:
- if k not in all_ret:
- all_ret[k] = []
- all_ret[k].append(ret[k])
-
- all_ret = {k : torch.cat(all_ret[k], 0) for k in all_ret}
- return all_ret
-
-def render(H, W, focal, chunk=1024*32, rays=None, c2w=None, ndc=True,
- near=0., far=1.,
- use_viewdirs=False, c2w_staticcam=None, img_idx=torch.Tensor(0),
- **kwargs):
- if c2w is not None:
- # special case to render full image
- rays_o, rays_d = get_rays(H, W, focal, c2w)
- else:
- # use provided ray batch
- rays_o, rays_d = rays
-
- if use_viewdirs:
- # provide ray directions as input
- viewdirs = rays_d
- if c2w_staticcam is not None:
- # special case to visualize effect of viewdirs
- rays_o, rays_d = get_rays(H, W, focal, c2w_staticcam)
- viewdirs = viewdirs / torch.norm(viewdirs, dim=-1, keepdim=True)
- viewdirs = torch.reshape(viewdirs, [-1,3]).float()
-
- sh = rays_d.shape # [..., 3]
- if ndc:
- # for forward facing scenes
- rays_o, rays_d = ndc_rays(H, W, focal, 1., rays_o, rays_d)
-
- # Create ray batch
- rays_o = torch.reshape(rays_o, [-1,3]).float()
- rays_d = torch.reshape(rays_d, [-1,3]).float()
- near, far = near * torch.ones_like(rays_d[...,:1]), far * torch.ones_like(rays_d[...,:1])
- rays = torch.cat([rays_o, rays_d, near, far], -1)
- if use_viewdirs:
- rays = torch.cat([rays, viewdirs], -1) # [1024, 11]
-
- # for NeRFW, we need to add frame index as input
- if img_idx.shape[0] != rays.shape[0]:
- img_idx = img_idx.repeat(rays.shape[0],1) # [1024, 1]
- rays = torch.cat([rays, img_idx], 1) # [1024, 12]
-
- # Render and reshape
- all_ret = batchify_rays(rays, chunk, **kwargs)
- for k in all_ret:
- k_sh = list(sh[:-1]) + list(all_ret[k].shape[1:])
- all_ret[k] = torch.reshape(all_ret[k], k_sh)
-
- k_extract = ['rgb_map', 'disp_map', 'acc_map']
- ret_list = [all_ret[k] for k in k_extract]
- ret_dict = {k : all_ret[k] for k in all_ret if k not in k_extract}
- return ret_list + [ret_dict]
-
-### This is for validation ###
-def render_path(args, render_poses, hwf, chunk, render_kwargs, gt_imgs=None, savedir=None, render_factor=0, single_gt_img=False, img_ids=torch.Tensor(0)):
- # render_kwargs['network_fn'].eval()
-
- H, W, focal = hwf
-
- if render_factor!=0:
- # Render downsampled for speed
- H = int(H//render_factor)
- W = int(W//render_factor)
- focal = focal/render_factor
-
- rgbs = []
- disps = []
- rgb0s = []
- psnr = []
-
- t = time.time()
- for i, c2w in enumerate(tqdm(render_poses)):
- t = time.time()
- rgb, disp, acc, extras = render(H, W, focal, chunk=chunk, c2w=c2w[:3,:4], img_idx=img_ids[i], **render_kwargs)
- rgbs.append(rgb.cpu().numpy())
- disps.append(disp.cpu().numpy())
- # rgb0s.append(extras['rgb0'].cpu().numpy())
- if i==0:
- print(rgb.shape, disp.shape)
-
- if gt_imgs is not None:
- if single_gt_img:
- p = -10. * np.log10(np.mean(np.square(rgb.cpu().numpy() - gt_imgs)))
- else:
- p = -10. * np.log10(np.mean(np.square(rgb.cpu().numpy() - gt_imgs[i])))
- psnr.append(p)#print(p)
-
- if savedir is not None:
- # rgb8_c = to8b(rgb0s[-1]) # save coarse img
- # filename = os.path.join(savedir, '{:03d}_coarse.png'.format(i))
- # imageio.imwrite(filename, rgb8_c)
-
- rgb8_f = to8b(rgbs[-1]) # save coarse+fine img
- filename = os.path.join(savedir, '{:03d}.png'.format(i))
- imageio.imwrite(filename, rgb8_f)
-
- ### Need validate code
- rgb_gt = to8b(gt_imgs[i]) # save GT img here
- filename = os.path.join(savedir, '{:03d}_GT.png'.format(i))
- imageio.imwrite(filename, rgb_gt)
-
- rgb_disp = to8b(disps[-1] / np.max(disps[-1])) # save GT img here
- filename = os.path.join(savedir, '{:03d}_disp.png'.format(i))
- imageio.imwrite(filename, rgb_disp)
-
- rgbs = np.stack(rgbs, 0)
- disps = np.stack(disps, 0)
- psnr = np.mean(psnr,0)
- print("Mean PSNR of this run is:", psnr)
- return rgbs, disps
-
-def render_test(args, train_dl, val_dl, hwf, start, render_kwargs_test, decoder_coarse=None, decoder_fine=None):
-
- ### Eval Training set result
- trainsavedir = os.path.join(args.basedir, args.expname, 'evaluate_train_{}_{:06d}'.format('test' if args.render_test else 'path', start))
- os.makedirs(trainsavedir, exist_ok=True)
- images_train = []
- poses_train = []
- index_train = []
- # views from validation set
- for img, pose, img_idx in train_dl:
- img_val = img.permute(0,2,3,1) # (1,240,360,3)
- pose_val = torch.zeros(1,4,4)
- pose_val[0,:3,:4] = pose.reshape(3,4)[:3,:4] # (1,3,4))
- pose_val[0,3,3] = 1.
- images_train.append(img_val)
- poses_train.append(pose_val)
- index_train.append(img_idx)
-
- images_train = torch.cat(images_train, dim=0).numpy()
- poses_train = torch.cat(poses_train, dim=0).to(device)
- index_train = torch.cat(index_train, dim=0).to(device)
- print('train poses shape', poses_train.shape)
-
- with torch.no_grad():
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- rgbs, disps = render_path(args, poses_train.to(device), hwf, args.chunk, render_kwargs_test, gt_imgs=images_train, savedir=trainsavedir, img_ids=index_train)
- torch.set_default_tensor_type('torch.FloatTensor')
- print('Saved train set')
- if args.render_video_train:
- print('Saving trainset as video', rgbs.shape, disps.shape)
- moviebase = os.path.join(args.basedir, args.expname, '{}_trainset_{:06d}_'.format(args.expname, start))
- imageio.mimwrite(moviebase + 'train_rgb.mp4', to8b(rgbs), fps=15, quality=8)
- imageio.mimwrite(moviebase + 'train_disp.mp4', to8b(disps / np.max(disps)), fps=15, quality=8)
- del images_train
- del poses_train
- # clean GPU memory after testing
- torch.cuda.empty_cache()
-
- ### Eval Validation set result
- testsavedir = os.path.join(args.basedir, args.expname, 'evaluate_val_{}_{:06d}'.format('test' if args.render_test else 'path', start))
- os.makedirs(testsavedir, exist_ok=True)
- images_val = []
- poses_val = []
- index_val = []
- # views from validation set
- for img, pose, img_idx in val_dl:
- img_val = img.permute(0,2,3,1) # (1,240,360,3)
- pose_val = torch.zeros(1,4,4)
- pose_val[0,:3,:4] = pose.reshape(3,4)[:3,:4] # (1,3,4))
- pose_val[0,3,3] = 1.
- images_val.append(img_val)
- poses_val.append(pose_val)
- index_val.append(img_idx)
-
- images_val = torch.cat(images_val, dim=0).numpy()
- poses_val = torch.cat(poses_val, dim=0).to(device)
- index_val = torch.cat(index_val, dim=0).to(device)
- print('test poses shape', poses_val.shape)
- with torch.no_grad():
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- rgbs, disps = render_path(args, poses_val.to(device), hwf, args.chunk, render_kwargs_test, gt_imgs=images_val, savedir=testsavedir, img_ids=index_val)
- torch.set_default_tensor_type('torch.FloatTensor')
- print('Saved test set')
- if args.render_video_test:
- print('Saving testset as video', rgbs.shape, disps.shape)
- moviebase = os.path.join(args.basedir, args.expname, '{}_test_{:06d}_'.format(args.expname, start))
- imageio.mimwrite(moviebase + 'test_rgb.mp4', to8b(rgbs), fps=15, quality=8)
- imageio.mimwrite(moviebase + 'test_disp.mp4', to8b(disps / np.max(disps)), fps=15, quality=8)
- del images_val
- del poses_val
- return
diff --git a/script/run_feature.py b/script/run_feature.py
deleted file mode 100755
index 7d5e62d..0000000
--- a/script/run_feature.py
+++ /dev/null
@@ -1,454 +0,0 @@
-import sys
-sys.path.append('../')
-import torch
-from torch import nn, optim
-from torchvision.utils import save_image
-import os, pdb
-from torchsummary import summary
-from dataset_loaders.load_7Scenes import load_7Scenes_dataloader
-from dataset_loaders.load_Cambridge import load_Cambridge_dataloader
-import os.path as osp
-import numpy as np
-from utils.utils import plot_features, save_image_saliancy, save_image_saliancy_single
-from utils.utils import freeze_bn_layer, freeze_bn_layer_train
-from models.nerfw import create_nerf
-from tqdm import tqdm
-from dm.callbacks import EarlyStopping
-from feature.dfnet import DFNet, DFNet_s
-# from feature.efficientnet import EfficientNetB3 as DFNet
-# from feature.efficientnet import EfficientNetB0 as DFNet
-from feature.misc import *
-from feature.options import config_parser
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-np.random.seed(0)
-torch.manual_seed(0)
-import random
-random.seed(0)
-
-def tmp_plot(target_in, rgb_in, features_target, features_rgb):
- '''
- print 1 pair of salient feature map
- '''
- print("for debug only...")
- pdb.set_trace()
- ### plot featues with pixel-wise addition
- save_image(target_in[1], './tmp/target_in.png')
- save_image(rgb_in[1], './tmp/rgb_in.png')
- save_image_saliancy(features_target[1], './tmp/target', True)
- save_image_saliancy(features_rgb[1], './tmp/rgb', True)
- ### plot featues seperately
- save_image(target_in[1], './tmp/target_in.png')
- save_image(rgb_in[1], './tmp/rgb_in.png')
- plot_features(features_target[:,1:2,...], './tmp/target', False)
- plot_features(features_rgb[:,1:2,...], './tmp/rgb', False)
- sys.exit()
-
-def tmp_plot2(target_in, rgb_in, features_target, features_rgb, i=0):
- '''
- print 1 pair of batch of salient feature map
- :param: target_in [B, 3, H, W]
- :param: rgb_in [B, 3, H, W]
- :param: features_target [B, C, H, W]
- :param: features_rgb [B, C, H, W]
- :param: frame index i of batch
- '''
- print("for debug only...")
- save_image(target_in[i], './tmp/target_in.png')
- save_image(rgb_in[i], './tmp/rgb_in.png')
- pdb.set_trace()
- features_t = features_target[i].clone()[:, None, :, :]
- features_r = features_rgb[i].clone()[:, None, :, :]
- save_image_saliancy(features_t, './tmp/target', True)
- save_image_saliancy(features_r, './tmp/rgb', True)
-
-def tmp_plot3(target_in, rgb_in, features_target, features_rgb, i=0):
- '''
- print 1 pair of 1 sample of salient feature map
- :param: target_in [B, 3, H, W]
- :param: rgb_in [B, 3, H, W]
- :param: features_target [B, C, H, W]
- :param: features_rgb [B, C, H, W]
- :param: frame index i of batch
- '''
- print("for debug only...")
- save_image(target_in[i], './tmp/target_in.png')
- save_image(rgb_in[i], './tmp/rgb_in.png')
- features_t = features_target[i].clone()[:, None, :, :]
- features_r = features_rgb[i].clone()[:, None, :, :]
- save_image_saliancy(features_t[0], './tmp/target', True)
- save_image_saliancy(features_r[0], './tmp/rgb', True)
-
-def lognuniform(low=-2, high=0, size=1, base=10):
- ''' sample from log uniform distribution between 0.01~1 '''
- return np.power(base, np.random.uniform(low, high, size))
-
-def getrelpose(pose1, pose2):
- ''' get relative pose from abs pose pose1 to abs pose pose2
- R^{v}_{gt} = R_v * R_gt.T
- :param: pose1 [B, 3, 4]
- :param: pose2 [B, 3, 4]
- return rel_pose [B, 3, 4]
- '''
- assert(pose1.shape == pose2.shape)
- rel_pose = pose1 - pose2 # compute translation term difference
- rel_pose[:,:3,:3] = pose2[:,:3,:3] @ torch.transpose(pose1[:,:3,:3], 1, 2) # compute rotation term difference
- return rel_pose
-
-parser = config_parser()
-args = parser.parse_args()
-
-def train_on_batch(args, targets, rgbs, poses, feat_model, dset_size, FeatureLoss, optimizer, hwf):
- ''' core training loop for featurenet'''
- feat_model.train()
- H, W, focal = hwf
- H, W = int(H), int(W)
- if args.freezeBN:
- feat_model = freeze_bn_layer_train(feat_model)
-
- train_loss_epoch = []
- select_inds = np.random.choice(dset_size, size=[dset_size], replace=False) # (N_rand,)
-
- batch_size=args.featurenet_batch_size # manual setting, use smaller batch size like featurenet_batch_size = 4 if OOM
- if dset_size % batch_size == 0:
- N_iters = dset_size//batch_size
- else:
- N_iters = dset_size//batch_size + 1
- i_batch = 0
-
- for i in range(0, N_iters):
- if i_batch + batch_size > dset_size:
- i_batch = 0
- break
- i_inds = select_inds[i_batch:i_batch+batch_size]
- i_batch = i_batch + batch_size
-
- # convert input shape to [B, 3, H, W]
- target_in = targets[i_inds].clone().permute(0,3,1,2).to(device)
- rgb_in = rgbs[i_inds].clone().permute(0,3,1,2).to(device)
- pose = poses[i_inds].clone().reshape(batch_size, 12).to(device)
- pose = torch.cat([pose, pose]) # double gt pose tensor
-
- features, predict_pose = feat_model(torch.cat([target_in, rgb_in]), True, upsampleH=H, upsampleW=W) # features: (1, [2, B, C, H, W])
-
- # get features_target and features_rgb
- if args.DFNet:
- features_target = features[0] # [3, B, C, H, W]
- features_rgb = features[1]
- else:
- features_target = features[0][0]
- features_rgb = features[0][1]
-
- # svd, seems not very benificial here, therefore removed
-
- if args.poselossonly:
- loss_pose = PoseLoss(args, predict_pose, pose, device) # target
- loss = loss_pose
- elif args.featurelossonly: # Not good. To be removed later
- loss_f = FeatureLoss(features_rgb, features_target)
- loss = loss_f
- else:
- loss_pose = PoseLoss(args, predict_pose, pose, device) # target
- if args.tripletloss:
- loss_f = triplet_loss_hard_negative_mining_plus(features_rgb, features_target, margin=args.triplet_margin)
- else:
- loss_f = FeatureLoss(features_rgb, features_target)
- loss = loss_pose + loss_f
-
- loss.backward()
- optimizer.step()
- optimizer.zero_grad()
- train_loss_epoch.append(loss.item())
- train_loss = np.mean(train_loss_epoch)
- return train_loss
-
-def train_on_batch_with_random_view_synthesis(args, targets, rgbs, poses, virtue_view, poses_perturb, feat_model, dset_size, FeatureLoss, optimizer, hwf, img_idxs, render_kwargs_test):
- ''' we implement random view synthesis for generating more views to help training posenet '''
- feat_model.train()
-
- H, W, focal = hwf
- H, W = int(H), int(W)
-
- if args.freezeBN:
- feat_model = freeze_bn_layer_train(feat_model)
-
- train_loss_epoch = []
-
- # random generate batch_size of idx
- select_inds = np.random.choice(dset_size, size=[dset_size], replace=False) # (N_rand,)
-
- batch_size=args.featurenet_batch_size # manual setting, use smaller batch size like featurenet_batch_size = 4 if OOM
- if dset_size % batch_size == 0:
- N_iters = dset_size//batch_size
- else:
- N_iters = dset_size//batch_size + 1
-
- i_batch = 0
- for i in range(0, N_iters):
- if i_batch + batch_size > dset_size:
- i_batch = 0
- break
- i_inds = select_inds[i_batch:i_batch+batch_size]
- i_batch = i_batch + batch_size
-
- # convert input shape to [B, 3, H, W]
- target_in = targets[i_inds].clone().permute(0,3,1,2).to(device)
- rgb_in = rgbs[i_inds].clone().permute(0,3,1,2).to(device)
- pose = poses[i_inds].clone().reshape(batch_size, 12).to(device)
- rgb_perturb = virtue_view[i_inds].clone().permute(0,3,1,2).to(device)
- pose_perturb = poses_perturb[i_inds].clone().reshape(batch_size, 12).to(device)
-
- # inference feature model for GT and nerf image
- pose = torch.cat([pose, pose]) # double gt pose tensor
- features, predict_pose = feat_model(torch.cat([target_in, rgb_in]), return_feature=True, upsampleH=H, upsampleW=W) # features: (1, [2, B, C, H, W])
-
- # get features_target and features_rgb
- if args.DFNet:
- features_target = features[0] # [3, B, C, H, W]
- features_rgb = features[1]
-
- loss_pose = PoseLoss(args, predict_pose, pose, device) # target
-
- if args.tripletloss:
- loss_f = triplet_loss_hard_negative_mining_plus(features_rgb, features_target, margin=args.triplet_margin)
- else:
- loss_f = FeatureLoss(features_rgb, features_target) # feature Maybe change to s2d-ce loss
-
- # inference model for RVS image
- _, virtue_pose = feat_model(rgb_perturb.to(device), False)
-
- # add relative pose loss here. TODO: This FeatureLoss is nn.MSE. Should be fixed later
- loss_pose_perturb = PoseLoss(args, virtue_pose, pose_perturb, device)
- loss = args.combine_loss_w[0]*loss_pose + args.combine_loss_w[1]*loss_f + args.combine_loss_w[2]*loss_pose_perturb
-
- loss.backward()
- optimizer.step()
- optimizer.zero_grad()
- train_loss_epoch.append(loss.item())
- train_loss = np.mean(train_loss_epoch)
- return train_loss
-
-def train_feature(args, train_dl, val_dl, test_dl, hwf, i_split, near, far):
-
- # # load pretrained PoseNet model
- if args.DFNet_s:
- feat_model = DFNet_s()
- else:
- feat_model = DFNet()
-
- if args.pretrain_model_path != '':
- print("load posenet from ", args.pretrain_model_path)
- feat_model.load_state_dict(torch.load(args.pretrain_model_path))
-
- # # Freeze BN to not updating gamma and beta
- if args.freezeBN:
- feat_model = freeze_bn_layer(feat_model)
-
- feat_model.to(device)
- # summary(feat_model, (3, 240, 427))
-
- # set optimizer
- optimizer = optim.Adam(feat_model.parameters(), lr=args.learning_rate)
- scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.95, patience=args.patience[1], verbose=True)
-
- # set callbacks parameters
- early_stopping = EarlyStopping(args, patience=args.patience[0], verbose=False)
-
- # loss function
- loss_func = nn.MSELoss(reduction='mean')
-
- i_train, i_val, i_test = i_split
- # Cast intrinsics to right types
- H, W, focal = hwf
- H, W = int(H), int(W)
- hwf = [H, W, focal]
-
- # Create log dir and copy the config file
- basedir = args.basedir
- expname = args.expname
- os.makedirs(os.path.join(basedir, expname), exist_ok=True)
- f = os.path.join(basedir, expname, 'args.txt')
- with open(f, 'w') as file:
- for arg in sorted(vars(args)):
- attr = getattr(args, arg)
- file.write('{} = {}\n'.format(arg, attr))
- if args.config is not None:
- f = os.path.join(basedir, expname, 'config.txt')
- with open(f, 'w') as file:
- file.write(open(args.config, 'r').read())
-
- # load NeRF
- _, render_kwargs_test, start, _, _ = create_nerf(args)
- global_step = start
-
- bds_dict = {
- 'near' : near,
- 'far' : far,
- }
- # render_kwargs_train.update(bds_dict)
- render_kwargs_test.update(bds_dict)
- if args.reduce_embedding==2:
- render_kwargs_test['i_epoch'] = start
-
- N_epoch = args.epochs + 1 # epoch
- print('Begin')
- print('TRAIN views are', i_train)
- print('TEST views are', i_test)
- print('VAL views are', i_val)
-
- world_setup_dict = {
- 'pose_scale' : train_dl.dataset.pose_scale,
- 'pose_scale2' : train_dl.dataset.pose_scale2,
- 'move_all_cam_vec' : train_dl.dataset.move_all_cam_vec,
- }
-
- if args.eval:
- feat_model.eval()
- ### testing
- # get_error_in_q(args, train_dl, feat_model, len(train_dl.dataset), device, batch_size=1)
- get_error_in_q(args, test_dl, feat_model, len(val_dl.dataset), device, batch_size=1)
- sys.exit()
-
- if args.render_feature_only:
- targets, rgbs, poses, img_idxs = render_nerfw_imgs(args, test_dl, hwf, device, render_kwargs_test, world_setup_dict)
- dset_size = poses.shape[0]
- feat_model.eval()
- # extract features
- for i in range(dset_size):
- target_in = targets[i:i+1].permute(0,3,1,2).to(device)
- rgb_in = rgbs[i:i+1].permute(0,3,1,2).to(device)
-
- features, _ = feat_model(torch.cat([target_in, rgb_in]), True, upsampleH=H, upsampleW=W)
- if args.DFNet:
- features_target = features[0] # [3, B, C, H, W]
- features_rgb = features[1]
-
- # save features
- save_i = 2 # saving feature index, save_i out of 128
- ft = features_target[0, None, :, save_i] # [1,1,H,W]
- fr = features_rgb[0, None, :, save_i] # [1,1,H,W]
-
- scene = 'shop_gap/'
- save_path = './tmp/'+scene
- save_path_t = './tmp/'+scene+'target/'
- save_path_r = './tmp/'+scene+'rgb/'
- if not os.path.isdir(save_path):
- os.mkdir(save_path)
- if not os.path.isdir(save_path_t):
- os.mkdir(save_path_t)
- if not os.path.isdir(save_path_r):
- os.mkdir(save_path_r)
- save_image_saliancy_single(ft, save_path_t + '%04d.png'%i, True)
- save_image_saliancy_single(fr, save_path_r + '%04d.png'%i, True)
-
- print("render features done")
- sys.exit()
-
-
- targets, rgbs, poses, img_idxs = render_nerfw_imgs(args, train_dl, hwf, device, render_kwargs_test, world_setup_dict)
-
- dset_size = len(train_dl.dataset)
- # clean GPU memory before testing, try to avoid OOM
- torch.cuda.empty_cache()
-
- model_log = tqdm(total=0, position=1, bar_format='{desc}')
- for epoch in tqdm(range(N_epoch), desc='epochs'):
-
- if args.random_view_synthesis:
- ### this is the implementation of RVS ###
- isRVS = epoch % args.rvs_refresh_rate == 0 # decide if to resynthesis new views
-
- if isRVS:
- # random sample virtual camera locations, todo:
- rand_trans = args.rvs_trans
- rand_rot = args.rvs_rotation
-
- # determine bounding box
- b_min = [poses[:,0,3].min()-args.d_max, poses[:,1,3].min()-args.d_max, poses[:,2,3].min()-args.d_max]
- b_max = [poses[:,0,3].max()+args.d_max, poses[:,1,3].max()+args.d_max, poses[:,2,3].max()+args.d_max]
-
- poses_perturb = poses.clone().numpy()
- for i in range(dset_size):
- poses_perturb[i] = perturb_single_render_pose(poses_perturb[i], rand_trans, rand_rot)
- for j in range(3):
- if poses_perturb[i,j,3] < b_min[j]:
- poses_perturb[i,j,3] = b_min[j]
- elif poses_perturb[i,j,3]> b_max[j]:
- poses_perturb[i,j,3] = b_max[j]
-
- poses_perturb = torch.Tensor(poses_perturb).to(device) # [B, 3, 4]
- tqdm.write("renders RVS...")
- virtue_view = render_virtual_imgs(args, poses_perturb, img_idxs, hwf, device, render_kwargs_test, world_setup_dict)
-
- train_loss = train_on_batch_with_random_view_synthesis(args, targets, rgbs, poses, virtue_view, poses_perturb, feat_model, dset_size, loss_func, optimizer, hwf, img_idxs, render_kwargs_test)
-
- else:
- train_loss = train_on_batch(args, targets, rgbs, poses, feat_model, dset_size, loss_func, optimizer, hwf)
-
- feat_model.eval()
- val_loss_epoch = []
- for data, pose, _ in val_dl:
- inputs = data.to(device)
- labels = pose.to(device)
-
- # pose loss
- _, predict = feat_model(inputs)
- loss = loss_func(predict, labels)
- val_loss_epoch.append(loss.item())
- val_loss = np.mean(val_loss_epoch)
-
- # reduce LR on plateau
- scheduler.step(val_loss)
-
- # logging
- tqdm.write('At epoch {0:6d} : train loss: {1:.4f}, val loss: {2:.4f}'.format(epoch, train_loss, val_loss))
-
- # check wether to early stop
- early_stopping(val_loss, feat_model, epoch=epoch, save_multiple=(not args.no_save_multiple), save_all=args.save_all_ckpt)
- if early_stopping.early_stop:
- print("Early stopping")
- break
-
- if args.featurelossonly:
- global_step += 1
- continue
-
- model_log.set_description_str(f'Best val loss: {early_stopping.val_loss_min:.4f}')
- if epoch % args.i_eval == 0:
- get_error_in_q(args, test_dl, feat_model, len(test_dl.dataset), device, batch_size=1)
- global_step += 1
-
- return
-
-def train():
-
- print(parser.format_values())
-
- # Load data
- if args.dataset_type == '7Scenes':
-
- train_dl, val_dl, test_dl, hwf, i_split, near, far = load_7Scenes_dataloader(args)
- near = near
- far = far
- print('NEAR FAR', near, far)
- train_feature(args, train_dl, val_dl, test_dl, hwf, i_split, near, far)
- return
-
- elif args.dataset_type == 'Cambridge':
-
- train_dl, val_dl, test_dl, hwf, i_split, near, far = load_Cambridge_dataloader(args)
- near = near
- far = far
-
- print('NEAR FAR', near, far)
- train_feature(args, train_dl, val_dl, test_dl, hwf, i_split, near, far)
- return
-
- else:
- print('Unknown dataset type', args.dataset_type, 'exiting')
- return
-
-if __name__ == "__main__":
-
- train()
\ No newline at end of file
diff --git a/script/run_nerf.py b/script/run_nerf.py
deleted file mode 100755
index 17d0950..0000000
--- a/script/run_nerf.py
+++ /dev/null
@@ -1,273 +0,0 @@
-import utils.set_sys_path
-import os, sys
-import numpy as np
-import imageio
-import json
-import random
-import time
-import torch
-import torch.nn as nn
-# from torch.utils.tensorboard import SummaryWriter
-from tqdm import tqdm, trange
-
-from models.ray_utils import *
-from models.nerfw import * # NeRF-w and NeRF-hist
-from models.options import config_parser
-from models.rendering import *
-from dataset_loaders.load_7Scenes import load_7Scenes_dataloader_NeRF
-from dataset_loaders.load_Cambridge import load_Cambridge_dataloader_NeRF
-
-# losses
-from models.losses import loss_dict
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-np.random.seed(0)
-torch.manual_seed(0)
-import random
-random.seed(0)
-
-parser = config_parser()
-args = parser.parse_args()
-
-def train_on_epoch_nerfw(args, train_dl, H, W, focal, N_rand, optimizer, loss_func, global_step, render_kwargs_train):
- for batch_idx, (target, pose, img_idx) in enumerate(train_dl):
- target = target[0].permute(1,2,0).to(device)
- pose = pose.reshape(3,4).to(device) # reshape to 3x4 rot matrix
- img_idx = img_idx.to(device)
-
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- if N_rand is not None:
- rays_o, rays_d = get_rays(H, W, focal, torch.Tensor(pose)) # (H, W, 3), (H, W, 3)
- coords = torch.stack(torch.meshgrid(torch.linspace(0, H-1, H), torch.linspace(0, W-1, W), indexing='ij'), -1) # (H, W, 2)
- coords = torch.reshape(coords, [-1,2]) # (H * W, 2)
- select_inds = np.random.choice(coords.shape[0], size=[N_rand], replace=False) # (N_rand,)
- select_coords = coords[select_inds].long() # (N_rand, 2)
- rays_o = rays_o[select_coords[:, 0], select_coords[:, 1]] # (N_rand, 3)
- rays_d = rays_d[select_coords[:, 0], select_coords[:, 1]] # (N_rand, 3)
- batch_rays = torch.stack([rays_o, rays_d], 0)
- target_s = target[select_coords[:, 0], select_coords[:, 1]] # (N_rand, 3)
-
- # ##### Core optimization loop #####
- rgb, disp, acc, extras = render(H, W, focal, chunk=args.chunk, rays=batch_rays, retraw=True, img_idx=img_idx, **render_kwargs_train)
- optimizer.zero_grad()
-
- # compute loss
- results = {}
- results['rgb_fine'] = rgb
- results['rgb_coarse'] = extras['rgb0']
- results['beta'] = extras['beta']
- results['transient_sigmas'] = extras['transient_sigmas']
-
- loss_d = loss_func(results, target_s)
- loss = sum(l for l in loss_d.values())
-
- with torch.no_grad():
- img_loss = img2mse(rgb, target_s)
- psnr = mse2psnr(img_loss)
- loss.backward()
- optimizer.step()
-
- # NOTE: IMPORTANT!
- ### update learning rate ###
- decay_rate = 0.1
- decay_steps = args.lrate_decay * 1000
- new_lrate = args.lrate * (decay_rate ** (global_step / decay_steps))
- for param_group in optimizer.param_groups:
- param_group['lr'] = new_lrate
- ################################
-
- torch.set_default_tensor_type('torch.FloatTensor')
- return loss, psnr
-
-def train_nerf(args, train_dl, val_dl, hwf, i_split, near, far, render_poses=None, render_img=None):
-
- i_train, i_val, i_test = i_split
- # Cast intrinsics to right types
- H, W, focal = hwf
- H, W = int(H), int(W)
- hwf = [H, W, focal]
-
- # Create log dir and copy the config file
- basedir = args.basedir
- expname = args.expname
- os.makedirs(os.path.join(basedir, expname), exist_ok=True)
- f = os.path.join(basedir, expname, 'args.txt')
- with open(f, 'w') as file:
- for arg in sorted(vars(args)):
- attr = getattr(args, arg)
- file.write('{} = {}\n'.format(arg, attr))
- if args.config is not None:
- f = os.path.join(basedir, expname, 'config.txt')
- with open(f, 'w') as file:
- file.write(open(args.config, 'r').read())
-
- # Create nerf model
- render_kwargs_train, render_kwargs_test, start, grad_vars, optimizer = create_nerf(args)
- global_step = start
-
- bds_dict = {
- 'near' : near,
- 'far' : far,
- }
- render_kwargs_train.update(bds_dict)
- render_kwargs_test.update(bds_dict)
- if args.reduce_embedding==2:
- render_kwargs_train['i_epoch'] = -1
- render_kwargs_test['i_epoch'] = -1
-
- if args.render_test:
- print('TRAIN views are', i_train)
- print('TEST views are', i_test)
- print('VAL views are', i_val)
- if args.reduce_embedding==2:
- render_kwargs_test['i_epoch'] = global_step
- render_test(args, train_dl, val_dl, hwf, start, render_kwargs_test)
- return
-
- # Prepare raybatch tensor if batching random rays
- N_rand = args.N_rand
- # use_batching = not args.no_batching
-
- N_epoch = args.epochs + 1 # epoch
- print('Begin')
- print('TRAIN views are', i_train)
- print('TEST views are', i_test)
- print('VAL views are', i_val)
-
-
- # loss function
- loss_func = loss_dict['nerfw'](coef=1)
-
- for i in trange(start, N_epoch):
- time0 = time.time()
- if args.reduce_embedding==2:
- render_kwargs_train['i_epoch'] = i
- loss, psnr = train_on_epoch_nerfw(args, train_dl, H, W, focal, N_rand, optimizer, loss_func, global_step, render_kwargs_train)
- dt = time.time()-time0
- ##### end #####
-
- # Rest is logging
- if i%args.i_weights==0 and i!=0:
- path = os.path.join(basedir, expname, '{:06d}.tar'.format(i))
- if args.N_importance > 0: # have fine sample network
- torch.save({
- 'global_step': global_step,
- 'network_fn_state_dict': render_kwargs_train['network_fn'].state_dict(),
- 'network_fine_state_dict': render_kwargs_train['network_fine'].state_dict(),
- 'embedding_a_state_dict': render_kwargs_train['embedding_a'].state_dict(),
- 'embedding_t_state_dict': render_kwargs_train['embedding_t'].state_dict(),
- 'optimizer_state_dict': optimizer.state_dict(),
- }, path)
- else:
- torch.save({
- 'global_step': global_step,
- 'network_fn_state_dict': render_kwargs_train['network_fn'].state_dict(),
- 'optimizer_state_dict': optimizer.state_dict(),
- }, path)
- print('Saved checkpoints at', path)
-
- if i%args.i_testset==0 and i > 0: # run thru all validation set
-
- # clean GPU memory before testing, try to avoid OOM
- torch.cuda.empty_cache()
-
- if args.reduce_embedding==2:
- render_kwargs_test['i_epoch'] = i
- trainsavedir = os.path.join(basedir, expname, 'trainset_{:06d}'.format(i))
- os.makedirs(trainsavedir, exist_ok=True)
- images_train = []
- poses_train = []
- index_train = []
- j_skip = 10 # save holdout view render result Trainset/j_skip
- # randomly choose some holdout views from training set
- for batch_idx, (img, pose, img_idx) in enumerate(train_dl):
- if batch_idx % j_skip != 0:
- continue
- img_val = img.permute(0,2,3,1) # (1,H,W,3)
- pose_val = torch.zeros(1,4,4)
- pose_val[0,:3,:4] = pose.reshape(3,4)[:3,:4] # (1,3,4))
- pose_val[0,3,3] = 1.
- images_train.append(img_val)
- poses_train.append(pose_val)
- index_train.append(img_idx)
- images_train = torch.cat(images_train, dim=0).numpy()
- poses_train = torch.cat(poses_train, dim=0).to(device)
- index_train = torch.cat(index_train, dim=0).to(device)
- print('train poses shape', poses_train.shape)
-
- with torch.no_grad():
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- render_path(args, poses_train, hwf, args.chunk, render_kwargs_test, gt_imgs=images_train, savedir=trainsavedir, img_ids=index_train)
- torch.set_default_tensor_type('torch.FloatTensor')
- print('Saved train set')
- del images_train
- del poses_train
-
- testsavedir = os.path.join(basedir, expname, 'testset_{:06d}'.format(i))
- os.makedirs(testsavedir, exist_ok=True)
- images_val = []
- poses_val = []
- index_val = []
- # views from validation set
- for img, pose, img_idx in val_dl:
- img_val = img.permute(0,2,3,1) # (1,H,W,3)
- pose_val = torch.zeros(1,4,4)
- pose_val[0,:3,:4] = pose.reshape(3,4)[:3,:4] # (1,3,4))
- pose_val[0,3,3] = 1.
- images_val.append(img_val)
- poses_val.append(pose_val)
- index_val.append(img_idx)
-
- images_val = torch.cat(images_val, dim=0).numpy()
- poses_val = torch.cat(poses_val, dim=0).to(device)
- index_val = torch.cat(index_val, dim=0).to(device)
- print('test poses shape', poses_val.shape)
-
- with torch.no_grad():
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- render_path(args, poses_val, hwf, args.chunk, render_kwargs_test, gt_imgs=images_val, savedir=testsavedir, img_ids=index_val)
- torch.set_default_tensor_type('torch.FloatTensor')
- print('Saved test set')
-
- # clean GPU memory after testing
- torch.cuda.empty_cache()
- del images_val
- del poses_val
-
- if i%args.i_print==0:
- tqdm.write(f"[TRAIN] Iter: {i} Loss: {loss.item()} PSNR: {psnr.item()}")
-
- global_step += 1
-
-def train():
-
- print(parser.format_values())
-
- # Load data
- if args.dataset_type == '7Scenes':
-
- train_dl, val_dl, hwf, i_split, bds, render_poses, render_img = load_7Scenes_dataloader_NeRF(args)
- near = bds[0]
- far = bds[1]
-
- print('NEAR FAR', near, far)
- train_nerf(args, train_dl, val_dl, hwf, i_split, near, far, render_poses, render_img)
- return
-
- elif args.dataset_type == 'Cambridge':
-
- train_dl, val_dl, hwf, i_split, bds, render_poses, render_img = load_Cambridge_dataloader_NeRF(args)
- near = bds[0]
- far = bds[1]
-
- print('NEAR FAR', near, far)
- train_nerf(args, train_dl, val_dl, hwf, i_split, near, far, render_poses, render_img)
- return
-
- else:
- print('Unknown dataset type', args.dataset_type, 'exiting')
- return
-
-if __name__=='__main__':
-
- train()
\ No newline at end of file
diff --git a/script/train.py b/script/train.py
deleted file mode 100755
index cd8323e..0000000
--- a/script/train.py
+++ /dev/null
@@ -1,169 +0,0 @@
-import utils.set_sys_path
-import os, sys
-import numpy as np
-import imageio
-import random
-import time
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torch.optim as optim
-from tqdm import tqdm, trange
-from torchsummary import summary
-# from torchinfo import summary
-import matplotlib.pyplot as plt
-
-from dm.pose_model import *
-from dm.direct_pose_model import *
-from dm.callbacks import EarlyStopping
-# from dm.prepare_data import prepare_data, load_dataset
-from dm.options import config_parser
-from models.rendering import render_path
-from models.nerfw import to8b
-from dataset_loaders.load_7Scenes import load_7Scenes_dataloader
-from dataset_loaders.load_Cambridge import load_Cambridge_dataloader
-from utils.utils import freeze_bn_layer
-from feature.direct_feature_matching import train_feature_matching
-# import torch.onnx
-
-parser = config_parser()
-args = parser.parse_args()
-device = torch.device('cuda:0') # this is really controlled in train.sh
-
-def render_test(args, train_dl, val_dl, hwf, start, model, device, render_kwargs_test):
- model.eval()
-
- # ### Eval Training set result
- if args.render_video_train:
- images_train = []
- poses_train = []
- # views from train set
- for img, pose in train_dl:
- predict_pose = inference_pose_regression(args, img, device, model)
- device_cpu = torch.device('cpu')
- predict_pose = predict_pose.to(device_cpu) # put predict pose back to cpu
-
- img_val = img.permute(0,2,3,1) # (1,240,320,3)
- pose_val = torch.zeros(1,4,4)
- pose_val[0,:3,:4] = predict_pose.reshape(3,4)[:3,:4] # (1,3,4))
- pose_val[0,3,3] = 1.
- images_train.append(img_val)
- poses_train.append(pose_val)
-
- images_train = torch.cat(images_train, dim=0).numpy()
- poses_train = torch.cat(poses_train, dim=0)
- print('train poses shape', poses_train.shape)
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- with torch.no_grad():
- rgbs, disps = render_path(poses_train.to(device), hwf, args.chunk, render_kwargs_test, gt_imgs=images_train, savedir=None)
- torch.set_default_tensor_type('torch.FloatTensor')
- print('Saving trainset as video', rgbs.shape, disps.shape)
- moviebase = os.path.join(args.basedir, args.model_name, '{}_trainset_{:06d}_'.format(args.model_name, start))
- imageio.mimwrite(moviebase + 'train_rgb.mp4', to8b(rgbs), fps=15, quality=8)
- imageio.mimwrite(moviebase + 'train_disp.mp4', to8b(disps / np.max(disps)), fps=15, quality=8)
-
- ### Eval Validation set result
- if args.render_video_test:
- images_val = []
- poses_val = []
- # views from val set
- for img, pose in val_dl:
- predict_pose = inference_pose_regression(args, img, device, model)
- device_cpu = torch.device('cpu')
- predict_pose = predict_pose.to(device_cpu) # put predict pose back to cpu
-
- img_val = img.permute(0,2,3,1) # (1,240,360,3)
- pose_val = torch.zeros(1,4,4)
- pose_val[0,:3,:4] = predict_pose.reshape(3,4)[:3,:4] # (1,3,4))
- pose_val[0,3,3] = 1.
- images_val.append(img_val)
- poses_val.append(pose_val)
-
- images_val = torch.cat(images_val, dim=0).numpy()
- poses_val = torch.cat(poses_val, dim=0)
- print('test poses shape', poses_val.shape)
- torch.set_default_tensor_type('torch.cuda.FloatTensor')
- with torch.no_grad():
- rgbs, disps = render_path(poses_val.to(device), hwf, args.chunk, render_kwargs_test, gt_imgs=images_val, savedir=None)
- torch.set_default_tensor_type('torch.FloatTensor')
- print('Saving testset as video', rgbs.shape, disps.shape)
- moviebase = os.path.join(args.basedir, args.model_name, '{}_test_{:06d}_'.format(args.model_name, start))
- imageio.mimwrite(moviebase + 'test_rgb.mp4', to8b(rgbs), fps=15, quality=8)
- imageio.mimwrite(moviebase + 'test_disp.mp4', to8b(disps / np.max(disps)), fps=15, quality=8)
- return
-
-def train():
- print(parser.format_values())
- # Load data
- if args.dataset_type == '7Scenes':
- train_dl, val_dl, test_dl, hwf, i_split, near, far = load_7Scenes_dataloader(args)
- elif args.dataset_type == 'Cambridge':
- train_dl, val_dl, test_dl, hwf, i_split, near, far = load_Cambridge_dataloader(args)
- else:
- print("please choose dataset_type: 7Scenes or Cambridge, exiting...")
- sys.exit()
-
- ### pose regression module
- if args.pretrain_model_path == '':
- print('training PoseNet from scratch')
- model = FeatureNet()
- else:
- # load pretrained FeatureNet model
- model = load_exisiting_model(args)
-
- if args.freezeBN:
- model = freeze_bn_layer(model)
- model.to(device)
-
- ### feature extraction module
- # load pretrained FeatureNet model
- if args.pretrain_featurenet_path == '':
- print('load FeatureNet from PoseNet')
- feat_model = load_exisiting_model(args)
- else:
- feat_model = load_exisiting_model(args, isFeatureNet=True)
-
- feat_model.eval()
- feat_model.to(device)
-
- # set optimizer
- optimizer = optim.Adam(model.parameters(), lr=args.learning_rate) #weight_decay=weight_decay, **kwargs
-
- # set callbacks parameters
- early_stopping = EarlyStopping(args, patience=args.patience[0], verbose=False)
-
- # start training
- if args.dataset_type == '7Scenes':
- train_feature_matching(args, model, feat_model, optimizer, i_split, hwf, near, far, device, early_stopping, train_dl=train_dl, val_dl=val_dl, test_dl=test_dl)
- elif args.dataset_type == 'Cambridge':
- train_feature_matching(args, model, feat_model, optimizer, i_split, hwf, near, far, device, early_stopping, train_dl=train_dl, val_dl=val_dl, test_dl=test_dl)
-
-def eval():
- print(parser.format_values())
- # Load data
- if args.dataset_type == '7Scenes':
- train_dl, val_dl, test_dl, hwf, i_split, near, far = load_7Scenes_dataloader(args)
- elif args.dataset_type == 'Cambridge':
- train_dl, val_dl, test_dl, hwf, i_split, near, far = load_Cambridge_dataloader(args)
- else:
- print("please choose dataset_type: 7Scenes or Cambridge, exiting...")
- sys.exit()
-
- # load pretrained PoseNet model
- model = load_exisiting_model(args)
- if args.freezeBN:
- model = freeze_bn_layer(model)
- model.to(device)
-
- print(len(val_dl.dataset))
-
- get_error_in_q(args, val_dl, model, len(val_dl.dataset), device, batch_size=1)
-
-if __name__ == '__main__':
- if args.eval:
- torch.manual_seed(0)
- random.seed(0)
- np.random.seed(0)
- eval()
- else:
- train()
diff --git a/script/utils/set_sys_path.py b/script/utils/set_sys_path.py
deleted file mode 100755
index 601085b..0000000
--- a/script/utils/set_sys_path.py
+++ /dev/null
@@ -1,2 +0,0 @@
-import sys
-sys.path.append('../')
\ No newline at end of file
diff --git a/script/utils/utils.py b/script/utils/utils.py
deleted file mode 100755
index f51d76d..0000000
--- a/script/utils/utils.py
+++ /dev/null
@@ -1,394 +0,0 @@
-'''
-helper functions to train robust feature extractors
-'''
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import datetime
-from torchvision.utils import make_grid
-import matplotlib.pyplot as plt
-import pdb
-from PIL import Image
-from torchvision.utils import save_image
-from math import pi
-import cv2
-# from pykalman import KalmanFilter
-
-def freeze_bn_layer(model):
- ''' freeze bn layer by not require grad but still behave differently when model.train() vs. model.eval() '''
- print("Freezing BatchNorm Layers...")
- for module in model.modules():
- if isinstance(module, nn.BatchNorm2d):
- # print("this is a BN layer:", module)
- if hasattr(module, 'weight'):
- module.weight.requires_grad_(False)
- if hasattr(module, 'bias'):
- module.bias.requires_grad_(False)
- return model
-
-def freeze_bn_layer_train(model):
- ''' set batchnorm to eval()
- it is useful to align train and testing result
- '''
- # model.train()
- # print("Freezing BatchNorm Layers...")
- for module in model.modules():
- if isinstance(module, nn.BatchNorm2d):
- module.eval()
- return model
-
-def save_image_saliancy(tensor, path, normalize: bool = False, scale_each: bool = False,):
- """
- Modification based on TORCHVISION.UTILS
- ::param: tensor (batch, channel, H, W)
- """
- # grid = make_grid(tensor.detach(), normalize=normalize, scale_each=scale_each, nrow=32)
- grid = make_grid(tensor.detach(), normalize=normalize, scale_each=scale_each, nrow=6)
- # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
- ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
- fig = plt.figure()
- plt.imshow(ndarr[:,:,0], cmap='jet') # viridis, plasma
- plt.axis('off')
- fig.savefig(path, bbox_inches='tight',dpi=fig.dpi,pad_inches=0.0)
- plt.close()
-
-def save_image_saliancy_single(tensor, path, normalize: bool = False, scale_each: bool = False,):
- """
- Modification based on TORCHVISION.UTILS, save single feature map
- ::param: tensor (batch, channel, H, W)
- """
- # grid = make_grid(tensor.detach(), normalize=normalize, scale_each=scale_each, nrow=32)
- grid = make_grid(tensor.detach(), normalize=normalize, scale_each=scale_each, nrow=1)
- # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
- ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
- fig = plt.figure()
- # plt.imshow(ndarr[:,:,0], cmap='plasma') # viridis, jet
- plt.imshow(ndarr[:,:,0], cmap='jet') # viridis, jet
- plt.axis('off')
- fig.savefig(path, bbox_inches='tight',dpi=fig.dpi,pad_inches=0.0)
- plt.close()
-
-def print_feature_examples(features, path):
- """
- print feature maps
- ::param: features
- """
- kwargs = {'normalize' : True, } # 'scale_each' : True
-
- for i in range(len(features)):
- fn = path + '{}.png'.format(i)
- # save_image(features[i].permute(1,0,2,3), fn, **kwargs)
- save_image_saliancy(features[i].permute(1,0,2,3), fn, normalize=True)
- # pdb.set_trace()
- ###
-
-def plot_features(features, path='f', isList=True):
- """
- print feature maps
- :param features: (3, [batch, H, W]) or [3, batch, H, W]
- :param path: save image path
- :param isList: wether the features is an list
- :return:
- """
- kwargs = {'normalize' : True, } # 'scale_each' : True
-
- if isList:
- dim = features[0].dim()
- else:
- dim = features.dim()
- assert(dim==3 or dim==4)
-
- if dim==4 and isList:
- print_feature_examples(features, path)
- elif dim==4 and (isList==False):
- fn = path
- lvl, b, H, W = features.shape
- for i in range(features.shape[0]):
- fn = path + '{}.png'.format(i)
- save_image_saliancy(features[i][None,...].permute(1,0,2,3).cpu(), fn, normalize=True)
-
- # # concat everything
- # features = features.reshape([-1, H, W])
- # # save_image(features[None,...].permute(1,0,2,3).cpu(), fn, **kwargs)
- # save_image_saliancy(features[None,...].permute(1,0,2,3).cpu(), fn, normalize=True)
-
- elif dim==3 and isList: # print all images in the list
- for i in range(len(features)):
- fn = path + '{}.png'.format(i)
- # save_image(features[i][None,...].permute(1,0,2,3).cpu(), fn, **kwargs)
- save_image_saliancy(features[i][None,...].permute(1,0,2,3).cpu(), fn, normalize=True)
- elif dim==3 and (isList==False):
- fn = path
- save_image_saliancy(features[None,...].permute(1,0,2,3).cpu(), fn, normalize=True)
-
-def sample_homography_np(
- shape, shift=0, perspective=True, scaling=True, rotation=True, translation=True,
- n_scales=5, n_angles=25, scaling_amplitude=0.1, perspective_amplitude_x=0.1,
- perspective_amplitude_y=0.1, patch_ratio=0.5, max_angle=pi/2,
- allow_artifacts=False, translation_overflow=0.):
- """Sample a random valid homography.
-
- Computes the homography transformation between a random patch in the original image
- and a warped projection with the same image size.
- As in `tf.contrib.image.transform`, it maps the output point (warped patch) to a
- transformed input point (original patch).
- The original patch, which is initialized with a simple half-size centered crop, is
- iteratively projected, scaled, rotated and translated.
-
- Arguments:
- shape: A rank-2 `Tensor` specifying the height and width of the original image.
- perspective: A boolean that enables the perspective and affine transformations.
- scaling: A boolean that enables the random scaling of the patch.
- rotation: A boolean that enables the random rotation of the patch.
- translation: A boolean that enables the random translation of the patch.
- n_scales: The number of tentative scales that are sampled when scaling.
- n_angles: The number of tentatives angles that are sampled when rotating.
- scaling_amplitude: Controls the amount of scale.
- perspective_amplitude_x: Controls the perspective effect in x direction.
- perspective_amplitude_y: Controls the perspective effect in y direction.
- patch_ratio: Controls the size of the patches used to create the homography. (like crop size)
- max_angle: Maximum angle used in rotations.
- allow_artifacts: A boolean that enables artifacts when applying the homography.
- translation_overflow: Amount of border artifacts caused by translation.
-
- Returns:
- A `Tensor` of shape `[1, 8]` corresponding to the flattened homography transform.
- """
-
- # print("debugging")
-
-
- # Corners of the output image
- pts1 = np.stack([[0., 0.], [0., 1.], [1., 1.], [1., 0.]], axis=0)
- # Corners of the input patch
- margin = (1 - patch_ratio) / 2
- pts2 = margin + np.array([[0, 0], [0, patch_ratio],
- [patch_ratio, patch_ratio], [patch_ratio, 0]])
-
- from numpy.random import normal
- from numpy.random import uniform
- from scipy.stats import truncnorm
-
- # Random perspective and affine perturbations
- # lower, upper = 0, 2
- std_trunc = 2
- # pdb.set_trace()
- if perspective:
- if not allow_artifacts:
- perspective_amplitude_x = min(perspective_amplitude_x, margin)
- perspective_amplitude_y = min(perspective_amplitude_y, margin)
- perspective_displacement = truncnorm(-1*std_trunc, std_trunc, loc=0, scale=perspective_amplitude_y/2).rvs(1)
- h_displacement_left = truncnorm(-1*std_trunc, std_trunc, loc=0, scale=perspective_amplitude_x/2).rvs(1)
- h_displacement_right = truncnorm(-1*std_trunc, std_trunc, loc=0, scale=perspective_amplitude_x/2).rvs(1)
- pts2 += np.array([[h_displacement_left, perspective_displacement],
- [h_displacement_left, -perspective_displacement],
- [h_displacement_right, perspective_displacement],
- [h_displacement_right, -perspective_displacement]]).squeeze()
-
- # Random scaling
- # sample several scales, check collision with borders, randomly pick a valid one
- if scaling:
- scales = truncnorm(-1*std_trunc, std_trunc, loc=1, scale=scaling_amplitude/2).rvs(n_scales)
- scales = np.concatenate((np.array([1]), scales), axis=0)
-
- center = np.mean(pts2, axis=0, keepdims=True)
- scaled = (pts2 - center)[np.newaxis, :, :] * scales[:, np.newaxis, np.newaxis] + center
- if allow_artifacts:
- valid = np.arange(n_scales) # all scales are valid except scale=1
- else:
- valid = (scaled >= 0.) * (scaled < 1.)
- valid = valid.prod(axis=1).prod(axis=1)
- valid = np.where(valid)[0]
- idx = valid[np.random.randint(valid.shape[0], size=1)].squeeze().astype(int)
- pts2 = scaled[idx,:,:]
-
- # Random translation
- if translation:
- # pdb.set_trace()
- t_min, t_max = np.min(pts2, axis=0), np.min(1 - pts2, axis=0)
- if allow_artifacts:
- t_min += translation_overflow
- t_max += translation_overflow
- pts2 += np.array([uniform(-t_min[0], t_max[0],1), uniform(-t_min[1], t_max[1], 1)]).T
-
- # Random rotation
- # sample several rotations, check collision with borders, randomly pick a valid one
- if rotation:
- angles = np.linspace(-max_angle, max_angle, num=n_angles)
- angles = np.concatenate((angles, np.array([0.])), axis=0) # in case no rotation is valid
- center = np.mean(pts2, axis=0, keepdims=True)
- rot_mat = np.reshape(np.stack([np.cos(angles), -np.sin(angles), np.sin(angles),
- np.cos(angles)], axis=1), [-1, 2, 2])
- rotated = np.matmul( (pts2 - center)[np.newaxis,:,:], rot_mat) + center
- if allow_artifacts:
- valid = np.arange(n_angles) # all scales are valid except scale=1
- else: # find multiple valid option and choose the valid one
- valid = (rotated >= 0.) * (rotated < 1.)
- valid = valid.prod(axis=1).prod(axis=1)
- valid = np.where(valid)[0]
- idx = valid[np.random.randint(valid.shape[0], size=1)].squeeze().astype(int)
- pts2 = rotated[idx,:,:]
-
- # Rescale to actual size
- shape = shape[::-1] # different convention [y, x]
- pts1 *= shape[np.newaxis,:]
- pts2 *= shape[np.newaxis,:]
-
- homography = cv2.getPerspectiveTransform(np.float32(pts1+shift), np.float32(pts2+shift))
- return homography
-
-def warp_points(points, homographies, device='cpu'):
- """
- Warp a list of points with the given homography.
-
- Arguments:
- points: list of N points, shape (N, 2(x, y))).
- homography: batched or not (shapes (B, 3, 3) and (...) respectively).
-
- Returns: a Tensor of shape (N, 2) or (B, N, 2(x, y)) (depending on whether the homography
- is batched) containing the new coordinates of the warped points.
-
- """
- # expand points len to (x, y, 1)
- no_batches = len(homographies.shape) == 2
- homographies = homographies.unsqueeze(0) if no_batches else homographies
-
- batch_size = homographies.shape[0]
- points = torch.cat((points.float(), torch.ones((points.shape[0], 1)).to(device)), dim=1)
- points = points.to(device)
- homographies = homographies.view(batch_size*3,3)
-
- warped_points = homographies@points.transpose(0,1)
-
- # normalize the points
- warped_points = warped_points.view([batch_size, 3, -1])
- warped_points = warped_points.transpose(2, 1)
- warped_points = warped_points[:, :, :2] / warped_points[:, :, 2:]
- return warped_points[0,:,:] if no_batches else warped_points
-
-def inv_warp_image_batch(img, mat_homo_inv, device='cpu', mode='bilinear'):
- '''
- Inverse warp images in batch
-
- :param img:
- batch of images
- tensor [batch_size, 1, H, W]
- :param mat_homo_inv:
- batch of homography matrices
- tensor [batch_size, 3, 3]
- :param device:
- GPU device or CPU
- :return:
- batch of warped images
- tensor [batch_size, 1, H, W]
- '''
- # compute inverse warped points
- if len(img.shape) == 2 or len(img.shape) == 3:
- img = img.view(1,1,img.shape[0], img.shape[1])
- if len(mat_homo_inv.shape) == 2:
- mat_homo_inv = mat_homo_inv.view(1,3,3)
-
- Batch, channel, H, W = img.shape
- coor_cells = torch.stack(torch.meshgrid(torch.linspace(-1, 1, W), torch.linspace(-1, 1, H), indexing='ij'), dim=2)
- coor_cells = coor_cells.transpose(0, 1)
- coor_cells = coor_cells.to(device)
- coor_cells = coor_cells.contiguous()
-
- src_pixel_coords = warp_points(coor_cells.view([-1, 2]), mat_homo_inv, device)
- src_pixel_coords = src_pixel_coords.view([Batch, H, W, 2])
- src_pixel_coords = src_pixel_coords.float()
-
- warped_img = F.grid_sample(img, src_pixel_coords, mode=mode, align_corners=True)
- return warped_img
-
-def compute_valid_mask(image_shape, inv_homography, device='cpu', erosion_radius=0):
- """
- Compute a boolean mask of the valid pixels resulting from an homography applied to
- an image of a given shape. Pixels that are False correspond to bordering artifacts.
- A margin can be discarded using erosion.
-
- Arguments:
- input_shape: Tensor of rank 2 representing the image shape, i.e. `[H, W]`.
- homography: Tensor of shape (B, 8) or (8,), where B is the batch size.
- `erosion_radius: radius of the margin to be discarded.
-
- Returns: a Tensor of type `tf.int32` and shape (H, W).
- """
-
- if inv_homography.dim() == 2:
- inv_homography = inv_homography.view(-1, 3, 3)
- batch_size = inv_homography.shape[0]
- mask = torch.ones(batch_size, 1, image_shape[0], image_shape[1]).to(device)
- mask = inv_warp_image_batch(mask, inv_homography, device=device, mode='nearest')
- mask = mask.view(batch_size, image_shape[0], image_shape[1])
- mask = mask.cpu().numpy()
- if erosion_radius > 0:
- kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (erosion_radius*2,)*2)
- for i in range(batch_size):
- mask[i, :, :] = cv2.erode(mask[i, :, :], kernel, iterations=1)
-
- return torch.tensor(mask).to(device)
-
-def Kalman1D(observations,damping=1):
- # To return the smoothed time series data
- observation_covariance = damping
- initial_value_guess = observations[0]
- transition_matrix = 1
- transition_covariance = 0.1
- initial_value_guess
- kf = KalmanFilter(
- initial_state_mean=initial_value_guess,
- initial_state_covariance=observation_covariance,
- observation_covariance=observation_covariance,
- transition_covariance=transition_covariance,
- transition_matrices=transition_matrix
- )
- pred_state, state_cov = kf.smooth(observations)
- return pred_state
-
-def Kalman3D(observations,damping=1):
- '''
- In:
- observation: Nx3
- Out:
- pred_state: Nx3
- '''
- # To return the smoothed time series data
- observation_covariance = damping
- transition_matrix = 1
- transition_covariance = 0.1
- initial_value_guess_x = observations[0,0]
- initial_value_guess_y = observations[0,1] # ?
- initial_value_guess_z = observations[0,2] # ?
-
- # perform 1D smooth for each axis
- kfx = KalmanFilter(
- initial_state_mean=initial_value_guess_x,
- initial_state_covariance=observation_covariance,
- observation_covariance=observation_covariance,
- transition_covariance=transition_covariance,
- transition_matrices=transition_matrix
- )
- pred_state_x, state_cov_x = kfx.smooth(observations[:, 0])
-
- kfy = KalmanFilter(
- initial_state_mean=initial_value_guess_y,
- initial_state_covariance=observation_covariance,
- observation_covariance=observation_covariance,
- transition_covariance=transition_covariance,
- transition_matrices=transition_matrix
- )
- pred_state_y, state_cov_y = kfy.smooth(observations[:, 1])
-
- kfz = KalmanFilter(
- initial_state_mean=initial_value_guess_z,
- initial_state_covariance=observation_covariance,
- observation_covariance=observation_covariance,
- transition_covariance=transition_covariance,
- transition_matrices=transition_matrix
- )
- pred_state_z, state_cov_z = kfy.smooth(observations[:, 2])
-
- pred_state = np.concatenate((pred_state_x, pred_state_y, pred_state_z), axis=1)
- return pred_state
\ No newline at end of file
diff --git a/style.css b/style.css
new file mode 100755
index 0000000..b4ae3f8
--- /dev/null
+++ b/style.css
@@ -0,0 +1,45 @@
+h1{
+ font-size: 35px;
+}
+
+h2{
+ font-size: 28px;
+}
+
+pre{
+ margin-bottom: 0;
+}
+
+.title {
+ padding-top: 3rem;
+ text-align: center;
+}
+
+.bmvc_label{
+ text-align: center;
+ font-size: 20px;
+}
+
+.authors{
+ font-size: 20px;
+ text-align: center;
+}
+
+.vid{
+ display: flex;
+ flex-direction: column;
+ align-items: center;
+ position: relative;
+ margin: auto;
+}
+
+/*.sizer {
+ width: 70%;
+ position: relative;
+ margin: auto;
+
+}
+
+.embed-responsive {
+ height:50%;
+}*/
\ No newline at end of file
diff --git a/supp_video.mp4 b/supp_video.mp4
new file mode 100755
index 0000000..e21851a
Binary files /dev/null and b/supp_video.mp4 differ