Official implementation of "S-VolSDF: Sparse Multi-View Stereo Regularization of Neural Implicit Surfaces" (ICCV 2023)
from torch.utils.data import Dataset
import numpy as np
import os, cv2, time
from PIL import Image
from datasets.data_io import *
import copy
from helpers.utils import load_K_Rt_from_P, glob_imgs, read_img
from volsdf.datasets.scene_dataset import scan2hash
from helpers.help import logger

s_h, s_w = 0, 0
class MVSDataset(Dataset):
    def __init__(self, datapath, listfile, mode, nviews, data_dir, ndepths=192, interval_scale=1.06, **kwargs):
        super(MVSDataset, self).__init__()
        self.datapath = datapath
        self.listfile = listfile
        self.mode = mode
        self.nviews = nviews
        self.nviews_max = 5
        self.ndepths = ndepths
        self.interval_scale = interval_scale
        self.data_dir = data_dir
        self._max_h, self._max_w = kwargs["max_h"], kwargs["max_w"]
        self.hparams = kwargs.get("args", dict())
        self.trains_i = kwargs.get("trains_i", None)
        self.fix_wh = False

        # check args
        logger.debug(f"max views used in MVS={self.nviews_max}")
        assert len(listfile) == 1 # one scan at a time
        assert self.mode == "test"
        if self.data_dir != 'DTU': assert interval_scale == 1
        assert self.data_dir in ['BlendedMVS', 'DTU']
        assert self.trains_i is not None
        
        self.meta_from_idr(listfile[0], self.data_dir)
        self.metas = self.build_list()
        assert self.trains_i == [x[1] for x in self.metas]

    def meta_from_idr(self, scan, data_dir):
        """
        read metadata in IDR format, see https://github.com/lioryariv/idr
        camera matrix, scale matrix, image paths
        """
        scan_id = scan[4:]
        
        instance_dir = os.path.join(self.hparams.data_dir_root, data_dir, 'scan{0}'.format(scan_id))
        image_dir = '{0}/image'.format(instance_dir)
        cam_file = '{0}/cameras.npz'.format(instance_dir)
        if not os.path.exists(cam_file) and int(scan_id) < 200:
            cam_file = os.path.join(self.hparams.data_dir_root, data_dir, 'scan114', 'cameras.npz')
        assert os.path.exists(image_dir), f"{image_dir} is empty"
        assert os.path.exists(cam_file), f"{cam_file} is empty"

        self.image_paths_idr = sorted(glob_imgs(image_dir))
        n_images = len(self.image_paths_idr)

        camera_dict = np.load(cam_file)
        scale_mats = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(n_images)]
        world_mats = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(n_images)]

        self.intrinsics_idr = []
        self.pose_idr = []

        if scan == 'scan5': # scan5only
            # scale_mat for scan5 is wrong, set it to 1 instead
            for scale_mat, world_mat in zip(scale_mats, world_mats):
                P = world_mat @ scale_mat
                intrinsics, pose = load_K_Rt_from_P(None, P[:3, :4])
                self.intrinsics_idr.append(intrinsics)
                self.pose_idr.append(pose)
            self.scale_mat = None
            self.scale_factor = 1.0
            self._scale_mvs = scale_mats[0][0,0]
        else:
            for scale_mat, world_mat in zip(scale_mats, world_mats):
                intrinsics, pose = load_K_Rt_from_P(None, world_mat[:3, :4])
                self.intrinsics_idr.append(intrinsics)
                self.pose_idr.append(pose)
            self.scale_mat = scale_mats[0] # first image of the scene
            self.scale_factor = scale_mats[0][0,0]

    def build_list(self):
        """
        get  (reference view, a list of source views with order from pair.txt)
        """
        metas = []
        scans = self.listfile

        interval_scale_dict = {}
        # scans
        for scan in scans:
            # determine the interval scale of each scene. default is 1.06
            if isinstance(self.interval_scale, float):
                interval_scale_dict[scan] = self.interval_scale
            else:
                interval_scale_dict[scan] = self.interval_scale[scan]

            if self.data_dir == 'DTU':
                pair_file = "{}/pair.txt".format(scan)
                if not os.path.exists(os.path.join(self.datapath, pair_file)): pair_file = "scan1/pair.txt"
            elif self.data_dir == 'BlendedMVS':
                pair_file = "{}/cams/pair.txt".format(scan2hash(scan))
            else:
                raise NotImplementedError
            if self.data_dir in ['DTU', 'BlendedMVS', None]:
                assert os.path.exists(os.path.join(self.datapath, pair_file))     
                # read the pair file
                with open(os.path.join(self.datapath, pair_file)) as f:
                    num_viewpoint = int(f.readline())
                    # viewpoints
                    for view_idx in range(num_viewpoint):
                        ref_view = int(f.readline().rstrip())
                        src_views = [int(x) for x in f.readline().rstrip().split()[1::2]]
                        # filter by no src view and fill to nviews
                        if len(src_views) > 0:
                            if ref_view not in self.trains_i:
                                continue
                            metas.append((scan, ref_view, src_views, scan))

        metas = [metas[[x[1] for x in metas].index(y)] for y in self.trains_i]
        logger.debug(f"metadata from pair.txt: {metas}")
        self.interval_scale = interval_scale_dict
        logger.debug(f"len_metas: {len(metas)}, interval_scale: {self.interval_scale}")
        return metas

    def __len__(self):
        return len(self.metas)

    def read_cam_file(self, filename, interval_scale):
        with open(filename) as f:
            lines = f.readlines()
            lines = [line.rstrip() for line in lines]
        # extrinsics: line [1,5), 4x4 matrix
        extrinsics = np.fromstring(' '.join(lines[1:5]), dtype=np.float32, sep=' ').reshape((4, 4))
        # intrinsics: line [7-10), 3x3 matrix
        intrinsics = np.fromstring(' '.join(lines[7:10]), dtype=np.float32, sep=' ').reshape((3, 3))
        intrinsics[:2, :] /= 4.0
        # depth_min & depth_interval: line 11
        depth_min = float(lines[11].split()[0])
        depth_interval = float(lines[11].split()[1])

        if self.data_dir == 'BlendedMVS':
            depth_max = float(lines[11].split()[-1])
            depth_interval = float(depth_max - depth_min) / self.ndepths
            return intrinsics, extrinsics, depth_min, depth_interval
        elif len(lines[11].split()) >= 3:
            # num_depth != 192 (default value)
            num_depth = lines[11].split()[2]
            depth_max = depth_min + int(float(num_depth)) * depth_interval
            depth_interval = (depth_max - depth_min) / self.ndepths

        depth_interval *= interval_scale

        return intrinsics, extrinsics, depth_min, depth_interval

    def scale_mvs_input(self, img, intrinsics, max_w, max_h, base=32):
        intrinsics_resize = copy.deepcopy(intrinsics)

        h, w = img.shape[:2]
        if h != max_h or w != max_w:
            scale = 1.0 * max_h / h
            if scale * w > max_w:
                scale = 1.0 * max_w / w
            new_w, new_h = scale * w // base * base, scale * h // base * base
        else:
            new_w, new_h = 1.0 * w // base * base, 1.0 * h // base * base

        scale_w = 1.0 * new_w / w
        scale_h = 1.0 * new_h / h
        intrinsics_resize[0, :] *= scale_w
        intrinsics_resize[1, :] *= scale_h

        img_resize = cv2.resize(img, (int(new_w), int(new_h)), interpolation=cv2.INTER_CUBIC)

        return img_resize, intrinsics_resize

    def __getitem__(self, idx):
        global s_h, s_w
        meta = self.metas[idx]
        scan, ref_view, src_views, scene_name = meta
        if self.trains_i is not None:
            _srcs = [x for x in src_views if x in self.trains_i]
            view_ids = [ref_view] + _srcs
            view_ids += [x for x in self.trains_i if x not in view_ids]
            assert ref_view not in _srcs and set(view_ids)==set(self.trains_i)
        else:
            # use only the reference view and first nviews-1 source views
            view_ids = [ref_view] + src_views[:self.nviews - 1]

        if len(view_ids) > self.nviews_max: view_ids = view_ids[:self.nviews_max]

        imgs = []
        depth_values = None
        proj_matrices = []

        for i, vid in enumerate(view_ids):
            if self.data_dir == 'BlendedMVS':
                proj_mat_filename = os.path.join(self.datapath, '{}/cams/{:0>8}_cam.txt'.format(scan2hash(scan), vid))
                _intrinsics, _extrinsics, depth_min, depth_interval = self.read_cam_file(proj_mat_filename, interval_scale=
                                                                                    self.interval_scale[scene_name])
                # adjust depth range for scan 4, 5
                if scan == 'scan5': # scan5only
                    # scale_mat for scan5 is wrong, set it to 1 instead
                    depth_min, depth_interval = depth_min/self._scale_mvs, depth_interval/self._scale_mvs
                if scan in ['scan4', 'scan5']:
                    depth_max = depth_min + self.ndepths * depth_interval
                    depth_max = min(depth_max, depth_min * 2.197)
                    depth_interval = float(depth_max - depth_min) / self.ndepths
            elif self.data_dir == 'DTU':
                depth_min = 425
                depth_interval = 2.5 * self.interval_scale[scene_name]
            else:
                raise NotImplementedError

            intrinsics = copy.deepcopy(self.intrinsics_idr[vid][:3, :3])
            intrinsics[:2, :] /= 4.0
            extrinsics = np.linalg.inv(self.pose_idr[vid])

            img_filename = os.path.join(self.image_paths_idr[vid])
            img = read_img(img_filename)

            # interpolate the images -> bigger
            # Feat_extract -> (288, 384) (576, 768) (1152, 1536)
            if self.hparams.get("x2_mvsres", False):
                _s_hw = 1536 / self._max_w
                assert self._max_w * _s_hw == 1536 and self._max_h * _s_hw == 1152 
                img, intrinsics = self.scale_mvs_input(img, intrinsics, self._max_w, self._max_h, base=1)
                img, intrinsics = self.scale_mvs_input(img, intrinsics, 1536, 1152)
            else:
                # scale input
                img, intrinsics = self.scale_mvs_input(img, intrinsics, self._max_w, self._max_h)

            # resize to standard height or width
            c_h, c_w = img.shape[:2]
            # assume all images have same size
            if i != 0: assert (c_h == s_h) and (c_w == s_w) 

            imgs.append(img)

            # extrinsics, intrinsics
            proj_mat = np.zeros(shape=(2, 4, 4), dtype=np.float32)
            proj_mat[0, :4, :4] = extrinsics
            proj_mat[1, :3, :3] = intrinsics
            proj_matrices.append(proj_mat)

            if i == 0:  # reference view
                depth_values = np.arange(depth_min, depth_interval * (self.ndepths - 0.5) + depth_min, depth_interval,
                                         dtype=np.float32)
                cam_near_far = np.array([depth_min, depth_interval, self.ndepths, depth_interval * self.ndepths + depth_min])
                s_h, s_w = img.shape[:2]

        # all
        imgs = np.stack(imgs).transpose([0, 3, 1, 2])

        proj_matrices = np.stack(proj_matrices)
        stage2_pjmats = proj_matrices.copy()
        stage2_pjmats[:, 1, :2, :] = proj_matrices[:, 1, :2, :] * 2
        stage3_pjmats = proj_matrices.copy()
        stage3_pjmats[:, 1, :2, :] = proj_matrices[:, 1, :2, :] * 4
        proj_matrices_ms = {
            "stage1": proj_matrices,
            "stage2": stage2_pjmats,
            "stage3": stage3_pjmats
        }

        return {"imgs": imgs[:,:3]*imgs[:,3:] if imgs.shape[1]==4 else imgs,
                "masks": imgs[:,3:] if imgs.shape[1]==4 else np.ones_like(imgs[:,:1]),
                "proj_matrices": proj_matrices_ms,
                "depth_values": depth_values,
                "cam_near_far": cam_near_far,
                "filename": scan + '/{}/' + '{:0>8}'.format(view_ids[0]) + "{}",
                }

datasets/general_eval.py
cvlab-stonybrook/s-volsdf
import argparse
import GPUtil
import os
import gc
from pyhocon import ConfigFactory
import torch
import numpy as np
from PIL import Image
from tqdm import tqdm
from skimage.metrics import structural_similarity

from lpips_tensorflow import lpips_tf
import tensorflow as tf
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

from datasets.data_io import read_pfm, save_pfm
import volsdf.utils.general as utils
import volsdf.utils.plots as plt
from volsdf.datasets.scene_dataset import get_trains_ids, get_eval_ids
from helpers.help import logger

torch.backends.cudnn.benchmark = True
torch.set_default_dtype(torch.float32)
torch.set_num_threads(1)

def evaluate(**kwargs):
    conf = ConfigFactory.parse_file(kwargs['conf'])
    exps_folder_name = kwargs['exps_folder_name']
    evals_folder_name = kwargs['evals_folder_name']

    root = './'
    expname = kwargs['expname']
    scan_id = kwargs['scan_id'] if kwargs['scan_id'] != -1 else conf.get_int('dataset.scan_id', default=-1)
    if scan_id != -1:
        expname = expname + '_{0}'.format(scan_id)
    else:
        scan_id = conf.get_string('dataset.object', default='')

    if kwargs['ckpt_dir'] == '' and kwargs['timestamp'] == 'latest':
        if os.path.exists(os.path.join(root, kwargs['exps_folder_name'], expname)):
            timestamps = os.listdir(os.path.join(root, kwargs['exps_folder_name'], expname))
            if (len(timestamps)) == 0:
                print('WRONG EXP FOLDER')
                exit()
            # self.timestamp = sorted(timestamps)[-1]
            timestamp = None
            for t in sorted(timestamps):
                if os.path.exists(os.path.join(root, kwargs['exps_folder_name'], expname, t, 'checkpoints',
                                               'ModelParameters', str(kwargs['checkpoint']) + ".pth")):
                    timestamp = t
            if timestamp is None:
                print('NO GOOD TIMSTAMP')
                exit()
        else:
            print('WRONG EXP FOLDER')
            exit()
    else:
        timestamp = kwargs['timestamp']

    utils.mkdir_ifnotexists(os.path.join(root, evals_folder_name))
    expdir = os.path.join(root, exps_folder_name, expname)
    evaldir = os.path.join(root, evals_folder_name, expname)
    utils.mkdir_ifnotexists(evaldir)

    dataset_conf = conf.get_config('dataset')
    if kwargs['scan_id'] != -1:
        dataset_conf['scan_id'] = kwargs['scan_id']
    dataset_conf['num_views'] = -1 # all images
    dataset_conf['data_dir_root'] = opt.data_dir_root
    eval_dataset = utils.get_class(conf.get_string('train.dataset_class'))(**dataset_conf)

    conf_model = conf.get_config('model')
    model = utils.get_class(conf.get_string('train.model_class'))(conf=conf_model)

    # settings for camera optimization
    scale_mat = eval_dataset.get_scale_mat()

    if opt.eval_rendering:
        eval_dataloader = torch.utils.data.DataLoader(eval_dataset,
                                                      batch_size=1,
                                                      num_workers=0,
                                                      shuffle=False,
                                                      collate_fn=eval_dataset.collate_fn
                                                      )
        total_pixels = eval_dataset.total_pixels
        img_res = eval_dataset.img_res
        logger.info(f'img res = {img_res}')

    if kwargs['ckpt_dir'] != '':
        old_checkpnts_dir = os.path.join(root, kwargs['ckpt_dir'], 'checkpoints')
    else:
        old_checkpnts_dir = os.path.join(expdir, timestamp, 'checkpoints')

    logger.info(f'load model from: {old_checkpnts_dir}')
    if opt.result_from != 'None':
        epoch = 0
        # use the latest epoch's rendering results
        for renderdir in os.listdir(evaldir):
            if renderdir.startswith('rendering_'):
                epoch = max(epoch, int(renderdir.replace('rendering_', '')))
    else:
        saved_model_state = torch.load(os.path.join(old_checkpnts_dir, 'ModelParameters', str(kwargs['checkpoint']) + ".pth"))
        model.load_state_dict(saved_model_state["model_state_dict"])
        epoch = saved_model_state['epoch']

    ####################################################################################################################
    model.cuda()
    model.eval()

    if opt.eval_mesh:
        with torch.no_grad():
            if 'bmvs' in opt.conf:
                mesh = plt.get_surface_high_res_mesh(
                    sdf=lambda x: model.implicit_network(x)[:, 0],
                    resolution=kwargs['resolution'],
                    grid_boundary=conf.get_list('plot.grid_boundary'),
                    level=conf.get_int('plot.level', default=0),
                    take_components = type(scan_id) is not str
                )
            elif 'dtu' in opt.conf:
                bb_dict = np.load(os.path.join(opt.data_dir_root, 'DTU/bbs.npz'))
                if scan_id == 82:
                    grid_params = bb_dict[str(83)]
                elif scan_id in [21, 34, 38]:
                    grid_params = bb_dict[str(24)]
                else:
                    grid_params = bb_dict[str(scan_id)]

                mesh = plt.get_surface_by_grid(
                    grid_params=grid_params,
                    sdf=lambda x: model.implicit_network(x)[:, 0],
                    resolution=kwargs['resolution'],
                    level=conf.get_int('plot.level', default=0),
                    higher_res=True
                )
            else:
                raise NotImplementedError

            # Transform to world coordinates
            mesh.apply_transform(scale_mat)

            # Taking the biggest connected component
            components = mesh.split(only_watertight=False)
            areas = np.array([c.area for c in components], dtype=np.float32)
            mesh_clean = components[areas.argmax()]

            mesh_folder = '{0}/mesh_{1}'.format(evaldir, epoch)
            utils.mkdir_ifnotexists(mesh_folder)
            mesh_clean.export('{0}/scan{1}.ply'.format(mesh_folder, scan_id), 'ply')

            del mesh_clean, components, areas, mesh
            gc.collect()
            torch.cuda.empty_cache()

    if opt.eval_rendering:
        images_dir = '{0}/rendering_{1}'.format(evaldir, epoch)
        logger.info(f"rendered images dir: {images_dir}")
        utils.mkdir_ifnotexists(images_dir)
        os.makedirs(os.path.join(images_dir, 'depth_est'), exist_ok=True)

        if 'dtu' in opt.conf:
            test_idx = get_eval_ids('DTU', scan_id=None)
            train_idx = get_trains_ids('DTU', scan=None, num_views=3)
        elif 'bmvs' in opt.conf:
            test_idx = get_eval_ids('BlendedMVS', scan_id)
            train_idx = get_trains_ids('BlendedMVS', f'scan{scan_id}', num_views=3)
            assert test_idx == [i for i in test_idx if i not in train_idx]
        else:
            raise NotImplementedError

        test_idx += train_idx[:3]
        logger.info(f"{len(test_idx)} images (including train)")

        psnrs, ssims, lpipss = [], [], []
        pred_ph = tf.placeholder(tf.float32)
        gt_ph = tf.placeholder(tf.float32)
        distance_t = lpips_tf.lpips(pred_ph, gt_ph, model='net-lin', net='vgg')
        for data_index, (indices, model_input, ground_truth) in enumerate(eval_dataloader):
            if indices not in test_idx: continue

            model_input["intrinsics"] = model_input["intrinsics"].cuda()
            model_input["uv"] = model_input["uv"].cuda()
            model_input['pose'] = model_input['pose'].cuda()

            # Already have results -> evaluate
            if opt.result_from != 'None':
                if indices in train_idx: continue

                if opt.result_from == 'blend':
                    pred_img = Image.open('{0}/eval_blend_{1}.png'.format(images_dir,'%03d' % indices[0]))
                elif opt.result_from == 'default':
                    pred_img = Image.open('{0}/eval_{1}.png'.format(images_dir,'%03d' % indices[0]))
                else:
                    raise NotImplementedError

                rgb_pred = np.array(pred_img, dtype=np.float32).reshape(-1, 3) / 255.
                rgb_pred = torch.from_numpy(rgb_pred)
                mask = ground_truth['mask'].reshape(-1, 3)
                mask_bin = (mask == 1.)
                rgb_fg = ground_truth['rgb'].reshape(-1, 3) * mask + (1 - mask)
                rgb_hat_fg = rgb_pred * mask + (1 - mask)
                rgb_fg = rgb_fg.reshape(img_res+[3,]).cpu().numpy() # (HW, 3) -> (H, W, 3)
                rgb_hat_fg = rgb_hat_fg.reshape(img_res+[3,]).cpu().numpy()
                mse = torch.mean((rgb_pred - ground_truth['rgb'].reshape(-1,3))[mask_bin] ** 2)
                psnr_masked = -10. * torch.log(mse) / torch.log(torch.Tensor([10.]))
                ssim_masked = float(structural_similarity(rgb_hat_fg, rgb_fg, multichannel=True))
                with tf.Session() as session:
                    lpips_masked = session.run(distance_t, feed_dict={pred_ph: rgb_hat_fg[None, ...], gt_ph: rgb_fg[None, ...]})[0]
                ssims.append(ssim_masked)
                psnrs.append(psnr_masked)
                lpipss.append(lpips_masked)
            
            # Otherwise, render and save RGBD for evaluation
            else:
                split = utils.split_input(model_input, total_pixels, n_pixels=opt.split_n_pixels)
                res = []
                for s in tqdm(split, ncols=60):
                    out = model(s)
                    res.append({
                        'rgb_values': out['rgb_values'].detach(),
                        'normal_map': out['normal_map'].detach(),
                        'depth_values': out['depth_values'].detach(),
                        'weights': out['weights'].detach(),
                    })

                batch_size = ground_truth['rgb'].shape[0]
                model_outputs = utils.merge_output(res, total_pixels, batch_size)

                # RGB
                rgb_eval = model_outputs['rgb_values']
                rgb_eval = rgb_eval.reshape(batch_size, total_pixels, 3)
                rgb_eval = plt.lin2img(rgb_eval, img_res).detach().cpu().numpy()[0]
                rgb_eval = rgb_eval.transpose(1, 2, 0)

                # Depth
                depth_map = model_outputs['depth_values'].reshape(batch_size, total_pixels)
                depth_map = plt.lin2img(depth_map[..., None], img_res).detach().cpu().squeeze()
                # save depth maps for image-based rendering
                depth_map_np = depth_map.numpy() * eval_dataset.scale_factor # (576, 768) float32
                save_pfm('{0}/depth_est/{1}.pfm'.format(images_dir, '%08d' % indices[0]), depth_map_np)

                # Save in png
                ## RGB
                img = Image.fromarray((rgb_eval * 255).astype(np.uint8))
                img.save('{0}/eval_{1}.png'.format(images_dir,'%03d' % indices[0]))
                
                ## Normal
                normal_eval = model_outputs['normal_map']
                normal_eval = normal_eval.reshape(batch_size, total_pixels, 3)
                normal_eval = (normal_eval + 1.) / 2.
                normal_eval = plt.lin2img(normal_eval, img_res).detach().cpu().numpy()[0]
                normal_eval = normal_eval.transpose(1, 2, 0)
                normal_eval = Image.fromarray((normal_eval * 255).astype(np.uint8))
                normal_eval.save('{0}/normal_{1}.png'.format(images_dir, '%03d' % indices[0]))

                ## Depth
                acc = model_outputs['weights'].sum(1).reshape(batch_size, total_pixels)
                acc = plt.lin2img(acc[..., None], img_res).detach().cpu().numpy().squeeze()
                depth_map = plt.visualize_depth(depth_map, acc)
                depth_map = Image.fromarray((depth_map * 255).astype(np.uint8))
                depth_map.save('{0}/dep_{1}.png'.format(images_dir, '%03d' % indices[0]))

                torch.cuda.empty_cache()

        del model
        gc.collect()
        torch.cuda.empty_cache()

        if opt.result_from != 'None':
            psnrs = np.array(psnrs).astype(np.float64)
            ssims = np.array(ssims).astype(np.float64)
            lpipss = np.array(lpipss).astype(np.float64)
            print(f"SCAN {scan_id}:")
            print("    psnr mean = {0}, std {1}".format("%.4f" % psnrs.mean(), "%.4f" % psnrs.std()))
            print("    ssim mean = {0}, std {1}".format("%.4f" % ssims.mean(), "%.4f" % ssims.std()))
            print("    lpips mean = {0}, std {1}".format("%.4f" % lpipss.mean(), "%.4f" % lpipss.std()))
            
            return psnrs, ssims, lpipss
        
    return np.array([0]), np.array([0]), np.array([0])

if __name__ == '__main__':
    ## prepare data
    # python eval_vsdf.py --conf dtu  --data_dir_root data_s_volsdf --scan_ids 106 --eval_rendering --gpu 0 --expname ours --exps_folder exps_vsdf --evals_folder exps_result # --eval_mesh
    # python eval_vsdf.py --conf bmvs --data_dir_root data_s_volsdf --scan_ids 4   --eval_mesh --eval_rendering --gpu 1 --expname ours --exps_folder exps_vsdf --evals_folder exps_result # --eval_mesh

    ## evaluate
    # python eval_vsdf.py --conf dtu  --data_dir_root data_s_volsdf --eval_rendering --gpu 0 --expname ours --exps_folder exps_vsdf --evals_folder exps_result --result_from blend
    # python eval_vsdf.py --conf bmvs --data_dir_root data_s_volsdf --eval_rendering --gpu 1 --expname ours --exps_folder exps_vsdf --evals_folder exps_result --result_from blend

    parser = argparse.ArgumentParser()
    parser.add_argument('--conf', type=str, default='dtu')
    parser.add_argument('--data_dir_root', type=str, default='data_s_volsdf', help='GT data dir')
    parser.add_argument('--eval_mesh', default=False, action="store_true", help='extract mesh via marching cube')
    parser.add_argument('--eval_rendering', default=False, action="store_true", help='If set, evaluate rendering quality.')
    parser.add_argument('--result_from', default='None', type=str, choices=['None', 'default', 'blend'])
    parser.add_argument('--expname', type=str, default='ours', help='The experiment name to be evaluated.')
    parser.add_argument('--exps_folder', type=str, default='exps_vsdf', help='The experiments folder name for train.')
    parser.add_argument('--evals_folder', type=str, default='exps_result', help='The evaluation folder name (a new folder).')
    parser.add_argument('--gpu', type=str, default='auto', help='GPU to use')
    parser.add_argument('--timestamp', default='latest', type=str, help='The experiemnt timestamp to test.')
    parser.add_argument('--checkpoint', default='latest',type=str, help='The trained model checkpoint to test')
    parser.add_argument('--ckpt_dir', default='',type=str)
    parser.add_argument('--scan_ids', nargs='+', type=int, default=None, help='e.g. --scan_ids 12 34 56')
    parser.add_argument('--resolution', default=512, type=int, help='Grid resolution for marching cube, set as 400 if not enough GPU')
    parser.add_argument('--split_n_pixels', default=512, type=int)
    opt = parser.parse_args()

    # configs
    opt.conf = f'./config/confs/{opt.conf}.conf'
    opt.eval_rendering = (opt.result_from != 'None') or opt.eval_rendering
    if opt.scan_ids is None:
        if 'dtu' in opt.conf:
            opt.scan_ids = [21, 24, 34, 37, 38, 40, 82, 106, 110, 114, 118]
        elif 'bmvs' in opt.conf:
            opt.scan_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9]
        else:
            raise NotImplementedError
    assert opt.result_from in ['None', 'default', 'blend']
    assert os.path.exists(opt.exps_folder)
    logger.warning(f'result folder = {opt.evals_folder}')
    logger.warning(f'volsdf ckpt folder = {opt.exps_folder}/{opt.expname}')
    logger.info(f'scan_ids = {opt.scan_ids}')

    # GPU
    if opt.gpu == "auto":
        deviceIDs = GPUtil.getAvailable(order='memory', limit=1, maxLoad=0.9, maxMemory=0.9, includeNan=False, excludeID=[], excludeUUID=[])
        gpu = deviceIDs[0]
    else:
        gpu = opt.gpu
    os.environ["CUDA_VISIBLE_DEVICES"] = '{0}'.format(gpu)

    # eval
    psnr_all, ssim_all, lpips_all = [], [], []
    for scan_id in opt.scan_ids:
        logger.info(f'scan_id = {scan_id}')
        psnr_i, ssim_i, lpips_i = evaluate(
                conf=opt.conf,
                expname=opt.expname,
                exps_folder_name=opt.exps_folder,
                evals_folder_name=opt.evals_folder,
                timestamp=opt.timestamp,
                checkpoint=opt.checkpoint,
                scan_id=int(scan_id),
                resolution=opt.resolution,
                ckpt_dir=opt.ckpt_dir,
                )
        psnr_all.append(psnr_i.tolist())
        ssim_all.append(ssim_i.tolist())
        lpips_all.append(lpips_i.tolist())

    if opt.result_from != 'None':
        print("FINAL metric: ")
        print(f"    psnr = {np.mean([np.mean(_) for _ in psnr_all]):.4f}")
        print(f"    ssim = {np.mean([np.mean(_) for _ in ssim_all]):.4f}")
        print(f"    lpips = {np.mean([np.mean(_) for _ in lpips_all]):.4f}")
eval_vsdf.py
cvlab-stonybrook/s-volsdf
# adapted from: https://github.com/jzhangbs/DTUeval-python

import numpy as np
import open3d as o3d
from tqdm import tqdm
import argparse
import os
import glob
import sklearn.neighbors as skln
import trimesh

def scan2hash(scan):
    """
    for BlendedMVS dataset
    """
    scan2hash_dict ={
        'scan1': '5a3ca9cb270f0e3f14d0eddb',
        'scan2': '5a6464143d809f1d8208c43c',
        'scan3': '5ab85f1dac4291329b17cb50',
        'scan4': '5b4933abf2b5f44e95de482a',
        'scan5': '5b22269758e2823a67a3bd03',
        'scan6': '5c0d13b795da9479e12e2ee9',
        'scan7': '5c1af2e2bee9a723c963d019',
        'scan8': '5c1dbf200843bc542d8ef8c4',
        'scan9': '5c34300a73a8df509add216d',
    }
    return scan2hash_dict[scan]

def apply_transform(vertices, matrix):
    """
    Transform mesh by a homogeneous transformation matrix.

    Does the bookkeeping to avoid recomputing things so this function
    should be used rather than directly modifying self.vertices
    if possible.

    Parameters
    ------------
    matrix : (4, 4) float
        Homogeneous transformation matrix
    """
    # get c-order float64 matrix
    matrix = np.asanyarray(
        matrix, order='C', dtype=np.float64)

    # only support homogeneous transformations
    if matrix.shape != (4, 4):
        raise ValueError('Transformation matrix must be (4, 4)!')

    # new vertex positions
    new_vertices = trimesh.transformations.transform_points(
        vertices,
        matrix=matrix)
    return new_vertices

def get_scales(data_dir_root):
    data_dir = 'DTU'
    scan_id = 114
    instance_dir = os.path.join(data_dir_root, data_dir, 'scan{0}'.format(scan_id))
    cam_file = '{0}/cameras.npz'.format(instance_dir)
    camera_dict = np.load(cam_file)
    assert camera_dict['scale_mat_0'].astype(np.float32)[0,0] == camera_dict['scale_mat_1'].astype(np.float32)[0,0]
    DTU_scale = camera_dict['scale_mat_0'].astype(np.float64)[0,0]
    print('DTU_scale:', DTU_scale)

    BMVS_scale = {}
    relative_scale = {}
    data_dir = 'BlendedMVS'
    for scan_id in range(1, 10):
        instance_dir = os.path.join(data_dir_root, data_dir, 'scan{0}'.format(scan_id))
        cam_file = '{0}/cameras.npz'.format(instance_dir)
        camera_dict = np.load(cam_file)
        assert camera_dict['scale_mat_0'].astype(np.float32)[0,0] == camera_dict['scale_mat_1'].astype(np.float32)[0,0]
        BMVS_scale[scan_id] = camera_dict['scale_mat_0'].astype(np.float64)[0,0]
        relative_scale[scan_id] = BMVS_scale[scan_id] / DTU_scale
    print('BMVS_scale:')
    print(BMVS_scale)
    print('relative_scale:')
    print(relative_scale)

def write_vis_pcd(file, points, colors):
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    pcd.colors = o3d.utility.Vector3dVector(colors)
    o3d.io.write_point_cloud(file, pcd)

if __name__ == "__main__":
    # python evals/eval_bmvs.py --data_dir_root data_s_volsdf --datadir exps_mvs --scan 4
    # python evals/eval_bmvs.py --data_dir_root data_s_volsdf --datadir exps_mvs > exps_mvs/bmvs_chamfer.txt

    """
    sample = 100K
        See VolSDF, "To measure the Chamfer l1 distance we used 100K random point samples from each surface."

    First, we set point cloud scale to be the same as the DTU dataset
        relative_scale = cam_scale_mat[scan_id] / cam_scale_DTU

    Then, we follow DTU dataset settings
        max_dist = 20, radius = 0.2
    """

    parser = argparse.ArgumentParser()
    parser.add_argument('--sample', type=int, default=100000)
    parser.add_argument('--scan', type=int, default=-1)
    parser.add_argument('--datadir', type=str, default='', help='pred point cloud')
    parser.add_argument('--dataset_dir', type=str, default='bmvs/dataset_textured_meshes', help='GT mesh')
    parser.add_argument('--data_dir_root', type=str, default='data_s_volsdf', help='GT data dir')
    parser.add_argument('--save_gt', action='store_true')
    parser.add_argument('-ve', '--visualize_error', action='store_true')
    parser.add_argument('--no_crop', action='store_true', help='NOT [eval only above the ground plane & using object masks]')
    
    args = parser.parse_args()

    gt_pcd_dir = os.path.join(args.data_dir_root, 'BlendedMVS', 'stl') # GT point cloud
    relative_scale = {1: 0.0010051393651899145, 2: 0.0015733906993148704, 3: 0.0012326845045689896, 4: 0.0015294108512811993, 5: 0.007349738091050388, 6: 0.01192223325424887, 7: 0.001284409757598681, 8: 0.0014762879597404273, 9: 0.022978406132555827}
    scans = [1, 2, 3, 4, 5, 6, 7, 8, 9]
    if args.scan in scans:
        scans = [args.scan]

    print("ply_name, chamfer(mm)")
    for scan in scans:
        pbar = tqdm(total=6)

        pbar.set_description('read pred pcd')
        pcd_file = os.path.join(args.datadir, 'mvsnet{:0>3}_l3.ply'.format(scan))
        data_pcd_o3d = o3d.io.read_point_cloud(pcd_file)
        data_pcd = np.asarray(data_pcd_o3d.points).astype('float32') # (N, 3)

        if scan == 5: # scan5only
            # scale_mat for scan5 is wrong, set it to 1 instead
            instance_dir = os.path.join(args.data_dir_root, 'BlendedMVS', f'scan{scan}')
            cam_file = '{0}/cameras.npz'.format(instance_dir)
            scale_mat = np.load(cam_file)['scale_mat_0']
            data_pcd = apply_transform(data_pcd, scale_mat)

        pbar.update(1)
        pbar.set_description('read gt mesh')

        gt_pointcloud_path = os.path.join(gt_pcd_dir, f'scan{scan}.ply')
        if args.save_gt:
            hash_name = scan2hash(f'scan{scan}')
            gt_dir = os.path.join(args.dataset_dir, hash_name, 'textured_mesh')

            # Find all '.obj' files in the current directory
            gt_files = glob.glob(os.path.join(gt_dir, '*.obj'))

            pbar.update(1)
            pbar.set_description('sample pcd from mesh')

            vertices_lis, triangles_lis = [], []
            tri_idx = 0
            # Load the first file as the initial mesh
            mesh = o3d.io.read_triangle_mesh(gt_files[0])
            vertices = np.asarray(mesh.vertices)
            triangles = np.asarray(mesh.triangles) #  0.. 731
            tri_idx += vertices.shape[0]
            vertices_lis.append(vertices)
            triangles_lis.append(triangles)

            # # Loop through the remaining files and add them to the mesh
            for file in gt_files[1:]:
                mesh = o3d.io.read_triangle_mesh(file)
                vertices = np.asarray(mesh.vertices)
                triangles = np.asarray(mesh.triangles) + tri_idx 
                tri_idx += vertices.shape[0]
                vertices_lis.append(vertices)
                triangles_lis.append(triangles)

            vertices = np.concatenate(vertices_lis, axis=0)
            triangles = np.concatenate(triangles_lis, axis=0)

            # GT pcd
            gt_mesh = o3d.geometry.TriangleMesh()
            gt_mesh.vertices = o3d.utility.Vector3dVector(vertices)
            gt_mesh.triangles = o3d.utility.Vector3iVector(triangles)
            sampled_pcd = gt_mesh.sample_points_uniformly(number_of_points=args.sample)
            gt_pointcloud = sampled_pcd
            pbar.set_description(f'{args.sample}')

            o3d.io.write_point_cloud(gt_pointcloud_path, gt_pointcloud)
            continue
        else:
            if not args.no_crop:
                gt_pointcloud_path = os.path.join(gt_pcd_dir, f'scan{scan}_crop.ply')
            assert os.path.exists(gt_pointcloud_path)
            gt_pointcloud = o3d.io.read_point_cloud(gt_pointcloud_path)
            gt_pcd = np.asarray(gt_pointcloud.points).astype('float32')
            pbar.set_description(f'{gt_pcd.shape[0]}')

        # ----------------------chamfer distance---------------------- #
        max_dist = 20
        radius = 0.2
        # x' = a*x + b
        # |x1'-x2'| = |(a*x1+b) - (a*x2+b)| = a|x1-x2|
        # |x1'/a-x2'/a| = |(x1+b/a) - (x2+b/a)| = |x1-x2|
        gt_pcd /= relative_scale[scan]
        data_pcd /= relative_scale[scan]

        pbar.update(1)
        pbar.set_description('random shuffle pcd index')
        shuffle_rng = np.random.default_rng()
        shuffle_rng.shuffle(data_pcd, axis=0)

        nn_engine = skln.NearestNeighbors(n_neighbors=1, radius=radius, algorithm='kd_tree', n_jobs=-1)
        # Find the closest neighbor within the radius of a point.

        pbar.update(1)
        pbar.set_description('compute data2stl')
        try:
            nn_engine.fit(gt_pcd)
            dist_d2s, idx_d2s = nn_engine.kneighbors(data_pcd, n_neighbors=1, return_distance=True)
        except:
            continue
        mean_d2s = dist_d2s[dist_d2s < max_dist].mean()

        pbar.update(1)
        pbar.set_description('compute stl2data')
        try:
            nn_engine.fit(data_pcd)
            dist_s2d, idx_s2d = nn_engine.kneighbors(gt_pcd, n_neighbors=1, return_distance=True)
        except:
            continue
        mean_s2d = dist_s2d[dist_s2d < max_dist].mean()

        pbar.update(1)
        if args.visualize_error:
            # Green: Errors larger than threshold (20)
            # White to Red: Errors counted in the reported statistics
            pbar.set_description('visualize error')
            vis_out_dir = os.path.join(args.datadir, 'result')
            os.makedirs(vis_out_dir, exist_ok=True)
            vis_dist = 10
            R = np.array([[1,0,0]], dtype=np.float64)
            G = np.array([[0,1,0]], dtype=np.float64)
            B = np.array([[0,0,1]], dtype=np.float64)
            W = np.array([[1,1,1]], dtype=np.float64)
            data_color = np.tile(B, (data_pcd.shape[0], 1))
            data_alpha = dist_d2s.clip(max=vis_dist) / vis_dist
            data_color = R * data_alpha + W * (1-data_alpha)
            data_color[dist_d2s[:,0] >= max_dist] = G
            write_vis_pcd(f'{vis_out_dir}/{scan}_d2s.ply', data_pcd, data_color)
            stl_color = np.tile(B, (gt_pcd.shape[0], 1))
            stl_alpha = dist_s2d.clip(max=vis_dist) / vis_dist
            stl_color = R * stl_alpha + W * (1-stl_alpha)
            stl_color[dist_s2d[:,0] >= max_dist] = G
            write_vis_pcd(f'{vis_out_dir}/{scan}_s2d.ply', gt_pcd, stl_color)

        pbar.update(1)
        pbar.set_description('done')
        pbar.close()
        over_all = (mean_d2s + mean_s2d) / 2
        print('scan{:0>3} {:.2f} {:.2f} {:.2f}'.format(scan, mean_d2s, mean_s2d, over_all))
evals/eval_bmvs.py
cvlab-stonybrook/s-volsdf
# adapted from: https://github.com/jzhangbs/DTUeval-python

import numpy as np
import open3d as o3d
import sklearn.neighbors as skln
from tqdm import tqdm
from scipy.io import loadmat
import multiprocessing as mp
import argparse
import os

def trun_n_d(n,d):
    return int(n*10**d)/10**d

def sample_single_tri(input_):
    n1, n2, v1, v2, tri_vert = input_
    c = np.mgrid[:n1+1, :n2+1]
    c += 0.5
    c[0] /= max(n1, 1e-7)
    c[1] /= max(n2, 1e-7)
    c = np.transpose(c, (1,2,0))
    k = c[c.sum(axis=-1) < 1]  # m2
    q = v1 * k[:,:1] + v2 * k[:,1:] + tri_vert
    return q

def write_vis_pcd(file, points, colors):
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(points)
    pcd.colors = o3d.utility.Vector3dVector(colors)
    o3d.io.write_point_cloud(file, pcd)

if __name__ == '__main__':
    # python evals/eval_dtu.py --data_dir_root data_s_volsdf --datadir exps_mvs --scan 106
    # python evals/eval_dtu.py --data_dir_root data_s_volsdf --datadir exps_mvs > exps_mvs/dtu_chamfer.txt

    mp.freeze_support()

    parser = argparse.ArgumentParser()
    parser.add_argument('--datadir', type=str, default='', help='pred point cloud')
    parser.add_argument('--data_dir_root', type=str, default='data_s_volsdf', help='GT data dir')
    parser.add_argument('--scan', type=int, default=-1)
    parser.add_argument('--mode', type=str, default='pcd', choices=['mesh', 'pcd'])
    parser.add_argument('--downsample_density', type=float, default=0.2)
    parser.add_argument('--patch_size', type=float, default=60)
    parser.add_argument('--max_dist', type=float, default=20)
    parser.add_argument('--visualize_threshold', type=float, default=10)
    parser.add_argument('-ve', '--visualize_error', action='store_true')
    args = parser.parse_args()

    dataset_dir = os.path.join(args.data_dir_root, 'DTU', 'DTU_MVS_Data')
    scans = [21, 34, 38, 82, 24, 37, 40, 106, 110, 114, 118]
    if args.scan in scans:
        scans = [args.scan]
    
    results = []
    print("ply_name, accuracy(mm), completeness(mm), overall(mm)")
    for scan in scans:
        args.scan = scan
        pred_datapath = os.path.join(args.datadir, 'mvsnet{:0>3}_l3.ply'.format(args.scan))
        #-------------------------------------#
        thresh = args.downsample_density
        if args.mode == 'mesh':
            pbar = tqdm(total=9)
            pbar.set_description('read data mesh')
            data_mesh = o3d.io.read_triangle_mesh(pred_datapath)

            vertices = np.asarray(data_mesh.vertices)
            triangles = np.asarray(data_mesh.triangles)
            tri_vert = vertices[triangles]

            pbar.update(1)
            pbar.set_description('sample pcd from mesh')
            v1 = tri_vert[:,1] - tri_vert[:,0]
            v2 = tri_vert[:,2] - tri_vert[:,0]
            l1 = np.linalg.norm(v1, axis=-1, keepdims=True)
            l2 = np.linalg.norm(v2, axis=-1, keepdims=True)
            area2 = np.linalg.norm(np.cross(v1, v2), axis=-1, keepdims=True)
            non_zero_area = (area2 > 0)[:,0]
            l1, l2, area2, v1, v2, tri_vert = [
                arr[non_zero_area] for arr in [l1, l2, area2, v1, v2, tri_vert]
            ]
            thr = thresh * np.sqrt(l1 * l2 / area2)
            n1 = np.floor(l1 / thr)
            n2 = np.floor(l2 / thr)

            with mp.Pool() as mp_pool:
                new_pts = mp_pool.map(sample_single_tri, ((n1[i,0], n2[i,0], v1[i:i+1], v2[i:i+1], tri_vert[i:i+1,0]) for i in range(len(n1))), chunksize=1024)

            new_pts = np.concatenate(new_pts, axis=0)
            data_pcd = np.concatenate([vertices, new_pts], axis=0)
        
        elif args.mode == 'pcd':
            pbar = tqdm(total=8)
            pbar.set_description('read data pcd')
            data_pcd_o3d = o3d.io.read_point_cloud(pred_datapath)
            data_pcd = np.asarray(data_pcd_o3d.points)

        pbar.update(1)
        pbar.set_description('random shuffle pcd index')
        shuffle_rng = np.random.default_rng()
        shuffle_rng.shuffle(data_pcd, axis=0)

        pbar.update(1)
        pbar.set_description('downsample pcd')
        nn_engine = skln.NearestNeighbors(n_neighbors=1, radius=thresh, algorithm='kd_tree', n_jobs=-1)
        try:
            nn_engine.fit(data_pcd)
        except:
            continue
        rnn_idxs = nn_engine.radius_neighbors(data_pcd, radius=thresh, return_distance=False)
        mask = np.ones(data_pcd.shape[0], dtype=np.bool_)
        for curr, idxs in enumerate(rnn_idxs):
            if mask[curr]:
                mask[idxs] = 0
                mask[curr] = 1
        data_down = data_pcd[mask]

        pbar.update(1)
        pbar.set_description('masking data pcd')
        obs_mask_file = loadmat(f'{dataset_dir}/ObsMask/ObsMask{args.scan}_10.mat')
        ObsMask, BB, Res = [obs_mask_file[attr] for attr in ['ObsMask', 'BB', 'Res']]
        BB = BB.astype(np.float32)

        patch = args.patch_size
        inbound = ((data_down >= BB[:1]-patch) & (data_down < BB[1:]+patch*2)).sum(axis=-1) ==3
        data_in = data_down[inbound]

        data_grid = np.around((data_in - BB[:1]) / Res).astype(np.int32)
        grid_inbound = ((data_grid >= 0) & (data_grid < np.expand_dims(ObsMask.shape, 0))).sum(axis=-1) ==3
        data_grid_in = data_grid[grid_inbound]
        in_obs = ObsMask[data_grid_in[:,0], data_grid_in[:,1], data_grid_in[:,2]].astype(np.bool_)
        data_in_obs = data_in[grid_inbound][in_obs]

        pbar.update(1)
        pbar.set_description('read STL pcd')
        stl_pcd = o3d.io.read_point_cloud(f'{dataset_dir}/Points/stl/stl{args.scan:03}_total.ply')
        stl = np.asarray(stl_pcd.points)

        try:
            pbar.update(1)
            pbar.set_description('compute data2stl')
            nn_engine.fit(stl)
            dist_d2s, idx_d2s = nn_engine.kneighbors(data_in_obs, n_neighbors=1, return_distance=True)
            max_dist = args.max_dist
            mean_d2s = dist_d2s[dist_d2s < max_dist].mean()

            pbar.update(1)
            pbar.set_description('compute stl2data')
            if args.scan == 82:
                ground_plane = loadmat(f'{dataset_dir}/ObsMask/Plane{83}.mat')['P']
            else:
                ground_plane = loadmat(f'{dataset_dir}/ObsMask/Plane{args.scan}.mat')['P']
        except:
            pbar.close()
            mean_d2s, mean_s2d, over_all = 10000.,10000., 10000.
            print('scan{:0>3} {:.2f} {:.2f} {:.2f}'.format(args.scan, trun_n_d(mean_d2s,2), trun_n_d(mean_s2d,2), trun_n_d(over_all,2)))
            results.append([mean_d2s, mean_s2d, over_all])
            continue

        stl_hom = np.concatenate([stl, np.ones_like(stl[:,:1])], -1)
        above = (ground_plane.reshape((1,4)) * stl_hom).sum(-1) > 0
        stl_above = stl[above]

        nn_engine.fit(data_in)
        dist_s2d, idx_s2d = nn_engine.kneighbors(stl_above, n_neighbors=1, return_distance=True)
        mean_s2d = dist_s2d[dist_s2d < max_dist].mean()

        pbar.update(1)
        if args.visualize_error:
            pbar.set_description('visualize error')
            vis_out_dir = os.path.join(args.datadir, 'result')
            os.makedirs(vis_out_dir, exist_ok=True)
            vis_dist = args.visualize_threshold
            R = np.array([[1,0,0]], dtype=np.float64)
            G = np.array([[0,1,0]], dtype=np.float64)
            B = np.array([[0,0,1]], dtype=np.float64)
            W = np.array([[1,1,1]], dtype=np.float64)
            data_color = np.tile(B, (data_down.shape[0], 1))
            data_alpha = dist_d2s.clip(max=vis_dist) / vis_dist
            data_color[ np.where(inbound)[0][grid_inbound][in_obs] ] = R * data_alpha + W * (1-data_alpha)
            data_color[ np.where(inbound)[0][grid_inbound][in_obs][dist_d2s[:,0] >= max_dist] ] = G
            write_vis_pcd(f'{vis_out_dir}/vis_{args.scan:03}_d2s.ply', data_down, data_color)
            stl_color = np.tile(B, (stl.shape[0], 1))
            stl_alpha = dist_s2d.clip(max=vis_dist) / vis_dist
            stl_color[ np.where(above)[0] ] = R * stl_alpha + W * (1-stl_alpha)
            stl_color[ np.where(above)[0][dist_s2d[:,0] >= max_dist] ] = G
            write_vis_pcd(f'{vis_out_dir}/vis_{args.scan:03}_s2d.ply', stl, stl_color)

        pbar.update(1)
        pbar.set_description('done')
        pbar.close()
        over_all = (mean_d2s + mean_s2d) / 2
        print('scan{:0>3} {:.2f} {:.2f} {:.2f}'.format(args.scan, trun_n_d(mean_d2s,2), trun_n_d(mean_s2d,2), trun_n_d(over_all,2)))
        results.append([mean_d2s, mean_s2d, over_all])
    results = np.array(results)
    results = results.mean(0)
    print('mean_err {:.3f} {:.3f} {:.3f}'.format(results[0], results[1], results[2]))
evals/eval_dtu.py
cvlab-stonybrook/s-volsdf
from loguru import logger
import sys
logger.remove()  # All configured handlers are removed
logger.add(sys.stderr, format="<level>{level: <8}</level> | - <level>{message}</level>", level="INFO")
# logger.add(sys.stderr, format="<level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>", level="DEBUG") 
# logger.add(sys.stderr, format="<level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>", level="INFO") 

from omegaconf import OmegaConf
import os 
import GPUtil

def check_CUDA_OOM(gpu_id, min_cuda_memory=12):
    import nvidia_smi
    nvidia_smi.nvmlInit()
    handle = nvidia_smi.nvmlDeviceGetHandleByIndex(gpu_id)
    info = nvidia_smi.nvmlDeviceGetMemoryInfo(handle)
    prevent_oom = info.free/1024/1024/1024 < min_cuda_memory
    nvidia_smi.nvmlShutdown()
    return prevent_oom

def run_help(args):
    if args.gpu == "auto":
        deviceIDs = GPUtil.getAvailable(order='memory', limit=1, maxLoad=1.0, maxMemory=0.8, includeNan=False, excludeID=[], excludeUUID=[])
        args.gpu = deviceIDs[0]
    logger.info(f"gpu -> {args.gpu}")
    os.environ["CUDA_VISIBLE_DEVICES"] = '{0}'.format(args.gpu)

    if not args.prevent_oom:
        try:
            args.prevent_oom = check_CUDA_OOM(gpu_id=args.gpu, min_cuda_memory=12)
        except:
            args.prevent_oom = False
            logger.warning(f"nvidia-smi is down.")

    args.vol.dataset.img_res = [args.max_h, args.max_w]
    args.vol.dataset.num_views = args.num_view
    if args.vol.dataset.data_dir != 'DTU':
        args.interval_scale = 1.0

    logger.debug(f"filter dist -> {args.filter_dist}, diff -> {args.filter_diff}")
    logger.debug(f"interval_scale for MVS -> {args.interval_scale}")
    logger.debug(f"img size -> {args.vol.dataset.img_res}")
    logger.debug(f"num_view -> {args.vol.dataset.num_views}")
    logger.debug(f"depth_inter_r > {args.depth_inter_r}")
    logger.debug(f"n_depth -> {args.ndepths}")
    
    # some imoprtant args
    assert len([int(nd) for nd in args.ndepths.split(",") if nd]) == 3
    assert args.depth_inter_r == "1,0.5,0.5"
    assert args.ndepths == "192,32,8"
    assert args.use_nerf_d == [1,0,0]
    assert args.x2_mvsres 
    assert args.vol.dataset.data_dir in ["BlendedMVS", "DTU"]

    os.makedirs(args.outdir, exist_ok=True)
    if not args.filter_only:
        # save global args for all scans
        with open(os.path.join(args.outdir, 'all_scans.yaml'), "w") as f:
            OmegaConf.save(args, f)
    logger.info('{0}'.format(' '.join(sys.argv)))
    
    return args 
helpers/help.py
cvlab-stonybrook/s-volsdf
import numpy as np
import torch, random
import cv2
from typing import List, Union, Tuple, Dict
import os
import json
import torch.nn as nn
from glob import glob
from PIL import Image

#---------------for runner-------------------#
# read intrinsics and extrinsics
def read_camera_parameters(filename):
    with open(filename) as f:
        lines = f.readlines()
        lines = [line.rstrip() for line in lines]
    # extrinsics: line [1,5), 4x4 matrix
    extrinsics = np.fromstring(' '.join(lines[1:5]), dtype=np.float32, sep=' ').reshape((4, 4))
    # intrinsics: line [7-10), 3x3 matrix
    intrinsics = np.fromstring(' '.join(lines[7:10]), dtype=np.float32, sep=' ').reshape((3, 3))
    return intrinsics, extrinsics

# read an image
def read_img(filename):
    img = Image.open(filename)
    # scale 0~255 to 0~1
    np_img = np.array(img, dtype=np.float32) / 255.
    return np_img

# read a binary mask
def read_mask(filename):
    return read_img(filename) > 0.5

# save a binary mask
def save_mask(filename, mask):
    assert mask.dtype == np.bool
    mask = mask.astype(np.uint8) * 255
    Image.fromarray(mask).save(filename)

# read a pair file, [(ref_view1, [src_view1-1, ...]), (ref_view2, [src_view2-1, ...]), ...]
def read_pair_file(filename):
    data = []
    with open(filename) as f:
        num_viewpoint = int(f.readline())
        # 49 viewpoints
        for view_idx in range(num_viewpoint):
            ref_view = int(f.readline().rstrip())
            src_views = [int(x) for x in f.readline().rstrip().split()[1::2]]
            if len(src_views) > 0:
                data.append((ref_view, src_views))
    return data

# write camera matrix to file
def write_cam(file, cam, cam_near_far=None):
    f = open(file, "w")
    f.write('extrinsic\n')
    for i in range(0, 4):
        for j in range(0, 4):
            f.write(str(cam[0][i][j]) + ' ')
        f.write('\n')
    f.write('\n')

    f.write('intrinsic\n')
    for i in range(0, 3):
        for j in range(0, 3):
            f.write(str(cam[1][i][j]) + ' ')
        f.write('\n')

    if cam_near_far is not None:
        f.write(f'\n{cam_near_far[0]:.4f} {cam_near_far[1]:.4f} {cam_near_far[2]:.4f} {cam_near_far[3]:.4f}\n')
    else:
        f.write('\n' + str(cam[1][3][0]) + ' ' + str(cam[1][3][1]) + ' ' + str(cam[1][3][2]) + ' ' + str(cam[1][3][3]) + '\n')

    f.close()

# project the reference point cloud into the source view, then project back
def reproject_with_depth(depth_ref, intrinsics_ref, extrinsics_ref, depth_src, intrinsics_src, extrinsics_src):
    width, height = depth_ref.shape[1], depth_ref.shape[0]
    ## step1. project reference pixels to the source view
    # reference view x, y
    x_ref, y_ref = np.meshgrid(np.arange(0, width), np.arange(0, height))
    x_ref, y_ref = x_ref.reshape([-1]), y_ref.reshape([-1])
    # reference 3D space
    xyz_ref = np.matmul(np.linalg.inv(intrinsics_ref),
                        np.vstack((x_ref, y_ref, np.ones_like(x_ref))) * depth_ref.reshape([-1]))
    # source 3D space
    xyz_src = np.matmul(np.matmul(extrinsics_src, np.linalg.inv(extrinsics_ref)),
                        np.vstack((xyz_ref, np.ones_like(x_ref))))[:3]
    # source view x, y
    K_xyz_src = np.matmul(intrinsics_src, xyz_src)
    xy_src = K_xyz_src[:2] / K_xyz_src[2:3]

    ## step2. reproject the source view points with source view depth estimation
    # find the depth estimation of the source view
    x_src = xy_src[0].reshape([height, width]).astype(np.float32)
    y_src = xy_src[1].reshape([height, width]).astype(np.float32)
    sampled_depth_src = cv2.remap(depth_src, x_src, y_src, interpolation=cv2.INTER_LINEAR)
    # mask = sampled_depth_src > 0

    # source 3D space
    # NOTE that we should use sampled source-view depth_here to project back
    xyz_src = np.matmul(np.linalg.inv(intrinsics_src),
                        np.vstack((xy_src, np.ones_like(x_ref))) * sampled_depth_src.reshape([-1]))
    # reference 3D space
    xyz_reprojected = np.matmul(np.matmul(extrinsics_ref, np.linalg.inv(extrinsics_src)),
                                np.vstack((xyz_src, np.ones_like(x_ref))))[:3]
    # source view x, y, depth
    depth_reprojected = xyz_reprojected[2].reshape([height, width]).astype(np.float32)
    K_xyz_reprojected = np.matmul(intrinsics_ref, xyz_reprojected)
    xy_reprojected = K_xyz_reprojected[:2] / K_xyz_reprojected[2:3]
    x_reprojected = xy_reprojected[0].reshape([height, width]).astype(np.float32)
    y_reprojected = xy_reprojected[1].reshape([height, width]).astype(np.float32)

    return depth_reprojected, x_reprojected, y_reprojected, x_src, y_src

def check_geometric_consistency(depth_ref, intrinsics_ref, extrinsics_ref, depth_src, intrinsics_src, extrinsics_src, filter_dist=1, filter_diff=0.01):
    width, height = depth_ref.shape[1], depth_ref.shape[0]
    x_ref, y_ref = np.meshgrid(np.arange(0, width), np.arange(0, height))
    depth_reprojected, x2d_reprojected, y2d_reprojected, x2d_src, y2d_src = reproject_with_depth(depth_ref, intrinsics_ref, extrinsics_ref,
                                                     depth_src, intrinsics_src, extrinsics_src)
    # check |p_reproj-p_1| < 1
    dist = np.sqrt((x2d_reprojected - x_ref) ** 2 + (y2d_reprojected - y_ref) ** 2)

    # check |d_reproj-d_1| / d_1 < 0.01
    depth_diff = np.abs(depth_reprojected - depth_ref)
    relative_depth_diff = depth_diff / depth_ref # */0 -> inf

    mask = np.logical_and(dist < filter_dist, relative_depth_diff < filter_diff)
    # mask = np.logical_and(dist < 1, relative_depth_diff < 0.01)
    depth_reprojected[~mask] = 0

    return mask, depth_reprojected, x2d_src, y2d_src
#----------------------------------#

def glob_imgs(path):
    imgs = []
    for ext in ['*.png', '*.jpg', '*.JPEG', '*.JPG']:
        imgs.extend(glob(os.path.join(path, ext)))
    return imgs

def load_K_Rt_from_P(filename, P=None):
    """
    return np.array
        intrinsic (4,4)
        pose (4,4)
    """
    if P is None:
        lines = open(filename).read().splitlines()
        if len(lines) == 4:
            lines = lines[1:]
        lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
        P = np.asarray(lines).astype(np.float32).squeeze()

    out = cv2.decomposeProjectionMatrix(P)
    K = out[0]
    R = out[1]
    t = out[2]

    K = K/K[2,2]
    intrinsics = np.eye(4)
    intrinsics[:3, :3] = K

    pose = np.eye(4, dtype=np.float32)
    pose[:3, :3] = R.transpose()
    pose[:3,3] = (t[:3] / t[3])[:,0]

    return intrinsics, pose

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 boolean_string(s):
    if s not in {'False', 'True'}:
        raise ValueError('Not a valid boolean string')
    return s == 'True'

class NanError(Exception):
    pass

def recursive_apply(obj: Union[List, Dict], func):
    assert type(obj) == dict or type(obj) == list
    idx_iter = obj if type(obj) == dict else range(len(obj))
    for k in idx_iter:
        if type(obj[k]) == list or type(obj[k]) == dict:
            recursive_apply(obj[k], func)
        else:
            obj[k] = func(obj[k])

def visualize_depth(depth, mask=None, depth_min=None, depth_max=None, direct=False):
    """Visualize the depth map with colormap.
       Rescales the values so that depth_min and depth_max map to 0 and 1,
       respectively.
    """
    # if not direct:
    #     depth = 1.0 / (depth + 1e-6)
    invalid_mask = np.logical_or(np.isnan(depth), np.logical_not(np.isfinite(depth)))
    if mask is not None:
        invalid_mask += np.logical_not(mask)
    if depth_min is None:
        depth_min = np.percentile(depth[np.logical_not(invalid_mask)], 5)
    if depth_max is None:
        depth_max = np.percentile(depth[np.logical_not(invalid_mask)], 95)
    depth[depth < depth_min] = depth_min
    depth[depth > depth_max] = depth_max
    depth[invalid_mask] = depth_max

    depth_scaled = (depth - depth_min) / (depth_max - depth_min)
    depth_scaled_uint8 = np.uint8(depth_scaled * 255)
    if not direct:
        depth_scaled_uint8 = 255 - depth_scaled_uint8
        depth_color = cv2.applyColorMap(depth_scaled_uint8, cv2.COLORMAP_JET)
        depth_color[invalid_mask, :] = 0
    else:
        depth_color = depth_scaled_uint8
        depth_color[invalid_mask] = 0
    return depth_color


def save_model_vis(obj, save_dir: str, job_name: str, global_step: int, max_keep: int):
    os.makedirs(os.path.join(save_dir, job_name), exist_ok=True)
    record_file = os.path.join(save_dir, job_name, 'record')
    cktp_file = os.path.join(save_dir, job_name, f'{global_step}.tar')
    if not os.path.exists(record_file):
        with open(record_file, 'w+') as f:
            json.dump([], f)
    with open(record_file, 'r') as f:
        record = json.load(f)
    record.append(global_step)
    if len(record) > max_keep:
        old = record[0]
        record = record[1:]
        os.remove(os.path.join(save_dir, job_name, f'{old}.tar'))
    torch.save(obj, cktp_file)
    with open(record_file, 'w') as f:
        json.dump(record, f)


def load_model_vis(model: nn.Module, load_path: str, load_step: int):
    model.load_state_dict(torch.load(load_path)['model'])
    return 0


# print arguments
def print_args(args):
    print("################################  args  ################################")
    for k, v in args.__dict__.items():
        print("{0: <10}\t{1: <30}\t{2: <20}".format(k, str(v), str(type(v))))
    print("########################################################################")


# torch.no_grad warpper for functions
def make_nograd_func(func):
    def wrapper(*f_args, **f_kwargs):
        with torch.no_grad():
            ret = func(*f_args, **f_kwargs)
        return ret

    return wrapper


# convert a function into recursive style to handle nested dict/list/tuple variables
def make_recursive_func(func):
    def wrapper(vars):
        if isinstance(vars, list):
            return [wrapper(x) for x in vars]
        elif isinstance(vars, tuple):
            return tuple([wrapper(x) for x in vars])
        elif isinstance(vars, dict):
            return {k: wrapper(v) for k, v in vars.items()}
        else:
            return func(vars)

    return wrapper


@make_recursive_func
def tensor2float(vars):
    if isinstance(vars, float):
        return vars
    elif isinstance(vars, torch.Tensor):
        return vars.data.item()
    else:
        raise NotImplementedError("invalid input type {} for tensor2float".format(type(vars)))


@make_recursive_func
def tensor2numpy(vars):
    if isinstance(vars, np.ndarray):
        return vars
    elif isinstance(vars, torch.Tensor):
        return vars.detach().cpu().numpy().copy()
    else:
        raise NotImplementedError("invalid input type {} for tensor2numpy".format(type(vars)))


@make_recursive_func
def tocuda(vars):
    if isinstance(vars, torch.Tensor):
        return vars.to(torch.device("cuda"))
    elif isinstance(vars, str):
        return vars
    elif vars is None:
        return None
    else:
        raise NotImplementedError("invalid input type {} for tensor2numpy".format(type(vars)))

@make_recursive_func
def tocpu(vars):
    if isinstance(vars, torch.Tensor):
        return vars.to(torch.device("cpu"))
    elif isinstance(vars, str):
        return vars
    elif vars is None:
        return None
    else:
        raise NotImplementedError("invalid input type {} for tensor2numpy".format(type(vars)))

def set_random_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
helpers/utils.py
cvlab-stonybrook/s-volsdf
import os
import sys

import tensorflow as tf
from six.moves import urllib

_URL = 'http://rail.eecs.berkeley.edu/models/lpips'


def _download(url, output_dir):
    """Downloads the `url` file into `output_dir`.

    Modified from https://github.com/tensorflow/models/blob/master/research/slim/datasets/dataset_utils.py
    """
    filename = url.split('/')[-1]
    filepath = os.path.join(output_dir, filename)

    def _progress(count, block_size, total_size):
        sys.stdout.write('\r>> Downloading %s %.1f%%' % (
            filename, float(count * block_size) / float(total_size) * 100.0))
        sys.stdout.flush()

    filepath, _ = urllib.request.urlretrieve(url, filepath, _progress)
    print()
    statinfo = os.stat(filepath)
    print('Successfully downloaded', filename, statinfo.st_size, 'bytes.')


def lpips(input0, input1, model='net-lin', net='alex', version=0.1):
    """
    Learned Perceptual Image Patch Similarity (LPIPS) metric.

    Args:
        input0: An image tensor of shape `[..., height, width, channels]`,
            with values in [0, 1].
        input1: An image tensor of shape `[..., height, width, channels]`,
            with values in [0, 1].

    Returns:
        The Learned Perceptual Image Patch Similarity (LPIPS) distance.

    Reference:
        Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, Oliver Wang.
        The Unreasonable Effectiveness of Deep Features as a Perceptual Metric.
        In CVPR, 2018.
    """
    # flatten the leading dimensions
    batch_shape = tf.shape(input0)[:-3]
    input0 = tf.reshape(input0, tf.concat([[-1], tf.shape(input0)[-3:]], axis=0))
    input1 = tf.reshape(input1, tf.concat([[-1], tf.shape(input1)[-3:]], axis=0))
    # NHWC to NCHW
    input0 = tf.transpose(input0, [0, 3, 1, 2])
    input1 = tf.transpose(input1, [0, 3, 1, 2])
    # normalize to [-1, 1]
    input0 = input0 * 2.0 - 1.0
    input1 = input1 * 2.0 - 1.0

    input0_name, input1_name = '0:0', '1:0'

    default_graph = tf.compat.v1.get_default_graph()
    producer_version = default_graph.graph_def_versions.producer

    cache_dir = os.path.expanduser('~/.lpips')
    os.makedirs(cache_dir, exist_ok=True)
    # files to try. try a specific producer version, but fallback to the version-less version (latest). Note from Haoyu: identical from md5sum.
    pb_fnames = [
        '%s_%s_v%s_%d.pb' % (model, net, version, producer_version),
        '%s_%s_v%s.pb' % (model, net, version),
    ]
    for pb_fname in pb_fnames:
        if not os.path.isfile(os.path.join(cache_dir, pb_fname)):
            try:
                _download(os.path.join(_URL, pb_fname), cache_dir)
            except urllib.error.HTTPError:
                pass
        if os.path.isfile(os.path.join(cache_dir, pb_fname)):
            break

    with open(os.path.join(cache_dir, pb_fname), 'rb') as f:
        graph_def = tf.compat.v1.GraphDef()
        graph_def.ParseFromString(f.read())
        _ = tf.import_graph_def(graph_def,
                                input_map={input0_name: input0, input1_name: input1})
        distance, = default_graph.get_operations()[-1].outputs

    if distance.shape.ndims == 4:
        distance = tf.squeeze(distance, axis=[-3, -2, -1])
    # reshape the leading dimensions
    distance = tf.reshape(distance, batch_shape)
    return distance

lpips_tensorflow/lpips_tf.py
cvlab-stonybrook/s-volsdf
import torch
import torch.nn as nn
import torch.nn.functional as F
from helpers.help import logger

Align_Corners_Range = False

def init_bn(module):
    if module.weight is not None:
        nn.init.ones_(module.weight)
    if module.bias is not None:
        nn.init.zeros_(module.bias)
    return


def init_uniform(module, init_method):
    if module.weight is not None:
        if init_method == "kaiming":
            nn.init.kaiming_uniform_(module.weight)
        elif init_method == "xavier":
            nn.init.xavier_uniform_(module.weight)
    return

class Conv2d(nn.Module):
    """Applies a 2D convolution (optionally with batch normalization and relu activation)
    over an input signal composed of several input planes.
    Attributes:
        conv (nn.Module): convolution module
        bn (nn.Module): batch normalization module
        relu (bool): whether to activate by relu
    Notes:
        Default momentum for batch normalization is set to be 0.01,
    """

    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 relu=True, bn=True, bn_momentum=0.1, init_method="xavier", **kwargs):
        super(Conv2d, self).__init__()

        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride,
                              bias=(not bn), **kwargs)
        self.kernel_size = kernel_size
        self.stride = stride
        self.bn = nn.BatchNorm2d(out_channels, momentum=bn_momentum) if bn else None
        self.relu = relu

        # assert init_method in ["kaiming", "xavier"]
        # self.init_weights(init_method)

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu:
            x = F.relu(x, inplace=True)
        return x

    def init_weights(self, init_method):
        """default initialization"""
        init_uniform(self.conv, init_method)
        if self.bn is not None:
            init_bn(self.bn)


class Deconv2d(nn.Module):
    """Applies a 2D deconvolution (optionally with batch normalization and relu activation)
       over an input signal composed of several input planes.
       Attributes:
           conv (nn.Module): convolution module
           bn (nn.Module): batch normalization module
           relu (bool): whether to activate by relu
       Notes:
           Default momentum for batch normalization is set to be 0.01,
       """

    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 relu=True, bn=True, bn_momentum=0.1, init_method="xavier", **kwargs):
        super(Deconv2d, self).__init__()
        self.out_channels = out_channels
        assert stride in [1, 2]
        self.stride = stride

        self.conv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride,
                                       bias=(not bn), **kwargs)
        self.bn = nn.BatchNorm2d(out_channels, momentum=bn_momentum) if bn else None
        self.relu = relu

        # assert init_method in ["kaiming", "xavier"]
        # self.init_weights(init_method)

    def forward(self, x):
        y = self.conv(x)
        if self.stride == 2:
            h, w = list(x.size())[2:]
            y = y[:, :, :2 * h, :2 * w].contiguous()
        if self.bn is not None:
            x = self.bn(y)
        if self.relu:
            x = F.relu(x, inplace=True)
        return x

    def init_weights(self, init_method):
        """default initialization"""
        init_uniform(self.conv, init_method)
        if self.bn is not None:
            init_bn(self.bn)

class Conv3d(nn.Module):
    """Applies a 3D convolution (optionally with batch normalization and relu activation)
    over an input signal composed of several input planes.
    Attributes:
        conv (nn.Module): convolution module
        bn (nn.Module): batch normalization module
        relu (bool): whether to activate by relu
    Notes:
        Default momentum for batch normalization is set to be 0.01,
    """

    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1,
                 relu=True, bn=True, bn_momentum=0.1, init_method="xavier", **kwargs):
        super(Conv3d, self).__init__()
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        assert stride in [1, 2]
        self.stride = stride

        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride,
                              bias=(not bn), **kwargs)
        self.bn = nn.BatchNorm3d(out_channels, momentum=bn_momentum) if bn else None
        self.relu = relu

        # assert init_method in ["kaiming", "xavier"]
        # self.init_weights(init_method)

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu:
            x = F.relu(x, inplace=True)
        return x

    def init_weights(self, init_method):
        """default initialization"""
        init_uniform(self.conv, init_method)
        if self.bn is not None:
            init_bn(self.bn)

class Deconv3d(nn.Module):
    """Applies a 3D deconvolution (optionally with batch normalization and relu activation)
       over an input signal composed of several input planes.
       Attributes:
           conv (nn.Module): convolution module
           bn (nn.Module): batch normalization module
           relu (bool): whether to activate by relu
       Notes:
           Default momentum for batch normalization is set to be 0.01,
       """

    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1,
                 relu=True, bn=True, bn_momentum=0.1, init_method="xavier", **kwargs):
        super(Deconv3d, self).__init__()
        self.out_channels = out_channels
        assert stride in [1, 2]
        self.stride = stride

        self.conv = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride,
                                       bias=(not bn), **kwargs)
        self.bn = nn.BatchNorm3d(out_channels, momentum=bn_momentum) if bn else None
        self.relu = relu

        # assert init_method in ["kaiming", "xavier"]
        # self.init_weights(init_method)

    def forward(self, x):
        y = self.conv(x)
        if self.bn is not None:
            x = self.bn(y)
        if self.relu:
            x = F.relu(x, inplace=True)
        return x

    def init_weights(self, init_method):
        """default initialization"""
        init_uniform(self.conv, init_method)
        if self.bn is not None:
            init_bn(self.bn)



class ConvBnReLU(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1):
        super(ConvBnReLU, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=pad, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        return F.relu(self.bn(self.conv(x)), inplace=True)


class ConvBn(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1):
        super(ConvBn, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=pad, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        return self.bn(self.conv(x))


class ConvBnReLU3D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1):
        super(ConvBnReLU3D, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=pad, bias=False)
        self.bn = nn.BatchNorm3d(out_channels)

    def forward(self, x):
        return F.relu(self.bn(self.conv(x)), inplace=True)


class ConvBn3D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1):
        super(ConvBn3D, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=pad, bias=False)
        self.bn = nn.BatchNorm3d(out_channels)

    def forward(self, x):
        return self.bn(self.conv(x))


class BasicBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride, downsample=None):
        super(BasicBlock, self).__init__()

        self.conv1 = ConvBnReLU(in_channels, out_channels, kernel_size=3, stride=stride, pad=1)
        self.conv2 = ConvBn(out_channels, out_channels, kernel_size=3, stride=1, pad=1)

        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        if self.downsample is not None:
            x = self.downsample(x)
        out += x
        return out


class Hourglass3d(nn.Module):
    def __init__(self, channels):
        super(Hourglass3d, self).__init__()

        self.conv1a = ConvBnReLU3D(channels, channels * 2, kernel_size=3, stride=2, pad=1)
        self.conv1b = ConvBnReLU3D(channels * 2, channels * 2, kernel_size=3, stride=1, pad=1)

        self.conv2a = ConvBnReLU3D(channels * 2, channels * 4, kernel_size=3, stride=2, pad=1)
        self.conv2b = ConvBnReLU3D(channels * 4, channels * 4, kernel_size=3, stride=1, pad=1)

        self.dconv2 = nn.Sequential(
            nn.ConvTranspose3d(channels * 4, channels * 2, kernel_size=3, padding=1, output_padding=1, stride=2,
                               bias=False),
            nn.BatchNorm3d(channels * 2))

        self.dconv1 = nn.Sequential(
            nn.ConvTranspose3d(channels * 2, channels, kernel_size=3, padding=1, output_padding=1, stride=2,
                               bias=False),
            nn.BatchNorm3d(channels))

        self.redir1 = ConvBn3D(channels, channels, kernel_size=1, stride=1, pad=0)
        self.redir2 = ConvBn3D(channels * 2, channels * 2, kernel_size=1, stride=1, pad=0)

    def forward(self, x):
        conv1 = self.conv1b(self.conv1a(x))
        conv2 = self.conv2b(self.conv2a(conv1))
        dconv2 = F.relu(self.dconv2(conv2) + self.redir2(conv1), inplace=True)
        dconv1 = F.relu(self.dconv1(dconv2) + self.redir1(x), inplace=True)
        return dconv1


def homo_warping(src_fea, src_proj, ref_proj, depth_values):
    # src_fea: [B, C, H, W]
    # src_proj: [B, 4, 4]
    # ref_proj: [B, 4, 4]
    # depth_values: [B, Ndepth] o [B, Ndepth, H, W]
    # out: [B, C, Ndepth, H, W]
    batch, channels = src_fea.shape[0], src_fea.shape[1]
    num_depth = depth_values.shape[1]
    height, width = src_fea.shape[2], src_fea.shape[3]

    with torch.no_grad():
        proj = torch.matmul(src_proj, torch.inverse(ref_proj))
        rot = proj[:, :3, :3]  # [B,3,3]
        trans = proj[:, :3, 3:4]  # [B,3,1]

        y, x = torch.meshgrid([torch.arange(0, height, dtype=torch.float32, device=src_fea.device),
                               torch.arange(0, width, dtype=torch.float32, device=src_fea.device)])
        y, x = y.contiguous(), x.contiguous()
        y, x = y.view(height * width), x.view(height * width)
        xyz = torch.stack((x, y, torch.ones_like(x)))  # [3, H*W]
        xyz = torch.unsqueeze(xyz, 0).repeat(batch, 1, 1)  # [B, 3, H*W]
        rot_xyz = torch.matmul(rot, xyz)  # [B, 3, H*W]
        rot_depth_xyz = rot_xyz.unsqueeze(2).repeat(1, 1, num_depth, 1) * depth_values.view(batch, 1, num_depth,
                                                                                            -1)  # [B, 3, Ndepth, H*W]
        proj_xyz = rot_depth_xyz + trans.view(batch, 3, 1, 1)  # [B, 3, Ndepth, H*W]
        proj_xy = proj_xyz[:, :2, :, :] / proj_xyz[:, 2:3, :, :]  # [B, 2, Ndepth, H*W]
        proj_x_normalized = proj_xy[:, 0, :, :] / ((width - 1) / 2) - 1
        proj_y_normalized = proj_xy[:, 1, :, :] / ((height - 1) / 2) - 1
        proj_xy = torch.stack((proj_x_normalized, proj_y_normalized), dim=3)  # [B, Ndepth, H*W, 2]
        grid = proj_xy

    warped_src_fea = F.grid_sample(src_fea, grid.view(batch, num_depth * height, width, 2), mode='bilinear',
                                   padding_mode='zeros', align_corners = False)
    warped_src_fea = warped_src_fea.view(batch, channels, num_depth, height, width)

    return warped_src_fea

class DeConv2dFuse(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, relu=True, bn=True,
                 bn_momentum=0.1):
        super(DeConv2dFuse, self).__init__()

        self.deconv = Deconv2d(in_channels, out_channels, kernel_size, stride=2, padding=1, output_padding=1,
                               bn=True, relu=relu, bn_momentum=bn_momentum)

        self.conv = Conv2d(2*out_channels, out_channels, kernel_size, stride=1, padding=1,
                           bn=bn, relu=relu, bn_momentum=bn_momentum)

        # assert init_method in ["kaiming", "xavier"]
        # self.init_weights(init_method)

    def forward(self, x_pre, x):
        x = self.deconv(x)
        x = torch.cat((x, x_pre), dim=1)
        x = self.conv(x)
        return x


class FeatureNet(nn.Module):
    def __init__(self, base_channels, num_stage=3, stride=4, arch_mode="unet"):
        super(FeatureNet, self).__init__()
        assert arch_mode in ["unet", "fpn"], print("mode must be in 'unet' or 'fpn', but get:{}".format(arch_mode))
        logger.debug("*************feature extraction arch mode:{}****************".format(arch_mode))
        self.arch_mode = arch_mode
        self.stride = stride
        self.base_channels = base_channels
        self.num_stage = num_stage

        self.conv0 = nn.Sequential(
            Conv2d(3, base_channels, 3, 1, padding=1),
            Conv2d(base_channels, base_channels, 3, 1, padding=1),
        )

        self.conv1 = nn.Sequential(
            Conv2d(base_channels, base_channels * 2, 5, stride=2, padding=2),
            Conv2d(base_channels * 2, base_channels * 2, 3, 1, padding=1),
            Conv2d(base_channels * 2, base_channels * 2, 3, 1, padding=1),
        )

        self.conv2 = nn.Sequential(
            Conv2d(base_channels * 2, base_channels * 4, 5, stride=2, padding=2),
            Conv2d(base_channels * 4, base_channels * 4, 3, 1, padding=1),
            Conv2d(base_channels * 4, base_channels * 4, 3, 1, padding=1),
        )

        self.out1 = nn.Conv2d(base_channels * 4, base_channels * 4, 1, bias=False)
        self.out_channels = [4 * base_channels]

        if self.arch_mode == 'unet':
            if num_stage == 3:
                self.deconv1 = DeConv2dFuse(base_channels * 4, base_channels * 2, 3)
                self.deconv2 = DeConv2dFuse(base_channels * 2, base_channels, 3)

                self.out2 = nn.Conv2d(base_channels * 2, base_channels * 2, 1, bias=False)
                self.out3 = nn.Conv2d(base_channels, base_channels, 1, bias=False)
                self.out_channels.append(2 * base_channels)
                self.out_channels.append(base_channels)

            elif num_stage == 2:
                self.deconv1 = DeConv2dFuse(base_channels * 4, base_channels * 2, 3)

                self.out2 = nn.Conv2d(base_channels * 2, base_channels * 2, 1, bias=False)
                self.out_channels.append(2 * base_channels)
        elif self.arch_mode == "fpn":
            final_chs = base_channels * 4
            if num_stage == 3:
                self.inner1 = nn.Conv2d(base_channels * 2, final_chs, 1, bias=True)
                self.inner2 = nn.Conv2d(base_channels * 1, final_chs, 1, bias=True)

                self.out2 = nn.Conv2d(final_chs, base_channels * 2, 3, padding=1, bias=False)
                self.out3 = nn.Conv2d(final_chs, base_channels, 3, padding=1, bias=False)
                self.out_channels.append(base_channels * 2)
                self.out_channels.append(base_channels)

            elif num_stage == 2:
                self.inner1 = nn.Conv2d(base_channels * 2, final_chs, 1, bias=True)

                self.out2 = nn.Conv2d(final_chs, base_channels, 3, padding=1, bias=False)
                self.out_channels.append(base_channels)

    def forward(self, x):
        conv0 = self.conv0(x)
        conv1 = self.conv1(conv0)
        conv2 = self.conv2(conv1)

        intra_feat = conv2
        outputs = {}
        out = self.out1(intra_feat)
        outputs["stage1"] = out
        if self.arch_mode == "unet":
            if self.num_stage == 3:
                intra_feat = self.deconv1(conv1, intra_feat)
                out = self.out2(intra_feat)
                outputs["stage2"] = out

                intra_feat = self.deconv2(conv0, intra_feat)
                out = self.out3(intra_feat)
                outputs["stage3"] = out

            elif self.num_stage == 2:
                intra_feat = self.deconv1(conv1, intra_feat)
                out = self.out2(intra_feat)
                outputs["stage2"] = out

        elif self.arch_mode == "fpn":
            if self.num_stage == 3:
                intra_feat = F.interpolate(intra_feat, scale_factor=2, mode="nearest") + self.inner1(conv1)
                out = self.out2(intra_feat)
                outputs["stage2"] = out

                intra_feat = F.interpolate(intra_feat, scale_factor=2, mode="nearest") + self.inner2(conv0)
                out = self.out3(intra_feat)
                outputs["stage3"] = out

            elif self.num_stage == 2:
                intra_feat = F.interpolate(intra_feat, scale_factor=2, mode="nearest") + self.inner1(conv1)
                out = self.out2(intra_feat)
                outputs["stage2"] = out

        return outputs

class CostRegNet(nn.Module):
    def __init__(self, in_channels, base_channels):
        super(CostRegNet, self).__init__()
        self.conv0 = Conv3d(in_channels, base_channels, padding=1)

        self.conv1 = Conv3d(base_channels, base_channels * 2, stride=2, padding=1)
        self.conv2 = Conv3d(base_channels * 2, base_channels * 2, padding=1)

        self.conv3 = Conv3d(base_channels * 2, base_channels * 4, stride=2, padding=1)
        self.conv4 = Conv3d(base_channels * 4, base_channels * 4, padding=1)

        self.conv5 = Conv3d(base_channels * 4, base_channels * 8, stride=2, padding=1)
        self.conv6 = Conv3d(base_channels * 8, base_channels * 8, padding=1)

        self.conv7 = Deconv3d(base_channels * 8, base_channels * 4, stride=2, padding=1, output_padding=1)

        self.conv9 = Deconv3d(base_channels * 4, base_channels * 2, stride=2, padding=1, output_padding=1)

        self.conv11 = Deconv3d(base_channels * 2, base_channels * 1, stride=2, padding=1, output_padding=1)

        self.prob = nn.Conv3d(base_channels, 1, 3, stride=1, padding=1, bias=False)

    def forward(self, x):
        conv0 = self.conv0(x)
        conv2 = self.conv2(self.conv1(conv0))
        conv4 = self.conv4(self.conv3(conv2))
        x = self.conv6(self.conv5(conv4))
        x = conv4 + self.conv7(x)
        x = conv2 + self.conv9(x)
        x = conv0 + self.conv11(x)
        x = self.prob(x)
        return x

class RefineNet(nn.Module):
    def __init__(self):
        super(RefineNet, self).__init__()
        self.conv1 = ConvBnReLU(4, 32)
        self.conv2 = ConvBnReLU(32, 32)
        self.conv3 = ConvBnReLU(32, 32)
        self.res = ConvBnReLU(32, 1)

    def forward(self, img, depth_init):
        concat = F.cat((img, depth_init), dim=1)
        depth_residual = self.res(self.conv3(self.conv2(self.conv1(concat))))
        depth_refined = depth_init + depth_residual
        return depth_refined


def depth_regression(p, depth_values):
    if depth_values.dim() <= 2:
        # print("regression dim <= 2")
        depth_values = depth_values.view(*depth_values.shape, 1, 1)
    depth = torch.sum(p * depth_values, 1)

    return depth

def cas_mvsnet_loss(inputs, depth_gt_ms, mask_ms, **kwargs):
    depth_loss_weights = kwargs.get("dlossw", None)

    total_loss = torch.tensor(0.0, dtype=torch.float32, device=mask_ms["stage1"].device, requires_grad=False)

    for (stage_inputs, stage_key) in [(inputs[k], k) for k in inputs.keys() if "stage" in k]:
        depth_est = stage_inputs["depth"]
        depth_gt = depth_gt_ms[stage_key]
        mask = mask_ms[stage_key]
        mask = mask > 0.5

        depth_loss = F.smooth_l1_loss(depth_est[mask], depth_gt[mask], reduction='mean')

        if depth_loss_weights is not None:
            stage_idx = int(stage_key.replace("stage", "")) - 1
            total_loss += depth_loss_weights[stage_idx] * depth_loss
        else:
            total_loss += 1.0 * depth_loss

    return total_loss, depth_loss

# https://github.com/alibaba/cascade-stereo/blob/678d7d0504b638c202fa183b3f4c4fd045328b24/CasMVSNet/models/module.py#L532
def get_cur_depth_range_samples(cur_depth, ndepth, depth_inteval_pixel, shape, max_depth=192.0, min_depth=0.0):
    #shape, (B, H, W)
    #cur_depth: (B, H, W)
    #return depth_range_values: (B, D, H, W)
    cur_depth_min = (cur_depth - ndepth / 2 * depth_inteval_pixel)  # (B, H, W)
    cur_depth_max = (cur_depth + ndepth / 2 * depth_inteval_pixel)
    # cur_depth_min = (cur_depth - ndepth / 2 * depth_inteval_pixel).clamp(min=0.0)   #(B, H, W)
    # cur_depth_max = (cur_depth_min + (ndepth - 1) * depth_inteval_pixel).clamp(max=max_depth)

    assert cur_depth.shape == torch.Size(shape), "cur_depth:{}, input shape:{}".format(cur_depth.shape, shape)
    new_interval = (cur_depth_max - cur_depth_min) / (ndepth - 1)  # (B, H, W)

    depth_range_samples = cur_depth_min.unsqueeze(1) + (torch.arange(0, ndepth, device=cur_depth.device,
                                                                  dtype=cur_depth.dtype,
                                                                  requires_grad=False).reshape(1, -1, 1,
                                                                                               1) * new_interval.unsqueeze(1))

    return depth_range_samples

def get_depth_range_samples_inverse(cur_depth, ndepth, depth_inteval_pixel, device, dtype, shape,
                           max_depth=192.0, min_depth=0.0):
    # stage-1
    if cur_depth.dim() == 2: # (B, D)
        cur_depth_min = cur_depth[:, 0]  # (B,)
        cur_depth_max = cur_depth[:, -1]
        # inverse sampling
        z_steps = torch.linspace(0, 1, ndepth, device=device)[None, :] # (1, D)
        depth_range_samples = 1. / (1./cur_depth_min[:, None] * (1.-z_steps) + 1./cur_depth_max[:, None] * z_steps) # (B, D)
        depth_range_samples = depth_range_samples.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, shape[1], shape[2]) # (B, D, H, W)
    # stage 2,3
    else:
        # cur_depth # (B, H, W)
        # assert cur_depth.shape == torch.Size(shape), "cur_depth:{}, input shape:{}".format(cur_depth.shape, shape)

        # [d-R/2, d+R/2]
        depth_range_samples = get_cur_depth_range_samples(cur_depth, ndepth, depth_inteval_pixel, shape, max_depth, min_depth)

        # [?, ?+R]
        """
        let R = ndepth*depth_inteval_pixel
        to make sure cur_depth ~ z_step=0.5
        min = 0.5 * (cur - R + sqrt(R^2 + cur^2))
        """
        # R = ndepth * depth_inteval_pixel (-10%)
        # _tmp_sum = R ** 2 + cur_depth ** 2 
        # cur_depth_min = 0.5 * (cur_depth - R + torch.sqrt(_tmp_sum)) 
        # cur_depth_max = cur_depth_min + R 
        # # inverse sampling
        # z_steps = torch.linspace(0, 1, ndepth, device=device)[None, :, None, None] # (1, D, 1, 1)
        # depth_range_samples = 1. / (1./cur_depth_min[:, None, :, :] * (1.-z_steps) + 1./cur_depth_max[:, None, :, :] * z_steps) # (B, D, H, W)

        # [d-R/2, d+R/2] + inverse (-25%)
        # cur_depth_min = (cur_depth - ndepth / 2 * depth_inteval_pixel)  # (B, H, W)
        # cur_depth_max = (cur_depth + ndepth / 2 * depth_inteval_pixel)
        # # inverse sampling
        # z_steps = torch.linspace(0, 1, ndepth, device=device)[None, :, None, None] # (1, D, 1, 1)
        # depth_range_samples = 1. / (1./cur_depth_min[:, None, :, :] * (1.-z_steps) + 1./cur_depth_max[:, None, :, :] * z_steps) # (B, D, H, W)

    return depth_range_samples

def get_depth_range_samples(cur_depth, ndepth, depth_inteval_pixel, device, dtype, shape,
                           max_depth=192.0, min_depth=0.0):
    #shape: (B, H, W)
    #cur_depth: (B, H, W) or (B, D)
    #return depth_range_samples: (B, D, H, W)
    if cur_depth.dim() == 2:
        cur_depth_min = cur_depth[:, 0]  # (B,)
        cur_depth_max = cur_depth[:, -1]
        new_interval = (cur_depth_max - cur_depth_min) / (ndepth - 1)  # (B, )

        depth_range_samples = cur_depth_min.unsqueeze(1) + (torch.arange(0, ndepth, device=device, dtype=dtype,
                                                                       requires_grad=False).reshape(1, -1) * new_interval.unsqueeze(1)) #(B, D)
        depth_range_samples = depth_range_samples.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, shape[1], shape[2]) #(B, D, H, W)
    else:
        depth_range_samples = get_cur_depth_range_samples(cur_depth, ndepth, depth_inteval_pixel, shape, max_depth, min_depth)

    return depth_range_samples

class DepthNet(nn.Module):
    def __init__(self):
        super(DepthNet, self).__init__()

    def forward(self, features, proj_matrices, depth_values, num_depth, cost_regularization, prob_volume_init=None, prevent_oom=False):
        proj_matrices = torch.unbind(proj_matrices, 1)
        assert len(features) == len(proj_matrices), "Different number of images and projection matrices"
        assert depth_values.shape[1] == num_depth, "depth_values.shape[1]:{}  num_depth:{}".format(depth_values.shapep[1], num_depth)
        num_views = len(features)

        # step 1. feature extraction
        # in: images; out: 32-channel feature maps
        ref_feature, src_features = features[0], features[1:]
        ref_proj, src_projs = proj_matrices[0], proj_matrices[1:]

        # step 2. differentiable homograph, build cost volume
        if prevent_oom:
            ref_volume = ref_feature.cpu().unsqueeze(2).repeat(1, 1, num_depth, 1, 1)
        else:
            ref_volume = ref_feature.unsqueeze(2).repeat(1, 1, num_depth, 1, 1)
        volume_sum = ref_volume
        volume_sq_sum = ref_volume ** 2
        del ref_volume
        for src_fea, src_proj in zip(src_features, src_projs):
            #warpped features
            src_proj_new = src_proj[:, 0].clone()
            src_proj_new[:, :3, :4] = torch.matmul(src_proj[:, 1, :3, :3], src_proj[:, 0, :3, :4])
            ref_proj_new = ref_proj[:, 0].clone()
            ref_proj_new[:, :3, :4] = torch.matmul(ref_proj[:, 1, :3, :3], ref_proj[:, 0, :3, :4])
            warped_volume = homo_warping(src_fea, src_proj_new, ref_proj_new, depth_values)
            if prevent_oom:
                warped_volume = warped_volume.cpu()
                torch.cuda.empty_cache()
            # warped_volume = homo_warping(src_fea, src_proj[:, 2], ref_proj[:, 2], depth_values)

            if self.training:
                volume_sum = volume_sum + warped_volume
                volume_sq_sum = volume_sq_sum + warped_volume ** 2
            else:
                # TODO: this is only a temporal solution to save memory, better way?
                volume_sum += warped_volume
                volume_sq_sum += warped_volume.pow_(2)  # the memory of warped_volume has been modified
            del warped_volume
        # aggregate multiple feature volumes by variance
        volume_variance = volume_sq_sum.div_(num_views).sub_(volume_sum.div_(num_views).pow_(2))
        volume_variance = volume_variance.cuda() # OOM

        # step 3. cost volume regularization
        cost_reg = cost_regularization(volume_variance) # [1, 1, 48, 88, 120]
        # cost_reg = F.upsample(cost_reg, [num_depth * 4, img_height, img_width], mode='trilinear')
        torch.cuda.empty_cache() # OOM
        prob_volume_pre = cost_reg.squeeze(1) # [1, 48, 88, 120]

        if prob_volume_init is not None:
            prob_volume_pre += prob_volume_init

        prob_volume = F.softmax(prob_volume_pre, dim=1)
        depth = depth_regression(prob_volume, depth_values=depth_values)

        with torch.no_grad():
            # photometric confidence
            prob_volume_sum4 = 4 * F.avg_pool3d(F.pad(prob_volume.unsqueeze(1), pad=(0, 0, 0, 0, 1, 2)), (4, 1, 1), stride=1, padding=0).squeeze(1)
            depth_index = depth_regression(prob_volume, depth_values=torch.arange(num_depth, device=prob_volume.device, dtype=torch.float)).long()
            depth_index = depth_index.clamp(min=0, max=num_depth-1)
            photometric_confidence = torch.gather(prob_volume_sum4, 1, depth_index.unsqueeze(1)).squeeze(1)

        return {"depth": depth,  "photometric_confidence": photometric_confidence, "prob_volume": prob_volume, "depth_values": depth_values}


class CascadeMVSNet(nn.Module):
    def __init__(self, refine=False, ndepths=[48, 32, 8], depth_interals_ratio=[4, 2, 1], share_cr=False,
                 grad_method="detach", arch_mode="fpn", cr_base_chs=[8, 8, 8]):
        super(CascadeMVSNet, self).__init__()
        self.refine = refine
        self.share_cr = share_cr
        self.ndepths = ndepths
        self.depth_interals_ratio = depth_interals_ratio
        self.grad_method = grad_method
        self.arch_mode = arch_mode
        self.cr_base_chs = cr_base_chs
        self.num_stage = len(ndepths)
        logger.debug("**********ndepths:{}, depth_intervals_ratio:{},  grad:{}, chs:{}************".format(ndepths, depth_interals_ratio, self.grad_method, self.cr_base_chs))

        assert len(ndepths) == len(depth_interals_ratio)

        self.stage_infos = {
            "stage1":{
                "scale": 4.0,
            },
            "stage2": {
                "scale": 2.0,
            },
            "stage3": {
                "scale": 1.0,
            }
        }

        self.feature = FeatureNet(base_channels=8, stride=4, num_stage=self.num_stage, arch_mode=self.arch_mode)
        if self.share_cr:
            self.cost_regularization = CostRegNet(in_channels=self.feature.out_channels, base_channels=8)
        else:
            self.cost_regularization = nn.ModuleList([CostRegNet(in_channels=self.feature.out_channels[i],
                                                                 base_channels=self.cr_base_chs[i])
                                                      for i in range(self.num_stage)])
        if self.refine:
            self.refine_network = RefineNet()
        self.DepthNet = DepthNet()

    @torch.no_grad()
    def forward(self, stage_idx, sample_cuda,
            features, extra, outputs, int_r, depth=None, prevent_oom=False, inverse_depth=False):
        
        imgs, proj_matrices, depth_values = sample_cuda["imgs"], sample_cuda["proj_matrices"], sample_cuda["depth_values"]
        
        # step 0.
        if depth is None:
            depth = outputs['depth'] if stage_idx > 0 else None
        outputs = {} if outputs is None else outputs
        depth_min = float(depth_values[0, 0].cpu().numpy())
        depth_max = float(depth_values[0, -1].cpu().numpy())
        depth_interval = (depth_max - depth_min) / depth_values.size(1)

        # step 1. feature extraction
        img = imgs[:, imgs.size(1)-1]

        # print("*********************stage{}*********************".format(stage_idx + 1))
        #stage feature, proj_mats, scales
        features_stage = [feat["stage{}".format(stage_idx + 1)] for feat in features]
        proj_matrices_stage = proj_matrices["stage{}".format(stage_idx + 1)]
        stage_scale = self.stage_infos["stage{}".format(stage_idx + 1)]["scale"]

        if depth is not None:
            if self.grad_method == "detach":
                cur_depth = depth.detach() # (1, 288, 384)
            else:
                cur_depth = depth
            cur_depth = F.interpolate(cur_depth.unsqueeze(1),
                                            [img.shape[2], img.shape[3]], mode='bilinear',
                                            align_corners=Align_Corners_Range).squeeze(1) # (1, H, W)
        else:
            cur_depth = depth_values

        get_depth_range_ = get_depth_range_samples_inverse if inverse_depth else get_depth_range_samples
        depth_range_samples = get_depth_range_(cur_depth=cur_depth,
                                                    ndepth=self.ndepths[stage_idx],
                                                    depth_inteval_pixel=int_r * depth_interval,
                                                    dtype=img[0].dtype,
                                                    device=img[0].device,
                                                    shape=[img.shape[0], img.shape[2], img.shape[3]],
                                                    max_depth=depth_max,
                                                    min_depth=depth_min) # (1, D, H, W)

        depth_values = F.interpolate(depth_range_samples.unsqueeze(1),
                            [self.ndepths[stage_idx], img.shape[2]//int(stage_scale), img.shape[3]//int(stage_scale)], mode='trilinear',
                            align_corners=Align_Corners_Range).squeeze(1) # (1, D, Ht, Wt)

        outputs_stage = self.DepthNet(features_stage, proj_matrices_stage,
                                        depth_values=depth_values,
                                        num_depth=self.ndepths[stage_idx],
                                        cost_regularization=self.cost_regularization if self.share_cr else self.cost_regularization[stage_idx],
                                        prevent_oom=prevent_oom)

        outputs["stage{}".format(stage_idx + 1)] = outputs_stage
        outputs.update(outputs_stage)

        return outputs, None
models/CasMVSNet.py
cvlab-stonybrook/s-volsdf
import copy
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import einops

'''
- We provide two different positional encoding methods as shown below.
- You can easily switch different pos-enc in the __init__() function of FMT.
- In our experiments, PositionEncodingSuperGule usually cost more GPU memory.
'''
from .position_encoding import PositionEncodingSuperGule, PositionEncodingSine


class LinearAttention(nn.Module):
    def __init__(self, eps=1e-6):
        super(LinearAttention, self).__init__()
        self.feature_map = lambda x: torch.nn.functional.elu(x) + 1
        self.eps = eps

    def forward(self, queries, keys, values):
        Q = self.feature_map(queries)
        K = self.feature_map(keys)

        # Compute the KV matrix, namely the dot product of keys and values so
        # that we never explicitly compute the attention matrix and thus
        # decrease the complexity
        KV = torch.einsum("nshd,nshm->nhmd", K, values)

        # Compute the normalizer
        Z = 1/(torch.einsum("nlhd,nhd->nlh", Q, K.sum(dim=1))+self.eps)

        # Finally compute and return the new values
        V = torch.einsum("nlhd,nhmd,nlh->nlhm", Q, KV, Z)

        return V.contiguous()


class AttentionLayer(nn.Module):
    def __init__(self, attention, d_model, n_heads, d_keys=None,
                 d_values=None):
        super(AttentionLayer, self).__init__()

        # Fill d_keys and d_values
        d_keys = d_keys or (d_model//n_heads)
        d_values = d_values or (d_model//n_heads)

        self.inner_attention = attention
        self.query_projection = nn.Linear(d_model, d_keys * n_heads)
        self.key_projection = nn.Linear(d_model, d_keys * n_heads)
        self.value_projection = nn.Linear(d_model, d_values * n_heads)
        self.out_projection = nn.Linear(d_values * n_heads, d_model)
        self.n_heads = n_heads

    def forward(self, queries, keys, values):
        # Extract the dimensions into local variables
        N, L, _ = queries.shape
        _, S, _ = keys.shape
        H = self.n_heads

        # Project the queries/keys/values
        queries = self.query_projection(queries).view(N, L, H, -1)
        keys = self.key_projection(keys).view(N, S, H, -1)
        values = self.value_projection(values).view(N, S, H, -1)

        # Compute the attention
        new_values = self.inner_attention(
            queries,
            keys,
            values,
        ).view(N, L, -1)

        # Project the output and return
        return self.out_projection(new_values)


class EncoderLayer(nn.Module):
    def __init__(self, d_model, n_heads, d_keys=None, d_values=None, d_ff=None, dropout=0.0,
                 activation="relu"):
        super(EncoderLayer, self).__init__()

        d_keys = d_keys or (d_model//n_heads)
        inner_attention = LinearAttention()
        attention = AttentionLayer(inner_attention, d_model, n_heads, d_keys, d_values)

        d_ff = d_ff or 2 * d_model
        self.attention = attention
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        self.activation = getattr(F, activation)

    def forward(self, x, source):
        # Normalize the masks
        N = x.shape[0]
        L = x.shape[1]

        # Run self attention and add it to the input
        x = x + self.dropout(self.attention(
            x, source, source,
        ))

        # Run the fully connected part of the layer
        y = x = self.norm1(x)
        y = self.dropout(self.activation(self.linear1(y)))
        y = self.dropout(self.linear2(y))

        return self.norm2(x+y)


class FMT(nn.Module):
    def __init__(self, config):
        super(FMT, self).__init__()

        self.d_model = config['d_model']
        self.nhead = config['nhead']
        self.layer_names = config['layer_names']
        encoder_layer = EncoderLayer(config['d_model'], config['nhead'])
        self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(len(self.layer_names))])
        self._reset_parameters()

        # self.pos_encoding = PositionEncodingSuperGule(config['d_model'])
        self.pos_encoding = PositionEncodingSine(config['d_model'])

    def _reset_parameters(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, ref_feature=None, src_feature=None, feat="ref"):
        """
        Args:
            ref_feature(torch.Tensor): [N, C, H, W]
            src_feature(torch.Tensor): [N, C, H, W]
        """

        assert ref_feature is not None

        if feat == "ref": # only self attention layer

            assert self.d_model == ref_feature.size(1)
            _, _, H, _ = ref_feature.shape

            ref_feature = einops.rearrange(self.pos_encoding(ref_feature), 'n c h w -> n (h w) c')

            ref_feature_list = []
            for layer, name in zip(self.layers, self.layer_names): # every self attention layer
                if name == 'self':
                    ref_feature = layer(ref_feature, ref_feature)
                    ref_feature_list.append(einops.rearrange(ref_feature, 'n (h w) c -> n c h w', h=H))
            return ref_feature_list

        elif feat == "src":

            assert self.d_model == ref_feature[0].size(1)
            _, _, H, _ = ref_feature[0].shape

            ref_feature = [einops.rearrange(_, 'n c h w -> n (h w) c') for _ in ref_feature]

            src_feature = einops.rearrange(self.pos_encoding(src_feature), 'n c h w -> n (h w) c')

            for i, (layer, name) in enumerate(zip(self.layers, self.layer_names)):
                if name == 'self':
                    src_feature = layer(src_feature, src_feature)
                elif name == 'cross':
                    src_feature = layer(src_feature, ref_feature[i // 2])
                else:
                    raise KeyError
            return einops.rearrange(src_feature, 'n (h w) c -> n c h w', h=H)
        else:
            raise ValueError("Wrong feature name")



class FMT_with_pathway(nn.Module):
    def __init__(self,
            base_channels=8,
            FMT_config={
                'd_model': 32,
                'nhead': 8,
                'layer_names': ['self', 'cross'] * 4}):

        super(FMT_with_pathway, self).__init__()

        self.FMT = FMT(FMT_config)

        self.dim_reduction_1 = nn.Conv2d(base_channels * 4, base_channels * 2, 1, bias=False)
        self.dim_reduction_2 = nn.Conv2d(base_channels * 2, base_channels * 1, 1, bias=False)

        self.smooth_1 = nn.Conv2d(base_channels * 2, base_channels * 2, 3, padding=1, bias=False)
        self.smooth_2 = nn.Conv2d(base_channels * 1, base_channels * 1, 3, padding=1, bias=False)

    def _upsample_add(self, x, y):
        """_upsample_add. Upsample and add two feature maps.

        :param x: top feature map to be upsampled.
        :param y: lateral feature map.
        """

        _, _, H, W = y.size()
        return F.interpolate(x, size=(H, W), mode='bilinear') + y


    def forward(self, features):
        """forward.

        :param features: multiple views and multiple stages features
        """

        for nview_idx, feature_multi_stages in enumerate(features):
            if nview_idx == 0: # ref view
                ref_fea_t_list = self.FMT(feature_multi_stages["stage1"].clone(), feat="ref")
                feature_multi_stages["stage1"] = ref_fea_t_list[-1]
                feature_multi_stages["stage2"] = self.smooth_1(self._upsample_add(self.dim_reduction_1(feature_multi_stages["stage1"]), feature_multi_stages["stage2"]))
                feature_multi_stages["stage3"] = self.smooth_2(self._upsample_add(self.dim_reduction_2(feature_multi_stages["stage2"]), feature_multi_stages["stage3"]))

            else: # src view
                feature_multi_stages["stage1"] = self.FMT([_.clone() for _ in ref_fea_t_list], feature_multi_stages["stage1"].clone(), feat="src")
                feature_multi_stages["stage2"] = self.smooth_1(self._upsample_add(self.dim_reduction_1(feature_multi_stages["stage1"]), feature_multi_stages["stage2"]))
                feature_multi_stages["stage3"] = self.smooth_2(self._upsample_add(self.dim_reduction_2(feature_multi_stages["stage2"]), feature_multi_stages["stage3"]))

        return features

models/FMT.py
cvlab-stonybrook/s-volsdf
import torch
import torch.nn as nn
import torch.nn.functional as F
from .module import *
from .FMT import FMT_with_pathway
import math
from helpers.utils import tocuda
from helpers.help import logger

Align_Corners_Range = False

class PixelwiseNet(nn.Module):

    def __init__(self):

        super(PixelwiseNet, self).__init__()
        self.conv0 = ConvBnReLU3D(in_channels=1, out_channels=16, kernel_size=1, stride=1, pad=0)
        self.conv1 = ConvBnReLU3D(in_channels=16, out_channels=8, kernel_size=1, stride=1, pad=0)
        self.conv2 = nn.Conv3d(in_channels=8, out_channels=1, kernel_size=1, stride=1, padding=0)
        self.output = nn.Sigmoid()

    def forward(self, x1):
        """forward.

        :param x1: [B, 1, D, H, W]
        """

        x1 = self.conv2(self.conv1(self.conv0(x1))).squeeze(1) # [B, D, H, W]
        output = self.output(x1)
        output = torch.max(output, dim=1, keepdim=True)[0] # [B, 1, H ,W]

        return output


class DepthNet(nn.Module):
    def __init__(self):
        super(DepthNet, self).__init__()
        self.pixel_wise_net = PixelwiseNet()

    def forward(self, features, proj_matrices, depth_values, num_depth, cost_regularization, prob_volume_init=None, view_weights=None):
        """forward.

        :param stage_idx: int, index of stage, [1, 2, 3], stage_1 corresponds to lowest image resolution
        :param features: torch.Tensor, TODO: [B, C, H, W]
        :param proj_matrices: torch.Tensor,
        :param depth_values: torch.Tensor, TODO: [B, D, H, W]
        :param num_depth: int, Ndepth
        :param cost_regularization: nn.Module, regularization network
        :param prob_volume_init:
        :param view_weights: pixel wise view weights for src views
        """
        proj_matrices = torch.unbind(proj_matrices, 1)
        assert len(features) == len(proj_matrices), "Different number of images and projection matrices"
        assert depth_values.shape[1] == num_depth, "depth_values.shape[1]:{}  num_depth:{}".format(depth_values.shapep[1], num_depth)

        # step 1. feature extraction
        ref_feature, src_features = features[0], features[1:] # [B, C, H, W]
        ref_proj, src_projs = proj_matrices[0], proj_matrices[1:] # [B, 2, 4, 4]

        # step 2. differentiable homograph, build cost volume
        if view_weights == None:
            view_weight_list = []

        similarity_sum = 0
        pixel_wise_weight_sum = 1e-5

        for i, (src_fea, src_proj) in enumerate(zip(src_features, src_projs)): # src_fea: [B, C, H, W]
            src_proj_new = src_proj[:, 0].clone() # [B, 4, 4]
            src_proj_new[:, :3, :4] = torch.matmul(src_proj[:, 1, :3, :3], src_proj[:, 0, :3, :4])
            ref_proj_new = ref_proj[:, 0].clone() # [B, 4, 4]
            ref_proj_new[:, :3, :4] = torch.matmul(ref_proj[:, 1, :3, :3], ref_proj[:, 0, :3, :4])
            warped_volume = homo_warping(src_fea, src_proj_new, ref_proj_new, depth_values)
            similarity = (warped_volume * ref_feature.unsqueeze(2)).mean(1, keepdim=True)

            if view_weights == None:
                view_weight = self.pixel_wise_net(similarity) # [B, 1, H, W]
                view_weight_list.append(view_weight)
            else:
                view_weight = view_weights[:, i:i+1]

            if self.training:
                similarity_sum = similarity_sum + similarity * view_weight.unsqueeze(1) # [B, 1, D, H, W]
                pixel_wise_weight_sum = pixel_wise_weight_sum + view_weight.unsqueeze(1) # [B, 1, 1, H, W]
            else:
                # TODO: this is only a temporal solution to save memory, better way?
                similarity_sum += similarity * view_weight.unsqueeze(1)
                pixel_wise_weight_sum += view_weight.unsqueeze(1)

            del warped_volume
        # aggregate multiple similarity across all the source views
        similarity = similarity_sum.div_(pixel_wise_weight_sum) # [B, 1, D, H, W]

        # step 3. cost volume regularization
        cost_reg = cost_regularization(similarity)
        prob_volume_pre = cost_reg.squeeze(1)

        if prob_volume_init is not None:
            prob_volume_pre += prob_volume_init

        prob_volume = torch.exp(F.log_softmax(prob_volume_pre, dim=1))
        depth = depth_wta(prob_volume, depth_values=depth_values)
        # # entropy
        # prob_volume = torch.clamp(prob_volume, 1e-8, 1.-1e-8)
        # entropy = (-prob_volume*torch.log(prob_volume) - (1-prob_volume)*torch.log(1-prob_volume))/(-math.log(0.5))
        # entropy = torch.sum(entropy, axis=1) # (1, H, W)
        # entropy = 1. - entropy / entropy.max()

        with torch.no_grad():
            photometric_confidence = torch.max(prob_volume, dim=1)[0]
        if view_weights == None:
            view_weights = torch.cat(view_weight_list, dim=1) # [B, Nview, H, W]
            # "view_weight":view_weight[0]
            return {"depth": depth,  "photometric_confidence": photometric_confidence, "prob_volume": prob_volume, "depth_values": depth_values}, view_weights.detach()
        else:
            return {"depth": depth,  "photometric_confidence": photometric_confidence, "prob_volume": prob_volume, "depth_values": depth_values}


class TransMVSNet(nn.Module):
    def __init__(self, refine=False, ndepths=[48, 32, 8], depth_interals_ratio=[4, 2, 1], share_cr=False,
            grad_method="detach", arch_mode="fpn", cr_base_chs=[8, 8, 8]):
        super(TransMVSNet, self).__init__()
        self.refine = refine
        self.share_cr = share_cr
        self.ndepths = ndepths
        self.depth_interals_ratio = depth_interals_ratio
        self.grad_method = grad_method
        self.arch_mode = arch_mode
        self.cr_base_chs = cr_base_chs
        self.num_stage = len(ndepths)
        logger.debug("**********netphs:{}, depth_intervals_ratio:{},  grad:{}, chs:{}************".format(ndepths, depth_interals_ratio, self.grad_method, self.cr_base_chs))

        assert len(ndepths) == len(depth_interals_ratio)

        self.stage_infos = {
                "stage1":{
                    "scale": 4.0,
                    },
                "stage2": {
                    "scale": 2.0,
                    },
                "stage3": {
                    "scale": 1.0,
                    }
                }

        self.feature = FeatureNet(base_channels=8)

        self.FMT_with_pathway = FMT_with_pathway()

        if self.share_cr:
            self.cost_regularization = CostRegNet(in_channels=1, base_channels=8)
        else:
            self.cost_regularization = nn.ModuleList([CostRegNet(in_channels=1, base_channels=self.cr_base_chs[i])
                for i in range(self.num_stage)])
        if self.refine:
            self.refine_network = RefineNet()

        self.DepthNet = DepthNet()

    @torch.no_grad()
    def forward(self, stage_idx, sample_cuda,
            features, extra, outputs, int_r, depth=None, inverse_depth=False, prevent_oom=False):

        view_weights = extra
        imgs, proj_matrices, depth_values = sample_cuda["imgs"], sample_cuda["proj_matrices"], sample_cuda["depth_values"]
       
        # step 0.
        if depth is None:
            depth = outputs['depth'] if stage_idx > 0 else None
        outputs = {} if outputs is None else outputs
        depth_min = float(depth_values[0, 0].cpu().numpy())
        depth_max = float(depth_values[0, -1].cpu().numpy())
        depth_interval = (depth_max - depth_min) / depth_values.size(1)

        # step 1. feature extraction
        img = imgs[:, imgs.size(1)-1]

        # for stage_idx in range(self.num_stage):
        features_stage = [feat["stage{}".format(stage_idx + 1)] for feat in features]
        proj_matrices_stage = proj_matrices["stage{}".format(stage_idx + 1)]
        stage_scale = self.stage_infos["stage{}".format(stage_idx + 1)]["scale"]

        Using_inverse_d = False

        if depth is not None:
            if self.grad_method == "detach":
                cur_depth = depth.detach()
            else:
                cur_depth = depth
            cur_depth = F.interpolate(cur_depth.unsqueeze(1),
                    [img.shape[2], img.shape[3]], mode='bilinear',
                    align_corners=Align_Corners_Range).squeeze(1)
        else:
            cur_depth = depth_values

        # [B, D, H, W]
        depth_range_samples = get_depth_range_samples(cur_depth=cur_depth,
                ndepth=self.ndepths[stage_idx],
                depth_inteval_pixel=int_r * depth_interval,
                dtype=img[0].dtype,
                device=img[0].device,
                shape=[img.shape[0], img.shape[2], img.shape[3]],
                max_depth=depth_max,
                min_depth=depth_min,
                use_inverse_depth=inverse_depth)

        if stage_idx + 1 > 1: # for stage 2 and 3
            view_weights = F.interpolate(view_weights, scale_factor=2, mode="nearest")

        if view_weights == None: # stage 1
            outputs_stage, view_weights = self.DepthNet(
                    features_stage,
                    proj_matrices_stage,
                    depth_values=F.interpolate(depth_range_samples.unsqueeze(1), [self.ndepths[stage_idx], img.shape[2]//int(stage_scale), img.shape[3]//int(stage_scale)], mode='trilinear', align_corners=Align_Corners_Range).squeeze(1),
                    num_depth=self.ndepths[stage_idx],
                    cost_regularization=self.cost_regularization if self.share_cr else self.cost_regularization[stage_idx], view_weights=view_weights)
        else:
            outputs_stage = self.DepthNet(
                    features_stage,
                    proj_matrices_stage,
                    depth_values=F.interpolate(depth_range_samples.unsqueeze(1), [self.ndepths[stage_idx], img.shape[2]//int(stage_scale), img.shape[3]//int(stage_scale)], mode='trilinear', align_corners=Align_Corners_Range).squeeze(1),
                    num_depth=self.ndepths[stage_idx],
                    cost_regularization=self.cost_regularization if self.share_cr else self.cost_regularization[stage_idx], view_weights=view_weights)

        wta_index_map = torch.argmax(outputs_stage['prob_volume'], dim=1, keepdim=True).type(torch.long)
        depth = torch.gather(outputs_stage['depth_values'], 1, wta_index_map).squeeze(1)
        outputs_stage['depth'] = depth

        outputs["stage{}".format(stage_idx + 1)] = outputs_stage
        outputs.update(outputs_stage)

        return outputs, view_weights
models/TransMVSNet.py
cvlab-stonybrook/s-volsdf
from models.TransMVSNet import TransMVSNet
from models.CasMVSNet import CascadeMVSNet
models/__init__.py
cvlab-stonybrook/s-volsdf
#!/usr/bin/env python
from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import math
import torch
from torch import nn
from torch.autograd import Function
from torch.nn.modules.utils import _pair
from torch.autograd.function import once_differentiable
from torchvision.ops import DeformConv2d, deform_conv2d


class DCNv2(nn.Module):

    def __init__(self, in_channels, out_channels,
                 kernel_size, stride, padding, dilation=1, deformable_groups=1):
        super(DCNv2, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = _pair(kernel_size)
        self.stride = _pair(stride)
        self.padding = _pair(padding)
        self.dilation = _pair(dilation)
        self.deformable_groups = deformable_groups

        self.weight = nn.Parameter(torch.Tensor(
            out_channels, in_channels, *self.kernel_size))
        self.bias = nn.Parameter(torch.Tensor(out_channels))
        self.reset_parameters()

    def reset_parameters(self):
        n = self.in_channels
        for k in self.kernel_size:
            n *= k
        stdv = 1. / math.sqrt(n)
        self.weight.data.uniform_(-stdv, stdv)
        self.bias.data.zero_()



class DCN(DCNv2):

    def __init__(self, in_channels, out_channels,
                 kernel_size, stride, padding,
                 dilation=1, deformable_groups=1,bias=True):
        super(DCN, self).__init__(in_channels, out_channels,
                                  kernel_size, stride, padding, dilation, deformable_groups)

        channels_ = self.deformable_groups * 3 * self.kernel_size[0] * self.kernel_size[1]
        self.conv_offset_mask = nn.Conv2d(self.in_channels,
                                          channels_,
                                          kernel_size=self.kernel_size,
                                          stride=self.stride,
                                          padding=self.padding,
                                          bias=True)
        if bias==False:
            self.bias = None
        self.init_offset()

    def init_offset(self):
        self.conv_offset_mask.weight.data.zero_()
        self.conv_offset_mask.bias.data.zero_()

    def forward(self, input):
        out = self.conv_offset_mask(input)
        o1, o2, mask = torch.chunk(out, 3, dim=1)
        offset = torch.cat((o1, o2), dim=1)
        mask = torch.sigmoid(mask)
        return deform_conv2d(
            input,
            offset,
            self.weight,
            self.bias,
            self.stride,
            self.padding,
            self.dilation,
            mask=mask
        )

models/dcn.py
cvlab-stonybrook/s-volsdf
# for TransMVSNet
import torch
import torch.nn as nn
import torch.nn.functional as F
from .dcn import DCN


def init_bn(module):
    if module.weight is not None:
        nn.init.ones_(module.weight)
    if module.bias is not None:
        nn.init.zeros_(module.bias)
    return


def init_uniform(module, init_method):
    if module.weight is not None:
        if init_method == "kaiming":
            nn.init.kaiming_uniform_(module.weight)
        elif init_method == "xavier":
            nn.init.xavier_uniform_(module.weight)
    return


class Conv2d(nn.Module):
    """Applies a 2D convolution (optionally with batch normalization and relu activation)
    over an input signal composed of several input planes.

    Attributes:
        conv (nn.Module): convolution module
        bn (nn.Module): batch normalization module
        relu (bool): whether to activate by relu

    Notes:
        Default momentum for batch normalization is set to be 0.01,

    """

    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 relu=True, bn=True, bn_momentum=0.1, init_method="xavier", **kwargs):
        super(Conv2d, self).__init__()

        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride,
                              bias=(not bn), **kwargs)
        self.kernel_size = kernel_size
        self.stride = stride
        self.bn = nn.BatchNorm2d(out_channels, momentum=bn_momentum) if bn else None
        self.relu = relu

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu:
            x = F.relu(x, inplace=True)
        return x

    def init_weights(self, init_method):
        """default initialization"""
        init_uniform(self.conv, init_method)
        if self.bn is not None:
            init_bn(self.bn)


class Deconv2d(nn.Module):
    """Applies a 2D deconvolution (optionally with batch normalization and relu activation)
       over an input signal composed of several input planes.

       Attributes:
           conv (nn.Module): convolution module
           bn (nn.Module): batch normalization module
           relu (bool): whether to activate by relu

       Notes:
           Default momentum for batch normalization is set to be 0.01,

       """

    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 relu=True, bn=True, bn_momentum=0.1, init_method="xavier", **kwargs):
        super(Deconv2d, self).__init__()
        self.out_channels = out_channels
        assert stride in [1, 2]
        self.stride = stride

        self.conv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride,
                                       bias=(not bn), **kwargs)
        self.bn = nn.BatchNorm2d(out_channels, momentum=bn_momentum) if bn else None
        self.relu = relu

    def forward(self, x):
        y = self.conv(x)
        if self.stride == 2:
            h, w = list(x.size())[2:]
            y = y[:, :, :2 * h, :2 * w].contiguous()
        if self.bn is not None:
            x = self.bn(y)
        if self.relu:
            x = F.relu(x, inplace=True)
        return x

    def init_weights(self, init_method):
        """default initialization"""
        init_uniform(self.conv, init_method)
        if self.bn is not None:
            init_bn(self.bn)


class Conv3d(nn.Module):
    """Applies a 3D convolution (optionally with batch normalization and relu activation)
    over an input signal composed of several input planes.

    Attributes:
        conv (nn.Module): convolution module
        bn (nn.Module): batch normalization module
        relu (bool): whether to activate by relu

    Notes:
        Default momentum for batch normalization is set to be 0.01,

    """

    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1,
                 relu=True, bn=True, bn_momentum=0.1, init_method="xavier", **kwargs):
        super(Conv3d, self).__init__()
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        assert stride in [1, 2]
        self.stride = stride

        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride,
                              bias=(not bn), **kwargs)
        self.bn = nn.BatchNorm3d(out_channels, momentum=bn_momentum) if bn else None
        self.relu = relu

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu:
            x = F.relu(x, inplace=True)
        return x

    def init_weights(self, init_method):
        """default initialization"""
        init_uniform(self.conv, init_method)
        if self.bn is not None:
            init_bn(self.bn)


class Deconv3d(nn.Module):
    """Applies a 3D deconvolution (optionally with batch normalization and relu activation)
       over an input signal composed of several input planes.

       Attributes:
           conv (nn.Module): convolution module
           bn (nn.Module): batch normalization module
           relu (bool): whether to activate by relu

       Notes:
           Default momentum for batch normalization is set to be 0.01,

       """

    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1,
                 relu=True, bn=True, bn_momentum=0.1, init_method="xavier", **kwargs):
        super(Deconv3d, self).__init__()
        self.out_channels = out_channels
        assert stride in [1, 2]
        self.stride = stride

        self.conv = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride,
                                       bias=(not bn), **kwargs)
        self.bn = nn.BatchNorm3d(out_channels, momentum=bn_momentum) if bn else None
        self.relu = relu

        # assert init_method in ["kaiming", "xavier"]
        # self.init_weights(init_method)

    def forward(self, x):
        y = self.conv(x)
        if self.bn is not None:
            x = self.bn(y)
        if self.relu:
            x = F.relu(x, inplace=True)
        return x

    def init_weights(self, init_method):
        """default initialization"""
        init_uniform(self.conv, init_method)
        if self.bn is not None:
            init_bn(self.bn)


class ConvBnReLU(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1):
        super(ConvBnReLU, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=pad, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        return F.relu(self.bn(self.conv(x)), inplace=True)


class ConvBn(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1):
        super(ConvBn, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride, padding=pad, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        return self.bn(self.conv(x))


class ConvBnReLU3D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1):
        super(ConvBnReLU3D, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=pad, bias=False)
        self.bn = nn.BatchNorm3d(out_channels)

    def forward(self, x):
        return F.relu(self.bn(self.conv(x)), inplace=True)


class ConvBn3D(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, pad=1):
        super(ConvBn3D, self).__init__()
        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride, padding=pad, bias=False)
        self.bn = nn.BatchNorm3d(out_channels)

    def forward(self, x):
        return self.bn(self.conv(x))


class BasicBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride, downsample=None):
        super(BasicBlock, self).__init__()

        self.conv1 = ConvBnReLU(in_channels, out_channels, kernel_size=3, stride=stride, pad=1)
        self.conv2 = ConvBn(out_channels, out_channels, kernel_size=3, stride=1, pad=1)

        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        out = self.conv1(x)
        out = self.conv2(out)
        if self.downsample is not None:
            x = self.downsample(x)
        out += x
        return out


class Hourglass3d(nn.Module):
    def __init__(self, channels):
        super(Hourglass3d, self).__init__()

        self.conv1a = ConvBnReLU3D(channels, channels * 2, kernel_size=3, stride=2, pad=1)
        self.conv1b = ConvBnReLU3D(channels * 2, channels * 2, kernel_size=3, stride=1, pad=1)

        self.conv2a = ConvBnReLU3D(channels * 2, channels * 4, kernel_size=3, stride=2, pad=1)
        self.conv2b = ConvBnReLU3D(channels * 4, channels * 4, kernel_size=3, stride=1, pad=1)

        self.dconv2 = nn.Sequential(
            nn.ConvTranspose3d(channels * 4, channels * 2, kernel_size=3, padding=1, output_padding=1, stride=2,
                               bias=False),
            nn.BatchNorm3d(channels * 2))

        self.dconv1 = nn.Sequential(
            nn.ConvTranspose3d(channels * 2, channels, kernel_size=3, padding=1, output_padding=1, stride=2,
                               bias=False),
            nn.BatchNorm3d(channels))

        self.redir1 = ConvBn3D(channels, channels, kernel_size=1, stride=1, pad=0)
        self.redir2 = ConvBn3D(channels * 2, channels * 2, kernel_size=1, stride=1, pad=0)

    def forward(self, x):
        conv1 = self.conv1b(self.conv1a(x))
        conv2 = self.conv2b(self.conv2a(conv1))
        dconv2 = F.relu(self.dconv2(conv2) + self.redir2(conv1), inplace=True)
        dconv1 = F.relu(self.dconv1(dconv2) + self.redir1(x), inplace=True)
        return dconv1


def homo_warping(src_fea, src_proj, ref_proj, depth_values):
    # src_fea: [B, C, H, W]
    # src_proj: [B, 4, 4]
    # ref_proj: [B, 4, 4]
    # depth_values: [B, Ndepth] or [B, Ndepth, H, W]
    # out: [B, C, Ndepth, H, W]
    batch, channels = src_fea.shape[0], src_fea.shape[1]
    num_depth = depth_values.shape[1]
    height, width = src_fea.shape[2], src_fea.shape[3]

    with torch.no_grad():
        proj = torch.matmul(src_proj, torch.inverse(ref_proj))
        rot = proj[:, :3, :3]  # [B,3,3]
        trans = proj[:, :3, 3:4]  # [B,3,1]

        y, x = torch.meshgrid([torch.arange(0, height, dtype=torch.float32, device=src_fea.device),
                               torch.arange(0, width, dtype=torch.float32, device=src_fea.device)])
        y, x = y.contiguous(), x.contiguous()
        y, x = y.view(height * width), x.view(height * width)
        xyz = torch.stack((x, y, torch.ones_like(x)))  # [3, H*W]
        xyz = torch.unsqueeze(xyz, 0).repeat(batch, 1, 1)  # [B, 3, H*W]
        rot_xyz = torch.matmul(rot, xyz)  # [B, 3, H*W]
        rot_depth_xyz = rot_xyz.unsqueeze(2).repeat(1, 1, num_depth, 1) * depth_values.view(batch, 1, num_depth,
                                                                                            -1)  # [B, 3, Ndepth, H*W]
        proj_xyz = rot_depth_xyz + trans.view(batch, 3, 1, 1)  # [B, 3, Ndepth, H*W]
        invalid = (proj_xyz[:, 2:3, :, :]<1e-6).squeeze(1) # [B, Ndepth, H*W]
        proj_xy = proj_xyz[:, :2, :, :] / (proj_xyz[:, 2:3, :, :])  # [B, 2, Ndepth, H*W]
        proj_x_normalized = proj_xy[:, 0, :, :] / ((width - 1) / 2) - 1
        proj_x_normalized[invalid] = -99.
        proj_y_normalized = proj_xy[:, 1, :, :] / ((height - 1) / 2) - 1
        proj_y_normalized[invalid] = -99.
        invalid = invalid | (proj_x_normalized<-1.001) | (proj_x_normalized>1.001) | (proj_y_normalized<-1.001) | (proj_y_normalized>1.001)
        proj_xy = torch.stack((proj_x_normalized, proj_y_normalized), dim=3)  # [B, Ndepth, H*W, 2]
        grid = proj_xy

    warped_src_fea = F.grid_sample(src_fea, grid.view(batch, num_depth * height, width, 2), mode='bilinear',
                                   padding_mode='zeros', align_corners=True)
    warped_src_fea = warped_src_fea.view(batch, channels, num_depth, height, width)

    return warped_src_fea


class DeConv2dFuse(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, relu=True, bn=True,
                 bn_momentum=0.1):
        super(DeConv2dFuse, self).__init__()

        self.deconv = Deconv2d(in_channels, out_channels, kernel_size, stride=2, padding=1, output_padding=1,
                               bn=True, relu=relu, bn_momentum=bn_momentum)

        self.conv = Conv2d(2*out_channels, out_channels, kernel_size, stride=1, padding=1,
                           bn=bn, relu=relu, bn_momentum=bn_momentum)

    def forward(self, x_pre, x):
        x = self.deconv(x)
        x = torch.cat((x, x_pre), dim=1)
        x = self.conv(x)
        return x


class FeatureNet(nn.Module):
    def __init__(self, base_channels):
        super(FeatureNet, self).__init__()
        self.base_channels = base_channels

        self.conv0 = nn.Sequential(
                Conv2d(3, base_channels, 3, 1, padding=1),
                Conv2d(base_channels, base_channels, 3, 1, padding=1))

        self.conv1 = nn.Sequential(
                Conv2d(base_channels, base_channels * 2, 5, stride=2, padding=2),
                Conv2d(base_channels * 2, base_channels * 2, 3, 1, padding=1),
                Conv2d(base_channels * 2, base_channels * 2, 3, 1, padding=1))

        self.conv2 = nn.Sequential(
                Conv2d(base_channels * 2, base_channels * 4, 5, stride=2, padding=2),
                Conv2d(base_channels * 4, base_channels * 4, 3, 1, padding=1),
                Conv2d(base_channels * 4, base_channels * 4, 3, 1, padding=1))

        self.out1 = nn.Sequential(
                Conv2d(base_channels * 4, base_channels * 4, 1),
                DCN(in_channels=base_channels * 4, out_channels=base_channels * 4, kernel_size=3, stride=1, padding=1),
                nn.BatchNorm2d(base_channels * 4),
                nn.ReLU(inplace=True),
                DCN(in_channels=base_channels * 4, out_channels=base_channels * 4, kernel_size=3,stride=1, padding=1),
                nn.BatchNorm2d(base_channels * 4),
                nn.ReLU(inplace=True),
                DCN(in_channels=base_channels * 4, out_channels=base_channels * 4, kernel_size=3,stride=1, padding=1))


        final_chs = base_channels * 4
        self.inner1 = nn.Conv2d(base_channels * 2, final_chs, 1, bias=True)
        self.inner2 = nn.Conv2d(base_channels * 1, final_chs, 1, bias=True)

        self.out2 = nn.Sequential(
                Conv2d(final_chs, final_chs, 3,1,padding=1),
                DCN(in_channels=final_chs, out_channels=final_chs,kernel_size=3, stride=1, padding=1),
                nn.BatchNorm2d(final_chs),
                nn.ReLU(inplace=True),
                DCN(in_channels=final_chs, out_channels=final_chs,kernel_size=3, stride=1, padding=1),
                nn.BatchNorm2d(final_chs),
                nn.ReLU(inplace=True),
                DCN(in_channels=final_chs, out_channels=base_channels * 2,kernel_size=3, stride=1, padding=1),
                                  )
        self.out3 = nn.Sequential(
                Conv2d(final_chs, final_chs, 3, 1, padding=1),
                DCN(in_channels=final_chs, out_channels=final_chs, kernel_size=3, stride=1,padding=1),
                nn.BatchNorm2d(final_chs),
                nn.ReLU(inplace=True),
                DCN(in_channels=final_chs, out_channels=final_chs, kernel_size=3, stride=1,padding=1),
                nn.BatchNorm2d(final_chs),
                nn.ReLU(inplace=True),
                DCN(in_channels=final_chs, out_channels=base_channels, kernel_size=3,stride=1, padding=1))

        self.out_channels = [4 * base_channels, base_channels * 2, base_channels]

    def forward(self, x):
        """forward.

        :param x: [B, C, H, W]
        :return outputs: stage1 [B, 32， 128， 160], stage2 [B, 16, 256, 320], stage3 [B, 8, 512, 640]
        """
        conv0 = self.conv0(x)
        conv1 = self.conv1(conv0)
        conv2 = self.conv2(conv1)

        intra_feat = conv2
        outputs = {}
        out = self.out1(intra_feat)
        outputs["stage1"] = out
        intra_feat = F.interpolate(intra_feat, scale_factor=2, mode="nearest") + self.inner1(conv1)
        out = self.out2(intra_feat)
        outputs["stage2"] = out

        intra_feat = F.interpolate(intra_feat, scale_factor=2, mode="nearest") + self.inner2(conv0)
        out = self.out3(intra_feat)
        outputs["stage3"] = out

        return outputs


class CostRegNet(nn.Module):
    def __init__(self, in_channels, base_channels):
        super(CostRegNet, self).__init__()
        self.conv0 = Conv3d(in_channels, base_channels, padding=1)

        self.conv1 = Conv3d(base_channels, base_channels * 2, stride=2, padding=1)
        self.conv2 = Conv3d(base_channels * 2, base_channels * 2, padding=1)

        self.conv3 = Conv3d(base_channels * 2, base_channels * 4, stride=2, padding=1)
        self.conv4 = Conv3d(base_channels * 4, base_channels * 4, padding=1)

        self.conv5 = Conv3d(base_channels * 4, base_channels * 8, stride=2, padding=1)
        self.conv6 = Conv3d(base_channels * 8, base_channels * 8, padding=1)

        self.conv7 = Deconv3d(base_channels * 8, base_channels * 4, stride=2, padding=1, output_padding=1)

        self.conv9 = Deconv3d(base_channels * 4, base_channels * 2, stride=2, padding=1, output_padding=1)

        self.conv11 = Deconv3d(base_channels * 2, base_channels * 1, stride=2, padding=1, output_padding=1)

        self.prob = nn.Conv3d(base_channels, 1, 3, stride=1, padding=1, bias=False)

    def forward(self, x):
        conv0 = self.conv0(x)
        conv2 = self.conv2(self.conv1(conv0))
        conv4 = self.conv4(self.conv3(conv2))
        x = self.conv6(self.conv5(conv4))
        x = conv4 + self.conv7(x)
        x = conv2 + self.conv9(x)
        x = conv0 + self.conv11(x)
        x = self.prob(x)
        return x


class RefineNet(nn.Module):
    def __init__(self):
        super(RefineNet, self).__init__()
        self.conv1 = ConvBnReLU(4, 32)
        self.conv2 = ConvBnReLU(32, 32)
        self.conv3 = ConvBnReLU(32, 32)
        self.res = ConvBnReLU(32, 1)

    def forward(self, img, depth_init):
        concat = F.cat((img, depth_init), dim=1)
        depth_residual = self.res(self.conv3(self.conv2(self.conv1(concat))))
        depth_refined = depth_init + depth_residual
        return depth_refined


def depth_wta(p, depth_values):
    '''Winner take all.'''
    wta_index_map = torch.argmax(p, dim=1, keepdim=True).type(torch.long)
    wta_depth_map = torch.gather(depth_values, 1, wta_index_map).squeeze(1)
    return wta_depth_map


def get_cur_depth_range_samples(cur_depth, ndepth, depth_inteval_pixel, shape, max_depth=192.0, min_depth=0.0):
    cur_depth_min = (cur_depth - ndepth / 2 * depth_inteval_pixel)  # (B, H, W)
    cur_depth_max = (cur_depth + ndepth / 2 * depth_inteval_pixel)
    assert cur_depth.shape == torch.Size(shape), "cur_depth:{}, input shape:{}".format(cur_depth.shape, shape)
    new_interval = (cur_depth_max - cur_depth_min) / (ndepth - 1)  # (B, H, W)
    depth_range_samples = cur_depth_min.unsqueeze(1) + (torch.arange(0, ndepth, device=cur_depth.device,
                                                                  dtype=cur_depth.dtype,
                                                                  requires_grad=False).reshape(1, -1, 1, 1) * new_interval.unsqueeze(1))

    return depth_range_samples


def get_depth_range_samples(cur_depth, ndepth, depth_inteval_pixel, device, dtype, shape,
                           max_depth=192.0, min_depth=0.0, use_inverse_depth=False):
    if cur_depth.dim() == 2:
        cur_depth_min = cur_depth[:, 0]  # (B,)
        cur_depth_max = cur_depth[:, -1]
        if use_inverse_depth is False:
            new_interval = (cur_depth_max - cur_depth_min) / (ndepth - 1)  # (B, )  Shouldn't cal this if we use inverse depth
            depth_range_samples = cur_depth_min.unsqueeze(1) + (torch.arange(0, ndepth, device=device, dtype=dtype,
                                                                        requires_grad=False).reshape(1, -1) * new_interval.unsqueeze(1)) #(B, D)
            depth_range_samples = depth_range_samples.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, shape[1], shape[2]) #(B, D, H, W)
        else:
            # inverse sampling
            z_steps = torch.linspace(0, 1, ndepth, device=device)[None, :] # (1, D)
            depth_range_samples = 1/(1/cur_depth_min[:, None] * (1-z_steps) + 1/cur_depth_max[:, None] * z_steps) # (B, D)
            depth_range_samples = depth_range_samples.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, shape[1], shape[2]) #(B, D, H, W)
    else:
        depth_range_samples = get_cur_depth_range_samples(cur_depth, ndepth, depth_inteval_pixel, shape, max_depth, min_depth)

    return depth_range_samples


models/module.py
cvlab-stonybrook/s-volsdf
import torch.nn as nn
import torch
import math


class PositionEncodingSuperGule(nn.Module):
    def __init__(self,d_model):
        super().__init__()
        self.d_model=d_model
        self.kenc = KeypointEncoder(d_model, [32, 64])

    def forward(self,x):
        # x : N,C,H,W
        y_position = torch.ones((x.shape[2], x.shape[3])).cumsum(0).float().unsqueeze(0).to(x)
        x_position = torch.ones((x.shape[2], x.shape[3])).cumsum(1).float().unsqueeze(0).to(x)
        xy_position = torch.cat([x_position, y_position]) - 1
        xy_position = xy_position.view(2, -1).permute(1, 0).repeat(x.shape[0], 1, 1)
        xy_position_n = normalize_keypoints(xy_position, x.shape)
        ret = x + self.kenc(xy_position_n).view(x.shape)
        return ret


class PositionEncodingSine(nn.Module):
    """
    This is a sinusoidal position encoding that generalized to 2-dimensional images
    """

    def __init__(self, d_model, max_shape=(600, 600), temp_bug_fix=True):
        """
        Args:
            max_shape (tuple): for 1/4 featmap, the max length of 600 corresponds to 2400 pixels
            temp_bug_fix (bool): As noted in this [issue](https://github.com/zju3dv/LoFTR/issues/41),
                the original implementation of LoFTR includes a bug in the pos-enc impl, which has little impact
                on the final performance. For now, we keep both impls for backward compatability.
                We will remove the buggy impl after re-training all variants of our released models.
        """
        super().__init__()

        pe = torch.zeros((d_model, *max_shape))
        y_position = torch.ones(max_shape).cumsum(0).float().unsqueeze(0)
        x_position = torch.ones(max_shape).cumsum(1).float().unsqueeze(0)
        if temp_bug_fix:
            div_term = torch.exp(torch.arange(0, d_model//2, 2).float() * (-math.log(10000.0) / (d_model//2)))
        else:  # a buggy implementation (for backward compatability only)
            div_term = torch.exp(torch.arange(0, d_model//2, 2).float() * (-math.log(10000.0) / d_model//2))
        div_term = div_term[:, None, None]  # [C//4, 1, 1]
        pe[0::4, :, :] = torch.sin(x_position * div_term)
        pe[1::4, :, :] = torch.cos(x_position * div_term)
        pe[2::4, :, :] = torch.sin(y_position * div_term)
        pe[3::4, :, :] = torch.cos(y_position * div_term)

        self.register_buffer('pe', pe.unsqueeze(0), persistent=False)  # [1, C, H, W]
        # self.register_buffer('pe11', pe.unsqueeze(0))  # [1, C, H, W]

    def forward(self, x):
        """
        Args:
            x: [N, C, H, W]
        """
        return x + self.pe[:, :, :x.size(2), :x.size(3)]


def MLP(channels: list, do_bn=True):
    """ Multi-layer perceptron """
    n = len(channels)
    layers = []
    for i in range(1, n):
        layers.append(
            nn.Conv1d(channels[i - 1], channels[i], kernel_size=1, bias=True))
        if i < (n-1):
            if do_bn:
                layers.append(nn.BatchNorm1d(channels[i]))
            layers.append(nn.ReLU())
    return nn.Sequential(*layers)


def normalize_keypoints(kpts, image_shape):
    """ Normalize keypoints locations based on image image_shape"""
    _, _, height, width = image_shape
    one = kpts.new_tensor(1)
    size = torch.stack([one*width, one*height])[None]
    center = size / 2
    scaling = size.max(1, keepdim=True).values * 0.7
    return (kpts - center[:, None, :]) / scaling[:, None, :]


class KeypointEncoder(nn.Module):
    """ Joint encoding of visual appearance and location using MLPs"""
    def __init__(self, feature_dim, layers):
        super().__init__()

        self.encoder = MLP([2] + layers + [feature_dim])
        nn.init.constant_(self.encoder[-1].bias, 0.0)

    def forward(self, kpts):
        inputs = kpts.transpose(1, 2)
        return self.encoder(inputs)


models/position_encoding.py
cvlab-stonybrook/s-volsdf
import torch
import torch.nn as nn
import torch.nn.functional as F

eps = 1e-12

def homo_warping(src_fea, src_proj, ref_proj, depth_values):
    # src_fea: [B, C, H, W]
    # src_proj: [B, 4, 4]
    # ref_proj: [B, 4, 4]
    # depth_values: [B, Ndepth] o [B, Ndepth, H, W]
    # out: [B, C, Ndepth, H, W]
    batch, channels = src_fea.shape[0], src_fea.shape[1]
    num_depth = depth_values.shape[1]
    height, width = src_fea.shape[2], src_fea.shape[3]

    with torch.no_grad():
        proj = torch.matmul(src_proj, torch.inverse(ref_proj))
        rot = proj[:, :3, :3]  # [B,3,3]
        trans = proj[:, :3, 3:4]  # [B,3,1]

        y, x = torch.meshgrid([torch.arange(0, height, dtype=torch.float32, device=src_fea.device),
                               torch.arange(0, width, dtype=torch.float32, device=src_fea.device)])
        y, x = y.contiguous(), x.contiguous()
        y, x = y.view(height * width), x.view(height * width)
        xyz = torch.stack((x, y, torch.ones_like(x)))  # [3, H*W]
        xyz = torch.unsqueeze(xyz, 0).repeat(batch, 1, 1)  # [B, 3, H*W]
        rot_xyz = torch.matmul(rot, xyz)  # [B, 3, H*W]
        rot_depth_xyz = rot_xyz.unsqueeze(2).repeat(1, 1, num_depth, 1) * depth_values.view(batch, 1, num_depth,
                                                                                            -1)  # [B, 3, Ndepth, H*W]
        proj_xyz = rot_depth_xyz + trans.view(batch, 3, 1, 1)  # [B, 3, Ndepth, H*W]
        proj_xy = proj_xyz[:, :2, :, :] / proj_xyz[:, 2:3, :, :]  # [B, 2, Ndepth, H*W]
        proj_x_normalized = proj_xy[:, 0, :, :] / ((width - 1) / 2) - 1
        proj_y_normalized = proj_xy[:, 1, :, :] / ((height - 1) / 2) - 1
        proj_xy = torch.stack((proj_x_normalized, proj_y_normalized), dim=3)  # [B, Ndepth, H*W, 2]
        grid = proj_xy

    warped_src_fea = F.grid_sample(src_fea, grid.view(batch, num_depth * height, width, 2), mode='bilinear',
                                   padding_mode='zeros', align_corners = False)
    warped_src_fea = warped_src_fea.view(batch, channels, num_depth, height, width)

    return warped_src_fea

def uncertainty_aware_samples(cur_depth, exp_var, ndepth, device, dtype, shape, inverse_depth=False):
    if cur_depth.dim() == 2:
        #must be the first stage
        cur_depth_min = cur_depth[:, 0]  # (B,)
        cur_depth_max = cur_depth[:, -1]
        if inverse_depth:
            # inverse sampling
            z_steps = torch.linspace(0, 1, ndepth, device=device)[None, :] # (1, D)
            depth_range_samples = 1/(1/cur_depth_min[:, None] * (1-z_steps) + 1/cur_depth_max[:, None] * z_steps) # (B, D)
        else:
            new_interval = (cur_depth_max - cur_depth_min) / (ndepth - 1)  # (B, )
            depth_range_samples = cur_depth_min.unsqueeze(1) + (torch.arange(0, ndepth, device=device, dtype=dtype,
                                                                        requires_grad=False).reshape(1, -1) * new_interval.unsqueeze(1)) # (B, D)
        depth_range_samples = depth_range_samples.unsqueeze(-1).unsqueeze(-1).repeat(1, 1, shape[1], shape[2]) # (B, D, H, W)
    else:
        low_bound = -torch.min(cur_depth, exp_var)
        high_bound = exp_var

        # assert exp_var.min() >= 0, exp_var.min()
        assert ndepth > 1

        step = (high_bound - low_bound) / (float(ndepth) - 1)
        new_samps = []
        for i in range(int(ndepth)):
            new_samps.append(cur_depth + low_bound + step * i + eps)

        depth_range_samples = torch.cat(new_samps, 1)
        # assert depth_range_samples.min() >= 0, depth_range_samples.min()
    return depth_range_samples

def depth_regression(p, depth_values):
    if depth_values.dim() <= 2:
        # print("regression dim <= 2")
        depth_values = depth_values.view(*depth_values.shape, 1, 1)
    depth = torch.sum(p * depth_values, 1)
    return depth

class Conv2dUnit(nn.Module):
    """Applies a 2D convolution (optionally with batch normalization and relu activation)
    over an input signal composed of several input planes.

    Attributes:
        conv (nn.Module): convolution module
        bn (nn.Module): batch normalization module
        relu (bool): whether to activate by relu

    Notes:
        Default momentum for batch normalization is set to be 0.01,

    """

    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 relu=True, bn=True, bn_momentum=0.1, **kwargs):
        super(Conv2dUnit, self).__init__()

        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride,
                              bias=(not bn), **kwargs)
        self.kernel_size = kernel_size
        self.stride = stride
        self.bn = nn.BatchNorm2d(out_channels, momentum=bn_momentum) if bn else None
        self.relu = relu

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu:
            x = F.relu(x, inplace=True)
        return x

class Deconv2dUnit(nn.Module):
    """Applies a 2D deconvolution (optionally with batch normalization and relu activation)
       over an input signal composed of several input planes.

       Attributes:
           conv (nn.Module): convolution module
           bn (nn.Module): batch normalization module
           relu (bool): whether to activate by relu

       Notes:
           Default momentum for batch normalization is set to be 0.01,

       """

    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 relu=True, bn=True, bn_momentum=0.1, **kwargs):
        super(Deconv2dUnit, self).__init__()
        self.out_channels = out_channels
        assert stride in [1, 2]
        self.stride = stride

        self.conv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride=stride,
                                       bias=(not bn), **kwargs)
        self.bn = nn.BatchNorm2d(out_channels, momentum=bn_momentum) if bn else None
        self.relu = relu

    def forward(self, x):
        y = self.conv(x)
        if self.stride == 2:
            h, w = list(x.size())[2:]
            y = y[:, :, :2 * h, :2 * w].contiguous()
        if self.bn is not None:
            x = self.bn(y)
        if self.relu:
            x = F.relu(x, inplace=True)
        return x

class Conv3dUnit(nn.Module):
    """Applies a 3D convolution (optionally with batch normalization and relu activation)
    over an input signal composed of several input planes.

    Attributes:
        conv (nn.Module): convolution module
        bn (nn.Module): batch normalization module
        relu (bool): whether to activate by relu

    Notes:
        Default momentum for batch normalization is set to be 0.01,

    """

    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1,
                 relu=True, bn=True, bn_momentum=0.1, init_method="xavier", **kwargs):
        super(Conv3dUnit, self).__init__()
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        assert stride in [1, 2]
        self.stride = stride

        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride=stride,
                              bias=(not bn), **kwargs)
        self.bn = nn.BatchNorm3d(out_channels, momentum=bn_momentum) if bn else None
        self.relu = relu

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu:
            x = F.relu(x, inplace=True)
        return x

class Deconv3dUnit(nn.Module):
    """Applies a 3D deconvolution (optionally with batch normalization and relu activation)
       over an input signal composed of several input planes.

       Attributes:
           conv (nn.Module): convolution module
           bn (nn.Module): batch normalization module
           relu (bool): whether to activate by relu

       Notes:
           Default momentum for batch normalization is set to be 0.01,

       """

    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1,
                 relu=True, bn=True, bn_momentum=0.1, init_method="xavier", **kwargs):
        super(Deconv3dUnit, self).__init__()
        self.out_channels = out_channels
        assert stride in [1, 2]
        self.stride = stride

        self.conv = nn.ConvTranspose3d(in_channels, out_channels, kernel_size, stride=stride,
                                       bias=(not bn), **kwargs)
        self.bn = nn.BatchNorm3d(out_channels, momentum=bn_momentum) if bn else None
        self.relu = relu

    def forward(self, x):
        y = self.conv(x)
        if self.bn is not None:
            x = self.bn(y)
        if self.relu:
            x = F.relu(x, inplace=True)
        return x

class Deconv2dBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, relu=True, bn=True,
                 bn_momentum=0.1):
        super(Deconv2dBlock, self).__init__()

        self.deconv = Deconv2dUnit(in_channels, out_channels, kernel_size, stride=2, padding=1, output_padding=1,
                                   bn=True, relu=relu, bn_momentum=bn_momentum)

        self.conv = Conv2dUnit(2 * out_channels, out_channels, kernel_size, stride=1, padding=1,
                               bn=bn, relu=relu, bn_momentum=bn_momentum)

    def forward(self, x_pre, x):
        x = self.deconv(x)
        x = torch.cat((x, x_pre), dim=1)
        x = self.conv(x)
        return x

class FeatExtNet(nn.Module):
    def __init__(self, base_channels, num_stage=3,):
        super(FeatExtNet, self).__init__()

        self.base_channels = base_channels
        self.num_stage = num_stage

        self.conv0 = nn.Sequential(
            Conv2dUnit(3, base_channels, 3, 1, padding=1),
            Conv2dUnit(base_channels, base_channels, 3, 1, padding=1),
        )

        self.conv1 = nn.Sequential(
            Conv2dUnit(base_channels, base_channels * 2, 5, stride=2, padding=2),
            Conv2dUnit(base_channels * 2, base_channels * 2, 3, 1, padding=1),
            Conv2dUnit(base_channels * 2, base_channels * 2, 3, 1, padding=1),
        )

        self.conv2 = nn.Sequential(
            Conv2dUnit(base_channels * 2, base_channels * 4, 5, stride=2, padding=2),
            Conv2dUnit(base_channels * 4, base_channels * 4, 3, 1, padding=1),
            Conv2dUnit(base_channels * 4, base_channels * 4, 3, 1, padding=1),
        )

        self.out1 = nn.Conv2d(base_channels * 4, base_channels * 4, 1, bias=False)
        self.out_channels = [4 * base_channels]

        if num_stage == 3:
            self.deconv1 = Deconv2dBlock(base_channels * 4, base_channels * 2, 3)
            self.deconv2 = Deconv2dBlock(base_channels * 2, base_channels, 3)

            self.out2 = nn.Conv2d(base_channels * 2, base_channels * 2, 1, bias=False)
            self.out3 = nn.Conv2d(base_channels, base_channels, 1, bias=False)
            self.out_channels.append(2 * base_channels)
            self.out_channels.append(base_channels)

        elif num_stage == 2:
            self.deconv1 = Deconv2dBlock(base_channels * 4, base_channels * 2, 3)

            self.out2 = nn.Conv2d(base_channels * 2, base_channels * 2, 1, bias=False)
            self.out_channels.append(2 * base_channels)

    def forward(self, x):
        conv0 = self.conv0(x)
        conv1 = self.conv1(conv0)
        conv2 = self.conv2(conv1)
        intra_feat = conv2
        outputs = {}
        out = self.out1(intra_feat)

        outputs["stage1"] = out
        if self.num_stage == 3:
            intra_feat = self.deconv1(conv1, intra_feat)
            out = self.out2(intra_feat)
            outputs["stage2"] = out

            intra_feat = self.deconv2(conv0, intra_feat)
            out = self.out3(intra_feat)
            outputs["stage3"] = out

        elif self.num_stage == 2:
            intra_feat = self.deconv1(conv1, intra_feat)
            out = self.out2(intra_feat)
            outputs["stage2"] = out

        return outputs

class CostRegNet(nn.Module):
    def __init__(self, in_channels, base_channels):
        super(CostRegNet, self).__init__()
        self.conv0 = Conv3dUnit(in_channels, base_channels, padding=1)

        self.conv1 = Conv3dUnit(base_channels, base_channels * 2, stride=2, padding=1)
        self.conv2 = Conv3dUnit(base_channels * 2, base_channels * 2, padding=1)

        self.conv3 = Conv3dUnit(base_channels * 2, base_channels * 4, stride=2, padding=1)
        self.conv4 = Conv3dUnit(base_channels * 4, base_channels * 4, padding=1)

        self.conv5 = Conv3dUnit(base_channels * 4, base_channels * 8, stride=2, padding=1)
        self.conv6 = Conv3dUnit(base_channels * 8, base_channels * 8, padding=1)

        self.deconv7 = Deconv3dUnit(base_channels * 8, base_channels * 4, stride=2, padding=1, output_padding=1)

        self.deconv8 = Deconv3dUnit(base_channels * 4, base_channels * 2, stride=2, padding=1, output_padding=1)

        self.deconv9 = Deconv3dUnit(base_channels * 2, base_channels * 1, stride=2, padding=1, output_padding=1)

        self.prob = nn.Conv3d(base_channels, 1, 3, stride=1, padding=1, bias=False)

    def forward(self, x):
        conv0 = self.conv0(x)
        conv2 = self.conv2(self.conv1(conv0))
        conv4 = self.conv4(self.conv3(conv2))
        x = self.conv6(self.conv5(conv4))
        x = conv4 + self.deconv7(x)
        x = conv2 + self.deconv8(x)
        x = conv0 + self.deconv9(x)
        x = self.prob(x)
        return x


def compute_depth(feats, proj_mats, depth_samps, cost_reg, lamb, is_training=False):
    '''

    :param feats: [(B, C, H, W), ] * num_views
    :param proj_mats: [()]
    :param depth_samps:
    :param cost_reg:
    :param lamb:
    :return:
    '''

    proj_mats = torch.unbind(proj_mats, 1)
    num_views = len(feats)
    num_depth = depth_samps.shape[1]

    assert len(proj_mats) == num_views, "Different number of images and projection matrices"

    ref_feat, src_feats = feats[0], feats[1:]
    ref_proj, src_projs = proj_mats[0], proj_mats[1:]

    ref_volume = ref_feat.unsqueeze(2).repeat(1, 1, num_depth, 1, 1)
    volume_sum = ref_volume
    volume_sq_sum = ref_volume ** 2
    del ref_volume

    #todo optimize impl
    for src_fea, src_proj in zip(src_feats, src_projs):
        src_proj_new = src_proj[:, 0].clone()
        src_proj_new[:, :3, :4] = torch.matmul(src_proj[:, 1, :3, :3], src_proj[:, 0, :3, :4])

        ref_proj_new = ref_proj[:, 0].clone()
        ref_proj_new[:, :3, :4] = torch.matmul(ref_proj[:, 1, :3, :3], ref_proj[:, 0, :3, :4])
        warped_volume = homo_warping(src_fea, src_proj_new, ref_proj_new, depth_samps)

        if is_training:
            volume_sum = volume_sum + warped_volume
            volume_sq_sum = volume_sq_sum + warped_volume ** 2
        else:
            volume_sum += warped_volume
            volume_sq_sum += warped_volume.pow_(2) #in_place method
        del warped_volume
    volume_variance = volume_sq_sum.div_(num_views).sub_(volume_sum.div_(num_views).pow_(2))

    prob_volume_pre = cost_reg(volume_variance).squeeze(1)
    prob_volume = F.softmax(prob_volume_pre, dim=1)
    depth = depth_regression(prob_volume, depth_values=depth_samps)

    with torch.no_grad():
        prob_volume_sum4 = 4 * F.avg_pool3d(F.pad(prob_volume.unsqueeze(1), pad=(0, 0, 0, 0, 1, 2)), (4, 1, 1),
                                            stride=1, padding=0).squeeze(1)
        depth_index = depth_regression(prob_volume, depth_values=torch.arange(num_depth, device=prob_volume.device,
                                                                              dtype=torch.float)).long()
        depth_index = depth_index.clamp(min=0, max=num_depth - 1)
        prob_conf = torch.gather(prob_volume_sum4, 1, depth_index.unsqueeze(1)).squeeze(1)

    samp_variance = (depth_samps - depth.unsqueeze(1)) ** 2
    exp_variance = lamb * torch.sum(samp_variance * prob_volume, dim=1, keepdim=False) ** 0.5

    return {"depth": depth, "photometric_confidence": prob_conf, "prob_volume": prob_volume, 'variance': exp_variance, "depth_values": depth_samps}

class UCSNet(nn.Module):
    def __init__(self, lamb=1.5, stage_configs=[64, 32, 8], grad_method="detach", base_chs=[8, 8, 8], feat_ext_ch=8):
        super(UCSNet, self).__init__()

        self.stage_configs = stage_configs
        self.grad_method = grad_method
        self.base_chs = base_chs
        self.lamb = lamb
        self.num_stage = len(stage_configs)
        self.ds_ratio = {"stage1": 4.0,
                         "stage2": 2.0,
                         "stage3": 1.0
                         }

        self.feature_extraction = FeatExtNet(base_channels=feat_ext_ch, num_stage=self.num_stage,)

        self.cost_regularization = nn.ModuleList([CostRegNet(in_channels=self.feature_extraction.out_channels[i],
                                                             base_channels=self.base_chs[i]) for i in range(self.num_stage)])
    
    @torch.no_grad()
    def forward(self, stage_idx, sample_cuda,
            features, extra, outputs, int_r, depth=None, prevent_oom=False, inverse_depth=False):

        exp_var = extra
        imgs, proj_matrices, depth_values = sample_cuda["imgs"], sample_cuda["proj_matrices"], sample_cuda["depth_values"]
        
        # step 0.
        if depth is None:
            depth = outputs['depth'] if stage_idx > 0 else None
        outputs = {} if outputs is None else outputs
        depth_min = float(depth_values[0, 0].cpu().numpy())
        depth_max = float(depth_values[0, -1].cpu().numpy())
        depth_interval = (depth_max - depth_min) / depth_values.size(1)

        # step 1. feature extraction
        img = imgs[:, imgs.size(1)-1]

        # print("*********************stage{}*********************".format(stage_idx + 1))
        #stage feature, proj_mats, scales
        features_stage = [feat["stage{}".format(stage_idx + 1)] for feat in features]
        proj_matrices_stage = proj_matrices["stage{}".format(stage_idx + 1)]
        stage_scale = self.ds_ratio["stage{}".format(stage_idx + 1)]
        cur_h = img.shape[2] // int(stage_scale)
        cur_w = img.shape[3] // int(stage_scale)

        if depth is not None:
            if self.grad_method == "detach":
                cur_depth = depth.detach()
                exp_var = exp_var.detach()
            else:
                cur_depth = depth

            cur_depth = F.interpolate(cur_depth.unsqueeze(1),
                                            [cur_h, cur_w], mode='bilinear')
            exp_var = F.interpolate(exp_var.unsqueeze(1), [cur_h, cur_w], mode='bilinear')

        else:
            cur_depth = depth_values

        depth_range_samples = uncertainty_aware_samples(cur_depth=cur_depth,
                                                        exp_var=exp_var,
                                                        ndepth=self.stage_configs[stage_idx],
                                                        dtype=img[0].dtype,
                                                        device=img[0].device,
                                                        shape=[img.shape[0], cur_h, cur_w],
                                                        inverse_depth=inverse_depth)

        outputs_stage = compute_depth(features_stage, proj_matrices_stage,
                                        depth_samps=depth_range_samples,
                                        cost_reg=self.cost_regularization[stage_idx],
                                        lamb=self.lamb,
                                        is_training=self.training)

        depth = outputs_stage['depth']
        exp_var = outputs_stage['variance']

        outputs["stage{}".format(stage_idx + 1)] = outputs_stage
        outputs.update(outputs_stage)

        return outputs, exp_var
models/ucsnet.py
cvlab-stonybrook/s-volsdf
#--------------------------args & GPU-------------------------------------#
from omegaconf import DictConfig, OmegaConf
import hydra
from helpers.help import logger
from helpers.help import run_help

# init args
@hydra.main(version_base=None, config_path="config", config_name="ours")
def get_config(cfg: DictConfig):
    OmegaConf.set_struct(cfg, False)
    global args
    args = cfg
get_config()

# update args & set GPU
args = run_help(args)
#------------------------------------------------------------------------#

import gc, cv2, signal
import os
import time
import numpy as np
import torch
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from plyfile import PlyData, PlyElement
from multiprocessing import Pool
from functools import partial
from pathlib import Path
from skimage.morphology import binary_dilation, disk
from models.TransMVSNet import TransMVSNet
from models.CasMVSNet import CascadeMVSNet
from models.ucsnet import UCSNet
from helpers.utils import *
from datasets.data_io import read_pfm, save_pfm
from datasets.general_eval import MVSDataset
from volsdf.datasets.scene_dataset import get_trains_ids, get_eval_ids
from volsdf.vsdf import VolOpt

cudnn.benchmark = True
torch.set_printoptions(sci_mode=False, precision=4)
np.set_printoptions(suppress=True, precision=4)

def save_depth(testlist):
    for scene in testlist:
        logger.info(f"parameter adjust - {scene}")
        _sparse_weight, _inverse_depth = args.vol.loss.sparse_weight, args.inverse_depth

        # sparsity regularization, slightly better performance with following settings
        if args.vol.dataset.data_dir == 'DTU':
            if scene in ['scan37']:
                args.vol.loss.sparse_weight = 0.1
            elif scene in ['scan24']:
                args.vol.loss.sparse_weight = 0
        elif args.vol.dataset.data_dir == 'BlendedMVS':
            if scene in ['scan2', 'scan3', 'scan7', 'scan9']:
                args.vol.loss.sparse_weight = 0

        # unbounded scenes, use inverse depth sampling
        if args.vol.dataset.data_dir == 'BlendedMVS':
            if scene in ['scan1', 'scan2', 'scan5', 'scan6', 'scan8', 'scan9']:
                args.inverse_depth = True
                logger.info("    inverse_D=[True,False,False]")

        logger.info(f"    inverse depth={args.inverse_depth} sparse_weight={args.vol.loss.sparse_weight}")

        save_scene_depth(scene)

        args.vol.loss.sparse_weight, args.inverse_depth = _sparse_weight, _inverse_depth

# prepare data for image-based rendering
def create_scene(scene):
    os.makedirs(os.path.join(args.outdir, scene), exist_ok=True)

    # dataset, dataloader
    trains_i = get_trains_ids(args.vol.dataset.data_dir, scene, args.num_view)
    evals_i = get_eval_ids(args.vol.dataset.data_dir, int(scene[4:])) # only save cam, no need for images
    trains_i += evals_i
    mvs_datapath = os.path.join(args.data_dir_root, args.vol.dataset.data_dir, 'mvs_data')
    test_dataset = MVSDataset(mvs_datapath, [scene], "test", args.num_view, args.vol.dataset.data_dir, args.numdepth, args.interval_scale,
                            max_h=args.max_h, max_w=args.max_w, trains_i=trains_i, args=args)
    TestImgLoader = DataLoader(test_dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, drop_last=False)

    for batch_idx, sample in enumerate(TestImgLoader):
        filenames = sample["filename"]
        cams = sample["proj_matrices"]["stage{}".format(3)].numpy()
        imgs = sample["imgs"].numpy()
        cam_near_far = sample["cam_near_far"].numpy()[0]
        del sample

        # save depth maps and confidence maps
        for filename, cam, img in zip(filenames, cams, imgs):
            img = img[0]  #ref view
            cam = cam[0]  #ref cam
            id_ = int(filename.split('/')[-1][:8])
            cam_filename = os.path.join(args.outdir, filename.format('cams', '_cam.txt'))
            img_filename = os.path.join(args.outdir, filename.format('images', '.png'))
            os.makedirs(cam_filename.rsplit('/', 1)[0], exist_ok=True)
            os.makedirs(img_filename.rsplit('/', 1)[0], exist_ok=True)
            # cam
            write_cam(cam_filename, cam, cam_near_far)
            # img
            if id_ not in evals_i:
                img = np.clip(np.transpose(img, (1, 2, 0)) * 255, 0, 255).astype(np.uint8)
                img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
                cv2.imwrite(img_filename, img_bgr)

# run CasMVS model to save depth maps and confidence maps
def save_scene_depth(scene):
    # save args for individual scan (scene)
    os.makedirs(os.path.join(args.outdir, scene), exist_ok=True)
    with open(os.path.join(args.outdir, scene, 'args.yaml'), "w") as f:
        OmegaConf.save(args, f)

    # dataset, dataloader
    trains_i = get_trains_ids(args.vol.dataset.data_dir, scene, args.num_view)
    mvs_datapath = os.path.join(args.data_dir_root, args.vol.dataset.data_dir, 'mvs_data')
    test_dataset = MVSDataset(mvs_datapath, [scene], "test", args.num_view, args.vol.dataset.data_dir, args.numdepth, args.interval_scale,
                            max_h=args.max_h, max_w=args.max_w, trains_i=trains_i, args=args)
    TestImgLoader = DataLoader(test_dataset, args.batch_size, shuffle=False, num_workers=args.num_workers, drop_last=False)

    # Let's begin
    torch.cuda.empty_cache()

    # MVS model
    if args.mvs_model_name == 'casmvsnet':
        mvs_model_ckpt_path = os.path.join(args.data_dir_root, 'mvs_models', 'casmvsnet.ckpt')
        model = CascadeMVSNet(refine=False, ndepths=[int(nd) for nd in args.ndepths.split(",") if nd],
                              depth_interals_ratio=[float(d_i) for d_i in args.depth_inter_r.split(",") if d_i],
                              share_cr=args.share_cr,
                              cr_base_chs=[int(ch) for ch in args.cr_base_chs.split(",") if ch],
                              grad_method=args.grad_method)
    elif args.mvs_model_name == 'ucsnet':
        mvs_model_ckpt_path = os.path.join(args.data_dir_root, 'mvs_models', 'ucsnet.ckpt')
        model = UCSNet(stage_configs=[int(nd) for nd in args.ndepths.split(",") if nd], lamb=1.5)
    elif args.mvs_model_name == 'transmvsnet':
        mvs_model_ckpt_path = os.path.join(args.data_dir_root, 'mvs_models', 'model_dtu.ckpt')
        model = TransMVSNet(refine=False, ndepths=[int(nd) for nd in args.ndepths.split(",") if nd],
                            depth_interals_ratio=[float(d_i) for d_i in args.depth_inter_r.split(",") if d_i],
                            share_cr=args.share_cr,
                            cr_base_chs=[int(ch) for ch in args.cr_base_chs.split(",") if ch],
                            grad_method=args.grad_method)
    else:
        mvs_model_ckpt_path = None
        raise NotImplementedError

    # load MVS model
    logger.info("loading model {}".format(args.mvs_model_name))
    state_dict = torch.load(mvs_model_ckpt_path, map_location=torch.device("cpu"))
    model.load_state_dict(state_dict['model'], strict=True)
    model.cuda()
    model.eval()

    # set some args for easy use later
    use_nerf_d = [d_i for d_i in args.use_nerf_d]
    opt_stepNs = [d_i for d_i in args.opt_stepNs]
    ndepths = [int(nd) for nd in args.ndepths.split(",") if nd]

    logger.debug(", ".join([f"{id_}: {sample['filename'][0]}" for id_, sample in enumerate(TestImgLoader)]))

    # volume optimizer
    vol_opt = VolOpt(args=args,
                    batch_size=1,
                    is_continue=args.get('is_continue', False),
                    timestamp='latest',
                    checkpoint='latest',
                    scan=scene)
    vol_opt.trains_i = trains_i
    assert vol_opt.trains_i == vol_opt.train_dataset.trains_ids()

    #-----------------------------------------------------------#
    img_n = len(TestImgLoader)
    view_extra_samples, outs_samples = [None,]*img_n, [None,]*img_n
    depths, confis = [None,]*img_n, [None,]*img_n

    for stage_idx in range(3):
        int_r = None if args.mvs_model_name =='ucsnet' else model.depth_interals_ratio[stage_idx]

        # (a) ---------- cost volume ----------
        start_time = time.time()
        outs, view_extras = [None,]*img_n, [None,]*img_n
        for i, sample in enumerate(TestImgLoader):
            sample_cuda = tocuda(sample)
            with torch.no_grad():
                # step 1. feature extraction
                imgs = sample_cuda["imgs"]
                features = []
                for nview_idx in range(imgs.size(1)): #imgs shape (B, N, C, H, W)
                    img = imgs[:, nview_idx]
                    if args.mvs_model_name =='ucsnet':
                        features.append(model.feature_extraction(img))
                    else:
                        features.append(model.feature(img))
                if args.mvs_model_name == 'transmvsnet':
                    features = model.FMT_with_pathway(features)
                del imgs
                torch.cuda.empty_cache()
                # step 2. cost volume
                outs[i], view_extras[i] = model(stage_idx, 
                    sample_cuda, features=features, extra=view_extra_samples[i], outputs=outs_samples[i],
                    int_r=int_r, prevent_oom=args.prevent_oom, inverse_depth=args.inverse_depth)
                if args.prevent_oom:
                    outs[i], view_extras[i] = tocpu(outs[i]), tocpu(view_extras[i])
                del features
                torch.cuda.empty_cache()

        if args.prevent_oom:
            outs, view_extras = tocuda(outs), tocuda(view_extras)
        torch.cuda.empty_cache()
        logger.debug(f"time(gen cost volume)={time.time()-start_time:.2f}")
        if args.ablate:
            for i in range(img_n):
                outs_samples[i], view_extra_samples[i] = outs[i], view_extras[i]
            continue

        # (b) ----------Volume optimization ----------
        do_volopt = opt_stepNs[stage_idx] > 0 and use_nerf_d[stage_idx] > 0
        if do_volopt:
            # create multi-scale P - data loader
            vol_opt.gen_dataset(stage_idx)
            vol_opt.stg = stage_idx
            vol_opt.loss.set_stg(stage_idx)
            # mvs inputs: P, depth_range
            vol_opt.get_mvs_input(outs)

            # noise-tolerant optimization
            epoch = 0
            if opt_stepNs[stage_idx] > 1:
                epoch = vol_opt.run(opt_stepNs[stage_idx])

            # Vol render depth
            logger.info(f"render volsdf at {vol_opt.plots_dir} ..")
            for i, id_k in enumerate(trains_i):
                depths[i], confis[i] = vol_opt.render_mvs(id_k, epoch)

            # next stage depth
            logger.info(f"mvs_depth replaced by vol_depth at stg={stage_idx} in 0,1,2")
            for i, sample in enumerate(TestImgLoader):
                # assert use_nerf_d[stage_idx] == 1
                outs[i][f"stage{stage_idx+1}"]['depth'] = depths[i]
                outs[i]['depth'] = depths[i]
        
        # update result
        for i in range(img_n):
            outs_samples[i], view_extra_samples[i] = outs[i], view_extras[i]
        del outs, view_extras, sample_cuda
        torch.cuda.empty_cache()
    #-----------------------------------------------------------#
    for batch_idx, sample in enumerate(TestImgLoader):
        outputs = outs_samples[batch_idx]
        dep_min, dep_max = sample["depth_values"].numpy().min(), sample["depth_values"].numpy().max()
        outputs = tensor2numpy(outputs)
        filenames = sample["filename"]
        cams = sample["proj_matrices"]["stage{}".format(3)].numpy()
        imgs = sample["imgs"].numpy()
        del sample

        # save depth maps and confidence maps
        for filename, cam, img, depth_est, photometric_confidence, conf_1, conf_2 in zip(filenames, cams, imgs, \
                                                        outputs["depth"], outputs["photometric_confidence"], outputs['stage1']["photometric_confidence"], outputs['stage2']["photometric_confidence"]):
            #ref view
            img = img[0]
            cam = cam[0]

            H,W = depth_est.shape
            conf_1 = cv2.resize(conf_1, (W,H))
            conf_2 = cv2.resize(conf_2, (W,H))
            photometric_confidence = cv2.resize(photometric_confidence, (W,H))
            conf_final = conf_1 * conf_2 * photometric_confidence

            depth_filename = os.path.join(args.outdir, filename.format('depth_est', '.pfm'))
            confidence_filename = os.path.join(args.outdir, filename.format('confidence', '.pfm'))
            cam_filename = os.path.join(args.outdir, filename.format('cams', '_cam.txt'))
            img_filename = os.path.join(args.outdir, filename.format('images', '.jpg'))
            os.makedirs(depth_filename.rsplit('/', 1)[0], exist_ok=True)
            os.makedirs(confidence_filename.rsplit('/', 1)[0], exist_ok=True)
            os.makedirs(cam_filename.rsplit('/', 1)[0], exist_ok=True)
            os.makedirs(img_filename.rsplit('/', 1)[0], exist_ok=True)
            #save depth maps
            save_pfm(depth_filename, depth_est) # (1152, 1536) float32
            _depth_min, _depth_max = np.quantile(depth_est, 0.01), dep_max
            depth_color = visualize_depth(depth_est, depth_min=_depth_min, depth_max=_depth_max)
            cv2.imwrite(os.path.join(args.outdir, filename.format('depth_est', '.png')), depth_color)
            cv2.imwrite(os.path.join(args.outdir, filename.format('depth_est', '_1.png')), visualize_depth(outputs['stage1']["depth"][0], depth_min=_depth_min, depth_max=_depth_max))
            cv2.imwrite(os.path.join(args.outdir, filename.format('depth_est', '_2.png')), visualize_depth(outputs['stage2']["depth"][0], depth_min=_depth_min, depth_max=_depth_max))
            #save confidence maps
            save_pfm(confidence_filename, conf_final)
            cv2.imwrite(os.path.join(args.outdir, filename.format('confidence', '_final.png')),visualize_depth(conf_final, direct=True))
            #save cams, img
            write_cam(cam_filename, cam)
            img = np.clip(np.transpose(img, (1, 2, 0)) * 255, 0, 255).astype(np.uint8)
            img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
            cv2.imwrite(img_filename, img_bgr)

    del outs_samples
    torch.cuda.empty_cache()
    gc.collect()

def filter_depth(scan_folder, out_folder, plyfilename):
    trains_i = get_trains_ids(args.vol.dataset.data_dir, Path(scan_folder).name, args.num_view)
    pair_data = [
        (trains_i[i], [x for x in trains_i if x != trains_i[i]]) for i in range(len(trains_i))
    ]
    logger.debug(pair_data)

    vertexs, vertex_colors = [], []
    # for each reference view and the corresponding source views
    for ref_view, src_views in pair_data:
        # load the camera parameters
        ref_intrinsics, ref_extrinsics = read_camera_parameters(
            os.path.join(scan_folder, 'cams/{:0>8}_cam.txt'.format(ref_view)))
        # load the reference image
        ref_img = read_img(os.path.join(scan_folder, 'images/{:0>8}.jpg'.format(ref_view)))
        # load the estimated depth of the reference view
        ref_depth_est = read_pfm(os.path.join(out_folder, 'depth_est/{:0>8}.pfm'.format(ref_view)))[0]
        # load the photometric mask of the reference view
        confidence = read_pfm(os.path.join(out_folder, 'confidence/{:0>8}.pfm'.format(ref_view)))[0]
        assert ref_depth_est.shape == ref_img.shape[:2]
        photo_mask = confidence > args.conf

        all_srcview_depth_ests = []
        all_srcview_x = []
        all_srcview_y = []
        all_srcview_geomask = []

        # compute the geometric mask
        geo_mask_sum = 0
        for src_view in src_views:
            # camera parameters of the source view
            src_intrinsics, src_extrinsics = read_camera_parameters(
                os.path.join(scan_folder, 'cams/{:0>8}_cam.txt'.format(src_view)))
            # the estimated depth of the source view
            src_depth_est = read_pfm(os.path.join(out_folder, 'depth_est/{:0>8}.pfm'.format(src_view)))[0]
            assert src_depth_est.shape == ref_img.shape[:2]
            geo_mask, depth_reprojected, x2d_src, y2d_src = check_geometric_consistency(ref_depth_est, ref_intrinsics, ref_extrinsics,
                                                                      src_depth_est,
                                                                      src_intrinsics, src_extrinsics,
                                                                      args.filter_dist, args.filter_diff)
            geo_mask_sum += geo_mask.astype(np.int32)
            all_srcview_depth_ests.append(depth_reprojected)
            all_srcview_x.append(x2d_src)
            all_srcview_y.append(y2d_src)
            all_srcview_geomask.append(geo_mask)

        depth_est_averaged = (sum(all_srcview_depth_ests) + ref_depth_est) / (geo_mask_sum + 1)
        geo_mask = geo_mask_sum >= args.thres_view
        final_mask = np.logical_and(photo_mask, geo_mask)
        if args.eval_mask:
            eval_mask_dir = os.path.join(args.data_dir_root, args.vol.dataset.data_dir, 'eval_mask')
            if args.vol.dataset.data_dir == 'BlendedMVS':
                eval_mask_path = os.path.join(eval_mask_dir, out_folder.split('/')[-1], "mask/{:0>8}.png".format(ref_view))
            elif args.vol.dataset.data_dir == 'DTU':
                eval_mask_path = os.path.join(eval_mask_dir, out_folder.split('/')[-1], "mask/{:0>3}.png".format(ref_view))
                if not os.path.exists(eval_mask_path):  
                    eval_mask_path = os.path.join(eval_mask_dir, out_folder.split('/')[-1], "{:0>3}.png".format(ref_view))
            else:
                raise NotImplementedError
            assert os.path.exists(eval_mask_path)

            eval_mask = read_img(eval_mask_path)
            if len(eval_mask.shape) == 3:
                eval_mask = eval_mask[:, :, -1]
            eval_mask = binary_dilation(eval_mask, disk(12))
            eval_mask = eval_mask * 1.
            eval_mask = cv2.resize(eval_mask, geo_mask.shape[::-1])
            final_mask = np.logical_and(final_mask, eval_mask > 0.)

        os.makedirs(os.path.join(out_folder, "mask"), exist_ok=True)
        save_mask(os.path.join(out_folder, "mask/{:0>8}_photo.png".format(ref_view)), photo_mask)
        save_mask(os.path.join(out_folder, "mask/{:0>8}_geo.png".format(ref_view)), geo_mask)
        save_mask(os.path.join(out_folder, "mask/{:0>8}_final.png".format(ref_view)), final_mask)

        logger.info("processing {}, ref-view{:0>2}, photo/geo/final-mask:{:.3f}/{:.3f}/{:.3f}".format(scan_folder, ref_view,
                                                                                    photo_mask.mean(),
                                                                                    geo_mask.mean(), final_mask.mean()))

        height, width = depth_est_averaged.shape[:2]
        x, y = np.meshgrid(np.arange(0, width), np.arange(0, height))
        valid_points = final_mask
        logger.debug(f"valid_points {valid_points.mean():.3f}")
        x, y, depth = x[valid_points], y[valid_points], depth_est_averaged[valid_points]
        color = ref_img[valid_points]
        xyz_ref = np.matmul(np.linalg.inv(ref_intrinsics),
                            np.vstack((x, y, np.ones_like(x))) * depth)
        xyz_world = np.matmul(np.linalg.inv(ref_extrinsics),
                              np.vstack((xyz_ref, np.ones_like(x))))[:3]
        vertexs.append(xyz_world.transpose((1, 0)))
        vertex_colors.append((color * 255).astype(np.uint8))

    # final point cloud
    vertexs = np.concatenate(vertexs, axis=0)
    vertex_colors = np.concatenate(vertex_colors, axis=0)
    vertexs = np.array([tuple(v) for v in vertexs], dtype=[('x', 'f4'), ('y', 'f4'), ('z', 'f4')])
    vertex_colors = np.array([tuple(v) for v in vertex_colors], dtype=[('red', 'u1'), ('green', 'u1'), ('blue', 'u1')])
    vertex_all = np.empty(len(vertexs), vertexs.dtype.descr + vertex_colors.dtype.descr)
    for prop in vertexs.dtype.names:
        vertex_all[prop] = vertexs[prop]
    for prop in vertex_colors.dtype.names:
        vertex_all[prop] = vertex_colors[prop]
    el = PlyElement.describe(vertex_all, 'vertex')
    PlyData([el]).write(plyfilename)
    logger.info(f"saving the final MVS result to {plyfilename}")

def init_worker():
    '''
    Catch Ctrl+C signal to termiante workers
    '''
    signal.signal(signal.SIGINT, signal.SIG_IGN)

def pcd_filter_worker(scan):
    assert 'scan' in scan
    scan_id = int(scan[4:])
    save_name = 'mvsnet{:0>3}_l3.ply'.format(scan_id)
    scan_folder = os.path.join(args.outdir, scan)
    out_folder = os.path.join(args.outdir, scan)
    filter_depth(scan_folder, out_folder, os.path.join(args.outdir, save_name))

def pcd_filter(testlist, number_worker):

    partial_func = partial(pcd_filter_worker)

    p = Pool(number_worker, init_worker)
    try:
        p.map(partial_func, testlist)
    except KeyboardInterrupt:
        print("....\nCaught KeyboardInterrupt, terminating workers")
        p.terminate()
    else:
        p.close()
    p.join()


if __name__ == '__main__':
    if 'txt' in args.testlist:
        with open(args.testlist) as f:
            content = f.readlines()
            testlist = [line.rstrip() for line in content] # ["scan1", "scan2"]
    else:
        testlist = [x for x in args.testlist.replace(' ', '').split(',') if x] # ["scan1",]

    logger.warning(f"{testlist} {args.outdir} {args.exps_folder}")
    time.sleep(5)

    # prepare data for image-based rendering
    if args.get("create_scene", False):
        args.x2_mvsres = False
        assert args.num_view == 3
        for scene in testlist:
            create_scene(scene)
        exit()

    # step1. save all the depth maps and the masks in outputs directory
    if not args.filter_only:
        save_depth(testlist)
    
    # step2. filter saved depth maps with photometric confidence maps and geometric constraints:
    pcd_filter(testlist, args.num_worker)
runner.py
cvlab-stonybrook/s-volsdf
""" simple image-based rendering """

#--------------------------args-------------------------------------#
from helpers.help import logger
from omegaconf import DictConfig, OmegaConf
import hydra

@hydra.main(version_base=None, config_path="config", config_name="ours")
def get_config(cfg: DictConfig):
    OmegaConf.set_struct(cfg, False)
    global args
    args = cfg

get_config()
#------------------------------------------------------------------#

from PIL import Image
from pathlib import Path
import os
import numpy as np
import cv2
import copy
import torch
import torch.nn.functional as F
from scipy.special import softmax

from volsdf.datasets.scene_dataset import get_trains_ids, get_eval_ids
from datasets.data_io import read_pfm
from helpers.utils import read_camera_parameters, read_img, check_geometric_consistency

def lift(x, y, z, intrinsics):
    fx = intrinsics[:, 0, 0]
    fy = intrinsics[:, 1, 1]
    cx = intrinsics[:, 0, 2]
    cy = intrinsics[:, 1, 2]
    sk = intrinsics[:, 0, 1]

    x_lift = (x - cx.unsqueeze(-1) + cy.unsqueeze(-1)*sk.unsqueeze(-1)/fy.unsqueeze(-1) - sk.unsqueeze(-1)*y/fy.unsqueeze(-1)) / fx.unsqueeze(-1) * z
    y_lift = (y - cy.unsqueeze(-1)) / fy.unsqueeze(-1) * z

    # homogeneous
    return torch.stack((x_lift, y_lift, z, torch.ones_like(z)), dim=-1)

def get_camera_params(uv, pose, intrinsics):
    """
    uv: (B, N, 2)
    pose, intrinsics: (B, 4, 4)
    """
    uv = torch.from_numpy(uv[None, :, :]).float()
    pose = torch.from_numpy(pose[None, :, :]).float()
    intrinsics = torch.from_numpy(intrinsics[None, :, :]).float()

    if pose.shape[1] == 7:
        raise NotImplementedError
    else: # In case of pose matrix representation
        cam_loc = pose[:, :3, 3]
        p = pose

    batch_size, num_samples, _ = uv.shape

    depth = torch.ones((batch_size, num_samples))
    x_cam = uv[:, :, 0].view(batch_size, -1) # (B, N)
    y_cam = uv[:, :, 1].view(batch_size, -1) # (B, N)
    z_cam = depth.view(batch_size, -1) # (B, N)

    pixel_points_cam = lift(x_cam, y_cam, z_cam, intrinsics=intrinsics) # (B, N, 4)

    # permute for batch matrix product
    pixel_points_cam = pixel_points_cam.permute(0, 2, 1) # (B, 4, N)

    # world_coords = (p[:, :3, :3] @ pixel_points_cam[:, :3, :] + p[:, :3, 3:]).permute(0, 2, 1) # (B, N, 3)
    world_coords = (torch.bmm(p[:, :3, :3], pixel_points_cam[:, :3, :]) + p[:, :3, 3:]).permute(0, 2, 1) # (B, N, 3)
    # assert torch.abs(world_coords - torch.bmm(p, pixel_points_cam).permute(0, 2, 1)[:, :, :3]).mean() < 0.0004
    
    ray_dirs = world_coords - cam_loc[:, None, :]
    ray_dirs = F.normalize(ray_dirs, dim=2)

    return ray_dirs, cam_loc

def get_dir_loc(_intrinsics, extrinsics, hw):
    h, w = hw[0], hw[1]
    intrinsics = np.eye(4)
    intrinsics[:3, :3] = _intrinsics
    pose = np.linalg.inv(extrinsics)
    uv = np.mgrid[0:h, 0:w].astype(np.int32)
    uv = np.flip(uv, axis=0).copy() # (2, 576, 768)
    uv = uv.reshape(2, -1).transpose(1, 0) # (442368, 2)
    ray_dirs, cam_loc = get_camera_params(uv, pose, intrinsics)
    ray_dirs, cam_loc = ray_dirs.squeeze(), cam_loc.squeeze() # (442368, 3) # (3)
    ray_dirs = ray_dirs.reshape(h, w, 3)
    return ray_dirs.numpy(), cam_loc.numpy()

def get_lpIMG(img_A, num_levels=4, is_mask=False):
    # generate Gaussian pyramid for A,B and mask
    G = img_A.copy().astype("float")
    gpA = [G]
    for i in range(num_levels):
        G = cv2.pyrDown(G)  
        gpA.append(G)
    
    # generate Laplacian Pyramids for A,B and masks
    if is_mask:
        lpA = [gpA[num_levels-1]]
        for i in range(num_levels-2,-1,-1):
            GE = gpA[i]
            lpA.append(GE) 
    else:
        lpA = [gpA[num_levels-1]]
        for i in range(num_levels-1,0,-1):
            GE = cv2.pyrUp(gpA[i])  
            L = cv2.subtract(gpA[i-1],GE)
            lpA.append(L)

    return lpA

def Laplacian_Blending(imgs, masks, num_levels=4):
    # Implement Laplacian_blending
    # assume mask is float32 [0,1], it has the same size to img_A and img_B
    # the mask indicates which parts of img_A or img_B are blended together
    # num_levels is the number of levels in the pyramid
    assert imgs.shape == masks.shape

    lp_imgs = []
    for img_i in imgs:
        lp_img_i = get_lpIMG(img_i, num_levels=num_levels)
        lp_imgs.append(lp_img_i)
    
    lp_masks = []
    for mask_i in masks:
        lp_mask_i = get_lpIMG(mask_i, num_levels=num_levels, is_mask=True)
        lp_masks.append(lp_mask_i)
    
    # Now blend images according to mask in each level
    LS = []
    for i in range(num_levels):
        ls = 0
        for j in range(len(masks)):
            ls += lp_masks[j][i] * lp_imgs[j][i]
        LS.append(ls)

    # now reconstruct
    ls_ = LS[0]
    for i in range(1,num_levels):
      ls_ = cv2.pyrUp(ls_)
      ls_ = cv2.add(ls_, LS[i])

    return np.clip(ls_, 0.0, 1.0)

def image_based_render(scan_folder, out_folder):    
    trains_i = get_trains_ids(args.vol.dataset.data_dir, Path(scan_folder).name, args.num_view)
    evals_i = get_eval_ids(args.vol.dataset.data_dir, int(Path(scan_folder).name[4:]))
    logger.debug(f'trains_i {trains_i}')
    logger.debug(f'evals_i {evals_i}')

    pair_data = [(idx, trains_i) for idx in evals_i]
    # for each reference view and the corresponding source views
    for ref_view, src_views in pair_data:
        # load the camera parameters
        ref_intrinsics, ref_extrinsics = read_camera_parameters(
            os.path.join(scan_folder, 'cams/{:0>8}_cam.txt'.format(ref_view)))
        # load the estimated depth of the reference view
        pred_ref_img = read_img(os.path.join(out_folder, 'eval_{:0>3}.png'.format(ref_view)))
        ref_depth_est = read_pfm(os.path.join(out_folder, 'depth_est/{:0>8}.pfm'.format(ref_view)))[0]
        ref_dir, ref_loc = get_dir_loc(ref_intrinsics, ref_extrinsics, ref_depth_est.shape)

        all_srcview_depth_ests = []
        all_srcview_x = []
        all_srcview_y = []
        all_srcview_geomask = []

        weight_mask_sum = 0
        weight_masks = []
        sampled_img_srcs = []

        # 0. compute the geometric mask
        for src_view in src_views:
            # camera parameters of the source view
            src_intrinsics, src_extrinsics = read_camera_parameters(
                os.path.join(scan_folder, 'cams/{:0>8}_cam.txt'.format(src_view)))
            # the estimated depth of the source view
            src_img = read_img(os.path.join(scan_folder, 'images/{:0>8}.png'.format(src_view)))
            src_depth_est = read_pfm(os.path.join(out_folder, 'depth_est/{:0>8}.pfm'.format(src_view)))[0]
            assert src_depth_est.shape == ref_depth_est.shape
            geo_mask, depth_reprojected, x2d_src, y2d_src = check_geometric_consistency(ref_depth_est, ref_intrinsics, ref_extrinsics,
                                                                      src_depth_est,
                                                                      src_intrinsics, src_extrinsics,
                                                                      filter_dist=2,
                                                                      )
            # x2d_src, y2d_src (576, 768) src_img (576, 768, 3)
            sampled_img_src = cv2.remap(src_img, x2d_src, y2d_src, interpolation=cv2.INTER_CUBIC)
            sampled_img_srcs.append(sampled_img_src)

            # per point direction
            src_dir, src_loc = get_dir_loc(src_intrinsics, src_extrinsics, src_depth_est.shape)
            sampled_src_dir = cv2.remap(src_dir, x2d_src, y2d_src, interpolation=cv2.INTER_CUBIC)
            sampled_src_dir /= np.linalg.norm(sampled_src_dir, axis=2, keepdims=True) # nan will be replaced by 0 later
            cos_dir = (sampled_src_dir*ref_dir).sum(axis=2) # -1 .. 1

            weight_mask = cos_dir
            weight_mask = np.nan_to_num(weight_mask)
            weight_mask *= geo_mask.astype(np.int32)
            weight_masks.append(weight_mask)

            weight_mask_sum += weight_mask
            all_srcview_depth_ests.append(depth_reprojected)
            all_srcview_x.append(x2d_src)
            all_srcview_y.append(y2d_src)
            all_srcview_geomask.append(geo_mask)

        weight_mask = 0.2 * np.ones_like(ref_depth_est)
        weight_masks.append(weight_mask)
        sampled_img_srcs.append(pred_ref_img)

        weight_masks = np.stack(weight_masks)
        weight_masks = softmax(20 * weight_masks, axis=0)
        weight_masks = weight_masks[..., None].repeat(3, -1) # (N, H, W) -> （N, H, W, 3)
        sampled_img_srcs = np.stack(sampled_img_srcs)

        # 1. fill undefined pixels
        sampled_img_srcs_4lap = sampled_img_srcs * weight_masks + sampled_img_srcs[-1:] * (1-weight_masks)

        # 2. erode mask, so that when bluring, undefined pixels won't come in
        weight_masks_4lap = copy.deepcopy(weight_masks)
        kernel = np.ones((5, 5), np.uint8)
        for i in range(weight_masks_4lap.shape[0]-1):
            erode_mask = cv2.erode((weight_masks_4lap[i]>0.2)*1.0, kernel) * 1.0
            weight_masks_4lap[i] = erode_mask * weight_masks_4lap[i]
        weight_masks_4lap[-1] += 1e-2
        weight_masks_4lap /= weight_masks_4lap.sum(0, keepdims=True)

        # 3. laplacian blend
        blend_image = Laplacian_Blending(sampled_img_srcs_4lap, weight_masks_4lap, num_levels=4)
        Image.fromarray((blend_image * 255).astype(np.uint8)).save(
            os.path.join(out_folder, 'eval_blend_{:0>3}.png'.format(ref_view))
        )


if __name__ == '__main__':
    # python simple_ibr.py          testlist='config/lists/dtu.txt'  outdir=exps_ibr +evals_folder=exps_result
    # python simple_ibr.py vol=bmvs testlist='config/lists/bmvs.txt' outdir=exps_ibr +evals_folder=exps_result

    if 'txt' in args.testlist:
        with open(args.testlist) as f:
            content = f.readlines()
            testlist = [line.rstrip() for line in content] # ["scan1", "scan2"]
    else:
        testlist = [x for x in args.testlist.replace(' ', '').split(',') if x] # ["scan1",]
    scan_ids = [int(x[4:]) for x in testlist]
    logger.warning(scan_ids)

    for scan_id in scan_ids:
        # use the latest epoch's rendering results
        evaldir = f'{args.evals_folder}/{args.vol.train.expname}_{scan_id}' # exps_result/ours_106
        epoch = 0
        for renderdir in os.listdir(evaldir):
            if renderdir.startswith('rendering_'):
                epoch = max(epoch, int(renderdir.replace('rendering_', '')))
        
        out_folder = os.path.join(evaldir, f'rendering_{epoch}') # exps_result/ours_106/rendering_1562
        scan_folder = os.path.join(args.outdir, f'scan{scan_id}') # data_ibr/scan106
        assert os.path.exists(scan_folder) and os.path.exists(out_folder)
        logger.warning(f'use cam & src_imgs in {scan_folder}')
        logger.warning(f'     add new images to {out_folder}')

        image_based_render(scan_folder, out_folder)
simple_ibr.py
cvlab-stonybrook/s-volsdf
import os
import torch
import numpy as np
import random
import cv2
from pathlib import Path
from PIL import Image

import volsdf.utils.general as utils
from volsdf.utils import rend_util

def scan2hash(scan):
    """
    for BlendedMVS dataset
    """
    scan2hash_dict ={
        'scan1': '5a3ca9cb270f0e3f14d0eddb',
        'scan2': '5a6464143d809f1d8208c43c',
        'scan3': '5ab85f1dac4291329b17cb50',
        'scan4': '5b4933abf2b5f44e95de482a',
        'scan5': '5b22269758e2823a67a3bd03',
        'scan6': '5c0d13b795da9479e12e2ee9',
        'scan7': '5c1af2e2bee9a723c963d019',
        'scan8': '5c1dbf200843bc542d8ef8c4',
        'scan9': '5c34300a73a8df509add216d',
    }
    return scan2hash_dict[scan]

def get_trains_ids(data_dir, scan, num_views=0, for_interp=False):
    if num_views <= 0: raise NotImplementedError
    
    if num_views == 49:
        return list(range(49))

    if data_dir == 'DTU':
        train_ids_all = [25, 22, 28, 40, 44, 48, 0, 8, 13]
        return train_ids_all[:num_views]
        
    elif data_dir == 'BlendedMVS':
        assert num_views == 3
        bmvs_train_ids = {
            1: [9, 10, 55],
            2: [59, 9, 52], 
            3: [26, 27, 22], 
            4: [11, 39, 53], 
            5: [32, 42, 47], 
            6: [28, 34, 57],
            7: [5, 25, 2], 
            8: [16, 21, 33], 
            9: [16, 60, 10], 
        }
        bmvs_train_ids_4_interp = {
            1 : [9, 10, 55],
            2 : [59, 52, 9],
            3 : [27, 26, 22],
            4 : [11, 39, 53],
            5 : [42, 32, 47],
            6 : [28, 34, 57],
            7 : [5, 25, 2],
            8 : [16, 33, 21],
            9 : [16, 60, 10],
        }
        if for_interp:
            train_ids_all = bmvs_train_ids_4_interp[int(scan[4:])][:num_views]
        else:
            train_ids_all = bmvs_train_ids[int(scan[4:])][:num_views]
        return train_ids_all
    
    else:
        raise NotImplementedError

def get_eval_ids(data_dir, scan_id=None):
    if 'DTU' == data_dir:
        # from regnerf/pixelnerf
        train_idx = [25, 22, 28, 40, 44, 48, 0, 8, 13]
        exclude_idx = [3, 4, 5, 6, 7, 16, 17, 18, 19, 20, 21, 36, 37, 38, 39]
        test_idx = [i for i in range(49) if i not in train_idx + exclude_idx]
        return test_idx
    elif 'BlendedMVS' == data_dir:
        bmvs_test_ids = {1: [19, 35, 58, 0, 33, 37, 31, 20, 61, 22, 36, 13], 2: [14, 49, 37, 1, 27, 8, 12, 39, 65, 23, 71, 68], 3: [33, 0, 28, 11, 6, 7, 15, 25, 13, 31, 16, 1], 4: [30, 48, 9, 68, 50, 59, 23, 29, 0, 46, 2, 71], 5: [58, 55, 24, 57, 17, 16, 41, 44, 0, 13, 20, 26], 6: [55, 25, 13, 75, 73, 2, 88, 22, 10, 80, 40, 1], 7: [17, 10, 22, 8, 23, 0, 27, 6, 14, 13, 29, 7], 8: [27, 25, 47, 28, 10, 17, 30, 18, 8, 26, 24, 43], 9: [118, 18, 136, 92, 68, 89, 133, 83, 25, 65, 94, 59]}
        return bmvs_test_ids[scan_id][:12]
    else:
        raise NotImplementedError

def get_near_id(data_dir, scan_id, idx):
    bmvs_dict = {'scan1': {0: 9, 1: 9, 2: 9, 3: 10, 4: 9, 5: 9, 6: 9, 7: 55, 8: 9, 9: 9, 10: 10, 11: 9, 12: 10, 13: 10, 14: 9, 15: 9, 16: 9, 17: 9, 18: 10, 19: 55, 20: 55, 21: 9, 22: 9, 23: 10, 24: 10, 25: 10, 26: 10, 27: 9, 28: 10, 29: 9, 30: 10, 31: 55, 32: 55, 33: 9, 34: 9, 35: 55, 36: 9, 37: 9, 38: 55, 39: 10, 40: 10, 41: 10, 42: 10, 43: 10, 44: 10, 45: 9, 46: 10, 47: 9, 48: 9, 49: 9, 50: 9, 51: 9, 52: 10, 53: 10, 54: 10, 55: 55, 56: 9, 57: 55, 58: 9, 59: 9, 60: 9, 61: 10, 62: 10, 63: 9}, 'scan2': {0: 52, 1: 59, 2: 59, 3: 59, 4: 9, 5: 9, 6: 59, 7: 52, 8: 59, 9: 9, 10: 9, 11: 52, 12: 59, 13: 52, 14: 52, 15: 52, 16: 52, 17: 52, 18: 52, 19: 9, 20: 52, 21: 52, 22: 9, 23: 9, 24: 9, 25: 9, 26: 52, 27: 52, 28: 52, 29: 52, 30: 59, 31: 59, 32: 9, 33: 52, 34: 59, 35: 52, 36: 52, 37: 59, 38: 52, 39: 59, 40: 9, 41: 59, 42: 52, 43: 59, 44: 59, 45: 9, 46: 52, 47: 52, 48: 9, 49: 52, 50: 52, 51: 52, 52: 52, 53: 52, 54: 9, 55: 9, 56: 52, 57: 52, 58: 52, 59: 59, 60: 59, 61: 52, 62: 52, 63: 9, 64: 9, 65: 52, 66: 52, 67: 52, 68: 52, 69: 9, 70: 52, 71: 9, 72: 9, 73: 52}, 'scan3': {0: 27, 1: 22, 2: 26, 3: 26, 4: 27, 5: 27, 6: 22, 7: 27, 8: 22, 9: 22, 10: 22, 11: 26, 12: 22, 13: 22, 14: 27, 15: 26, 16: 26, 17: 26, 18: 27, 19: 27, 20: 22, 21: 27, 22: 22, 23: 27, 24: 22, 25: 26, 26: 26, 27: 27, 28: 26, 29: 26, 30: 27, 31: 26, 32: 27, 33: 27, 34: 26, 35: 27}, 'scan4': {0: 53, 1: 11, 2: 53, 3: 53, 4: 11, 5: 39, 6: 53, 7: 53, 8: 11, 9: 53, 10: 11, 11: 11, 12: 53, 13: 11, 14: 39, 15: 11, 16: 53, 17: 39, 18: 11, 19: 53, 20: 53, 21: 11, 22: 11, 23: 53, 24: 11, 25: 39, 26: 53, 27: 11, 28: 39, 29: 53, 30: 53, 31: 11, 32: 11, 33: 11, 34: 11, 35: 11, 36: 39, 37: 39, 38: 11, 39: 39, 40: 11, 41: 39, 42: 53, 43: 11, 44: 11, 45: 11, 46: 11, 47: 39, 48: 53, 49: 11, 50: 53, 51: 11, 52: 53, 53: 53, 54: 53, 55: 11, 56: 11, 57: 53, 58: 53, 59: 53, 60: 11, 61: 53, 62: 53, 63: 11, 64: 11, 65: 53, 66: 53, 67: 53, 68: 11, 69: 53, 70: 53, 71: 11, 72: 53, 73: 53, 74: 11, 75: 53, 76: 39, 77: 53, 78: 53, 79: 11, 80: 53, 81: 11, 82: 53}, 'scan5': {0: 47, 1: 47, 2: 47, 3: 42, 4: 47, 5: 32, 6: 42, 7: 32, 8: 42, 9: 32, 10: 32, 11: 32, 12: 42, 13: 42, 14: 42, 15: 42, 16: 47, 17: 42, 18: 32, 19: 42, 20: 32, 21: 42, 22: 47, 23: 47, 24: 32, 25: 47, 26: 47, 27: 47, 28: 47, 29: 42, 30: 47, 31: 32, 32: 32, 33: 42, 34: 42, 35: 42, 36: 42, 37: 42, 38: 32, 39: 47, 40: 42, 41: 47, 42: 42, 43: 42, 44: 47, 45: 42, 46: 32, 47: 47, 48: 47, 49: 47, 50: 42, 51: 47, 52: 42, 53: 42, 54: 42, 55: 42, 56: 42, 57: 42, 58: 32, 59: 42, 60: 42, 61: 32, 62: 42, 63: 32}, 'scan6': {0: 57, 1: 34, 2: 57, 3: 28, 4: 28, 5: 28, 6: 57, 7: 57, 8: 57, 9: 28, 10: 57, 11: 28, 12: 34, 13: 28, 14: 57, 15: 28, 16: 57, 17: 28, 18: 28, 19: 57, 20: 57, 21: 34, 22: 57, 23: 28, 24: 34, 25: 28, 26: 57, 27: 34, 28: 28, 29: 34, 30: 57, 31: 34, 32: 34, 33: 57, 34: 34, 35: 57, 36: 57, 37: 57, 38: 34, 39: 57, 40: 28, 41: 34, 42: 28, 43: 28, 44: 28, 45: 28, 46: 57, 47: 28, 48: 57, 49: 34, 50: 34, 51: 34, 52: 57, 53: 28, 54: 28, 55: 28, 56: 57, 57: 57, 58: 57, 59: 28, 60: 28, 61: 34, 62: 28, 63: 34, 64: 28, 65: 28, 66: 57, 67: 57, 68: 57, 69: 57, 70: 28, 71: 57, 72: 57, 73: 28, 74: 34, 75: 28, 76: 57, 77: 34, 78: 28, 79: 28, 80: 57, 81: 57, 82: 28, 83: 28, 84: 28, 85: 28, 86: 28, 87: 57, 88: 57, 89: 28}, 'scan7': {0: 5, 1: 2, 2: 2, 3: 25, 4: 2, 5: 5, 6: 5, 7: 5, 8: 5, 9: 2, 10: 25, 11: 25, 12: 5, 13: 5, 14: 5, 15: 2, 16: 2, 17: 2, 18: 2, 19: 2, 20: 2, 21: 2, 22: 2, 23: 2, 24: 5, 25: 25, 26: 5, 27: 2, 28: 5, 29: 5, 30: 5}, 'scan8': {0: 21, 1: 16, 2: 21, 3: 21, 4: 16, 5: 16, 6: 16, 7: 16, 8: 16, 9: 33, 10: 16, 11: 21, 12: 21, 13: 16, 14: 16, 15: 16, 16: 16, 17: 16, 18: 33, 19: 33, 20: 33, 21: 21, 22: 16, 23: 16, 24: 16, 25: 33, 26: 16, 27: 16, 28: 21, 29: 21, 30: 16, 31: 16, 32: 16, 33: 33, 34: 16, 35: 16, 36: 16, 37: 33, 38: 33, 39: 16, 40: 21, 41: 16, 42: 16, 43: 33, 44: 21, 45: 16, 46: 21, 47: 21, 48: 16}, 'scan9': {0: 60, 1: 10, 2: 16, 3: 10, 4: 16, 5: 10, 6: 16, 7: 16, 8: 10, 9: 16, 10: 10, 11: 16, 12: 10, 13: 10, 14: 16, 15: 16, 16: 16, 17: 16, 18: 16, 19: 16, 20: 10, 21: 16, 22: 16, 23: 16, 24: 16, 25: 60, 26: 16, 27: 16, 28: 16, 29: 16, 30: 60, 31: 16, 32: 16, 33: 10, 34: 10, 35: 16, 36: 10, 37: 16, 38: 60, 39: 16, 40: 16, 41: 16, 42: 16, 43: 16, 44: 10, 45: 16, 46: 10, 47: 10, 48: 10, 49: 10, 50: 10, 51: 16, 52: 16, 53: 16, 54: 16, 55: 10, 56: 60, 57: 16, 58: 16, 59: 10, 60: 60, 61: 10, 62: 10, 63: 16, 64: 16, 65: 60, 66: 16, 67: 16, 68: 16, 69: 16, 70: 10, 71: 16, 72: 16, 73: 10, 74: 10, 75: 60, 76: 16, 77: 16, 78: 10, 79: 16, 80: 10, 81: 10, 82: 16, 83: 60, 84: 10, 85: 16, 86: 16, 87: 16, 88: 16, 89: 16, 90: 16, 91: 10, 92: 16, 93: 16, 94: 16, 95: 10, 96: 16, 97: 16, 98: 16, 99: 16, 100: 10, 101: 10, 102: 16, 103: 16, 104: 60, 105: 10, 106: 16, 107: 16, 108: 16, 109: 10, 110: 16, 111: 10, 112: 16, 113: 60, 114: 10, 115: 60, 116: 16, 117: 10, 118: 60, 119: 10, 120: 16, 121: 16, 122: 10, 123: 16, 124: 16, 125: 16, 126: 10, 127: 16, 128: 60, 129: 10, 130: 60, 131: 10, 132: 16, 133: 60, 134: 10, 135: 16, 136: 60, 137: 10, 138: 16, 139: 10, 140: 16, 141: 60, 142: 10, 143: 16}}
    if 'BlendedMVS' == data_dir:
        return bmvs_dict[f'scan{scan_id}'][idx]
    else:
        raise NotImplementedError



class SceneDataset(torch.utils.data.Dataset):

    def __init__(self,
                 data_dir,
                 img_res,
                 scan_id=0,
                 num_views=-1,
                 data_dir_root=None,
                 ):

        self.data_dir, self.scan_id, self.num_views = data_dir, scan_id, num_views
        self.total_pixels = img_res[0] * img_res[1]
        self.img_res = img_res
        assert num_views in [-1, 3, 4, 5, 6, 9]
        self.mode, self.plot_id = 'train', 0
        self.sampling_idx = None
        self.use_pixel_centers = False

        instance_dir = os.path.join(data_dir_root, data_dir, 'scan{0}'.format(scan_id))
        image_dir = '{0}/image'.format(instance_dir)
        self.cam_file = '{0}/cameras.npz'.format(instance_dir)
        if not os.path.exists(self.cam_file) and int(scan_id) < 200:
            self.cam_file = os.path.join(data_dir_root, data_dir, 'scan114', 'cameras.npz')
        assert os.path.exists(image_dir), "Data directory is empty"
        assert os.path.exists(self.cam_file), "Data directory is empty"
        image_paths = sorted(utils.glob_imgs(image_dir))
        self.n_images = len(image_paths)
        camera_dict = np.load(self.cam_file)
        scale_mats = [camera_dict['scale_mat_%d' % idx].astype(np.float32) for idx in range(self.n_images)]
        world_mats = [camera_dict['world_mat_%d' % idx].astype(np.float32) for idx in range(self.n_images)]

        # low res
        rgb = rend_util.load_rgb(image_paths[0]) # (C, H, W)
        scale_h, scale_w = self.img_res[0]*1./rgb.shape[0], self.img_res[1]*1./rgb.shape[1]

        # mask
        mask_path = os.path.join(data_dir_root, data_dir, 'eval_mask')
        if data_dir == 'DTU':
            maskf_fn = lambda x: os.path.join(mask_path, f'scan{scan_id}', 'mask', f'{x:03d}.png')
            if not os.path.exists(maskf_fn(0)):
                maskf_fn = lambda x: os.path.join(mask_path, f'scan{scan_id}', f'{x:03d}.png')
        elif data_dir == 'BlendedMVS':
            maskf_fn = lambda x: os.path.join(mask_path, f'scan{scan_id}', 'mask', f'{x:08d}.png')
        else:
            raise NotImplementedError

        self.rgb_images = []
        self.rgb_smooth = []
        self.masks = []
        self.intrinsics_all = []
        self.pose_all = []
        self.scale_factor = scale_mats[0][0,0]
        if int(self.scan_id) == 5 and data_dir == 'BlendedMVS': # scan5only
            # scale_mat for scan5 is wrong, set it to 1 instead
            self.scale_factor = 1.0

        for id_, path in enumerate(image_paths):
            # K, pose
            scale_mat, world_mat = scale_mats[id_], world_mats[id_]
            P = world_mat @ scale_mat
            P = P[:3, :4]
            intrinsics, pose = rend_util.load_K_Rt_from_P(None, P)
            # pose - scale
            intrinsics[0, :] *= scale_w
            intrinsics[1, :] *= scale_h
            self.intrinsics_all.append(torch.from_numpy(intrinsics).float())
            self.pose_all.append(torch.from_numpy(pose).float())

            # rgb
            img = rend_util.load_rgb(path) # (H, W, 3)
            # rgb - scale
            if scale_h != 1 or scale_w != 1:
                img = cv2.resize(img, (self.img_res[1], self.img_res[0]), interpolation=cv2.INTER_CUBIC)
            rgb = img.transpose(2, 0, 1) # (3, H, W)
            rgb = rgb.reshape(3, -1).transpose(1, 0) # -> (3, H*W) -> (H*W, 3)
            self.rgb_images.append(torch.from_numpy(rgb).float())

            # rgb smooth
            img = cv2.GaussianBlur(img, (31,31), 90)
            rgb = img.transpose(2, 0, 1) # (3, H, W)
            rgb = rgb.reshape(3, -1).transpose(1, 0) # -> (3, H*W) -> (H*W, 3)
            self.rgb_smooth.append(torch.from_numpy(rgb).float())

            # mask
            if data_dir == 'DTU' and id_ in get_eval_ids(data_dir=data_dir) and scan_id not in [1, 4, 11, 13, 48]:
                fname = maskf_fn(id_)
                with open(fname, 'rb') as imgin:
                    mask_image = np.array(Image.open(imgin), dtype=np.float32)[:, :, :3] / 255.
                    mask_image = (mask_image == 1).astype(np.float32)
                if scale_h != 1 or scale_w != 1:
                    mask_image = cv2.resize(mask_image, (self.img_res[1], self.img_res[0]), cv2.INTER_NEAREST)
                    mask_image = (mask_image > 0.5).astype(np.float32)
                mask_image = mask_image.transpose(2, 0, 1) # (3, H, W)
                mask_image = mask_image.reshape(3, -1).transpose(1, 0) # -> (3, H*W) -> (H*W, 3)
                self.masks.append(torch.from_numpy(mask_image).float())
            elif data_dir == 'BlendedMVS':
                idxs_with_mask = get_eval_ids(data_dir=data_dir, scan_id=scan_id) + get_trains_ids(data_dir=data_dir, scan=f'scan{scan_id}', num_views=3)
                if id_ in idxs_with_mask:
                    fname = maskf_fn(id_)
                    assert os.path.exists(fname)
                    with open(fname, 'rb') as imgin:
                        mask_image = np.array(Image.open(imgin), dtype=np.float32)
                        assert len(mask_image.shape) == 3 and mask_image.shape[2] == 4
                        mask_image = mask_image[:, :, -1] * 1. / 255.
                        mask_image = np.stack([mask_image,mask_image,mask_image], -1)
                        mask_image = cv2.resize(mask_image, (self.img_res[1], self.img_res[0]), cv2.INTER_NEAREST)
                        mask_image = (mask_image > 0.5).astype(np.float32)
                        mask_image = mask_image.reshape(-1, 3)
                    self.masks.append(torch.from_numpy(mask_image).float())
                else:
                    self.masks.append(torch.ones_like(self.rgb_images[0]).float())
            else:
                self.masks.append(torch.ones_like(self.rgb_images[0]).float())
          
    def __len__(self):
        return self.n_images

    def __getitem__(self, idx):
        """
        select an image with random N rays/pixels
        """
        if self.num_views >= 1:
            train_ids = self.trains_ids()
            unobs_ids = [x for x in range(self.n_images) if x not in train_ids]
            if self.mode == 'train':
                idx = train_ids[random.randint(0, self.num_views - 1)]
            elif self.mode == 'plot':
                eval_ids = get_eval_ids(data_dir=self.data_dir, scan_id=self.scan_id)
                if len(eval_ids) == 0:
                    eval_ids = unobs_ids
                idx = eval_ids[self.plot_id]
                self.plot_id = (self.plot_id + 1) % len(eval_ids)

        uv = np.mgrid[0:self.img_res[0], 0:self.img_res[1]].astype(np.int32)
        uv = torch.from_numpy(np.flip(uv, axis=0).copy()).float()
        uv = uv.reshape(2, -1).transpose(1, 0)

        if self.use_pixel_centers:
            uv += 0.5

        sample = {
            "uv": uv,
            "intrinsics": self.intrinsics_all[idx],
            "pose": self.pose_all[idx],
        }
        if self.data_dir in ['BlendedMVS', ]:
            sample["near_pose"] = self.pose_all[get_near_id(data_dir=self.data_dir, scan_id=self.scan_id, idx=idx)]

        ground_truth = {
            "rgb": self.rgb_images[idx],
            "rgb_smooth": self.rgb_smooth[idx],
            "mask": self.masks[idx]
        }

        if self.sampling_idx is not None:
            ground_truth["rgb"] = self.rgb_images[idx][self.sampling_idx, :]
            ground_truth["rgb_smooth"] = self.rgb_smooth[idx][self.sampling_idx, :]
            sample["uv"] = uv[self.sampling_idx, :]

        return idx, sample, ground_truth

    def trains_ids(self):
        return get_trains_ids(data_dir=self.data_dir, scan=f'scan{self.scan_id}', num_views=self.num_views)

    def collate_fn(self, batch_list):
        # get list of dictionaries and returns input, ground_true as dictionary for all batch instances
        batch_list = zip(*batch_list)

        all_parsed = []
        for entry in batch_list:
            if type(entry[0]) is dict:
                # make them all into a new dict
                ret = {}
                for k in entry[0].keys():
                    ret[k] = torch.stack([obj[k] for obj in entry])
                all_parsed.append(ret)
            else:
                all_parsed.append(torch.LongTensor(entry))

        return tuple(all_parsed)

    def change_sampling_idx(self, sampling_size):
        if sampling_size == -1:
            self.sampling_idx = None
        else:
            self.sampling_idx = torch.randperm(self.total_pixels)[:sampling_size]

    def get_scale_mat(self):
        return np.load(self.cam_file)['scale_mat_0']

volsdf/datasets/scene_dataset.py
cvlab-stonybrook/s-volsdf
import torch.nn as nn
import torch


class Density(nn.Module):
    def __init__(self, params_init={}):
        super().__init__()
        for p in params_init:
            param = nn.Parameter(torch.tensor(params_init[p]))
            setattr(self, p, param)

    def forward(self, sdf, beta=None):
        return self.density_func(sdf, beta=beta)


class LaplaceDensity(Density):  # alpha * Laplace(loc=0, scale=beta).cdf(-sdf)
    def __init__(self, params_init={}, beta_min=0.0001):
        super().__init__(params_init=params_init)
        self.beta_min = torch.tensor(beta_min).cuda()

    def density_func(self, sdf, beta=None):
        if beta is None:
            beta = self.get_beta()

        alpha = 1 / beta
        return alpha * (0.5 + 0.5 * sdf.sign() * torch.expm1(-sdf.abs() / beta))

    def get_beta(self):
        beta = self.beta.abs() + self.beta_min
        return beta


class AbsDensity(Density):  # like NeRF++
    def density_func(self, sdf, beta=None):
        return torch.abs(sdf)


class SimpleDensity(Density):  # like NeRF
    def __init__(self, params_init={}, noise_std=1.0):
        super().__init__(params_init=params_init)
        self.noise_std = noise_std

    def density_func(self, sdf, beta=None):
        if self.training and self.noise_std > 0.0:
            noise = torch.randn(sdf.shape).cuda() * self.noise_std
            sdf = sdf + noise
        return torch.relu(sdf)

volsdf/model/density.py
cvlab-stonybrook/s-volsdf
import torch

""" Positional encoding embedding. Code was taken from https://github.com/bmild/nerf. """

class Embedder:
    def __init__(self, **kwargs):
        self.kwargs = kwargs
        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']
        N_freqs = self.kwargs['num_freqs']

        if self.kwargs['log_sampling']:
            freq_bands = 2. ** torch.linspace(0., max_freq, N_freqs)
        else:
            freq_bands = torch.linspace(2.**0., 2.**max_freq, N_freqs)

        for freq in freq_bands:
            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):
        return torch.cat([fn(inputs) for fn in self.embed_fns], -1)

def get_embedder(multires, input_dims=3):
    embed_kwargs = {
        'include_input': True,
        'input_dims': input_dims,
        'max_freq_log2': multires-1,
        'num_freqs': multires,
        'log_sampling': True,
        'periodic_fns': [torch.sin, torch.cos],
    }

    embedder_obj = Embedder(**embed_kwargs)
    def embed(x, eo=embedder_obj): return eo.embed(x)
    return embed, embedder_obj.out_dim

volsdf/model/embedder.py
cvlab-stonybrook/s-volsdf
import torch
from torch import nn
import volsdf.utils.general as utils
import math
import numpy as np
from helpers.help import logger

def anneal_linearly(t, val0, val1):
    if t >= 1:
        return val1
    elif t <= 0:
        return val0
    return val0 + (val1 - val0) * np.minimum(t, 1.)

class VolSDFLoss(nn.Module):
    def __init__(self, rgb_loss, 
            eikonal_weight, rgb_weight=1., mvs_weight=0., 
            sparse_weight=0., anneal_rgb=0,
            gce=1, confi=0):
        super().__init__()
        self.eikonal_weight = eikonal_weight
        self.rgb_weight = rgb_weight
        self.mvs_weight = mvs_weight
        self.sparse_weight = sparse_weight
        self.gce = gce
        self.anneal_rgb = anneal_rgb
        self.confi = confi
        logger.info(f"loss lambda: RGB_{rgb_weight} EK_{eikonal_weight} MVS_{mvs_weight} SP_{sparse_weight}")
        self.rgb_loss = utils.get_class(rgb_loss)(reduction='mean')

    def set_stg(self, stg):
        self.iter_step = 0
        if stg >= 1:
            self.anneal_rgb = 0
            self.sparse_weight = 0
            raise NotImplementedError

    def get_rgb_loss(self,rgb_values, rgb_gt, model_outputs=None, t=0):
        rgb_gt = rgb_gt.reshape(-1, 3)
        if t > 0:
            pi, pj = model_outputs['pi'], model_outputs['pj'] # (512, 98)
            confi = (pi*pj).sum(-1)
            loss = torch.abs(rgb_values-rgb_gt).mean(-1)
            loss = loss * (confi<t) # confi<0.01, uncertainty
            return loss.mean()
        rgb_loss = self.rgb_loss(rgb_values, rgb_gt)
        return rgb_loss

    def get_eikonal_loss(self, grad_theta):
        eikonal_loss = ((grad_theta.norm(2, dim=1) - 1) ** 2).mean()
        return eikonal_loss

    def get_mvs_loss(self, model_outputs):
        pi, pj = model_outputs['pi'], model_outputs['pj'] # (512, 98)
        w = model_outputs['weights'] # (512, 98)
        pw = pi * pj

        if self.gce == 1:
            loss = - pw * w
        elif self.gce == 0:
            loss = - pw * torch.log(w+1e-8)
        else:
            loss = - pw * w.detach() ** self.gce
            loss = loss * torch.log(w+1e-8)
        loss = loss.sum(1)
        loss = 1. * (pw.sum(1) > self.confi) * loss
        return loss.mean()

    def get_sparse_loss(self, model_outputs):
        pi, pj = model_outputs['pi'], model_outputs['pj'] # (512, 98)
        confi = (pi*pj).sum(-1)
        if 'depth_values_all' in model_outputs:
            dep = model_outputs['depth_values_all'].squeeze()
        else:
            dep = model_outputs['depth_values'].squeeze()
        loss = 1. / (dep + 1e-3)
        loss = loss * (confi<self.confi)
        return loss.mean()

    def forward(self, model_outputs, ground_truth):
        rgb_gt = ground_truth['rgb'].cuda()

        output = {}
        output['rgb_loss'] = self.get_rgb_loss(model_outputs['rgb_values'], rgb_gt)
        if 'grad_theta' in model_outputs:
            output['eikonal_loss'] = self.get_eikonal_loss(model_outputs['grad_theta'])
        else:
            output['eikonal_loss'] = torch.tensor(0.0).cuda().float()

        # generalized cross entropy loss
        if 'pi' in model_outputs and self.mvs_weight > 0:
            output['mvs_loss'] = self.get_mvs_loss(model_outputs)
        else:
            output['mvs_loss'] = torch.tensor(0.0).cuda().float()

        # sparsity regularization
        if 'pi' in model_outputs and self.sparse_weight > 0 and self.anneal_rgb > 0 and self.iter_step < self.anneal_rgb:
            output['sparse_loss'] = self.get_sparse_loss(model_outputs)
        else:
            output['sparse_loss'] = torch.tensor(0.0).cuda().float()

        anneal_sparse = 0
        if self.sparse_weight > 0 and self.anneal_rgb > 0 and self.iter_step < self.anneal_rgb:
            anneal_sparse = anneal_linearly(self.iter_step/self.anneal_rgb, 1.0, 0.)
            output['rgb_loss'] = self.get_rgb_loss(model_outputs['rgb_values'], ground_truth['rgb_smooth'].cuda(), model_outputs, t=1e-8)

        # total loss
        output['loss'] = self.rgb_weight * output['rgb_loss'] +\
            self.eikonal_weight * output['eikonal_loss'] +\
            self.mvs_weight * output['mvs_loss'] +\
            self.sparse_weight * anneal_sparse * output['sparse_loss']

        self.iter_step += 1

        return output
volsdf/model/loss.py
cvlab-stonybrook/s-volsdf
from loguru import logger
import torch.nn as nn
import numpy as np

from volsdf.utils import rend_util
from volsdf.model.embedder import *
from volsdf.model.density import LaplaceDensity
from volsdf.model.ray_sampler import ErrorBoundSampler

class ImplicitNetwork(nn.Module):
    def __init__(
            self,
            feature_vector_size,
            sdf_bounding_sphere,
            d_in,
            d_out,
            dims,
            geometric_init=True,
            bias=1.0,
            skip_in=(),
            weight_norm=True,
            multires=0,
            sphere_scale=1.0,
    ):
        super().__init__()

        self.sdf_bounding_sphere = sdf_bounding_sphere
        self.sphere_scale = sphere_scale
        dims = [d_in] + dims + [d_out + feature_vector_size]

        self.embed_fn = None
        if multires > 0:
            embed_fn, input_ch = get_embedder(multires, input_dims=d_in)
            self.embed_fn = embed_fn
            dims[0] = input_ch

        self.num_layers = len(dims)
        self.skip_in = skip_in

        for l in range(0, self.num_layers - 1):
            if l + 1 in self.skip_in:
                out_dim = dims[l + 1] - dims[0]
            else:
                out_dim = dims[l + 1]

            lin = nn.Linear(dims[l], out_dim)

            if geometric_init:
                if l == self.num_layers - 2:
                    torch.nn.init.normal_(lin.weight, mean=np.sqrt(np.pi) / np.sqrt(dims[l]), std=0.0001)
                    torch.nn.init.constant_(lin.bias, -bias)
                elif multires > 0 and l == 0:
                    torch.nn.init.constant_(lin.bias, 0.0)
                    torch.nn.init.constant_(lin.weight[:, 3:], 0.0)
                    torch.nn.init.normal_(lin.weight[:, :3], 0.0, np.sqrt(2) / np.sqrt(out_dim))
                elif multires > 0 and l in self.skip_in:
                    torch.nn.init.constant_(lin.bias, 0.0)
                    torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))
                    torch.nn.init.constant_(lin.weight[:, -(dims[0] - 3):], 0.0)
                else:
                    torch.nn.init.constant_(lin.bias, 0.0)
                    torch.nn.init.normal_(lin.weight, 0.0, np.sqrt(2) / np.sqrt(out_dim))

            if weight_norm:
                lin = nn.utils.weight_norm(lin)

            setattr(self, "lin" + str(l), lin)

        self.softplus = nn.Softplus(beta=100)

    def forward(self, input):
        if self.embed_fn is not None:
            input = self.embed_fn(input)

        x = input

        for l in range(0, self.num_layers - 1):
            lin = getattr(self, "lin" + str(l))

            if l in self.skip_in:
                x = torch.cat([x, input], 1) / np.sqrt(2)

            x = lin(x)

            if l < self.num_layers - 2:
                x = self.softplus(x)

        return x

    def gradient(self, x, with_sdf=False):
        x.requires_grad_(True)
        y = self.forward(x)[:,:1]
        d_output = torch.ones_like(y, requires_grad=False, device=y.device)
        gradients = torch.autograd.grad(
            outputs=y,
            inputs=x,
            grad_outputs=d_output,
            create_graph=True,
            retain_graph=True,
            only_inputs=True)[0]
        if with_sdf:
            return gradients, y
        return gradients

    def get_outputs(self, x):
        x.requires_grad_(True)
        output = self.forward(x)
        sdf = output[:,:1]
        ''' Clamping the SDF with the scene bounding sphere, so that all rays are eventually occluded '''
        if self.sdf_bounding_sphere > 0.0:
            sphere_sdf = self.sphere_scale * (self.sdf_bounding_sphere - x.norm(2,1, keepdim=True))
            sdf = torch.minimum(sdf, sphere_sdf)
        feature_vectors = output[:, 1:]
        d_output = torch.ones_like(sdf, requires_grad=False, device=sdf.device)
        gradients = torch.autograd.grad(
            outputs=sdf,
            inputs=x,
            grad_outputs=d_output,
            create_graph=True,
            retain_graph=True,
            only_inputs=True)[0]

        return sdf, feature_vectors, gradients

    def get_sdf_vals(self, x):
        sdf = self.forward(x)[:,:1]
        ''' Clamping the SDF with the scene bounding sphere, so that all rays are eventually occluded '''
        if self.sdf_bounding_sphere > 0.0:
            sphere_sdf = self.sphere_scale * (self.sdf_bounding_sphere - x.norm(2,1, keepdim=True))
            sdf = torch.minimum(sdf, sphere_sdf)
        return sdf


class RenderingNetwork(nn.Module):
    def __init__(
            self,
            feature_vector_size,
            mode,
            d_in,
            d_out,
            dims,
            weight_norm=True,
            multires_view=0,
    ):
        super().__init__()

        self.mode = mode
        dims = [d_in + feature_vector_size] + dims + [d_out]

        self.embedview_fn = None
        if multires_view > 0:
            embedview_fn, input_ch = get_embedder(multires_view)
            self.embedview_fn = embedview_fn
            dims[0] += (input_ch - 3)

        self.num_layers = len(dims)

        for l in range(0, self.num_layers - 1):
            out_dim = dims[l + 1]
            lin = nn.Linear(dims[l], out_dim)

            if weight_norm:
                lin = nn.utils.weight_norm(lin)

            setattr(self, "lin" + str(l), lin)

        self.relu = nn.ReLU()
        self.sigmoid = torch.nn.Sigmoid()

    def forward(self, points, normals, view_dirs, feature_vectors):
        if self.embedview_fn is not None:
            view_dirs = self.embedview_fn(view_dirs)

        if self.mode == 'idr':
            rendering_input = torch.cat([points, view_dirs, normals, feature_vectors], dim=-1)
        elif self.mode == 'nerf':
            rendering_input = torch.cat([view_dirs, feature_vectors], dim=-1)

        x = rendering_input

        for l in range(0, self.num_layers - 1):
            lin = getattr(self, "lin" + str(l))

            x = lin(x)

            if l < self.num_layers - 2:
                x = self.relu(x)

        x = self.sigmoid(x)
        return x

class VolSDFNetwork(nn.Module):
    def __init__(self, conf):
        super().__init__()
        self.feature_vector_size = conf.get_int('feature_vector_size')
        self.scene_bounding_sphere = conf.get_float('scene_bounding_sphere', default=1.0)
        self.white_bkgd = conf.get_bool('white_bkgd', default=False)
        self.bg_color = torch.tensor(conf.get_list("bg_color", default=[1.0, 1.0, 1.0])).float().cuda()

        self.implicit_network = ImplicitNetwork(self.feature_vector_size, 0.0 if self.white_bkgd else self.scene_bounding_sphere, **conf.get_config('implicit_network'))
        self.rendering_network = RenderingNetwork(self.feature_vector_size, **conf.get_config('rendering_network'))

        self.density = LaplaceDensity(**conf.get_config('density'))
        self.ray_sampler = ErrorBoundSampler(self.scene_bounding_sphere, **conf.get_config('ray_sampler'))

    def forward(self, input, fast=-1):
        # Parse model input
        intrinsics = input["intrinsics"]
        uv = input["uv"]
        pose = input["pose"]
        iter_step = input.get("iter_step", 1)

        ray_dirs, cam_loc = rend_util.get_camera_params(uv, pose, intrinsics)

        # we should use unnormalized ray direction for depth (sphere vs pinhole)
        ray_dirs_tmp, _ = rend_util.get_camera_params(uv, torch.eye(4).to(pose.device)[None], intrinsics)
        depth_scale = ray_dirs_tmp[0, :, 2:]

        batch_size, num_pixels, _ = ray_dirs.shape

        cam_loc = cam_loc.unsqueeze(1).repeat(1, num_pixels, 1).reshape(-1, 3)
        ray_dirs = ray_dirs.reshape(-1, 3)

        z_vals, z_samples_eik = self.ray_sampler.get_z_vals(ray_dirs, cam_loc, self, fast=fast, iter_step=iter_step) # (N_rays, N_samples)
        N_samples = z_vals.shape[1] # = 1 + 64 + 32 + 1

        points = cam_loc.unsqueeze(1) + z_vals.unsqueeze(2) * ray_dirs.unsqueeze(1) # (N_rays, N_samples, 3)
        points_flat = points.reshape(-1, 3)

        dirs = ray_dirs.unsqueeze(1).repeat(1,N_samples,1)
        dirs_flat = dirs.reshape(-1, 3)

        sdf, feature_vectors, gradients = self.implicit_network.get_outputs(points_flat)
        rgb_flat = self.rendering_network(points_flat, gradients, dirs_flat, feature_vectors)
        rgb = rgb_flat.reshape(-1, N_samples, 3)

        weights, dists = self.volume_rendering(z_vals, sdf) # (N_rays, N_samples)

        rgb_values = torch.sum(weights.unsqueeze(-1) * rgb, 1)

        depth_values = torch.sum(weights * z_vals, 1, keepdims=True) / (weights.sum(dim=1, keepdims=True) +1e-8)
        # we should scale rendered distance to depth along z direction
        depth_values = depth_scale * depth_values

        # white background assumption
        if self.white_bkgd:
            acc_map = torch.sum(weights, -1)
            rgb_values = rgb_values + (1. - acc_map[..., None]) * self.bg_color.unsqueeze(0)

        output = {
            'rgb_values': rgb_values,
            'depth_values': depth_values,
            'depth_vals': z_vals * depth_scale,
            'weights': weights,
            'xyz': points,
        }

        if self.training:
            # Sample points for the eikonal loss
            n_eik_points = batch_size * num_pixels
            eikonal_points = torch.empty(n_eik_points, 3).uniform_(-self.scene_bounding_sphere, self.scene_bounding_sphere).cuda()

            # add some of the near surface points
            eik_near_points = (cam_loc.unsqueeze(1) + z_samples_eik.unsqueeze(2) * ray_dirs.unsqueeze(1)).reshape(-1, 3)
            eikonal_points = torch.cat([eikonal_points, eik_near_points], 0)

            grad_theta = self.implicit_network.gradient(eikonal_points)
            output['grad_theta'] = grad_theta

        if not self.training:
            gradients = gradients.detach()
            normals = gradients / gradients.norm(2, -1, keepdim=True)
            normals = normals.reshape(-1, N_samples, 3)
            normal_map = torch.sum(weights.unsqueeze(-1) * normals, 1)

            output['normal_map'] = normal_map


        return output

    def volume_rendering(self, z_vals, sdf):
        density_flat = self.density(sdf)
        density = density_flat.reshape(-1, z_vals.shape[1])  # (batch_size * num_pixels) x N_samples

        dists = z_vals[:, 1:] - z_vals[:, :-1]
        dists = torch.cat([dists, torch.tensor([1e10]).cuda().unsqueeze(0).repeat(dists.shape[0], 1)], -1)

        # LOG SPACE
        free_energy = dists * density
        shifted_free_energy = torch.cat([torch.zeros(dists.shape[0], 1).cuda(), free_energy[:, :-1]], dim=-1)  # shift one step
        alpha = 1 - torch.exp(-free_energy)  # probability of it is not empty here
        transmittance = torch.exp(-torch.cumsum(shifted_free_energy, dim=-1))  # probability of everything is empty up to now
        weights = alpha * transmittance # probability of the ray hits something here

        return weights, dists

volsdf/model/network.py
cvlab-stonybrook/s-volsdf
import torch
import torch.nn as nn
import time

import volsdf.utils.general as utils
from volsdf.utils import rend_util
from volsdf.model.network import ImplicitNetwork, RenderingNetwork
from volsdf.model.density import LaplaceDensity, AbsDensity
from volsdf.model.ray_sampler import ErrorBoundSampler


"""
For modeling more complex backgrounds, we follow the inverted sphere parametrization from NeRF++ 
https://github.com/Kai-46/nerfplusplus 
"""


class VolSDFNetworkBG(nn.Module):
    def __init__(self, conf):
        super().__init__()
        self.feature_vector_size = conf.get_int('feature_vector_size')
        self.scene_bounding_sphere = conf.get_float('scene_bounding_sphere', default=1.0)

        # Foreground object's networks
        self.implicit_network = ImplicitNetwork(self.feature_vector_size, 0.0, **conf.get_config('implicit_network'))
        self.rendering_network = RenderingNetwork(self.feature_vector_size, **conf.get_config('rendering_network'))

        self.density = LaplaceDensity(**conf.get_config('density'))
        self.ray_sampler = ErrorBoundSampler(self.scene_bounding_sphere, inverse_sphere_bg=True, **conf.get_config('ray_sampler'))

        # Background's networks
        bg_feature_vector_size = conf.get_int('bg_network.feature_vector_size')
        self.bg_implicit_network = ImplicitNetwork(bg_feature_vector_size, 0.0, **conf.get_config('bg_network.implicit_network'))
        self.bg_rendering_network = RenderingNetwork(bg_feature_vector_size, **conf.get_config('bg_network.rendering_network'))
        self.bg_density = AbsDensity(**conf.get_config('bg_network.density', default={}))

    def forward(self, input, fast=-1):
        # Parse model input
        intrinsics = input["intrinsics"]
        uv = input["uv"]
        pose = input["pose"]
        iter_step = input.get("iter_step", 1)

        ray_dirs, cam_loc = rend_util.get_camera_params(uv, pose, intrinsics)

        # we should use unnormalized ray direction for depth (sphere vs pinhole)
        ray_dirs_tmp, _ = rend_util.get_camera_params(uv, torch.eye(4).to(pose.device)[None], intrinsics)
        depth_scale = ray_dirs_tmp[0, :, 2:]

        batch_size, num_pixels, _ = ray_dirs.shape

        cam_loc = cam_loc.unsqueeze(1).repeat(1, num_pixels, 1).reshape(-1, 3)
        ray_dirs = ray_dirs.reshape(-1, 3)

        z_vals, z_samples_eik = self.ray_sampler.get_z_vals(ray_dirs, cam_loc, self, fast=fast)

        z_vals, z_vals_bg = z_vals #  (z_vals, z_vals_inverse_sphere)
        z_max = z_vals[:,-1]
        z_vals = z_vals[:,:-1]
        N_samples = z_vals.shape[1]

        points = cam_loc.unsqueeze(1) + z_vals.unsqueeze(2) * ray_dirs.unsqueeze(1)
        points_flat = points.reshape(-1, 3)

        dirs = ray_dirs.unsqueeze(1).repeat(1,N_samples,1)
        dirs_flat = dirs.reshape(-1, 3)

        sdf, feature_vectors, gradients = self.implicit_network.get_outputs(points_flat)

        if not self.training:
            # use nearest viewing direction in training set, slightly better performance
            ray_dirs_near, _ = rend_util.get_camera_params(uv, input["near_pose"].cuda(), intrinsics)
            ray_dirs_near = ray_dirs_near.reshape(-1, 3)
            dirs_flat = ray_dirs_near.unsqueeze(1).repeat(1,N_samples,1).reshape(-1, 3)
        rgb_flat = self.rendering_network(points_flat, gradients, dirs_flat, feature_vectors)
        rgb = rgb_flat.reshape(-1, N_samples, 3)

        weights, bg_transmittance, dists = self.volume_rendering(z_vals, z_max, sdf)

        fg_rgb_values = torch.sum(weights.unsqueeze(-1) * rgb, 1)

        # Background rendering
        N_bg_samples = z_vals_bg.shape[1]
        z_vals_bg = torch.flip(z_vals_bg, dims=[-1, ])  # 1--->0

        bg_dirs = ray_dirs.unsqueeze(1).repeat(1,N_bg_samples,1)
        bg_locs = cam_loc.unsqueeze(1).repeat(1,N_bg_samples,1)

        bg_points, bg_depth_vals = self.depth2pts_outside(bg_locs, bg_dirs, z_vals_bg)  # [..., N_samples, 4]
        bg_points_flat = bg_points.reshape(-1, 4)
        bg_dirs_flat = bg_dirs.reshape(-1, 3)

        output = self.bg_implicit_network(bg_points_flat)
        bg_sdf = output[:,:1]
        bg_feature_vectors = output[:, 1:]
        if not self.training:
            bg_dirs_flat = ray_dirs_near.unsqueeze(1).repeat(1,N_bg_samples,1).reshape(-1, 3)
        bg_rgb_flat = self.bg_rendering_network(None, None, bg_dirs_flat, bg_feature_vectors)
        bg_rgb = bg_rgb_flat.reshape(-1, N_bg_samples, 3)

        bg_weights = self.bg_volume_rendering(z_vals_bg, bg_sdf)

        bg_rgb_values = torch.sum(bg_weights.unsqueeze(-1) * bg_rgb, 1)

        weights_all = torch.cat([weights, bg_transmittance[:, None]*bg_weights], 1)
        depth_vals_all = depth_scale * torch.cat([z_vals, bg_depth_vals], 1)
        depth_values_all = torch.sum(weights_all * depth_vals_all, 1, keepdims=True) / (weights_all.sum(dim=1, keepdims=True) +1e-8)

        depth_vals = z_vals * depth_scale
        depth_values = torch.sum(weights * depth_vals, 1, keepdims=True) / (weights.sum(dim=1, keepdims=True) +1e-8)

        # Composite foreground and background
        bg_rgb_values = bg_transmittance.unsqueeze(-1) * bg_rgb_values
        rgb_values = fg_rgb_values + bg_rgb_values

        output = {
            'rgb_values': rgb_values,
            'depth_values_all': depth_values_all,
            'depth_values': depth_values,
            'depth_vals': depth_vals,
            'weights': weights,
            'xyz': points.detach(),
        }

        if self.training:
            # Sample points for the eikonal loss
            n_eik_points = batch_size * num_pixels
            eikonal_points = torch.empty(n_eik_points, 3).uniform_(-self.scene_bounding_sphere, self.scene_bounding_sphere).cuda()

            # add some of the near surface points
            eik_near_points = (cam_loc.unsqueeze(1) + z_samples_eik.unsqueeze(2) * ray_dirs.unsqueeze(1)).reshape(-1, 3)
            eikonal_points = torch.cat([eikonal_points, eik_near_points], 0)

            grad_theta = self.implicit_network.gradient(eikonal_points)
            output['grad_theta'] = grad_theta

        if not self.training:
            gradients = gradients.detach()
            normals = gradients / gradients.norm(2, -1, keepdim=True)
            normals = normals.reshape(-1, N_samples, 3)
            normal_map = torch.sum(weights.unsqueeze(-1) * normals, 1)

            output['normal_map'] = normal_map

        return output

    def volume_rendering(self, z_vals, z_max, sdf):
        density_flat = self.density(sdf)
        density = density_flat.reshape(-1, z_vals.shape[1]) # (batch_size * num_pixels) x N_samples

        # included also the dist from the sphere intersection
        dists = z_vals[:, 1:] - z_vals[:, :-1]
        dists = torch.cat([dists, z_max.unsqueeze(-1) - z_vals[:, -1:]], -1)

        # LOG SPACE
        free_energy = dists * density
        shifted_free_energy = torch.cat([torch.zeros(dists.shape[0], 1).cuda(), free_energy], dim=-1)  # add 0 for transperancy 1 at t_0
        alpha = 1 - torch.exp(-free_energy)  # probability of it is not empty here
        transmittance = torch.exp(-torch.cumsum(shifted_free_energy, dim=-1))  # probability of everything is empty up to now
        fg_transmittance = transmittance[:, :-1]
        weights = alpha * fg_transmittance  # probability of the ray hits something here
        bg_transmittance = transmittance[:, -1]  # factor to be multiplied with the bg volume rendering

        return weights, bg_transmittance, dists

    def bg_volume_rendering(self, z_vals_bg, bg_sdf):
        bg_density_flat = self.bg_density(bg_sdf)
        bg_density = bg_density_flat.reshape(-1, z_vals_bg.shape[1]) # (batch_size * num_pixels) x N_samples

        bg_dists = z_vals_bg[:, :-1] - z_vals_bg[:, 1:]
        bg_dists = torch.cat([bg_dists, torch.tensor([1e10]).cuda().unsqueeze(0).repeat(bg_dists.shape[0], 1)], -1)

        # LOG SPACE
        bg_free_energy = bg_dists * bg_density
        bg_shifted_free_energy = torch.cat([torch.zeros(bg_dists.shape[0], 1).cuda(), bg_free_energy[:, :-1]], dim=-1)  # shift one step
        bg_alpha = 1 - torch.exp(-bg_free_energy)  # probability of it is not empty here
        bg_transmittance = torch.exp(-torch.cumsum(bg_shifted_free_energy, dim=-1))  # probability of everything is empty up to now
        bg_weights = bg_alpha * bg_transmittance # probability of the ray hits something here

        return bg_weights

    def depth2pts_outside(self, ray_o, ray_d, depth):

        '''
        ray_o, ray_d: [..., 3]
        depth: [...]; inverse of distance to sphere origin
        '''

        o_dot_d = torch.sum(ray_d * ray_o, dim=-1)
        under_sqrt = o_dot_d ** 2 - ((ray_o ** 2).sum(-1) - self.scene_bounding_sphere ** 2)
        d_sphere = torch.sqrt(under_sqrt) - o_dot_d
        p_sphere = ray_o + d_sphere.unsqueeze(-1) * ray_d
        p_mid = ray_o - o_dot_d.unsqueeze(-1) * ray_d
        p_mid_norm = torch.norm(p_mid, dim=-1)

        rot_axis = torch.cross(ray_o, p_sphere, dim=-1)
        rot_axis = rot_axis / torch.norm(rot_axis, dim=-1, keepdim=True)
        phi = torch.asin(p_mid_norm / self.scene_bounding_sphere)
        theta = torch.asin(p_mid_norm * depth)  # depth is inside [0, 1]
        rot_angle = (phi - theta).unsqueeze(-1)  # [..., 1]

        # now rotate p_sphere
        # Rodrigues formula: https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
        p_sphere_new = p_sphere * torch.cos(rot_angle) + \
                       torch.cross(rot_axis, p_sphere, dim=-1) * torch.sin(rot_angle) + \
                       rot_axis * torch.sum(rot_axis * p_sphere, dim=-1, keepdim=True) * (1. - torch.cos(rot_angle))
        p_sphere_new = p_sphere_new / torch.norm(p_sphere_new, dim=-1, keepdim=True)
        pts = torch.cat((p_sphere_new, depth.unsqueeze(-1)), dim=-1)

        # now calculate conventional depth
        d1 = - o_dot_d / torch.sum(ray_d * ray_d, dim=-1)
        ray_d_cos = 1. / torch.norm(ray_d, dim=-1)
        depth_real = 1. / (depth + 1e-6) * torch.cos(theta) * ray_d_cos + d1
        return pts, depth_real

volsdf/model/network_bg.py
cvlab-stonybrook/s-volsdf
import abc
import torch

from volsdf.utils import rend_util

class RaySampler(metaclass=abc.ABCMeta):
    def __init__(self,near, far):
        self.near = near
        self.far = far

    @abc.abstractmethod
    def get_z_vals(self, ray_dirs, cam_loc, model):
        pass

class UniformSampler(RaySampler):
    def __init__(self, scene_bounding_sphere, near, N_samples, take_sphere_intersection=False, far=-1):
        super().__init__(near, 2.0 * scene_bounding_sphere if far == -1 else far)  # default far is 2*R
        self.N_samples = N_samples
        self.scene_bounding_sphere = scene_bounding_sphere
        self.take_sphere_intersection = take_sphere_intersection

    def get_z_vals(self, ray_dirs, cam_loc, model, iter_step=None):
        if not self.take_sphere_intersection:
            near, far = self.near * torch.ones(ray_dirs.shape[0], 1).cuda(), self.far * torch.ones(ray_dirs.shape[0], 1).cuda()
        else:
            sphere_intersections = rend_util.get_sphere_intersections(cam_loc, ray_dirs, r=self.scene_bounding_sphere)
            near = self.near * torch.ones(ray_dirs.shape[0], 1).cuda()
            far = sphere_intersections[:,1:]

        t_vals = torch.linspace(0., 1., steps=self.N_samples).cuda()
        z_vals = near * (1. - t_vals) + far * (t_vals)

        if model.training:
            # get intervals between samples
            mids = .5 * (z_vals[..., 1:] + z_vals[..., :-1])
            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).cuda()

            z_vals = lower + (upper - lower) * t_rand

        return z_vals


class ErrorBoundSampler(RaySampler):
    def __init__(self, scene_bounding_sphere, near, N_samples, N_samples_eval, N_samples_extra,
                 eps, beta_iters, max_total_iters,
                 inverse_sphere_bg=False, N_samples_inverse_sphere=0, add_tiny=0.0):
        super().__init__(near, 2.0 * scene_bounding_sphere)
        self.N_samples = N_samples
        self.N_samples_eval = N_samples_eval
        self.uniform_sampler = UniformSampler(scene_bounding_sphere, near, N_samples_eval, take_sphere_intersection=inverse_sphere_bg)

        self.N_samples_extra = N_samples_extra

        self.eps = eps
        self.beta_iters = beta_iters
        self.max_total_iters = max_total_iters
        self.scene_bounding_sphere = scene_bounding_sphere
        self.add_tiny = add_tiny

        self.inverse_sphere_bg = inverse_sphere_bg
        if inverse_sphere_bg:
            self.inverse_sphere_sampler = UniformSampler(1.0, 0.0, N_samples_inverse_sphere, False, far=1.0)

    def get_z_vals(self, ray_dirs, cam_loc, model, fast=-1, iter_step=None):
        max_total_iters = fast if fast >= 0 else self.max_total_iters
        beta0 = model.density.get_beta().detach()

        # Start with uniform sampling
        z_vals = self.uniform_sampler.get_z_vals(ray_dirs, cam_loc, model, iter_step=iter_step)
        samples, samples_idx = z_vals, None

        # Get maximum beta from the upper bound (Lemma 2)
        dists = z_vals[:, 1:] - z_vals[:, :-1]
        bound = (1.0 / (4.0 * torch.log(torch.tensor(self.eps + 1.0)))) * (dists ** 2.).sum(-1)
        beta = torch.sqrt(bound)

        total_iters, not_converge = 0, True

        # Algorithm 1: soft ray tracing
        while not_converge and total_iters < max_total_iters:
            points = cam_loc.unsqueeze(1) + samples.unsqueeze(2) * ray_dirs.unsqueeze(1)
            points_flat = points.reshape(-1, 3)

            # Calculating the SDF only for the new sampled points
            with torch.no_grad():
                samples_sdf = model.implicit_network.get_sdf_vals(points_flat)
            if samples_idx is not None:
                sdf_merge = torch.cat([sdf.reshape(-1, z_vals.shape[1] - samples.shape[1]),
                                       samples_sdf.reshape(-1, samples.shape[1])], -1)
                sdf = torch.gather(sdf_merge, 1, samples_idx).reshape(-1, 1)
            else:
                sdf = samples_sdf


            # Calculating the bound d* (Theorem 1)
            d = sdf.reshape(z_vals.shape)
            dists = z_vals[:, 1:] - z_vals[:, :-1]
            a, b, c = dists, d[:, :-1].abs(), d[:, 1:].abs()
            first_cond = a.pow(2) + b.pow(2) <= c.pow(2)
            second_cond = a.pow(2) + c.pow(2) <= b.pow(2)
            d_star = torch.zeros(z_vals.shape[0], z_vals.shape[1] - 1).cuda()
            d_star[first_cond] = b[first_cond]
            d_star[second_cond] = c[second_cond]
            s = (a + b + c) / 2.0
            area_before_sqrt = s * (s - a) * (s - b) * (s - c)
            mask = ~first_cond & ~second_cond & (b + c - a > 0)
            d_star[mask] = (2.0 * torch.sqrt(area_before_sqrt[mask])) / (a[mask])
            d_star = (d[:, 1:].sign() * d[:, :-1].sign() == 1) * d_star  # Fixing the sign


            # Updating beta using line search
            curr_error = self.get_error_bound(beta0, model, sdf, z_vals, dists, d_star)
            beta[curr_error <= self.eps] = beta0
            beta_min, beta_max = beta0.unsqueeze(0).repeat(z_vals.shape[0]), beta
            for j in range(self.beta_iters):
                beta_mid = (beta_min + beta_max) / 2.
                curr_error = self.get_error_bound(beta_mid.unsqueeze(-1), model, sdf, z_vals, dists, d_star)
                beta_max[curr_error <= self.eps] = beta_mid[curr_error <= self.eps]
                beta_min[curr_error > self.eps] = beta_mid[curr_error > self.eps]
            beta = beta_max

            # Upsample more points
            density = model.density(sdf.reshape(z_vals.shape), beta=beta.unsqueeze(-1))
            dists = torch.cat([dists, torch.tensor([1e10]).cuda().unsqueeze(0).repeat(dists.shape[0], 1)], -1)
            free_energy = dists * density
            shifted_free_energy = torch.cat([torch.zeros(dists.shape[0], 1).cuda(), free_energy[:, :-1]], dim=-1)
            alpha = 1 - torch.exp(-free_energy)
            transmittance = torch.exp(-torch.cumsum(shifted_free_energy, dim=-1))
            weights = alpha * transmittance  # probability of the ray hits something here

            #  Check if we are done and this is the last sampling
            total_iters += 1
            not_converge = beta.max() > beta0

            if not_converge and total_iters < max_total_iters:
                ''' Sample more points proportional to the current error bound'''

                N = self.N_samples_eval

                bins = z_vals
                error_per_section = torch.exp(-d_star / beta.unsqueeze(-1)) * (dists[:,:-1] ** 2.) / (4 * beta.unsqueeze(-1) ** 2)
                error_integral = torch.cumsum(error_per_section, dim=-1)
                bound_opacity = (torch.clamp(torch.exp(error_integral),max=1.e6) - 1.0) * transmittance[:,:-1]

                pdf = bound_opacity + self.add_tiny
                pdf = pdf / torch.sum(pdf, -1, keepdim=True)
                cdf = torch.cumsum(pdf, -1)
                cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)

            else:
                ''' Sample the final sample set to be used in the volume rendering integral '''

                N = self.N_samples

                bins = z_vals
                pdf = weights[..., :-1]
                pdf = pdf + 1e-5  # prevent nans
                pdf = pdf / torch.sum(pdf, -1, keepdim=True)
                cdf = torch.cumsum(pdf, -1)
                cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], -1)  # (batch, len(bins))


            # Invert CDF
            if (not_converge and total_iters < max_total_iters) or (not model.training):
                u = torch.linspace(0., 1., steps=N).cuda().unsqueeze(0).repeat(cdf.shape[0], 1)
            else:
                u = torch.rand(list(cdf.shape[:-1]) + [N]).cuda()
            u = u.contiguous()

            inds = torch.searchsorted(cdf, 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])


            # Adding samples if we not converged
            if not_converge and total_iters < max_total_iters:
                z_vals, samples_idx = torch.sort(torch.cat([z_vals, samples], -1), -1)


        z_samples = samples

        near, far = self.near * torch.ones(ray_dirs.shape[0], 1).cuda(), self.far * torch.ones(ray_dirs.shape[0],1).cuda()
        if self.inverse_sphere_bg: # if inverse sphere then need to add the far sphere intersection
            far = rend_util.get_sphere_intersections(cam_loc, ray_dirs, r=self.scene_bounding_sphere)[:,1:]

        if self.N_samples_extra > 0:
            if model.training:
                sampling_idx = torch.randperm(z_vals.shape[1])[:self.N_samples_extra]
            else:
                sampling_idx = torch.linspace(0, z_vals.shape[1]-1, self.N_samples_extra).long()
            z_vals_extra = torch.cat([near, far, z_vals[:,sampling_idx]], -1)
        else:
            z_vals_extra = torch.cat([near, far], -1)

        z_vals, _ = torch.sort(torch.cat([z_samples, z_vals_extra], -1), -1)

        # add some of the near surface points
        idx = torch.randint(z_vals.shape[-1], (z_vals.shape[0],)).cuda()
        z_samples_eik = torch.gather(z_vals, 1, idx.unsqueeze(-1))

        if self.inverse_sphere_bg:
            z_vals_inverse_sphere = self.inverse_sphere_sampler.get_z_vals(ray_dirs, cam_loc, model)
            z_vals_inverse_sphere = z_vals_inverse_sphere * (1./self.scene_bounding_sphere)
            z_vals = (z_vals, z_vals_inverse_sphere)

        return z_vals, z_samples_eik

    def get_error_bound(self, beta, model, sdf, z_vals, dists, d_star):
        density = model.density(sdf.reshape(z_vals.shape), beta=beta)
        shifted_free_energy = torch.cat([torch.zeros(dists.shape[0], 1).cuda(), dists * density[:, :-1]], dim=-1)
        integral_estimation = torch.cumsum(shifted_free_energy, dim=-1)
        error_per_section = torch.exp(-d_star / beta) * (dists ** 2.) / (4 * beta ** 2)
        error_integral = torch.cumsum(error_per_section, dim=-1)
        bound_opacity = (torch.clamp(torch.exp(error_integral), max=1.e6) - 1.0) * torch.exp(-integral_estimation[:, :-1])

        return bound_opacity.max(-1)[0]



volsdf/model/ray_sampler.py
cvlab-stonybrook/s-volsdf
import os
from glob import glob
import torch
import numpy as np

def mkdir_ifnotexists(directory):
    if not os.path.exists(directory):
        os.makedirs(directory)

def get_class(kls):
    parts = kls.split('.')
    module = ".".join(parts[:-1])
    m = __import__(module)
    for comp in parts[1:]:
        m = getattr(m, comp)
    return m

def glob_imgs(path):
    imgs = []
    for ext in ['*.png', '*.jpg', '*.JPEG', '*.JPG']:
        imgs.extend(glob(os.path.join(path, ext)))
    return imgs

def split_input(model_input, total_pixels, n_pixels=10000):
    '''
     Split the input to fit Cuda memory for large resolution.
     Can decrease the value of n_pixels in case of cuda out of memory error.
     '''
    split = []
    for i, indx in enumerate(torch.split(torch.arange(total_pixels).cuda(), n_pixels, dim=0)):
        data = model_input.copy()
        data['uv'] = torch.index_select(model_input['uv'], 1, indx)
        if 'object_mask' in data:
            data['object_mask'] = torch.index_select(model_input['object_mask'], 1, indx)
        for key in ['rgb']:
            if key in data:
                data[key] = torch.index_select(model_input[key], 1, indx)
        split.append(data)
    return split

def merge_output(res, total_pixels, batch_size):
    ''' Merge the split output. '''

    model_outputs = {}
    for entry in res[0]:
        if res[0][entry] is None:
            continue
        if len(res[0][entry].shape) == 1:
            model_outputs[entry] = torch.cat([r[entry].reshape(batch_size, -1, 1) for r in res],
                                             1).reshape(batch_size * total_pixels)
        elif len(res[0][entry].shape) == 2:
            model_outputs[entry] = torch.cat([r[entry].reshape(batch_size, -1, r[entry].shape[-1]) for r in res],
                                             1).reshape(batch_size * total_pixels, -1)
        elif len(res[0][entry].shape) == 3:
            model_outputs[entry] = torch.cat([r[entry].reshape(batch_size, -1, r[entry].shape[-2], r[entry].shape[-1]) for r in res],
                                             1).reshape(batch_size * total_pixels, -1, res[0][entry].shape[-1])
        else:
            raise NotImplementedError

    return model_outputs
volsdf/utils/general.py
cvlab-stonybrook/s-volsdf
import plotly.graph_objs as go
import plotly.offline as offline
from plotly.subplots import make_subplots
import numpy as np
import torch
from skimage import measure
import torchvision
import trimesh
from PIL import Image
import matplotlib.pyplot as plt
import cv2
import math
import matplotlib.cm as cm

from volsdf.utils import rend_util

def stacked_plot(indices, plot_data, path, epoch, img_res, plot_nimgs, resolution, grid_boundary, level=0):
    cam_loc, cam_dir = rend_util.get_camera_for_plot(plot_data['pose'])

    # plot images (2+300+2+300, 2+400+2, 3) 0 ~ 255
    img = plot_images(plot_data['rgb_eval'], plot_data['rgb_gt'], path, epoch, plot_nimgs, img_res, indices, save=False)

    # plot normal maps (300, 400, 3) 0 ~ 255
    normal = plot_normal_maps(plot_data['normal_map'], path, epoch, plot_nimgs, img_res, indices, save=False)

    # plot depth maps (300, 400, 3) 0 ~ 255
    depth = plot_depth_maps(plot_data['depth_map'], plot_data['acc'], path, epoch, plot_nimgs, img_res, indices, save=False)

    stack = [img[2:2+img_res[0], 2:2+img_res[1]], depth, normal]
    stack = [torch.from_numpy(_).permute(2, 0, 1) for _ in stack]
    return stack

def plot(implicit_network, indices, plot_data, path, epoch, img_res, plot_nimgs, resolution, grid_boundary, level=0):

    if plot_data is not None:
        cam_loc, cam_dir = rend_util.get_camera_for_plot(plot_data['pose'])

        # plot images
        plot_images(plot_data['rgb_eval'], plot_data['rgb_gt'], path, epoch, plot_nimgs, img_res, indices)

        # plot normal maps
        plot_normal_maps(plot_data['normal_map'], path, epoch, plot_nimgs, img_res, indices)

        # plot depth maps
        plot_depth_maps(plot_data['depth_map'], plot_data['acc'], path, epoch, plot_nimgs, img_res, indices)

        # concat output images to single large image
        images = []
        for name in ["rendering", "depth", "normal"]:
            images.append(cv2.imread('{0}/{1}_{2}_{3}.png'.format(path, name, epoch, indices[0]))) 
        images = [images[0][2:2+img_res[0], :]] + images[1:]

        images = np.concatenate(images, axis=1)
        cv2.imwrite('{0}/merge_{1}_{2}.png'.format(path, epoch, indices[0]), images)

    del plot_data
    torch.cuda.empty_cache()

    surface_traces = get_surface_trace(path=path, 
                                        epoch=epoch,
                                        sdf=lambda x: implicit_network(x)[:, 0],
                                        resolution=resolution,
                                        grid_boundary=grid_boundary,
                                        level=level
                                        )
    del surface_traces
 

def get_surface_trace(path, epoch, sdf, resolution=100, grid_boundary=[-2.0, 2.0], return_mesh=False, level=0):
    grid = get_grid_uniform(resolution, grid_boundary)
    points = grid['grid_points']

    z = []
    for i, pnts in enumerate(torch.split(points, 100000, dim=0)):
        z.append(sdf(pnts).detach().cpu().numpy())
    z = np.concatenate(z, axis=0)

    if (not (np.min(z) > level or np.max(z) < level)):

        z = z.astype(np.float32)

        verts, faces, normals, values = measure.marching_cubes(
            volume=z.reshape(grid['xyz'][1].shape[0], grid['xyz'][0].shape[0],
                             grid['xyz'][2].shape[0]).transpose([1, 0, 2]),
            level=level,
            spacing=(grid['xyz'][0][2] - grid['xyz'][0][1],
                     grid['xyz'][0][2] - grid['xyz'][0][1],
                     grid['xyz'][0][2] - grid['xyz'][0][1]))

        verts = verts + np.array([grid['xyz'][0][0], grid['xyz'][1][0], grid['xyz'][2][0]])

        I, J, K = faces.transpose()

        traces = [go.Mesh3d(x=verts[:, 0], y=verts[:, 1], z=verts[:, 2],
                            i=I, j=J, k=K, name='implicit_surface',
                            color='#ffffff', opacity=1.0, flatshading=False,
                            lighting=dict(diffuse=1, ambient=0, specular=0),
                            lightposition=dict(x=0, y=0, z=-1), showlegend=True)]

        meshexport = trimesh.Trimesh(verts, faces, normals)
        meshexport.export('{0}/surface_{1}.ply'.format(path, epoch), 'ply')

        if return_mesh:
            return meshexport
        return traces
    return None

def get_surface_high_res_mesh(sdf, resolution=100, grid_boundary=[-2.0, 2.0], level=0, take_components=True):
    # get low res mesh to sample point cloud
    grid = get_grid_uniform(100, grid_boundary)
    z = []
    points = grid['grid_points']

    for i, pnts in enumerate(torch.split(points, 100000, dim=0)):
        z.append(sdf(pnts).detach().cpu().numpy())
    z = np.concatenate(z, axis=0)

    z = z.astype(np.float32)

    verts, faces, normals, values = measure.marching_cubes(
        volume=z.reshape(grid['xyz'][1].shape[0], grid['xyz'][0].shape[0],
                         grid['xyz'][2].shape[0]).transpose([1, 0, 2]),
        level=level,
        spacing=(grid['xyz'][0][2] - grid['xyz'][0][1],
                 grid['xyz'][0][2] - grid['xyz'][0][1],
                 grid['xyz'][0][2] - grid['xyz'][0][1]))

    verts = verts + np.array([grid['xyz'][0][0], grid['xyz'][1][0], grid['xyz'][2][0]])

    mesh_low_res = trimesh.Trimesh(verts, faces, normals)
    if take_components:
        components = mesh_low_res.split(only_watertight=False)
        areas = np.array([c.area for c in components], dtype=np.float)
        mesh_low_res = components[areas.argmax()]

    recon_pc = trimesh.sample.sample_surface(mesh_low_res, 10000)[0]
    recon_pc = torch.from_numpy(recon_pc).float().cuda()

    # Center and align the recon pc
    s_mean = recon_pc.mean(dim=0)
    s_cov = recon_pc - s_mean
    s_cov = torch.mm(s_cov.transpose(0, 1), s_cov)
    vecs = torch.view_as_real(torch.linalg.eig(s_cov)[1].transpose(0, 1))[:, :, 0]
    if torch.det(vecs) < 0:
        vecs = torch.mm(torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]]).cuda().float(), vecs)
    helper = torch.bmm(vecs.unsqueeze(0).repeat(recon_pc.shape[0], 1, 1),
                       (recon_pc - s_mean).unsqueeze(-1)).squeeze()

    grid_aligned = get_grid(helper.cpu(), resolution)

    grid_points = grid_aligned['grid_points']

    g = []
    for i, pnts in enumerate(torch.split(grid_points, 100000, dim=0)):
        g.append(torch.bmm(vecs.unsqueeze(0).repeat(pnts.shape[0], 1, 1).transpose(1, 2),
                           pnts.unsqueeze(-1)).squeeze() + s_mean)
    grid_points = torch.cat(g, dim=0)

    # MC to new grid
    points = grid_points
    z = []
    for i, pnts in enumerate(torch.split(points, 100000, dim=0)):
        z.append(sdf(pnts).detach().cpu().numpy())
    z = np.concatenate(z, axis=0)

    meshexport = None
    if (not (np.min(z) > level or np.max(z) < level)):

        z = z.astype(np.float32)

        verts, faces, normals, values = measure.marching_cubes(
            volume=z.reshape(grid_aligned['xyz'][1].shape[0], grid_aligned['xyz'][0].shape[0],
                             grid_aligned['xyz'][2].shape[0]).transpose([1, 0, 2]),
            level=level,
            spacing=(grid_aligned['xyz'][0][2] - grid_aligned['xyz'][0][1],
                     grid_aligned['xyz'][0][2] - grid_aligned['xyz'][0][1],
                     grid_aligned['xyz'][0][2] - grid_aligned['xyz'][0][1]))

        verts = torch.from_numpy(verts).cuda().float()
        verts = torch.bmm(vecs.unsqueeze(0).repeat(verts.shape[0], 1, 1).transpose(1, 2),
                   verts.unsqueeze(-1)).squeeze()
        verts = (verts + grid_points[0]).cpu().numpy()

        meshexport = trimesh.Trimesh(verts, faces, normals)

    return meshexport

def get_surface_by_grid(grid_params, sdf, resolution=100, level=0, higher_res=False, splitn=100000):
    grid_params = grid_params * [[1.5], [1.0]]

    # params = PLOT_DICT[scan_id]
    input_min = torch.tensor(grid_params[0]).float()
    input_max = torch.tensor(grid_params[1]).float()

    if higher_res:
        # get low res mesh to sample point cloud
        grid = get_grid(None, 100, input_min=input_min, input_max=input_max, eps=0.0)
        z = []
        points = grid['grid_points']

        for i, pnts in enumerate(torch.split(points, splitn, dim=0)):
            z.append(sdf(pnts).detach().cpu().numpy())
        z = np.concatenate(z, axis=0)

        z = z.astype(np.float32)

        verts, faces, normals, values = measure.marching_cubes(
            volume=z.reshape(grid['xyz'][1].shape[0], grid['xyz'][0].shape[0],
                             grid['xyz'][2].shape[0]).transpose([1, 0, 2]),
            level=level,
            spacing=(grid['xyz'][0][2] - grid['xyz'][0][1],
                     grid['xyz'][0][2] - grid['xyz'][0][1],
                     grid['xyz'][0][2] - grid['xyz'][0][1]))

        verts = verts + np.array([grid['xyz'][0][0], grid['xyz'][1][0], grid['xyz'][2][0]])

        mesh_low_res = trimesh.Trimesh(verts, faces, normals)
        components = mesh_low_res.split(only_watertight=False)
        areas = np.array([c.area for c in components], dtype=np.float)
        mesh_low_res = components[areas.argmax()]

        recon_pc = trimesh.sample.sample_surface(mesh_low_res, 10000)[0]
        recon_pc = torch.from_numpy(recon_pc).float().cuda()

        # Center and align the recon pc
        s_mean = recon_pc.mean(dim=0)
        s_cov = recon_pc - s_mean
        s_cov = torch.mm(s_cov.transpose(0, 1), s_cov)
        vecs = torch.view_as_real(torch.linalg.eig(s_cov)[1].transpose(0, 1))[:, :, 0]
        if torch.det(vecs) < 0:
            vecs = torch.mm(torch.tensor([[1, 0, 0], [0, 0, 1], [0, 1, 0]]).cuda().float(), vecs)
        helper = torch.bmm(vecs.unsqueeze(0).repeat(recon_pc.shape[0], 1, 1),
                           (recon_pc - s_mean).unsqueeze(-1)).squeeze()

        grid_aligned = get_grid(helper.cpu(), resolution, eps=0.01)
    else:
        grid_aligned = get_grid(None, resolution, input_min=input_min, input_max=input_max, eps=0.0)

    grid_points = grid_aligned['grid_points']

    if higher_res:
        g = []
        for i, pnts in enumerate(torch.split(grid_points, splitn, dim=0)):
            g.append(torch.bmm(vecs.unsqueeze(0).repeat(pnts.shape[0], 1, 1).transpose(1, 2),
                               pnts.unsqueeze(-1)).squeeze() + s_mean)
        grid_points = torch.cat(g, dim=0)

    # MC to new grid
    points = grid_points
    z = []
    for i, pnts in enumerate(torch.split(points, splitn, dim=0)):
        z.append(sdf(pnts).detach().cpu().numpy())
    z = np.concatenate(z, axis=0)

    meshexport = None
    if (not (np.min(z) > level or np.max(z) < level)):

        z = z.astype(np.float32)

        verts, faces, normals, values = measure.marching_cubes(
            volume=z.reshape(grid_aligned['xyz'][1].shape[0], grid_aligned['xyz'][0].shape[0],
                             grid_aligned['xyz'][2].shape[0]).transpose([1, 0, 2]),
            level=level,
            spacing=(grid_aligned['xyz'][0][2] - grid_aligned['xyz'][0][1],
                     grid_aligned['xyz'][0][2] - grid_aligned['xyz'][0][1],
                     grid_aligned['xyz'][0][2] - grid_aligned['xyz'][0][1]))

        if higher_res:
            verts = torch.from_numpy(verts).cuda().float()
            verts = torch.bmm(vecs.unsqueeze(0).repeat(verts.shape[0], 1, 1).transpose(1, 2),
                       verts.unsqueeze(-1)).squeeze()
            verts = (verts + grid_points[0]).cpu().numpy()
        else:
            verts = verts + np.array([grid_aligned['xyz'][0][0], grid_aligned['xyz'][1][0], grid_aligned['xyz'][2][0]])

        meshexport = trimesh.Trimesh(verts, faces, normals)

        # CUTTING MESH ACCORDING TO THE BOUNDING BOX
        if higher_res:
            bb = grid_params
            transformation = np.eye(4)
            transformation[:3, 3] = (bb[1,:] + bb[0,:])/2.
            bounding_box = trimesh.creation.box(extents=bb[1,:] - bb[0,:], transform=transformation)

            meshexport = meshexport.slice_plane(bounding_box.facets_origin, -bounding_box.facets_normal)

    return meshexport

def get_grid_uniform(resolution, grid_boundary=[-2.0, 2.0]):
    x = np.linspace(grid_boundary[0], grid_boundary[1], resolution)
    y = x
    z = x

    xx, yy, zz = np.meshgrid(x, y, z)
    grid_points = torch.tensor(np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T, dtype=torch.float)

    return {"grid_points": grid_points.cuda(),
            "shortest_axis_length": 2.0,
            "xyz": [x, y, z],
            "shortest_axis_index": 0}

def get_grid(points, resolution, input_min=None, input_max=None, eps=0.1):
    if input_min is None or input_max is None:
        input_min = torch.min(points, dim=0)[0].squeeze().numpy()
        input_max = torch.max(points, dim=0)[0].squeeze().numpy()

    bounding_box = input_max - input_min
    shortest_axis = np.argmin(bounding_box)
    if (shortest_axis == 0):
        x = np.linspace(input_min[shortest_axis] - eps,
                        input_max[shortest_axis] + eps, resolution)
        length = np.max(x) - np.min(x)
        y = np.arange(input_min[1] - eps, input_max[1] + length / (x.shape[0] - 1) + eps, length / (x.shape[0] - 1))
        z = np.arange(input_min[2] - eps, input_max[2] + length / (x.shape[0] - 1) + eps, length / (x.shape[0] - 1))
    elif (shortest_axis == 1):
        y = np.linspace(input_min[shortest_axis] - eps,
                        input_max[shortest_axis] + eps, resolution)
        length = np.max(y) - np.min(y)
        x = np.arange(input_min[0] - eps, input_max[0] + length / (y.shape[0] - 1) + eps, length / (y.shape[0] - 1))
        z = np.arange(input_min[2] - eps, input_max[2] + length / (y.shape[0] - 1) + eps, length / (y.shape[0] - 1))
    elif (shortest_axis == 2):
        z = np.linspace(input_min[shortest_axis] - eps,
                        input_max[shortest_axis] + eps, resolution)
        length = np.max(z) - np.min(z)
        x = np.arange(input_min[0] - eps, input_max[0] + length / (z.shape[0] - 1) + eps, length / (z.shape[0] - 1))
        y = np.arange(input_min[1] - eps, input_max[1] + length / (z.shape[0] - 1) + eps, length / (z.shape[0] - 1))

    xx, yy, zz = np.meshgrid(x, y, z)
    grid_points = torch.tensor(np.vstack([xx.ravel(), yy.ravel(), zz.ravel()]).T, dtype=torch.float).cuda()
    return {"grid_points": grid_points,
            "shortest_axis_length": length,
            "xyz": [x, y, z],
            "shortest_axis_index": shortest_axis}


def plot_depth_maps(depth_maps, acc, path, epoch, plot_nrow, img_res, indices, save=True):
    depth_maps_plot = lin2img(depth_maps[..., None], img_res)
    acc_plot = lin2img(acc[..., None], img_res)

    assert sum(depth_maps_plot.shape[:2]) + sum(acc_plot.shape[:2]) == 4
    tensor = visualize_depth(depth_maps_plot[0,0].cpu().numpy(), acc_plot[0,0].cpu().numpy())

    scale_factor = 255
    tensor = (tensor * scale_factor).astype(np.uint8)

    if not save:
        return tensor
    img = Image.fromarray(tensor)
    img.save('{0}/depth_{1}_{2}.png'.format(path, epoch, indices[0]))

def plot_normal_maps(normal_maps, path, epoch, plot_nrow, img_res, indices, save=True):
    normal_maps_plot = lin2img(normal_maps, img_res)

    tensor = torchvision.utils.make_grid(normal_maps_plot,
                                         scale_each=False,
                                         normalize=False,
                                         nrow=plot_nrow).cpu().detach().numpy()
    tensor = tensor.transpose(1, 2, 0)
    scale_factor = 255
    tensor = (tensor * scale_factor).astype(np.uint8)

    if not save:
        return tensor
    img = Image.fromarray(tensor)
    img.save('{0}/normal_{1}_{2}.png'.format(path, epoch, indices[0]))

def plot_images(rgb_points, ground_true, path, epoch, plot_nrow, img_res, indices, save=True):
    ground_true = ground_true.cpu()

    output_vs_gt = torch.cat((rgb_points, ground_true), dim=0)
    output_vs_gt_plot = lin2img(output_vs_gt, img_res)

    tensor = torchvision.utils.make_grid(output_vs_gt_plot,
                                         scale_each=False,
                                         normalize=False,
                                         nrow=plot_nrow).cpu().detach().numpy()

    tensor = tensor.transpose(1, 2, 0)
    scale_factor = 255
    tensor = (tensor * scale_factor).astype(np.uint8)

    if not save:
        return tensor
    img = Image.fromarray(tensor)
    img.save('{0}/rendering_{1}_{2}.png'.format(path, epoch, indices[0]))


def lin2img(tensor, img_res):
    batch_size, num_samples, channels = tensor.shape
    return tensor.permute(0, 2, 1).view(batch_size, channels, img_res[0], img_res[1])

def visualize_depth(x, acc, lo=None, hi=None):
    """Visualizes depth maps."""

    depth_curve_fn = lambda x: -np.log(x + np.finfo(np.float32).eps)
    return visualize_cmap(
        x, acc, cm.get_cmap('turbo'), curve_fn=depth_curve_fn, lo=lo, hi=hi)

def weighted_percentile(x, w, ps, assume_sorted=False):
    """Compute the weighted percentile(s) of a single vector."""
    x = x.reshape([-1])
    w = w.reshape([-1])
    if not assume_sorted:
        sortidx = np.argsort(x)
        x, w = x[sortidx], w[sortidx]
    acc_w = np.cumsum(w)
    return np.interp(np.array(ps) * (acc_w[-1] / 100), acc_w, x)

def matte(vis, acc, dark=0.8, light=1.0, width=8):
    """Set non-accumulated pixels to a Photoshop-esque checker pattern."""
    bg_mask = np.logical_xor(
        (np.arange(acc.shape[0]) % (2 * width) // width)[:, None],
        (np.arange(acc.shape[1]) % (2 * width) // width)[None, :])
    bg = np.where(bg_mask, light, dark)
    return vis * acc[:, :, None] + (bg * (1 - acc))[:, :, None]

def visualize_cmap(value,
                   weight,
                   colormap,
                   lo=None,
                   hi=None,
                   percentile=99.,
                   curve_fn=lambda x: x,
                   modulus=None,
                   matte_background=True):
    """Visualize a 2D image and a 2D weighting according to some colormap.
    Args:
        value: A 2D image.
        weight: A weight map, in [0, 1].
        colormap: A colormap function.
        lo: The lower bound to use when rendering, if None then use a percentile.
        hi: The upper bound to use when rendering, if None then use a percentile.
        percentile: What percentile of the value map to crop to when automatically
        generating `lo` and `hi`. Depends on `weight` as well as `value'.
        curve_fn: A curve function that gets applied to `value`, `lo`, and `hi`
        before the rest of visualization. Good choices: x, 1/(x+eps), log(x+eps).
        modulus: If not None, mod the normalized value by `modulus`. Use (0, 1]. If
        `modulus` is not None, `lo`, `hi` and `percentile` will have no effect.
        matte_background: If True, matte the image over a checkerboard.
    Returns:
        A colormap rendering.
    """
    # Identify the values that bound the middle of `value' according to `weight`.
    lo_auto, hi_auto = weighted_percentile(
        value, weight, [50 - percentile / 2, 50 + percentile / 2])

    # If `lo` or `hi` are None, use the automatically-computed bounds above.
    eps = np.finfo(np.float32).eps
    lo = lo or (lo_auto - eps)
    hi = hi or (hi_auto + eps)

    # Curve all values.
    value, lo, hi = [curve_fn(x) for x in [value, lo, hi]]

    # Wrap the values around if requested.
    if modulus:
        value = np.mod(value, modulus) / modulus
    else:
        # Otherwise, just scale to [0, 1].
        value = np.nan_to_num(
            np.clip((value - np.minimum(lo, hi)) / np.abs(hi - lo), 0, 1))

    if colormap:
        colorized = colormap(value)[:, :, :3] # 2D
    else:
        assert len(value.shape) == 3 and value.shape[-1] == 3
        colorized = value

    return matte(colorized, weight) if matte_background else colorized
volsdf/utils/plots.py
cvlab-stonybrook/s-volsdf
import numpy as np
import imageio
import skimage
import cv2
import torch
from torch.nn import functional as F

def find_closest_number(numbers, x):
    differences = np.abs(np.array(numbers) - x)
    closest_index = np.argmin(differences)
    closest_number = numbers[closest_index]
    return closest_number

def get_psnr(img1, img2, normalize_rgb=False):
    if normalize_rgb: # [-1,1] --> [0,1]
        img1 = (img1 + 1.) / 2.
        img2 = (img2 + 1. ) / 2.

    mse = torch.mean((img1 - img2) ** 2)
    psnr = -10. * torch.log(mse) / torch.log(torch.Tensor([10.]).cuda())

    return psnr


def load_rgb(path, normalize_rgb = False):
    img = imageio.imread(path)
    img = skimage.img_as_float32(img)

    if normalize_rgb: # [-1,1] --> [0,1]
        img -= 0.5
        img *= 2.
    # img = img.transpose(2, 0, 1)
    return img


def load_K_Rt_from_P(filename, P=None):
    if P is None:
        lines = open(filename).read().splitlines()
        if len(lines) == 4:
            lines = lines[1:]
        lines = [[x[0], x[1], x[2], x[3]] for x in (x.split(" ") for x in lines)]
        P = np.asarray(lines).astype(np.float32).squeeze()

    out = cv2.decomposeProjectionMatrix(P)
    K = out[0]
    R = out[1]
    t = out[2]

    K = K/K[2,2]
    intrinsics = np.eye(4)
    intrinsics[:3, :3] = K

    pose = np.eye(4, dtype=np.float32)
    pose[:3, :3] = R.transpose()
    pose[:3,3] = (t[:3] / t[3])[:,0]

    return intrinsics, pose


def get_camera_params(uv, pose, intrinsics):
    """
    uv: (B, N, 2)
    pose, intrinsics: (B, 4, 4)
    """
    if pose.shape[1] == 7: #In case of quaternion vector representation
        cam_loc = pose[:, 4:]
        R = quat_to_rot(pose[:,:4])
        p = torch.eye(4).repeat(pose.shape[0],1,1).cuda().float()
        p[:, :3, :3] = R
        p[:, :3, 3] = cam_loc
    else: # In case of pose matrix representation
        cam_loc = pose[:, :3, 3]
        p = pose

    batch_size, num_samples, _ = uv.shape

    depth = torch.ones((batch_size, num_samples)).cuda()
    x_cam = uv[:, :, 0].view(batch_size, -1) # (B, N)
    y_cam = uv[:, :, 1].view(batch_size, -1) # (B, N)
    z_cam = depth.view(batch_size, -1) # (B, N)

    pixel_points_cam = lift(x_cam, y_cam, z_cam, intrinsics=intrinsics) # (B, N, 4)

    # permute for batch matrix product
    pixel_points_cam = pixel_points_cam.permute(0, 2, 1) # (B, 4, N)

    # world_coords = (p[:, :3, :3] @ pixel_points_cam[:, :3, :] + p[:, :3, 3:]).permute(0, 2, 1) # (B, N, 3)
    world_coords = (torch.bmm(p[:, :3, :3], pixel_points_cam[:, :3, :]) + p[:, :3, 3:]).permute(0, 2, 1) # (B, N, 3)
    # assert torch.abs(world_coords - torch.bmm(p, pixel_points_cam).permute(0, 2, 1)[:, :, :3]).mean() < 0.0004
    # u, v = get_uv(world_coords, pose, intrinsics)
    
    ray_dirs = world_coords - cam_loc[:, None, :]
    ray_dirs = F.normalize(ray_dirs, dim=2)

    return ray_dirs, cam_loc

def get_uv(world_coords, pose, intrinsics):
    """
    uv: (B, N, 2)
    pose, intrinsics: (B, 4, 4)
    """
    if pose.shape[1] == 7: #In case of quaternion vector representation
        cam_loc = pose[:, 4:]
        R = quat_to_rot(pose[:,:4])
        p = torch.eye(4).repeat(pose.shape[0],1,1).cuda().float()
        p[:, :3, :3] = R
        p[:, :3, 3] = cam_loc
    else: # In case of pose matrix representation
        cam_loc = pose[:, :3, 3]
        p = pose

    intrinsics = intrinsics.cuda()
    p = p
    world_coords = world_coords
    fx = intrinsics[:, 0, 0].unsqueeze(-1) # (B, 1)
    fy = intrinsics[:, 1, 1].unsqueeze(-1)
    cx = intrinsics[:, 0, 2].unsqueeze(-1)
    cy = intrinsics[:, 1, 2].unsqueeze(-1)
    sk = intrinsics[:, 0, 1].unsqueeze(-1)

    # [(B, N, 3) - (B, 1, 3)] @ (1, 3, 3) -> (B, N, 3)
    pts_cam = (world_coords - p[:, :3, 3].view(-1, 1, 3)) @ p[:, :3, :3]
    x_lift = pts_cam[:, :, 0]
    y_lift = pts_cam[:, :, 1]
    z = pts_cam[:, :, 2] # (B, N)

    y = y_lift / z * fy + cy
    x = x_lift / z * fx + cx - cy*sk/fy + sk*y/fy

    return x, y

def get_camera_for_plot(pose):
    if pose.shape[1] == 7: #In case of quaternion vector representation
        cam_loc = pose[:, 4:].detach()
        R = quat_to_rot(pose[:,:4].detach())
    else: # In case of pose matrix representation
        cam_loc = pose[:, :3, 3]
        R = pose[:, :3, :3]
    cam_dir = R[:, :3, 2]
    return cam_loc, cam_dir


def lift(x, y, z, intrinsics):
    # parse intrinsics
    intrinsics = intrinsics.cuda()
    fx = intrinsics[:, 0, 0]
    fy = intrinsics[:, 1, 1]
    cx = intrinsics[:, 0, 2]
    cy = intrinsics[:, 1, 2]
    sk = intrinsics[:, 0, 1]

    x_lift = (x - cx.unsqueeze(-1) + cy.unsqueeze(-1)*sk.unsqueeze(-1)/fy.unsqueeze(-1) - sk.unsqueeze(-1)*y/fy.unsqueeze(-1)) / fx.unsqueeze(-1) * z
    y_lift = (y - cy.unsqueeze(-1)) / fy.unsqueeze(-1) * z

    # homogeneous
    return torch.stack((x_lift, y_lift, z, torch.ones_like(z).cuda()), dim=-1)


def quat_to_rot(q):
    batch_size, _ = q.shape
    q = F.normalize(q, dim=1)
    R = torch.ones((batch_size, 3,3)).cuda()
    qr=q[:,0]
    qi = q[:, 1]
    qj = q[:, 2]
    qk = q[:, 3]
    R[:, 0, 0]=1-2 * (qj**2 + qk**2)
    R[:, 0, 1] = 2 * (qj *qi -qk*qr)
    R[:, 0, 2] = 2 * (qi * qk + qr * qj)
    R[:, 1, 0] = 2 * (qj * qi + qk * qr)
    R[:, 1, 1] = 1-2 * (qi**2 + qk**2)
    R[:, 1, 2] = 2*(qj*qk - qi*qr)
    R[:, 2, 0] = 2 * (qk * qi-qj * qr)
    R[:, 2, 1] = 2 * (qj*qk + qi*qr)
    R[:, 2, 2] = 1-2 * (qi**2 + qj**2)
    return R


def rot_to_quat(R):
    batch_size, _,_ = R.shape
    q = torch.ones((batch_size, 4)).cuda()

    R00 = R[:, 0,0]
    R01 = R[:, 0, 1]
    R02 = R[:, 0, 2]
    R10 = R[:, 1, 0]
    R11 = R[:, 1, 1]
    R12 = R[:, 1, 2]
    R20 = R[:, 2, 0]
    R21 = R[:, 2, 1]
    R22 = R[:, 2, 2]

    q[:,0]=torch.sqrt(1.0+R00+R11+R22)/2
    q[:, 1]=(R21-R12)/(4*q[:,0])
    q[:, 2] = (R02 - R20) / (4 * q[:, 0])
    q[:, 3] = (R10 - R01) / (4 * q[:, 0])
    return q


def get_sphere_intersections(cam_loc, ray_directions, r = 1.0):
    # Input: n_rays x 3 ; n_rays x 3
    # Output: n_rays x 1, n_rays x 1 (close and far)

    ray_cam_dot = torch.bmm(ray_directions.view(-1, 1, 3),
                            cam_loc.view(-1, 3, 1)).squeeze(-1)
    under_sqrt = ray_cam_dot ** 2 - (cam_loc.norm(2, 1, keepdim=True) ** 2 - r ** 2)

    # sanity check
    if (under_sqrt <= 0).sum() > 0:
        print('BOUNDING SPHERE PROBLEM!')
        exit()

    sphere_intersections = torch.sqrt(under_sqrt) * torch.Tensor([-1, 1]).cuda().float() - ray_cam_dot
    sphere_intersections = sphere_intersections.clamp_min(0.0)

    return sphere_intersections

volsdf/utils/rend_util.py
cvlab-stonybrook/s-volsdf
from helpers.help import logger
import os
from datetime import datetime
from pyhocon import ConfigFactory
import copy
import itertools
from omegaconf import OmegaConf
import torch
from tqdm import tqdm
from torch.nn.functional import grid_sample
from torch.utils.tensorboard import SummaryWriter

import volsdf.utils.general as utils
import volsdf.utils.plots as plt
from volsdf.utils import rend_util
from volsdf.datasets.scene_dataset import SceneDataset

class VolOpt():
    def __init__(self, **kwargs):
        torch.set_default_dtype(torch.float32)
        torch.set_num_threads(1)

        # get configs
        self.hparams = copy.deepcopy(kwargs['args'])
        resolved = OmegaConf.to_container(self.hparams['vol'], resolve=True, throw_on_missing=True)
        self.conf = ConfigFactory.from_dict(resolved)
        self.batch_size = kwargs['batch_size']
        self.exps_folder_name = self.hparams.exps_folder

        root = './'
        self.expname = self.conf.get_string('train.expname')
        kwargs_scan_id = int(kwargs['scan'][4:])
        scan_id = kwargs_scan_id if kwargs_scan_id != -1 else self.conf.get_int('dataset.scan_id', default=-1)
        self.scan_id = scan_id
        if scan_id != -1:
            self.expname = self.expname + '_{0}'.format(scan_id)

        if kwargs['is_continue'] and kwargs['timestamp'] == 'latest':
            if os.path.exists(os.path.join(root, self.hparams.exps_folder, self.expname)):
                timestamps = os.listdir(os.path.join(root, self.hparams.exps_folder, self.expname))
                if (len(timestamps)) == 0:
                    is_continue = False
                    timestamp = None
                else:
                    timestamp = sorted(timestamps)[-1]
                    is_continue = True
            else:
                is_continue = False
                timestamp = None
        else:
            timestamp = kwargs['timestamp']
            is_continue = kwargs['is_continue']

        # create exps dirs
        utils.mkdir_ifnotexists(os.path.join(root, self.exps_folder_name))
        self.expdir = os.path.join(root, self.exps_folder_name, self.expname)
        utils.mkdir_ifnotexists(self.expdir)
        self.timestamp = '{:%Y_%m_%d_%H_%M_%S}'.format(datetime.now())
        utils.mkdir_ifnotexists(os.path.join(self.expdir, self.timestamp))
        self.plots_dir = os.path.join(self.expdir, self.timestamp, 'plots')
        utils.mkdir_ifnotexists(self.plots_dir)
        # create checkpoints dirs
        self.checkpoints_path = os.path.join(self.expdir, self.timestamp, 'checkpoints')
        utils.mkdir_ifnotexists(self.checkpoints_path)
        self.model_params_subdir = "ModelParameters"
        self.optimizer_params_subdir = "OptimizerParameters"
        utils.mkdir_ifnotexists(os.path.join(self.checkpoints_path, self.model_params_subdir))
        utils.mkdir_ifnotexists(os.path.join(self.checkpoints_path, self.optimizer_params_subdir))

        # save configs
        with open(os.path.join(self.expdir, self.timestamp, 'run.yaml'), "w") as f:
            OmegaConf.save(self.hparams, f)
        
        # dataset config
        logger.info('Loading data ...')
        dataset_conf = self.conf.get_config('dataset')
        if kwargs_scan_id != -1:
            dataset_conf['scan_id'] = kwargs_scan_id
        dataset_conf['data_dir_root'] = self.hparams.data_dir_root
        logger.info(f"    full resolution in VolOpt {dataset_conf['img_res']}")
        assert self.conf.get_string('train.dataset_class') == "volsdf.datasets.scene_dataset.SceneDataset"
        assert [self.hparams.max_h, self.hparams.max_w] == dataset_conf['img_res']

        # generate dataset
        self.data_confs = [copy.deepcopy(dataset_conf) for _ in range(3)]
        self.gen_dataset(stg = 2) # full resolution
        self.gen_plot_dataset()
        self.stg = 2
        self.ds_len = len(self.train_dataset) # number of training images

        # model
        conf_model = self.conf.get_config('model')
        self.model = utils.get_class(self.conf.get_string('train.model_class'))(conf=conf_model)
        if torch.cuda.is_available():
            self.model.cuda()

        # loss
        self.loss = utils.get_class(self.conf.get_string('train.loss_class'))(**self.conf.get_config('loss'))

        # optimizer
        self.lr = self.conf.get_float('train.learning_rate')
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)

        # load ckpt
        self.start_epoch = 0
        self.iter_step, self.total_step = 0, 0
        ckpt_dir = self.conf.get_string('train.ckpt_dir', '')
        if is_continue:
            self.load_from_dir(dir=os.path.join(self.expdir, timestamp), checkpoint=kwargs['checkpoint'])
        elif ckpt_dir != '':
            self.load_from_dir(dir=ckpt_dir, checkpoint='latest')

        # some parameters
        self.num_pixels = self.conf.get_int('train.num_pixels')
        self.plot_freq = self.conf.get_int('train.plot_freq')
        self.render_freq = self.conf.get_int('train.render_freq')
        self.checkpoint_freq = self.conf.get_int('train.checkpoint_freq', default=100)
        self.split_n_pixels = self.conf.get_int('train.split_n_pixels', default=10000)
        self.plot_conf = self.conf.get_config('plot')

        # logs
        self.writer = SummaryWriter(log_dir=os.path.join(self.plots_dir, 'logs'))

        # copy hparams to every module
        self.model.hparams = self.hparams
        self.loss.hparams = self.hparams

    def load_from_dir(self, dir, checkpoint='latest'):
        """
        e.g. dir = exps_debug/ours_114/2022_10_21_12_46_36
        """
        old_checkpnts_dir = os.path.join(dir, 'checkpoints')
        logger.info(f"Load from {old_checkpnts_dir} at {checkpoint}.pth")

        saved_model_state = torch.load(
            os.path.join(old_checkpnts_dir, 'ModelParameters', str(checkpoint) + ".pth"))
        self.model.load_state_dict(saved_model_state["model_state_dict"])

        self.start_epoch = saved_model_state['epoch']
        self.iter_step = saved_model_state['iter_step']

        data = torch.load(
            os.path.join(old_checkpnts_dir, 'OptimizerParameters', str(checkpoint) + ".pth"))
        self.optimizer.load_state_dict(data["optimizer_state_dict"])

    def gen_plot_dataset(self):
        data_conf = copy.deepcopy(self.conf.get_config('dataset'))
        data_conf['img_res'] = [int(_/4.) for _ in data_conf['img_res']]
        self.plot_dataset = SceneDataset(**data_conf)
        self.plot_dataloader = torch.utils.data.DataLoader(self.plot_dataset,
                                                        batch_size=self.conf.get_int('plot.plot_nimgs'),
                                                        shuffle=False,
                                                        collate_fn=self.plot_dataset.collate_fn
                                                        )

    def gen_dataset(self, stg):
        data_conf_stg = self.data_confs[stg]

        self.train_dataset = SceneDataset(**data_conf_stg)
        logger.info('Finish loading data. Data-set size: {0}'.format(len(self.train_dataset)))

        self.train_dataloader = torch.utils.data.DataLoader(self.train_dataset,
                                                            batch_size=self.batch_size,
                                                            shuffle=True,
                                                            collate_fn=self.train_dataset.collate_fn
                                                            )
        self.eval_dataloader = torch.utils.data.DataLoader(self.train_dataset,
                                                        batch_size=1,
                                                        shuffle=False,
                                                        collate_fn=self.train_dataset.collate_fn
                                                        )

        self.total_pixels = self.train_dataset.total_pixels
        self.img_res = self.train_dataset.img_res
        self.scale_factor = self.train_dataset.scale_factor
        self.n_batches = len(self.train_dataloader)

    def save_checkpoints(self, epoch, latest_only=False):
        torch.save(
            {"epoch": epoch, "model_state_dict": self.model.state_dict(), "iter_step": self.iter_step},
            os.path.join(self.checkpoints_path, self.model_params_subdir, "latest.pth"))
        torch.save(
            {"epoch": epoch, "optimizer_state_dict": self.optimizer.state_dict()},
            os.path.join(self.checkpoints_path, self.optimizer_params_subdir, "latest.pth"))

        if latest_only:
            return 0

        torch.save(
            {"epoch": epoch, "model_state_dict": self.model.state_dict(), "iter_step": self.iter_step},
            os.path.join(self.checkpoints_path, self.model_params_subdir, str(epoch) + ".pth"))
        torch.save(
            {"epoch": epoch, "optimizer_state_dict": self.optimizer.state_dict()},
            os.path.join(self.checkpoints_path, self.optimizer_params_subdir, str(epoch) + ".pth"))

    def train_step(self, batch, use_mvs=False):
        self.model.train()

        indices, model_input, ground_truth = batch
        model_input["intrinsics"] = model_input["intrinsics"].cuda()
        model_input["uv"] = model_input["uv"].cuda() # [1, 512, 2]           
        model_input['pose'] = model_input['pose'].cuda()
        model_input['iter_step'] = self.iter_step
        
        fast = 1
        model_outputs = self.model(model_input, fast=fast)
        if use_mvs:
            model_outputs['pj'], model_outputs['pi'], _ = self.cost_mapping(
                z_vals=model_outputs['depth_vals'], ts=indices, xyz_raw=model_outputs['xyz'])
        loss_output = self.loss(model_outputs, ground_truth)

        loss = loss_output['loss']

        self.optimizer.zero_grad()
        loss.backward()
        if self.hparams.grad_clip:
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
        self.on_after_backward()
        self.optimizer.step()

        psnr = rend_util.get_psnr(model_outputs['rgb_values'],
                                    ground_truth['rgb'].cuda().reshape(-1,3))
           
        if self.total_step % 50 == 0:
            for k, v in loss_output.items():
                self.writer.add_scalar('t/'+k, v, self.total_step)
            
            self.writer.add_scalar('t/beta', self.model.density.get_beta().item(), self.total_step)
            self.writer.add_scalar('t/alpha', 1. / self.model.density.get_beta().item(), self.total_step)
            self.writer.add_scalar('t/psnr', psnr.item(), self.total_step)

        self.train_dataset.change_sampling_idx(self.num_pixels)

        self.iter_step += 1
        self.total_step += 1

    def render_step(self, batch, epoch, dataset, fast=-1):
        self.model.eval()

        indices, model_input, ground_truth = batch
        model_input["intrinsics"] = model_input["intrinsics"].cuda()
        model_input["uv"] = model_input["uv"].cuda()
        model_input['pose'] = model_input['pose'].cuda()
        model_input['iter_step'] = self.iter_step

        split = utils.split_input(model_input, dataset.total_pixels, n_pixels=self.split_n_pixels)
        res = []
        for s in tqdm(split, leave=False):
            out = self.model(s, fast=fast)
            d = {'rgb_values': out['rgb_values'].detach().cpu(),
                    'normal_map': out['normal_map'].detach().cpu(),
                    'depth_values': out['depth_values'].detach().cpu(),
                    'depth_vals': out['depth_vals'].detach().cpu(),
                    'weights': out['weights'].detach().cpu(),
                    'xyz': out['xyz'].detach().cpu(),
                    }
            res.append(d)
        
        del out

        batch_size = ground_truth['rgb'].shape[0]
        model_outputs = utils.merge_output(res, dataset.total_pixels, batch_size)
        stack = []

        plot_data = self.get_plot_data(model_input, model_outputs, model_input['pose'], ground_truth['rgb'])

        depth_cuda = plt.lin2img(plot_data['depth_map'][..., None], dataset.img_res)
        depth_cuda = depth_cuda[0].cuda() * self.scale_factor
        acc = model_outputs['weights'].sum(1).reshape(depth_cuda.shape)
        depth_cuda[acc<0.2] = depth_cuda.max()

        mask = ground_truth['mask'].reshape(-1, 3)
        mask_bin = (mask == 1.)
        mse = torch.mean((model_outputs['rgb_values'] - ground_truth['rgb'].reshape(-1,3))[mask_bin] ** 2)
        psnr = -10. * torch.log(mse) / torch.log(torch.Tensor([10.]))
        self.writer.add_scalar('val/psnr', psnr.item(), self.total_step)

        stack += plt.stacked_plot(indices, plot_data, self.plots_dir, epoch, dataset.img_res, **self.plot_conf)
        stack[0][~mask_bin.reshape(dataset.img_res+[3,]).permute(2, 0, 1)] = 0
        stack = torch.stack(stack, dim=0) # (B, 3, H, W)
        self.writer.add_images('val/vis', stack, self.total_step)

        photo_conf = None

        self.total_step += 1

        return depth_cuda, photo_conf

    def get_plot_data(self, model_input, model_outputs, pose, rgb_gt):
        batch_size, num_samples, _ = rgb_gt.shape

        rgb_eval = model_outputs['rgb_values'].reshape(batch_size, num_samples, 3)
        normal_map = model_outputs['normal_map'].reshape(batch_size, num_samples, 3)
        normal_map = (normal_map + 1.) / 2.
      
        depth_map = model_outputs['depth_values'].reshape(batch_size, num_samples)
        acc = model_outputs['weights'].sum(1).reshape(batch_size, num_samples)
        
        plot_data = {
            'rgb_gt': rgb_gt,
            'pose': pose,
            'rgb_eval': rgb_eval,
            'normal_map': normal_map,
            'depth_map': depth_map,
            'acc': acc
        }

        return plot_data

    def render_mvs(self, id_k, epoch):
        self.train_dataset.mode = 'test'

        self.train_dataset.change_sampling_idx(-1)

        batch = next(itertools.islice(self.eval_dataloader, id_k, None))
        depth_cuda, depth_confi = self.render_step(batch, epoch, self.train_dataset, fast=-1)

        self.train_dataset.mode = 'train'

        return depth_cuda, depth_confi

    def run(self, opt_stepN=1e8):
        # set everything here
        start_iter_step = self.iter_step
        logger.info(f"train volsdf at {self.checkpoints_path} ..")
        logger.info(f"[NOW] total_step={self.total_step}, iter_step={self.iter_step} stg={self.stg} opt_stepN={opt_stepN}")
        
        # start pbar
        pbar = tqdm(total=opt_stepN-start_iter_step, desc="Train", ncols=60)

        # optimization: VolSDF uses 2000 epoch with 50 images, 2000*50 / 3 images -> ?
        epoch = self.start_epoch
        while True:
            if epoch % self.checkpoint_freq == 0:
                self.save_checkpoints(epoch)

            # render during training to check what's wrong
            early_render = (self.iter_step-start_iter_step <= 120*50) and epoch % (20*50//self.ds_len) == 0
            if epoch % self.render_freq == 0 or early_render:
                self.plot_dataset.change_sampling_idx(-1)
                self.plot_dataset.mode = 'plot'
                batch = next(iter(self.plot_dataloader))
                self.plot_dataset.mode = 'train'
                self.render_step(batch, epoch, self.plot_dataset, fast=-1)

                self.save_checkpoints(epoch, latest_only=True)
                torch.cuda.empty_cache()

            self.train_dataset.change_sampling_idx(self.num_pixels)

            for data_index, batch in enumerate(self.train_dataloader):
                self.train_step(batch, self.hparams.use_mvs)
                pbar.update(1)
            
            if self.iter_step-start_iter_step > opt_stepN:
                break

            epoch += 1

        # close pbar
        pbar.close()

        # save
        self.save_checkpoints(epoch)
        self.start_epoch = epoch

        return epoch

    def get_mvs_input(self, outs_samples):
        self.costs, self.z_mvs, self.bd_mvs = dict(), dict(), dict()
        for i in range(len(outs_samples)): # i = 0, 1, 2, id_k = 25, 22, 28
            prob_volume = outs_samples[i]['prob_volume']
            depth_values = outs_samples[i]['depth_values'] / self.scale_factor
            self.costs[i] = prob_volume # (1, 48, 288, 384)
            self.z_mvs[i] = depth_values # (1, 48, 288, 384)
            bd_mvs_i = depth_values[:, [0, -1], :, :] # (1, 2, 288, 384)
            bd_mvs_i[:, 0, :, :] = torch.minimum(bd_mvs_i[:, 0, :, :], torch.ones_like(bd_mvs_i[:, 0, :, :])*self.conf.get_float('model.scene_bounding_sphere'))
            self.bd_mvs[i] = bd_mvs_i

        del prob_volume, depth_values

    @torch.no_grad()
    def cost_mapping(self, z_vals, ts, xyz_raw):
        """
        xyz_w = p @ xy_lift
        x = x_lift/z * fx + cx + (y - cy)*sk/fy
        y = y_lift/z * fy + cy
        """
        results_cost_j = torch.zeros_like(z_vals, requires_grad=False) # (N=N_rays, D=N_samples)
        results_cost_mvs = torch.zeros_like(z_vals, requires_grad=False) # (N=N_rays, D=N_samples)
        valid_mask = torch.zeros_like(z_vals, requires_grad=False, dtype=bool)
        for i, id_k in enumerate(self.trains_i):
            same_view = ts[0] == id_k # (N_rays,)

            xyz_j = xyz_raw.detach().clone() # N, D, 3
            cost, z_mvs = self.costs[i].to(xyz_raw.device), self.z_mvs[i].to(xyz_raw.device) # (1, 48, H, W)
            _h, _w = self.train_dataset.img_res
            K = self.train_dataset.intrinsics_all[id_k] # (4, 4)
            c2w = self.train_dataset.pose_all[id_k].to(xyz_raw.device) # Ps[id_]
            # assert list(cost.shape[2:4]) == self.train_dataset.img_res # check K == cost.K
            c2w = c2w[:3] # (3, 4)
            fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2]
            sk = K[0, 1]         

            xyz_j -= c2w[:, 3].view(1, 1, 3) # 1, 1, 3 rays_o_j
            xyz_j = xyz_j @ c2w[:, :3] # N, D, 3
            # xyz_j[..., 1:] *= -1 # "right down front"
            # assert (xyz_j[...,2] - uvw[...,2]).sum() < 0.001 # (xyz_j[...,2] - z_vals).sum()
            # assert xyz_j[..., 2:].min() > 0
            xyz_j[..., :2] /= xyz_j[..., 2:] # N, D, 1
            xyz_j[..., 1] = xyz_j[..., 1] * fy + cy
            xyz_j[..., 0] = xyz_j[..., 0] * fx + cx + (xyz_j[..., 1] - cy)*sk/fy
            # assert (xyz_j[..., :2] - uvw[...,:2]).sum() < 0.5
            xyz_j[..., 0] = xyz_j[..., 0] / ((_w - 1) / 2) - 1 # K[0, 2]*2
            xyz_j[..., 1] = xyz_j[..., 1] / ((_h - 1) / 2) - 1 # K[1, 2]*2
            # interpolat near,far, then normalize depth
            Hj, Dj, _ = xyz_j.shape
            xyz_j = xyz_j.view(1,Hj,Dj,3) # 1, H', D', 3
            mvs_near, mvs_far = z_mvs[:,:1,:,:], z_mvs[:,-1:,:,:] # 1, C=1, H, W
            # 1, C=1, H, W    1, H', D', 2  ->  1, C=1, H', D'
            bound_hw = 1.001
            invalid_j = (xyz_j[..., 2]<1e-5)|(xyz_j[..., 0]>bound_hw)|(xyz_j[..., 0]<-bound_hw)|(xyz_j[..., 1]>bound_hw)|(xyz_j[..., 1]<-bound_hw) # 1, H', D'
            xyz_j[invalid_j, :] = -99.
            near_j = grid_sample(mvs_near, xyz_j[...,:2], mode='bilinear', padding_mode='zeros', align_corners=True)[:,0,:,:]
            far_j = grid_sample(mvs_far, xyz_j[...,:2], mode='bilinear', padding_mode='zeros', align_corners=True)[:,0,:,:]
            if self.hparams.inverse_depth and self.stg == 0:
                far_j[invalid_j] = 1e-8
                xyz_j[..., 2] = 2 * (1. - near_j/xyz_j[..., 2]) / (1. - near_j/far_j) - 1 # 1, H', D'
            elif self.hparams.inverse_depth and self.stg >= 1:
                raise NotImplementedError
            else:
                xyz_j[..., 2] = 2 * (xyz_j[..., 2]-near_j) / (far_j-near_j) - 1
            bound_z = 1.01
            invalid_j = (near_j<1e-5)|(far_j<1e-5)|(xyz_j[..., 2]>bound_z)|(xyz_j[..., 2]<-bound_z)|invalid_j # 1, H', D'
            xyz_j[invalid_j, :] = -99.
            # interpolate cost
            xyz_j = xyz_j.view(1,Hj,Dj,1,3) # 1, H', D', 1, 3
            xyz_j = xyz_j.permute(0,2,1,3,4) # 1, D', H', 1, 3
            # B=1,C=1,D=48,H=144,W=192    1, D', H', W'=1, 3   ->  1, C=1, D', H', W'=1
            cost_j = grid_sample(cost[None,...], xyz_j, mode='bilinear', padding_mode='zeros', align_corners=True)
            cost_j = cost_j.squeeze_().permute(1,0) # (N_rays, N_samples)
            # add to cost
            if same_view:
                results_cost_mvs = cost_j
            else:
                results_cost_j += cost_j
                # mask out points not inside at least 1 another view's frustum
                valid_mask = valid_mask | (~invalid_j[0, :, :])
        # mask out points..
        results_cost_mvs[~valid_mask] = 0.
        del cost_j, z_mvs, cost, xyz_j
        return results_cost_j, results_cost_mvs, valid_mask

    def on_after_backward(self) -> None:
        valid_gradients = True
        for name, param in self.model.named_parameters():
            if param.grad is not None:
                valid_gradients = not (torch.isnan(param.grad).any() or torch.isinf(param.grad).any())
                if not valid_gradients:
                    break

        if not valid_gradients:
            logger.warning(f'detected inf or nan values in gradients. not updating model parameters')
            self.optimizer.zero_grad()
volsdf/vsdf.py
cvlab-stonybrook/s-volsdf