Official Code for Bayes' Rays Paper
#####bayesrays init#####
bayesrays/__init__.py
BayesRays/BayesRays
# Copyright 2022 The Nerfstudio Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" Data parser for nerfstudio datasets in nerfbuster-changes branch. """

from __future__ import annotations

import math
import os
from dataclasses import dataclass, field
from pathlib import Path, PurePath
from typing import Optional, Tuple, Type

import numpy as np
import torch
from PIL import Image
from rich.console import Console
from typing_extensions import Literal

from nerfstudio.cameras import camera_utils
from nerfstudio.cameras.cameras import CAMERA_MODEL_TO_TYPE, Cameras, CameraType
from nerfstudio.data.dataparsers.base_dataparser import (
    DataParser,
    DataParserConfig,
    DataparserOutputs,
)
from nerfstudio.data.scene_box import SceneBox
from nerfstudio.utils.io import load_from_json

CONSOLE = Console(width=120)
MAX_AUTO_RESOLUTION = 1600


@dataclass
class NerfbusterDataparserConfig(DataParserConfig):
    """Nerfbuster dataset config"""

    _target: Type = field(default_factory=lambda: Nerfbuster)
    """target class to instantiate"""
    data: Path = Path()
    """Directory or explicit json file path specifying location of data."""
    scale_factor: float = 1.0
    """How much to scale the camera origins by."""
    downscale_factor: Optional[int] = None
    """How much to downscale images. If not set, images are chosen such that the max dimension is <1600px."""
    scene_scale: float = 1.0
    """How much to scale the region of interest by."""
    orientation_method: Literal["pca", "up", "vertical", "none"] = "up"
    """The method to use for orientation."""
    center_method: Literal["poses", "focus", "none"] = "poses"
    """The method to use to center the poses."""
    auto_scale_poses: bool = True
    """Whether to automatically scale the poses to fit in +/- 1 bounding box."""
    eval_mode: Literal["train-split-fraction", "eval-frame-index", "eval-interval"] = "train-split-fraction"
    """The method to use for splitting the dataset into train and eval."""
    train_split_fraction: float = 0.9
    """The fraction of images to use for training. The remaining images are for eval."""
    train_frame_indices: Tuple[int, ...] = (0,)
    """The index of the frames to use for train."""
    eval_frame_indices: Tuple[int, ...] = (1,)
    """The index of the frames to use for eval."""
    eval_interval: int = 8
    """The interval between frames to use for eval."""
    depth_unit_scale_factor: float = 1e-3
    """Scales the depth values to meters. Default value is 0.001 for a millimeter to meter conversion."""


@dataclass
class Nerfbuster(DataParser):
    """Nerfbuster DatasetParser"""

    config: NerfbusterDataParserConfig
    downscale_factor: Optional[int] = None

    def _generate_dataparser_outputs(self, split="train"):
        # pylint: disable=too-many-statements

        assert self.config.data.exists(), f"Data directory {self.config.data} does not exist."

        if self.config.data.suffix == ".json":
            meta = load_from_json(self.config.data)
            data_dir = self.config.data.parent
        else:
            meta = load_from_json(self.config.data / "transforms.json")
            data_dir = self.config.data

        image_filenames = []
        mask_filenames = []
        depth_filenames = []
        poses = []
        num_skipped_image_filenames = 0

        fx_fixed = "fl_x" in meta
        fy_fixed = "fl_y" in meta
        cx_fixed = "cx" in meta
        cy_fixed = "cy" in meta
        height_fixed = "h" in meta
        width_fixed = "w" in meta
        distort_fixed = False
        for distort_key in ["k1", "k2", "k3", "p1", "p2"]:
            if distort_key in meta:
                distort_fixed = True
                break
        fx = []
        fy = []
        cx = []
        cy = []
        height = []
        width = []
        distort = []

        # sort the frames by fname
        # they do this in mipnerf360 code
        fnames = []
        for frame in meta["frames"]:
            filepath = PurePath(frame["file_path"])
            fname = self._get_fname(filepath, data_dir)
            fnames.append(fname)
        inds = np.argsort(fnames)
        frames = [meta["frames"][ind] for ind in inds]

        for frame in frames:
            filepath = PurePath(frame["file_path"])
            fname = self._get_fname(filepath, data_dir)
            if not fname.exists():
                num_skipped_image_filenames += 1
                continue

            if not fx_fixed:
                assert "fl_x" in frame, "fx not specified in frame"
                fx.append(float(frame["fl_x"]))
            if not fy_fixed:
                assert "fl_y" in frame, "fy not specified in frame"
                fy.append(float(frame["fl_y"]))
            if not cx_fixed:
                assert "cx" in frame, "cx not specified in frame"
                cx.append(float(frame["cx"]))
            if not cy_fixed:
                assert "cy" in frame, "cy not specified in frame"
                cy.append(float(frame["cy"]))
            if not height_fixed:
                assert "h" in frame, "height not specified in frame"
                height.append(int(frame["h"]))
            if not width_fixed:
                assert "w" in frame, "width not specified in frame"
                width.append(int(frame["w"]))
            if not distort_fixed:
                distort.append(
                    camera_utils.get_distortion_params(
                        k1=float(frame["k1"]) if "k1" in frame else 0.0,
                        k2=float(frame["k2"]) if "k2" in frame else 0.0,
                        k3=float(frame["k3"]) if "k3" in frame else 0.0,
                        k4=float(frame["k4"]) if "k4" in frame else 0.0,
                        p1=float(frame["p1"]) if "p1" in frame else 0.0,
                        p2=float(frame["p2"]) if "p2" in frame else 0.0,
                    )
                )

            image_filenames.append(fname)
            poses.append(np.array(frame["transform_matrix"]))
            if "mask_path" in frame:
                mask_filepath = PurePath(frame["mask_path"])
                mask_fname = self._get_fname(
                    mask_filepath,
                    data_dir,
                    downsample_folder_prefix="masks_",
                )
                mask_filenames.append(mask_fname)

            if "depth_file_path" in frame:
                depth_filepath = PurePath(frame["depth_file_path"])
                depth_fname = self._get_fname(depth_filepath, data_dir, downsample_folder_prefix="depths_")
                depth_filenames.append(depth_fname)

        if num_skipped_image_filenames >= 0:
            CONSOLE.log(f"Skipping {num_skipped_image_filenames} files in dataset split {split}.")
        assert (
            len(image_filenames) != 0
        ), """
        No image files found. 
        You should check the file_paths in the transforms.json file to make sure they are correct.
        """
        assert len(mask_filenames) == 0 or (
            len(mask_filenames) == len(image_filenames)
        ), """
        Different number of image and mask filenames.
        You should check that mask_path is specified for every frame (or zero frames) in transforms.json.
        """
        assert len(depth_filenames) == 0 or (
            len(depth_filenames) == len(image_filenames)
        ), """
        Different number of image and depth filenames.
        You should check that depth_file_path is specified for every frame (or zero frames) in transforms.json.
        """

        if self.config.eval_mode == "train-split-fraction":
            has_split_files_spec = any(f"{split}_filenames" in meta for split in ("train", "val", "test"))
            if f"{split}_filenames" in meta:
                # Validate split first
                split_filenames = set(self._get_fname(PurePath(x), data_dir) for x in meta[f"{split}_filenames"])
                unmatched_filenames = split_filenames.difference(image_filenames)
                if unmatched_filenames:
                    raise RuntimeError(f"Some filenames for split {split} were not found: {unmatched_filenames}.")

                indices = [i for i, path in enumerate(image_filenames) if path in split_filenames]
                CONSOLE.log(f"[yellow] Dataset is overriding {split}_indices to {indices}")
                indices = np.array(indices, dtype=np.int32)
                i_train = indices
                i_eval = indices
            elif has_split_files_spec:
                raise RuntimeError(f"The dataset's list of filenames for split {split} is missing.")
            else:
                # filter image_filenames and poses based on train/eval split percentage
                num_images = len(image_filenames)
                num_train_images = math.ceil(num_images * self.config.train_split_fraction)
                num_eval_images = num_images - num_train_images
                i_all = np.arange(num_images)
                i_train = np.linspace(
                    0, num_images - 1, num_train_images, dtype=int
                )  # equally spaced training images starting and ending at 0 and num_images-1
                i_eval = np.setdiff1d(i_all, i_train)  # eval images are the remaining images
                assert len(i_eval) == num_eval_images
        elif self.config.eval_mode == "eval-frame-index":
            # keep around some metadata
            eval_frame_index_0_metadata = []
            # filter image_filenames and poses based on train and eval frame indices
            num_images = len(image_filenames)
            basenames = [os.path.basename(image_filename) for image_filename in image_filenames]
            i_all = np.arange(num_images)
            i_train = []
            i_eval = []
            for idx, basename in zip(i_all, basenames):
                # check the frame index
                if len(basename.split("_")) == 2:
                    frame_index = 0
                else:
                    frame_index = int(basename.split("_")[1])
                if frame_index in self.config.train_frame_indices:
                    i_train.append(idx)
                    if frame_index == 0:
                        # set 1 where frame_index is 0
                        eval_frame_index_0_metadata.append(1)
                    else:
                        eval_frame_index_0_metadata.append(0)
                if frame_index in self.config.eval_frame_indices:
                    i_eval.append(idx)
        elif self.config.eval_mode == "eval-interval":
            # filter image_filenames and poses based on a specified interval for eval
            # this chunk of code is very similar to the mipnerf360 code
            num_images = len(image_filenames)
            all_indices = np.arange(num_images)
            train_indices = all_indices[all_indices % self.config.eval_interval != 0]
            eval_indices = all_indices[all_indices % self.config.eval_interval == 0]
            i_train = train_indices
            i_eval = eval_indices
        else:
            raise ValueError(f"Unknown eval mode {self.config.eval_mode}")

        if split == "train":
            indices = i_train
        elif split in ["val", "test"]:
            indices = i_eval
        else:
            raise ValueError(f"Unknown dataparser split {split}")

        if "orientation_override" in meta:
            orientation_method = meta["orientation_override"]
            CONSOLE.log(f"[yellow] Dataset is overriding orientation method to {orientation_method}")
        else:
            orientation_method = self.config.orientation_method

        poses = torch.from_numpy(np.array(poses).astype(np.float32))
        poses, transform_matrix = camera_utils.auto_orient_and_center_poses(
            poses,
            method=orientation_method,
            center_method=self.config.center_method,
        )

        # Scale poses
        scale_factor = 1.0
        if self.config.auto_scale_poses:
            scale_factor /= float(torch.max(torch.abs(poses[:, :3, 3])))
        scale_factor *= self.config.scale_factor

        poses[:, :3, 3] *= scale_factor

        # Choose image_filenames and poses based on split, but after auto orient and scaling the poses.
        image_filenames = [image_filenames[i] for i in indices]
        mask_filenames = [mask_filenames[i] for i in indices] if len(mask_filenames) > 0 else []
        depth_filenames = [depth_filenames[i] for i in indices] if len(depth_filenames) > 0 else []
        poses = poses[indices]

        # in x,y,z order
        # assumes that the scene is centered at the origin
        aabb_scale = self.config.scene_scale
        scene_box = SceneBox(
            aabb=torch.tensor(
                [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32
            )
        )

        if "camera_model" in meta:
            camera_type = CAMERA_MODEL_TO_TYPE[meta["camera_model"]]
        else:
            camera_type = CameraType.PERSPECTIVE

        idx_tensor = torch.tensor(indices, dtype=torch.long)
        fx = float(meta["fl_x"]) if fx_fixed else torch.tensor(fx, dtype=torch.float32)[idx_tensor]
        fy = float(meta["fl_y"]) if fy_fixed else torch.tensor(fy, dtype=torch.float32)[idx_tensor]
        cx = float(meta["cx"]) if cx_fixed else torch.tensor(cx, dtype=torch.float32)[idx_tensor]
        cy = float(meta["cy"]) if cy_fixed else torch.tensor(cy, dtype=torch.float32)[idx_tensor]
        height = int(meta["h"]) if height_fixed else torch.tensor(height, dtype=torch.int32)[idx_tensor]
        width = int(meta["w"]) if width_fixed else torch.tensor(width, dtype=torch.int32)[idx_tensor]
        if distort_fixed:
            distortion_params = camera_utils.get_distortion_params(
                k1=float(meta["k1"]) if "k1" in meta else 0.0,
                k2=float(meta["k2"]) if "k2" in meta else 0.0,
                k3=float(meta["k3"]) if "k3" in meta else 0.0,
                k4=float(meta["k4"]) if "k4" in meta else 0.0,
                p1=float(meta["p1"]) if "p1" in meta else 0.0,
                p2=float(meta["p2"]) if "p2" in meta else 0.0,
            )
        else:
            distortion_params = torch.stack(distort, dim=0)[idx_tensor]

        cameras = Cameras(
            fx=fx,
            fy=fy,
            cx=cx,
            cy=cy,
            distortion_params=distortion_params,
            height=height,
            width=width,
            camera_to_worlds=poses[:, :3, :4],
            camera_type=camera_type,
        )

        assert self.downscale_factor is not None
        cameras.rescale_output_resolution(scaling_factor=1.0 / self.downscale_factor)

        metadata = {
            "depth_filenames": depth_filenames if len(depth_filenames) > 0 else None,
            "depth_unit_scale_factor": self.config.depth_unit_scale_factor,
        }

        if self.config.eval_mode == "eval-frame-index":
            metadata["eval_frame_index_0_metadata"] = torch.tensor(eval_frame_index_0_metadata, dtype=torch.long)

        if "applied_transform" in meta:
            applied_transform = torch.tensor(meta["applied_transform"], dtype=transform_matrix.dtype)
            transform_matrix = transform_matrix @ torch.cat(
                [applied_transform, torch.tensor([[0, 0, 0, 1]], dtype=transform_matrix.dtype)], 0
            )
        if "applied_scale" in meta:
            applied_scale = float(meta["applied_scale"])
            scale_factor *= applied_scale

        dataparser_outputs = DataparserOutputs(
            image_filenames=image_filenames,
            cameras=cameras,
            scene_box=scene_box,
            mask_filenames=mask_filenames if len(mask_filenames) > 0 else None,
            dataparser_scale=scale_factor,
            dataparser_transform=transform_matrix,
            metadata=metadata,
        )
        return dataparser_outputs

    def _get_fname(self, filepath: PurePath, data_dir: PurePath, downsample_folder_prefix="images_") -> Path:
        """Get the filename of the image file.
        downsample_folder_prefix can be used to point to auxiliary image data, e.g. masks

        filepath: the base file name of the transformations.
        data_dir: the directory of the data that contains the transform file
        downsample_folder_prefix: prefix of the newly generated downsampled images
        """

        if self.downscale_factor is None:
            if self.config.downscale_factor is None:
                test_img = Image.open(data_dir / filepath)
                h, w = test_img.size
                max_res = max(h, w)
                df = 0
                while True:
                    if (max_res / 2 ** (df)) < MAX_AUTO_RESOLUTION:
                        break
                    if not (data_dir / f"{downsample_folder_prefix}{2**(df+1)}" / filepath.name).exists():
                        break
                    df += 1

                self.downscale_factor = 2**df
                CONSOLE.log(f"Auto image downscale factor of {self.downscale_factor}")
            else:
                self.downscale_factor = self.config.downscale_factor

        if self.downscale_factor > 1:
            return data_dir / f"{downsample_folder_prefix}{self.downscale_factor}" / filepath.name
        return data_dir / filepath
bayesrays/dataparsers/nerfbusters/nb_dataparser.py
BayesRays/BayesRays
''' This file only serves the purpose of loading nerfbuster dataset for bayesrays and does not work for running nerfbuster mode. To run/test nerfbuster model please refer to nerfbuster repository '''

from nerfstudio.plugins.registry_dataparser import DataParserSpecification
from bayesrays.dataparsers.nerfbusters.nb_dataparser import NerfbusterDataparserConfig

nbDataparser = DataParserSpecification(config=NerfbusterDataparserConfig())
bayesrays/dataparsers/nerfbusters/nerfbuster_config.py
BayesRays/BayesRays
''' Loading a few forward-facing views of the whole dataset as train and the rest as test data for bayesrays '''
# view picks for LF based on CF-NeRF and for ScanNet based on NerfingMVS codebases.

from nerfstudio.plugins.registry_dataparser import DataParserSpecification
from bayesrays.dataparsers.sparse.sparse_nerfstudio_dataparser import SparseNsDataParserConfig

sparseNsDataparser = DataParserSpecification(config=SparseNsDataParserConfig())
bayesrays/dataparsers/sparse/sparse_nerfstudio_config.py
BayesRays/BayesRays

""" Same as nerfstudio dataparse but uses CF-NeRF and NerfingMVS train/test splits for LF and ScanNet datasets """

from __future__ import annotations

import math
from dataclasses import dataclass, field
from pathlib import Path
from typing import Literal, Optional, Type

import numpy as np
import torch
from PIL import Image

from nerfstudio.cameras import camera_utils
from nerfstudio.cameras.cameras import CAMERA_MODEL_TO_TYPE, Cameras, CameraType
from nerfstudio.data.dataparsers.base_dataparser import (
    DataParser,
    DataParserConfig,
    DataparserOutputs,
)
from nerfstudio.data.scene_box import SceneBox
from nerfstudio.utils.io import load_from_json
from nerfstudio.utils.rich_utils import CONSOLE
from numpy.random import RandomState


MAX_AUTO_RESOLUTION = 1600
seed = 12345


@dataclass
class SparseNsDataParserConfig(DataParserConfig):
    """Sparse Nerfstudio dataset config"""

    _target: Type = field(default_factory=lambda: SparseNerfstudio)
    """target class to instantiate"""
    data: Path = Path()
    """Directory or explicit json file path specifying location of data."""
    scale_factor: float = 1.0
    """How much to scale the camera origins by."""
    downscale_factor: Optional[int] = 2
    """How much to downscale images. If not set, images are chosen such that the max dimension is <1600px."""
    orientation_method: Literal["pca", "up", "vertical", "none"] = "up"
    """The method to use for orientation."""
    center_method: Literal["poses", "focus", "none"] = "poses"
    """The method to use to center the poses."""
    auto_scale_poses: bool = False
    """Whether to automatically scale the poses to fit in +/- 1 bounding box."""
    depth_unit_scale_factor: float = 1e-3
    """Scales the depth values to meters. Default value is 0.001 for a millimeter to meter conversion."""
    dataset_name: Literal["africa", "basket", "statue", "torch", "scene_079","scene_000","scene_158","scene_316"] = "africa"


@dataclass
class SparseNerfstudio(DataParser):
    """Sparse Nerfstudio DatasetParser"""

    config: SparseNsDataParserConfig
    downscale_factor: Optional[int] = None

    def _generate_dataparser_outputs(self, split="train"):
        assert self.config.data.exists(), f"Data directory {self.config.data} does not exist."

        if self.config.data.suffix == ".json":
            meta = load_from_json(self.config.data)
            data_dir = self.config.data.parent
        else:
            meta = load_from_json(self.config.data / "transforms.json")
            data_dir = self.config.data

        image_filenames = []
        mask_filenames = []
        depth_filenames = []
        poses = []

        fx_fixed = "fl_x" in meta
        fy_fixed = "fl_y" in meta
        cx_fixed = "cx" in meta
        cy_fixed = "cy" in meta
        height_fixed = "h" in meta
        width_fixed = "w" in meta
        distort_fixed = False
        for distort_key in ["k1", "k2", "k3", "p1", "p2"]:
            if distort_key in meta:
                distort_fixed = True
                break
        fx = []
        fy = []
        cx = []
        cy = []
        height = []
        width = []
        distort = []

        for frame in meta["frames"]:
            filepath = Path(frame["file_path"])
            fname = self._get_fname(filepath, data_dir)

            if not fx_fixed:
                assert "fl_x" in frame, "fx not specified in frame"
                fx.append(float(frame["fl_x"]))
            if not fy_fixed:
                assert "fl_y" in frame, "fy not specified in frame"
                fy.append(float(frame["fl_y"]))
            if not cx_fixed:
                assert "cx" in frame, "cx not specified in frame"
                cx.append(float(frame["cx"]))
            if not cy_fixed:
                assert "cy" in frame, "cy not specified in frame"
                cy.append(float(frame["cy"]))
            if not height_fixed:
                assert "h" in frame, "height not specified in frame"
                height.append(int(frame["h"]))
            if not width_fixed:
                assert "w" in frame, "width not specified in frame"
                width.append(int(frame["w"]))
            if not distort_fixed:
                distort.append(
                    camera_utils.get_distortion_params(
                        k1=float(frame["k1"]) if "k1" in frame else 0.0,
                        k2=float(frame["k2"]) if "k2" in frame else 0.0,
                        k3=float(frame["k3"]) if "k3" in frame else 0.0,
                        k4=float(frame["k4"]) if "k4" in frame else 0.0,
                        p1=float(frame["p1"]) if "p1" in frame else 0.0,
                        p2=float(frame["p2"]) if "p2" in frame else 0.0,
                    )
                )

            image_filenames.append(fname)
            poses.append(np.array(frame["transform_matrix"]))
            if "mask_path" in frame:
                mask_filepath = Path(frame["mask_path"])
                mask_fname = self._get_fname(
                    mask_filepath,
                    data_dir,
                    downsample_folder_prefix="masks_",
                )
                mask_filenames.append(mask_fname)

            if "depth_file_path" in frame:
                depth_filepath = Path(frame["depth_file_path"])
                depth_fname = self._get_fname(depth_filepath, data_dir, downsample_folder_prefix="depths_")
                depth_filenames.append(depth_fname)
        
        if "mask_path" in frame:
            sorted_lists = sorted(zip(image_filenames, mask_filenames))
            _, mask_filenames = zip(*sorted_lists)
        if "depth_file_path" in frame:
            sorted_lists = sorted(zip(image_filenames, depth_filenames))
            _, depth_filenames = zip(*sorted_lists)    
            
        sorted_lists = sorted(zip(image_filenames, poses))
        image_filenames, poses = zip(*sorted_lists)
        
        assert len(mask_filenames) == 0 or (
            len(mask_filenames) == len(image_filenames)
        ), """
        Different number of image and mask filenames.
        You should check that mask_path is specified for every frame (or zero frames) in transforms.json.
        """
        assert len(depth_filenames) == 0 or (
            len(depth_filenames) == len(image_filenames)
        ), """
        Different number of image and depth filenames.
        You should check that depth_file_path is specified for every frame (or zero frames) in transforms.json.
        """

        has_split_files_spec = any(f"{split}_filenames" in meta for split in ("train", "val", "test"))
        if f"{split}_filenames" in meta:
            # Validate split first
            split_filenames = set(self._get_fname(Path(x), data_dir) for x in meta[f"{split}_filenames"])
            unmatched_filenames = split_filenames.difference(image_filenames)
            if unmatched_filenames:
                raise RuntimeError(f"Some filenames for split {split} were not found: {unmatched_filenames}.")

            indices = [i for i, path in enumerate(image_filenames) if path in split_filenames]
            CONSOLE.log(f"[yellow] Dataset is overriding {split}_indices to {indices}")
            indices = np.array(indices, dtype=np.int32)
        elif has_split_files_spec:
            raise RuntimeError(f"The dataset's list of filenames for split {split} is missing.")
        else:
            # filter image_filenames and poses based on train/eval split percentage
            num_images = len(image_filenames)
            
            i_all = np.arange(num_images)
            
#train/test split from CF-NeRF https://github.com/poetrywanderer/CFNeRF/blob/66918a9748c137e1c0242c12be7aa6efa39ece06/run_nerf_uncertainty_NF.py#L750
            if self.config.dataset_name == 'basket':
                # 4 views
                i_train = list(np.arange(43,50,2))
                i_val = list(np.arange(42,50,2))
                self.config.scene_scale = 50. # half the far bound computed by colmap bounds 
            elif self.config.dataset_name == 'africa':
                # 5 views
                i_train = list(np.arange(5,14,2))
                i_val = list(np.arange(6,14,2))
                self.config.scene_scale = 36.

            elif self.config.dataset_name == 'statue':
                # 5 views
                i_train = list(np.arange(67,76,2))
                i_val = list(np.arange(68,76,2))
                self.config.scene_scale = 50.

            elif self.config.dataset_name == 'torch':
                # 5 views
                i_train = list(np.arange(8,17,2))
                i_val = list(np.arange(9,17,2))
                self.config.scene_scale = 72.
                
# train/test split from scannet scenes of https://github.com/weiyithu/NerfingMVS/tree/28511191239daf25cd8ded17e7fa21a68df54de1
            elif self.config.dataset_name == 'scene_079':
                i_val = [4, 12, 20, 28, 37]
                i_train = list([i for i in range(40) if i not in i_val])
                self.config.scene_scale = 150.
            elif self.config.dataset_name == 'scene_000':
                i_val = [4, 12, 20, 28, 36]
                i_train = list([i for i in range(40) if i not in i_val])
                self.config.scene_scale = 140.
            elif self.config.dataset_name == 'scene_316':
                i_val = [4, 12, 20, 28, 36]
                i_train = list([i for i in range(40) if i not in i_val])
                self.config.scene_scale = 100.
            elif self.config.dataset_name == 'scene_158':
                i_val = [4, 12, 20, 28, 36]
                i_train = list([i for i in range(40) if i not in i_val]) 
                self.config.scene_scale = 2000
            
            else:
                raise ValueError(f"Unknown scene {self.config.dataset_name}")
            if split == "train":
                indices = i_train
            elif split in ["val", "test"]:
                indices = i_val
            else:
                raise ValueError(f"Unknown dataparser split {split}")

        if "orientation_override" in meta:
            orientation_method = meta["orientation_override"]
            CONSOLE.log(f"[yellow] Dataset is overriding orientation method to {orientation_method}")
        else:
            orientation_method = self.config.orientation_method

        poses = torch.from_numpy(np.array(poses).astype(np.float32))
        poses, transform_matrix = camera_utils.auto_orient_and_center_poses(
            poses,
            method=orientation_method,
            center_method=self.config.center_method,
        )
        # Scale poses
        scale_factor = 1.0
        if self.config.auto_scale_poses:
            scale_factor /= float(torch.max(torch.abs(poses[:, :3, 3])))
        scale_factor *= self.config.scale_factor

        poses[:, :3, 3] *= scale_factor

        # Choose image_filenames and poses based on split, but after auto orient and scaling the poses.
        image_filenames = [image_filenames[i] for i in indices]
        mask_filenames = [mask_filenames[i] for i in indices] if len(mask_filenames) > 0 else []
        depth_filenames = [depth_filenames[i] for i in indices] if len(depth_filenames) > 0 else []

        idx_tensor = torch.tensor(indices, dtype=torch.long)
        poses = poses[idx_tensor]

        # in x,y,z order
        # assumes that the scene is centered at the origin
        aabb_scale = self.config.scene_scale
        scene_box = SceneBox(
            aabb=torch.tensor(
                [[-aabb_scale, -aabb_scale, -aabb_scale], [aabb_scale, aabb_scale, aabb_scale]], dtype=torch.float32
            )
        )

        if "camera_model" in meta:
            camera_type = CAMERA_MODEL_TO_TYPE[meta["camera_model"]]
        else:
            camera_type = CameraType.PERSPECTIVE

        fx = float(meta["fl_x"]) if fx_fixed else torch.tensor(fx, dtype=torch.float32)[idx_tensor]
        fy = float(meta["fl_y"]) if fy_fixed else torch.tensor(fy, dtype=torch.float32)[idx_tensor]
        cx = float(meta["cx"]) if cx_fixed else torch.tensor(cx, dtype=torch.float32)[idx_tensor]
        cy = float(meta["cy"]) if cy_fixed else torch.tensor(cy, dtype=torch.float32)[idx_tensor]
        height = int(meta["h"]) if height_fixed else torch.tensor(height, dtype=torch.int32)[idx_tensor]
        width = int(meta["w"]) if width_fixed else torch.tensor(width, dtype=torch.int32)[idx_tensor]
        if distort_fixed:
            distortion_params = camera_utils.get_distortion_params(
                k1=float(meta["k1"]) if "k1" in meta else 0.0,
                k2=float(meta["k2"]) if "k2" in meta else 0.0,
                k3=float(meta["k3"]) if "k3" in meta else 0.0,
                k4=float(meta["k4"]) if "k4" in meta else 0.0,
                p1=float(meta["p1"]) if "p1" in meta else 0.0,
                p2=float(meta["p2"]) if "p2" in meta else 0.0,
            )
        else:
            distortion_params = torch.stack(distort, dim=0)[idx_tensor]

        cameras = Cameras(
            fx=fx,
            fy=fy,
            cx=cx,
            cy=cy,
            distortion_params=distortion_params,
            height=height,
            width=width,
            camera_to_worlds=poses[:, :3, :4],
            camera_type=camera_type,
        )

        assert self.downscale_factor is not None
        cameras.rescale_output_resolution(scaling_factor=1.0 / self.downscale_factor)

        if "applied_transform" in meta:
            applied_transform = torch.tensor(meta["applied_transform"], dtype=transform_matrix.dtype)
            transform_matrix = transform_matrix @ torch.cat(
                [applied_transform, torch.tensor([[0, 0, 0, 1]], dtype=transform_matrix.dtype)], 0
            )
        if "applied_scale" in meta:
            applied_scale = float(meta["applied_scale"])
            scale_factor *= applied_scale
        dataparser_outputs = DataparserOutputs(
            image_filenames=image_filenames,
            cameras=cameras,
            scene_box=scene_box,
            mask_filenames=mask_filenames if len(mask_filenames) > 0 else None,
            dataparser_scale=scale_factor,
            dataparser_transform=transform_matrix,
            metadata={
                "depth_filenames": depth_filenames if len(depth_filenames) > 0 else None,
                "depth_unit_scale_factor": self.config.depth_unit_scale_factor,
            },
        )
        return dataparser_outputs

    def _get_fname(self, filepath: Path, data_dir: Path, downsample_folder_prefix="images_") -> Path:
        """Get the filename of the image file.
        downsample_folder_prefix can be used to point to auxiliary image data, e.g. masks

        filepath: the base file name of the transformations.
        data_dir: the directory of the data that contains the transform file
        downsample_folder_prefix: prefix of the newly generated downsampled images
        """

        if self.downscale_factor is None:
            if self.config.downscale_factor is None:
                test_img = Image.open(data_dir / filepath)
                h, w = test_img.size
                max_res = max(h, w)
                df = 0
                while True:
                    if (max_res / 2 ** (df)) < MAX_AUTO_RESOLUTION:
                        break
                    if not (data_dir / f"{downsample_folder_prefix}{2**(df+1)}" / filepath.name).exists():
                        break
                    df += 1

                self.downscale_factor = 2**df
                CONSOLE.log(f"Auto image downscale factor of {self.downscale_factor}")
            else:
                self.downscale_factor = self.config.downscale_factor

        if self.downscale_factor > 1:
            return data_dir / f"{downsample_folder_prefix}{self.downscale_factor}" / filepath.name
        return data_dir / filepath

bayesrays/dataparsers/sparse/sparse_nerfstudio_dataparser.py
BayesRays/BayesRays
import torch
import numpy as np

# original code from https://github.com/poetrywanderer/CF-NeRF/blob/66918a9748c137e1c0242c12be7aa6efa39ece06/run_nerf_helpers.py#L382

def ause(unc_vec, err_vec, err_type='rmse'):
    ratio_removed = np.linspace(0, 1, 100, endpoint=False)
    # Sort the error
    err_vec_sorted, _ = torch.sort(err_vec)

    # Calculate the error when removing a fraction pixels with error
    n_valid_pixels = len(err_vec)
    ause_err = []
    for r in ratio_removed:
        err_slice = err_vec_sorted[0:int((1-r)*n_valid_pixels)]
        if err_type == 'rmse':
            ause_err.append(torch.sqrt(err_slice.mean()).cpu().numpy())
        elif err_type == 'mae' or err_type == 'mse':
            ause_err.append(err_slice.mean().cpu().numpy())
       

    ###########################################

    # Sort by variance
    _, var_vec_sorted_idxs = torch.sort(unc_vec)
    # Sort error by variance
    err_vec_sorted_by_var = err_vec[var_vec_sorted_idxs]
    ause_err_by_var = []
    for r in ratio_removed:
        
        err_slice = err_vec_sorted_by_var[0:int((1 - r) * n_valid_pixels)]
        if err_type == 'rmse':
            ause_err_by_var.append(torch.sqrt(err_slice.mean()).cpu().numpy())
        elif err_type == 'mae'or err_type == 'mse':
            ause_err_by_var.append(err_slice.mean().cpu().numpy())
    
    #Normalize and append
    max_val = max(max(ause_err), max(ause_err_by_var))
    ause_err = ause_err / max_val
    ause_err = np.array(ause_err)
    
    ause_err_by_var = ause_err_by_var / max_val
    ause_err_by_var = np.array(ause_err_by_var)
    ause = np.trapz(ause_err_by_var - ause_err, ratio_removed)
    return ratio_removed, ause_err, ause_err_by_var, ause
bayesrays/metrics/ause.py
BayesRays/BayesRays
# source: https://github.com/nerfstudio-project/nerfstudio/blob/nerfbusters-changes/nerfstudio/utils/metrics.py

"""
This file contains NeRF metrics with masking capabilities.
"""

from abc import abstractmethod
from typing import Optional

import torch
from torch import Tensor
from torch import nn
from torchmetrics.functional import structural_similarity_index_measure
from torchtyping import TensorType
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity


def _valid_img(img: Tensor, normalize: bool):
    """check that input is a valid image to the network."""
    value_check = img.max() <= 1.0 and img.min() >= 0.0 if normalize else img.min() >= -1
    return img.ndim == 4 and img.shape[1] == 3 and value_check


class LearnedPerceptualImagePatchSimilarityWithMasking(LearnedPerceptualImagePatchSimilarity):
    """LearnedPerceptualImagePatchSimilarity module that will allow for masking capabilities."""

    def update(self, img1: Tensor, img2: Tensor) -> None:  # pylint: disable=arguments-differ
        """Update internal states with lpips score with masking."""

        # hardcode this to True for now to avoid touching a lot of the torchmetrics code
        self.net.spatial = True

        if not (_valid_img(img1, self.normalize) and _valid_img(img2, self.normalize)):
            raise ValueError(
                "Expected both input arguments to be normalized tensors with shape [N, 3, H, W]."
                f" Got input with shape {img1.shape} and {img2.shape} and values in range"
                f" {[img1.min(), img1.max()]} and {[img2.min(), img2.max()]} when all values are"
                f" expected to be in the {[0,1] if self.normalize else [-1,1]} range."
            )
        loss = self.net(img1, img2, normalize=self.normalize)
        # now loss is the shape [batch size, H, W]
        # we set loss to self.sum_scores to use the existing API from torchvision
        self.sum_scores = loss  # pylint: disable=attribute-defined-outside-init

    def compute(self) -> Tensor:
        """Compute final perceptual similarity metric."""
        # note that we don't use self.reduction anymore
        return self.sum_scores
    
class ImageMetricModule(nn.Module):
    """Computes image metrics with masking capabilities.
    We assume that the pred and target inputs are in the range [0, 1].
    """

    def __init__(self):
        super().__init__()
        self.populate_modules()

    def populate_modules(self):
        """Populates the modules that will be used to compute the metric."""

    @abstractmethod
    def forward(
        self,
        preds: TensorType["bs", 3, "H", "W"],
        target: TensorType["bs", 3, "H", "W"],
        mask: Optional[TensorType["bs", 1, "H", "W"]] = None,
    ) -> TensorType["bs"]:
        """Computes the metric.
        Args:
            preds: Predictions.
            target: Ground truth.
            mask: Mask to use to only compute the metrics where the mask is True.
        Returns:
            Metric value.
        """


class PSNRModule(ImageMetricModule):
    """Computes PSNR with masking capabilities."""

    def forward(
        self,
        preds: TensorType["bs", 3, "H", "W"],
        target: TensorType["bs", 3, "H", "W"],
        mask: Optional[TensorType["bs", 1, "H", "W"]] = None,
    ) -> TensorType["bs"]:
        bs, h, w = preds.shape[0], preds.shape[2], preds.shape[3]
        hw = h * w

        preds_reshaped = preds.view(bs, 3, hw)
        target_reshaped = target.view(bs, 3, hw)
        num = (preds_reshaped - target_reshaped) ** 2
        # the non-masked version
        if mask is None:
            den = hw
        else:
            mask_reshaped = mask.view(bs, 1, hw)
            num = num * mask_reshaped
            den = mask_reshaped.sum(-1)
        mse = num.sum(-1) / den
        psnr = 10 * torch.log10(1.0 / mse)
        psnr = psnr.mean(-1)
        return psnr


class SSIMModule(ImageMetricModule):
    """Computes PSNR with masking capabilities."""

    def forward(
        self,
        preds: TensorType["bs", 3, "H", "W"],
        target: TensorType["bs", 3, "H", "W"],
        mask: Optional[TensorType["bs", 1, "H", "W"]] = None,
    ) -> TensorType["bs"]:
        bs, h, w = preds.shape[0], preds.shape[2], preds.shape[3]
        hw = h * w

        _, ssim_image = structural_similarity_index_measure(
            preds=preds, target=target, reduction="none", data_range=1.0, return_full_image=True
        )
        ssim_image = ssim_image.mean(1)  # average over the channels
        assert ssim_image.shape == (bs, h, w)

        # the non-masked version
        if mask is None:
            ssim = ssim_image.view(bs, hw).mean(1)
            return ssim

        # the masked version
        ssim_reshaped = ssim_image.view(bs, hw)
        mask_reshaped = mask.view(bs, hw)
        den = mask_reshaped.sum(-1, keepdim=True)
        ssim = (ssim_reshaped * mask_reshaped / den).sum(-1)
        return ssim


class LPIPSModule(ImageMetricModule):
    """Computes LPIPS with masking capabilities."""

    def populate_modules(self):
        # by setting normalize=True, we assume that the pred and target inputs are in the range [0, 1]
        self.lpips_with_masking = LearnedPerceptualImagePatchSimilarityWithMasking(normalize=True)

    def forward(
        self,
        preds: TensorType["bs", 3, "H", "W"],
        target: TensorType["bs", 3, "H", "W"],
        mask: Optional[TensorType["bs", 1, "H", "W"]] = None,
    ) -> TensorType["bs"]:
        bs, h, w = preds.shape[0], preds.shape[2], preds.shape[3]
        hw = h * w

        with torch.no_grad():
            lpips_image = self.lpips_with_masking(preds, target)
        lpips_image = lpips_image.mean(1)  # average over the channels
        assert lpips_image.shape == (bs, h, w)

        # the non-masked version
        if mask is None:
            lpips = lpips_image.view(bs, hw).mean(1)
            return lpips

        # the masked version
        lpips_reshaped = lpips_image.view(bs, hw)
        mask_reshaped = mask.view(bs, hw)
        den = mask_reshaped.sum(-1, keepdim=True)
        lpips = (lpips_reshaped * mask_reshaped / den).sum(-1)
        return lpips
bayesrays/metrics/image_metrics.py
BayesRays/BayesRays
# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#!/usr/bin/env python
"""
eval.py
"""
from __future__ import annotations

from datetime import datetime
import json
import types
import torch
import mediapy as media
from dataclasses import dataclass
from time import time
from pathlib import Path
from typing import Optional
from rich.progress import (
    BarColumn,
    MofNCompleteColumn,
    Progress,
    TextColumn,
    TimeElapsedColumn,
)
import tyro
import numpy as np
import cv2

import nerfstudio
import pkg_resources

from nerfstudio.utils.eval_utils import eval_setup
from nerfstudio.utils.rich_utils import CONSOLE
from nerfstudio.utils import colormaps

from bayesrays.metrics.ause import ause
from bayesrays.metrics.image_metrics import PSNRModule, SSIMModule, LPIPSModule
from torchmetrics import MeanSquaredError


from PIL import Image
import matplotlib.pyplot as plt
from matplotlib.cm import inferno

from bayesrays.scripts.output_uncertainty import get_output_nerfacto_new, get_output_mipnerf_new, get_output_ngp_new, get_output_fn, get_output_nerfacto_all, get_uncertainty

def plot_errors(ratio_removed, ause_err, ause_err_by_var, err_type, scene_no, output_path): #AUSE plots, with oracle curve also visible
    plt.plot(ratio_removed, ause_err, '--')
    plt.plot(ratio_removed, ause_err_by_var, '-r')
    # plt.plot(ratio_removed, ause_err_by_var - ause_err, '-g') # uncomment for getting plots similar to the paper, without visible oracle curve
    path = output_path.parent / Path("plots") 
    path.mkdir(parents=True, exist_ok=True)
    plt.savefig(path/ Path('plot_'+err_type+'_'+str(scene_no)+'.png'))
    plt.figure()

def visualize_ranks(unc ,gt, colormap='jet'):
    flattened_unc = unc.flatten()
    flattened_gt = gt.flatten()
    
    # Find the ranks of the pixel values
    ranks_unc = np.argsort(np.argsort(flattened_unc)) 
    ranks_gt = np.argsort(np.argsort(flattened_gt)) 
    
    max_rank = max(np.max(ranks_unc),np.max(ranks_gt))
    
    cmap = plt.get_cmap(colormap, max_rank)
    
    # Normalize the ranks to the range [0, 1]
    normalized_ranks_unc = ranks_unc / max_rank
    normalized_ranks_gt = ranks_gt / max_rank
    
    # Apply the colormap to the normalized ranks
    colored_ranks_unc = cmap(normalized_ranks_unc)
    colored_ranks_gt = cmap(normalized_ranks_gt)
    
    colored_unc = colored_ranks_unc.reshape((*unc.shape,4))
    colored_gt = colored_ranks_gt.reshape((*gt.shape,4))
    
    return colored_unc, colored_gt

def get_filtered_image_metrics(self, 
                                 outputs: Dict[str, torch.Tensor],
                                 batch: Dict[str, torch.Tensor], thresh: torch.Tensor, add_nb_mask=False, visibility_mask: torch.Tensor=None) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
    
    image = batch["image"]
    rgb = outputs["rgb-"+"{:.2f}".format(thresh.item())]
    h,w = rgb.shape[0], rgb.shape[1]
    image = torch.clip(torch.tensor(media.resize_image(image, (h, w))).to(self.device),0.,1.)
    
    acc = outputs["accumulation-"+"{:.2f}".format(thresh.item())]
    depth = outputs["depth-"+"{:.2f}".format(thresh.item())]
    #implementing masked psnr,lpisp,ssim like https://github.com/ethanweber
    #/nerfbusters/blob/1f4240344ecff1313f6dfa7be5e06fe7d3e29154/scripts/launch_nerf.py#L258
    depth_mask = (depth < 2.0).float()
   
    if add_nb_mask:
        mask = depth_mask[...,0] * visibility_mask
        mask = mask[..., None].repeat(1, 1, 3)
    else:
        mask = depth_mask[...,0]
        mask = mask[..., None].repeat(1, 1, 3)
    image, rgb = image * mask, rgb * mask
    
    # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
    image = image.permute(2, 0, 1).unsqueeze(0)
    rgb = rgb.permute(2, 0, 1).unsqueeze(0)
    m = mask.permute(2, 0, 1).unsqueeze(0)[:, 0:1]

    psnr = float(self.psnr_module(rgb, image, m)[0])
    ssim = float(self.ssim_module(rgb, image, m)[0])
    lpips = float(self.lpips_module(rgb, image, m)[0])

    
    metrics_dict = {"psnr": float(psnr), "ssim": float(ssim)}  # type: ignore
    metrics_dict["lpips"] = float(lpips)
    if add_nb_mask:
        metrics_dict["coverage"] = float(mask[..., 0].sum() / visibility_mask.sum() * 100)
    else:    
        metrics_dict["coverage"] = float(mask[..., 0].sum()/(image.shape[-1]*image.shape[-2]) * 100)
    return metrics_dict

def get_average_filtered_image_metrics(self, step: Optional[int] = None):
    """Iterate over all the images in the eval dataset and get the average.
    From https://github.com/nerfstudio-project/nerfstudio/blob/main/nerfstudio/pipelines/base_pipeline.py#L342

    Returns:
        metrics_dict: dictionary of metrics
    """
    self.eval()
    num_images = len(self.datamanager.fixed_indices_eval_dataloader)

    # Override evaluation function
    self.model.get_image_metrics_and_images = types.MethodType(get_filtered_image_metrics, self.model)
    
    thresh_values = torch.linspace(0.1,1,10)
    self.model.thresh_range = thresh_values
    
    self.model.psnr_module = PSNRModule().to(self.device)
    self.model.ssim_module = SSIMModule().to(self.device)
    self.model.lpips_module = LPIPSModule().to(self.device)
    
    views = ["view"]
    psnr = [["psnr-"+"{:.2f}".format(i.item())] for i in thresh_values]
    lpips = [["lpisp-"+"{:.2f}".format(i.item())] for i in thresh_values]
    ssim = [["ssim-"+"{:.2f}".format(i.item())] for i in thresh_values]
    coverage = [["coverage-"+"{:.2f}".format(i.item())] for i in thresh_values]
    self.datamanager.fixed_indices_eval_dataloader.cameras.rescale_output_resolution(1./self.downscale_factor)
    with Progress(
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TimeElapsedColumn(),
        MofNCompleteColumn(),
        transient=True,
    ) as progress:
        task = progress.add_task("[green]Evaluating all eval images...", total=num_images)
        view_no = 0 
        
        for camera_ray_bundle, batch in self.datamanager.fixed_indices_eval_dataloader:
            
            if self.add_nb_mask:
                base_path = self.nb_mask_path
                pseudo_gt_visibility = media.read_image(str(base_path)+"/{:05d}.png".format(view_no))
                pseudo_gt_visibility = torch.from_numpy(pseudo_gt_visibility).long().to(self.device)
                pseudo_gt_visibility = (pseudo_gt_visibility[..., 0] >= 1).float()
            else:
                pseudo_gt_visibility = 1
            
            # time this the following line
            inner_start = time()
            height, width = camera_ray_bundle.shape
            num_rays = height * width
            
            outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
            for n,thresh in enumerate(thresh_values):
                metrics_dict = self.model.get_image_metrics_and_images(outputs, batch, thresh, self.add_nb_mask, pseudo_gt_visibility)
                psnr[n].append(float(metrics_dict["psnr"]))
                lpips[n].append(float(metrics_dict["lpips"]))
                ssim[n].append(float(metrics_dict["ssim"]))
                coverage[n].append(float(metrics_dict["coverage"]))
                
                print("view:", view_no, "-", str(thresh), "psnr:", psnr[n][-1], "coverage:", coverage[n][-1])

            views.append(str(view_no))    
            view_no +=1
            progress.advance(task)
            

    # average the metrics list
    metrics_dict = {}

    lists = [*psnr, *ssim, *lpips, *coverage]
    for l in lists:
        l.append(sum(l[1:])/view_no)
    views.append('average')
    lists = [views] + lists 
    self.train()
    return {}, lists



def get_image_metrics_and_images_unc(self, no:int,
                                 outputs: Dict[str, torch.Tensor],
                                 batch: Dict[str, torch.Tensor], 
                                 err_all: list,
                                 err_var_all: list,
                                 last: bool,
                                 eval_depth: bool) -> Tuple[Dict[str, float], Dict[str, torch.Tensor]]:
    """ From https://github.com/nerfstudio-project/nerfstudio/blob/main/nerfstudio/models/nerfacto.py#L357 """
    image = batch["image"].to(self.device)
    rgb = outputs["rgb"]

    unc = outputs["uncertainty"]

    acc = colormaps.apply_colormap(outputs["accumulation"])
    depth = colormaps.apply_depth_colormap(
        outputs["depth"],
        accumulation=outputs["accumulation"],
    )
    combined_rgb = torch.cat([image, rgb], dim=1)
    combined_acc = torch.cat([acc], dim=1)
    combined_depth = torch.cat([depth], dim=1)
    
    if eval_depth:
        depth = outputs["depth"].squeeze(-1)
        
        # load the calculated scale, to run evaluation on depth in same scale as GT depth
        a = float(np.loadtxt(str(self.dataset_path) +'/scale_parameters.txt', delimiter=','))
        depth = a * depth 
        depth_gt_dir = str(self.dataset_path) + '/depth_gt_{:02d}.npy'.format(no)
        depth_gt = np.load(depth_gt_dir)
        depth_gt = torch.tensor(depth_gt, device=depth.device)
        depth = depth/depth_gt.max()
        depth_gt = depth_gt/depth_gt.max()


        squared_error = ((depth_gt - depth) ** 2)
        absolute_error = (abs(depth_gt - depth))
        unc_flat = unc.flatten()
        absolute_error_flat = absolute_error.flatten()
        squared_error_flat = squared_error.flatten()

        ratio, err_mse, err_var_mse, ause_mse = ause(unc_flat, squared_error_flat, err_type='mse')
        plot_errors(ratio, err_mse, err_var_mse, 'mse', no, self.output_path)
        ratio, err_mae, err_var_mae, ause_mae =  ause(unc_flat, absolute_error_flat, err_type='mae')
        plot_errors(ratio, err_mae, err_var_mae, 'mae', no, self.output_path)
        ratio, err_rmse, err_var_rmse, ause_rmse =  ause(unc_flat, squared_error_flat, err_type='rmse')
        plot_errors(ratio, err_rmse, err_var_rmse, 'rmse', no, self.output_path)

        err_all[0] += err_mse
        err_all[1] += err_rmse
        err_all[2] += err_mae
        err_var_all[0] += err_var_mse
        err_var_all[1] += err_var_rmse
        err_var_all[2] += err_var_mae

        if last:
            ratio_all = ratio
            err_mse_all, err_rmse_all, err_mae_all = err_all[0]/(no+1), err_all[1]/(no+1), err_all[2]/(no+1)
            err_var_mse_all, err_var_rmse_all, err_var_mae_all = err_var_all[0]/(no+1), err_var_all[1]/(no+1), err_var_all[2]/(no+1)
            plot_errors(ratio_all, err_mse_all, err_var_mse_all, 'mse', 'all', self.output_path)
            plot_errors(ratio_all, err_rmse_all, err_var_rmse_all, 'rmse', 'all', self.output_path)
            plot_errors(ratio_all, err_mae_all, err_var_mae_all, 'mae', 'all', self.output_path)



        #for visualizaiton
        depth_img = torch.clip(depth, min=0., max=1.)
        absolute_error_img = torch.clip(absolute_error, min=0., max=1.)
    
    #save images
    path = self.output_path.parent / "plots" 
    path.mkdir(parents=True, exist_ok=True)
    if eval_depth:
        im = Image.fromarray((depth_gt.cpu().numpy()* 255).astype('uint8'))
        im.save(path / Path(str(no)+"_depth_gt.jpeg"))
        im = Image.fromarray((depth_img.cpu().numpy()* 255).astype('uint8'))
        im.save(path / Path(str(no)+"_depth.jpeg"))
        im = Image.fromarray(np.uint8(inferno(absolute_error_img.cpu().numpy()) * 255) )
        im.save(path / Path(str(no)+"_error.png"))

        im = Image.fromarray((image.cpu().numpy()* 255).astype('uint8'))
        im.save(path / Path(str(no)+"_gt_image.jpeg"))
        uu, errr = visualize_ranks(unc.squeeze(-1).cpu().numpy(), absolute_error.cpu().numpy())
        im = Image.fromarray(np.uint8(uu * 255))
        im.save(path / Path(str(no)+"_unc_colored.png"))

        im = Image.fromarray(np.uint8(errr * 255))
        im.save(path / Path(str(no)+"_error_colored.png"))

    
  
    im = Image.fromarray(np.uint8(inferno(unc.squeeze(-1).cpu().numpy()) * 255).astype('uint8'))
    im.save(path / Path(str(no)+"_unc.png"))
    im = Image.fromarray(np.uint8(rgb.cpu().numpy() * 255).astype('uint8'))
    im.save(path / Path(str(no)+"_rgb.jpeg"))
    
    
    
    # Switch images from [H, W, C] to [1, C, H, W] for metrics computations
    image = torch.moveaxis(image, -1, 0)[None, ...]
    rgb = torch.moveaxis(rgb, -1, 0)[None, ...]

    psnr = self.psnr(image, rgb)
    ssim = self.ssim(image, rgb)
    lpips = self.lpips(image, rgb)


    # all of these metrics will be logged as scalars
    metrics_dict = {"psnr": float(psnr.item()), "ssim": float(ssim)}  # type: ignore
    metrics_dict["lpips"] = float(lpips)
    if eval_depth:
        metrics_dict["ause_mse"] = float(ause_mse)
        metrics_dict["ause_mae"] = float(ause_mae)
        metrics_dict["ause_rmse"] = float(ause_rmse)
        metrics_dict["mse"] = float(squared_error.mean().item())
    
    images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth}
    if eval_depth:
        images_dict['err_all'] = err_all
        images_dict['err_var_all'] = err_var_all


    return metrics_dict, images_dict


def get_average_uncertainty_metrics(self, step: Optional[int] = None):
    """Iterate over all the images in the eval dataset and get the average.
    From https://github.com/nerfstudio-project/nerfstudio/blob/main/nerfstudio/pipelines/base_pipeline.py#L342

    Returns:
        metrics_dict: dictionary of metrics
    """
    self.eval()
    metrics_dict_list = []
    unc_list = []
    mse_list = []
    num_images = len(self.datamanager.fixed_indices_eval_dataloader)

    # Override evaluation function
    self.model.get_image_metrics_and_images = types.MethodType(get_image_metrics_and_images_unc, self.model)
    
    views = ["view"]
    psnr = ["psnr"]
    lpips = ["lpips"]
    ssim = ["ssim"]
    ause_mse = ["ause_mse"]
    ause_rmse = ["ause_rmse"]
    ause_mae = ["ause_mae"]
    mse = ["mse"]
    
    with Progress(
        TextColumn("[progress.description]{task.description}"),
        BarColumn(),
        TimeElapsedColumn(),
        MofNCompleteColumn(),
        transient=True,
    ) as progress:
        task = progress.add_task("[green]Evaluating all eval images...", total=num_images)
        view_no = 0
        err_all = [np.zeros(100),np.zeros(100),np.zeros(100)]
        err_var_all = [np.zeros(100),np.zeros(100),np.zeros(100)] 

        for camera_ray_bundle, batch in self.datamanager.fixed_indices_eval_dataloader:
            
            # time this the following line
            inner_start = time()
            height, width = camera_ray_bundle.shape
            num_rays = height * width
            outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
            metrics_dict, images_dict = self.model.get_image_metrics_and_images(view_no, outputs, batch, err_all, err_var_all, view_no == len(self.datamanager.fixed_indices_eval_dataloader)-1, self.eval_depth)
            if self.eval_depth:
                mse_list.append(metrics_dict["mse"])
                err_all = images_dict['err_all'] 
                err_var_all = images_dict['err_var_all']

            
            # TODO do this in a cleaner way
            views.append(str(view_no))
            psnr.append(float(metrics_dict["psnr"]))
            lpips.append(float(metrics_dict["lpips"]))
            ssim.append(float(metrics_dict["ssim"]))
            if self.eval_depth:
                ause_mse.append(float(metrics_dict["ause_mse"]))
                ause_mae.append(float(metrics_dict["ause_mae"]))
                ause_rmse.append(float(metrics_dict["ause_rmse"]))
                mse.append(float(metrics_dict["mse"]))

            assert "num_rays_per_sec" not in metrics_dict
            metrics_dict["num_rays_per_sec"] = num_rays / (time() - inner_start)
            fps_str = "fps"
            assert fps_str not in metrics_dict
            metrics_dict[fps_str] = metrics_dict["num_rays_per_sec"] / (height * width)
            metrics_dict_list.append(metrics_dict)
            view_no +=1
            progress.advance(task)
            
    # average the metrics list
    metrics_dict = {}
    for key in metrics_dict_list[0].keys():
        metrics_dict[key] = float(
            torch.mean(torch.tensor([metrics_dict[key] for metrics_dict in metrics_dict_list]))
        )
    
    if self.eval_depth:    
        lists = [psnr, ssim, lpips, ause_mse, ause_rmse, ause_mae, mse]
    else:
        lists = [psnr, ssim, lpips]
    for l in lists:
        l.append(sum(l[1:])/view_no)
    views.append('average')
   
    lists = [views] + lists
    self.train()
    return metrics_dict, lists


@dataclass
class ComputeMetrics:
    """Load a checkpoint, compute some metrics, and save it to a JSON file."""

    # Path to config YAML file.
    load_config: Path
    # Name of the output file.
    output_path: Path = Path("output.json")
    # Path to output video file.
    unc_path: Path = Path("unc.npy")
    # Render with filtering.
    dataset_path: Path = Path("./data")
    # dataset path    
    downscale_factor: float = 2.0
    filter_out: bool = False
    # filter floater results
    nb_mask: bool = False
    # add extra visibility mask to metrics like nerfbusters for fair evaluation
    visibility_path: Path = Path("./data/nerfbusters/aloe/visibility_masks")
    # if nb_mask set to true, specify the path to the masks
    eval_depth: bool = True
    #perform evaluation on depth error
        

    def main(self) -> None:
        """Main function."""
        
        if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
            config, pipeline, checkpoint_path, _ = eval_setup(self.load_config)
        else:
            config, pipeline, checkpoint_path = eval_setup(self.load_config)
        
        # Dynamic change of get_outputs method to include uncertainty
        self.device = pipeline.device
        pipeline.model.filter_out = self.filter_out
        pipeline.model.hessian = torch.tensor(np.load(str(self.unc_path))).to(self.device)
        pipeline.model.lod = np.log2(round(pipeline.model.hessian.shape[0]**(1/3))-1)
        pipeline.model.dataset_path = self.dataset_path
        pipeline.model.output_path = self.output_path
        pipeline.model.white_bg = False
        pipeline.model.black_bg = False
        pipeline.model.N =  4096*1000 #approx ray dataset size (train batch size x number of query iterations in uncertainty extraction step)
        pipeline.model.get_uncertainty = types.MethodType(get_uncertainty, pipeline.model)
        pipeline.expname = config.experiment_name
        pipeline.add_nb_mask = self.nb_mask
        pipeline.nb_mask_path = self.visibility_path
        pipeline.downscale_factor = self.downscale_factor
        pipeline.eval_depth = self.eval_depth
        if not self.filter_out:
            new_method = get_output_fn(pipeline.model)
            pipeline.model.get_outputs = types.MethodType(new_method, pipeline.model)
        else: #filter out metrics only implemented for nerfacto currently!
            pipeline.model.get_outputs = types.MethodType(get_output_nerfacto_all, pipeline.model)

        # Override evaluation function
        if self.filter_out:
            pipeline.get_average_eval_image_metrics = types.MethodType(get_average_filtered_image_metrics, pipeline)
        else:    
            pipeline.get_average_eval_image_metrics = types.MethodType(get_average_uncertainty_metrics, pipeline)

        assert self.output_path.suffix == ".json"
        metrics_dict, metric_lists = pipeline.get_average_eval_image_metrics()
        self.output_path.parent.mkdir(parents=True, exist_ok=True)
        
        
        # Get the output and define the names to save to
        benchmark_info = {
            "experiment_name": config.experiment_name,
            "method_name": config.method_name,
            "checkpoint": str(checkpoint_path),
            "results": metrics_dict,
        }
        # Save output to output file
        self.output_path.write_text(json.dumps(benchmark_info, indent=2), "utf8")
        
            
        # Save output to output file
        nb_filter = '_nb' if self.nb_mask else ''
        timestamp = datetime.now().timestamp()
        date_time = datetime.fromtimestamp(timestamp)
        str_date_time = date_time.strftime("%d-%m-%Y-%H%M%S")
        csv_path = str(self.output_path).split('.')[0] + '_' + config.experiment_name + '_'+ str_date_time + nb_filter +'.csv'
        
        np.savetxt(csv_path, [p for p in zip(*metric_lists)], delimiter=',', fmt='%s')
        CONSOLE.print(f"Saved results to: {self.output_path}")


def entrypoint():
    """Entrypoint for use with pyproject scripts."""
    tyro.extras.set_accent_color("bright_yellow")
    tyro.cli(ComputeMetrics).main()


if __name__ == "__main__":
    entrypoint()

# For sphinx docs


def get_parser_fn(): return tyro.extras.get_parser(ComputeMetrics)  # noqa

bayesrays/scripts/eval_uncertainty.py
BayesRays/BayesRays
import numpy as np
import torch
import nerfstudio
import pkg_resources
import nerfacc
from nerfstudio.models.nerfacto import NerfactoModel
from nerfstudio.models.instant_ngp import NGPModel
from nerfstudio.models.mipnerf import MipNerfModel
from nerfstudio.field_components.field_heads import FieldHeadNames
from bayesrays.utils.utils import normalize_point_coords, find_grid_indices
from nerfstudio.utils import colors
from nerfstudio.model_components.losses import (
    orientation_loss,
    pred_normal_loss
)



def get_uncertainty(self, points):
    aabb = self.scene_box.aabb.to(points.device)
    ## samples outside aabb will have 0 coeff and hence 0 uncertainty. To avoid problems with these samples we set zero_out=False
    inds, coeffs = find_grid_indices(points, aabb, self.field.spatial_distortion ,self.lod, points.device, zero_out=False)
    cfs_2 = (coeffs**2)/torch.sum((coeffs**2),dim=0, keepdim=True)
    uns = self.un[inds.long()] #[8,N]
    un_points = torch.sqrt(torch.sum((uns*cfs_2),dim=0)).unsqueeze(1)
    
    #for stability in volume rendering we use log uncertainty
    un_points = torch.log10(un_points+1e-12)
    un_points = un_points.view((points.shape[0], points.shape[1],1))
    return un_points

def get_output_nerfacto_new(self, ray_bundle):
    ''' reimplementation of get_output function from models because of lack of proper interface to outputs dict'''
#     original_outputs = self.__class__.get_outputs(self, ray_bundle)  # Call original get_outputs (this is slower than just copying the original method here)
    
    N = self.N
    reg_lambda = 1e-4 /( (2**self.lod)**3)
    H = self.hessian/N + reg_lambda
    self.un = 1/H
            
    max_uncertainty = 6 #approximate upper bound of the function log10(1/(x+lambda)) when lambda=1e-4/(256^3) and x is the hessian
    min_uncertainty = -3 #approximate lower bound of that function (cutting off at hessian = 1000)
    density_fns_new = []
    if self.filter_out:
        for i in self.density_fns:
            density_fns_new.append(lambda x, i=i: i(x) * (self.get_uncertainty(x)<= self.filter_thresh*max_uncertainty))
    else:
        density_fns_new = self.density_fns
    
    if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
        ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=density_fns_new)
    else:
        ray_samples,_, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=density_fns_new)
    field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals)
    points = ray_samples.frustums.get_positions()
    un_points = self.get_uncertainty(points)

    #get weights
    if self.filter_out:
        density = field_outputs[FieldHeadNames.DENSITY] * (un_points <= self.filter_thresh*max_uncertainty)
    else:
        density = field_outputs[FieldHeadNames.DENSITY]
    weights = ray_samples.get_weights(density)
    
    uncertainty = torch.sum(weights * un_points, dim=-2) 
    uncertainty += (1-torch.sum(weights,dim=-2)) * min_uncertainty #alpha blending
    
    #normalize into acceptable range for rendering
    uncertainty = torch.clip(uncertainty, min_uncertainty, max_uncertainty)
    uncertainty = (uncertainty-min_uncertainty)/(max_uncertainty-min_uncertainty)
    
    if self.white_bg:
        self.renderer_rgb.background_color=colors.WHITE 
    elif self.black_bg:
        self.renderer_rgb.background_color=colors.BLACK     
    rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
    depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)
    accumulation = self.renderer_accumulation(weights=weights)

    original_outputs = {
        "rgb": rgb,
        "accumulation": accumulation,
        "depth": depth,
    }
                                    
    original_outputs['uncertainty'] = uncertainty 
    if self.training:
        original_outputs["weights_list"] = weights_list
        original_outputs["ray_samples_list"] = ray_samples_list
    if self.config.predict_normals:
        normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
        pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)
        original_outputs["normals"] = self.normals_shader(normals)
        original_outputs["pred_normals"] = self.normals_shader(pred_normals)    

    if self.training and self.config.predict_normals:
        original_outputs["rendered_orientation_loss"] = orientation_loss(
            weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
        )

        original_outputs["rendered_pred_normal_loss"] = pred_normal_loss(
            weights.detach(),
            field_outputs[FieldHeadNames.NORMALS].detach(),
            field_outputs[FieldHeadNames.PRED_NORMALS],
        )

    for i in range(self.config.num_proposal_iterations):
        original_outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])


    return original_outputs

def get_output_ngp_new(self, ray_bundle):
    assert self.field is not None
    assert pkg_resources.get_distribution("nerfstudio").version >= "0.3.1"
    
    N = self.N  #approx ray dataset size
    reg_lambda = 1e-4 /( (2**self.lod)**3)
    H =   reg_lambda + self.hessian/N
    self.un = 1/H 
            
    max_uncertainty = 6 #approximate upper bound of the function log10(1/(x+lambda)) when lambda=1e-4/(256^3) and x is the hessian
    min_uncertainty = -3 #approximate lower bound of that function (cutting off at hessian = 1000)
    if self.filter_out:
        density_fn_new = lambda x: self.field.density_fn(x) * (self.get_uncertainty(x.unsqueeze(0)).squeeze(0)<= self.filter_thresh*max_uncertainty)
        self.sampler.density_fn = density_fn_new    
    num_rays = len(ray_bundle)
    if self.collider is not None:
        ray_bundle = self.collider(ray_bundle)
    with torch.no_grad():
        ray_samples, ray_indices = self.sampler(
            ray_bundle=ray_bundle,
            near_plane=self.config.near_plane,
            far_plane=self.config.far_plane,
            render_step_size=self.config.render_step_size*0.001,
            alpha_thre=self.config.alpha_thre,
            cone_angle=self.config.cone_angle,
        )

    field_outputs = self.field(ray_samples)
    points = ray_samples.frustums.get_positions()
    un_points = self.get_uncertainty(points.unsqueeze(0)).squeeze(0)
    if self.filter_out:
        density = field_outputs[FieldHeadNames.DENSITY] * (un_points <= self.filter_thresh*max_uncertainty)
    else:
        density = field_outputs[FieldHeadNames.DENSITY] 
    # accumulation
    packed_info = nerfacc.pack_info(ray_indices, num_rays)
    weights = nerfacc.render_weight_from_density(
        t_starts=ray_samples.frustums.starts[..., 0],
        t_ends=ray_samples.frustums.ends[..., 0],
        sigmas=density[..., 0],
        packed_info=packed_info,
    )[0]
        
    weights = weights[..., None]
    
    comp_uncertainty = nerfacc.accumulate_along_rays(
                weights[..., 0], values=un_points, ray_indices=ray_indices, n_rays=num_rays
            )
    accumulated_weight = nerfacc.accumulate_along_rays(
        weights[..., 0], values=None, ray_indices=ray_indices, n_rays=num_rays
    )
    uncertainty = comp_uncertainty + min_uncertainty * (1.0 - accumulated_weight)#alpha blending

    #normalize into acceptable range for rendering
    uncertainty = torch.clip(uncertainty, min_uncertainty, max_uncertainty)
    uncertainty = (uncertainty-min_uncertainty)/(max_uncertainty-min_uncertainty)
    
    if self.white_bg:
        self.renderer_rgb.background_color=colors.WHITE    
    elif self.black_bg:
        self.renderer_rgb.background_color=colors.BLACK  

    rgb = self.renderer_rgb(
        rgb=field_outputs[FieldHeadNames.RGB],
        weights=weights,
        ray_indices=ray_indices,
        num_rays=num_rays,
    )
    depth = self.renderer_depth(
        weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=num_rays
    )
    accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=num_rays)
    
    outputs = {
        "rgb": rgb,
        "accumulation": accumulation,
        "depth": depth,
        "uncertainty": uncertainty
    }
    return outputs

def get_output_mipnerf_new(self, ray_bundle):
    ''' reimplementation of get_output function from models because of lack of proper interface to outputs dict'''
    
    N = self.N  #approx ray dataset size
    reg_lambda = 1e-4 /( (2**self.lod)**3)
    H = self.hessian/N + reg_lambda
    self.un = 1/H
            
    max_uncertainty = 6 #approximate upper bound of the function log10(1/(x+lambda)) when lambda=1e-4/(256^3) and x is the hessian
    min_uncertainty = -3 #approximate lower bound of that function (cutting off at hessian = 1000)
    
    # uniform sampling
    ray_samples = self.sampler_uniform(ray_bundle)
        
    # First pass:
    field_outputs_coarse = self.field.forward(ray_samples)
    weights_coarse = ray_samples.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY])
    rgb_coarse = self.renderer_rgb(
        rgb=field_outputs_coarse[FieldHeadNames.RGB],
        weights=weights_coarse,
    )
    accumulation_coarse = self.renderer_accumulation(weights_coarse)
    depth_coarse = self.renderer_depth(weights_coarse, ray_samples)

    # pdf sampling
    ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples, weights_coarse)

    # Second pass:
    field_outputs_fine = self.field.forward(ray_samples_pdf)
        
    points = ray_samples_pdf.frustums.get_positions()
    un_points = self.get_uncertainty(points)
    
    if self.filter_out:
        density = field_outputs_fine[FieldHeadNames.DENSITY] * (un_points <= self.filter_thresh*max_uncertainty)
    else:
        density = field_outputs_fine[FieldHeadNames.DENSITY] 
    #get weights
    weights_fine = ray_samples_pdf.get_weights(density)
    rgb_fine = self.renderer_rgb(
        rgb=field_outputs_fine[FieldHeadNames.RGB],
        weights=weights_fine,
    )
    accumulation_fine = self.renderer_accumulation(weights_fine)
    depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf)
    
    
    uncertainty = torch.sum(weights_fine * un_points, dim=-2) 
    uncertainty += (1-torch.sum(weights_fine,dim=-2)) * min_uncertainty #alpha blending
    
    #normalize into acceptable range for rendering
    uncertainty = torch.clip(uncertainty, min_uncertainty, max_uncertainty)
    uncertainty = (uncertainty-min_uncertainty)/(max_uncertainty-min_uncertainty)
    
    if self.white_bg:
        self.renderer_rgb.background_color=colors.WHITE    
    elif self.black_bg:
        self.renderer_rgb.background_color=colors.BLACK        
   
    original_outputs = {
            "rgb_coarse": rgb_coarse,
            "rgb_fine": rgb_fine,
            "accumulation_coarse": accumulation_coarse,
            "accumulation_fine": accumulation_fine,
            "depth_coarse": depth_coarse,
            "depth_fine": depth_fine,
            "depth": depth_fine,
            "rgb": rgb_fine,
            "accumulation": accumulation_fine
        }
    original_outputs['uncertainty'] = uncertainty 

    return original_outputs

def get_output_fn(model):

    if isinstance(model, NerfactoModel):
        return get_output_nerfacto_new
    elif isinstance(model, NGPModel):
        return get_output_ngp_new
    elif isinstance(model, MipNerfModel):
        return get_output_mipnerf_new
    else:
        raise Exception("Sorry, this model is not currently supported.")

def get_output_nerfacto_all(self, ray_bundle):
    ''' reimplementation of get_output function from models for evaluation with different filter levels'''
    
    N = self.N  #approx ray dataset size
    reg_lambda = 1e-4 /( (2**self.lod)**3)
    H = self.hessian/N + reg_lambda
    self.un = 1/H
            
    max_uncertainty = 6 #approximate upper bound of the function log10(1/(x+lambda)) when lambda=1e-4/(64^3) and x is the hessian
    min_uncertainty = -3 #approximate lower bound of that function (cutting off at hessian = 1000)
    original_outputs={}
    for thresh in self.thresh_range:
        density_fns_new = []
        for i in self.density_fns:
            density_fns_new.append(lambda x, i=i: i(x) * (self.get_uncertainty(x)<= thresh*max_uncertainty))

        if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
            ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=density_fns_new)
        else:
            ray_samples,_, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=density_fns_new)
        field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals)
        points = ray_samples.frustums.get_positions()

        #get weights
        uns = self.get_uncertainty(ray_samples.frustums.get_positions())
        density = field_outputs[FieldHeadNames.DENSITY] * (uns <= thresh*max_uncertainty)
        weights = ray_samples.get_weights(density)

        rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
        depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)
        accumulation = self.renderer_accumulation(weights=weights)
        original_outputs["rgb-"+"{:.2f}".format(thresh.item())] = rgb
        original_outputs["accumulation-"+"{:.2f}".format(thresh.item())] = accumulation     
        original_outputs["depth-"+"{:.2f}".format(thresh.item())] = depth     

    return original_outputs


bayesrays/scripts/output_uncertainty.py
BayesRays/BayesRays
#This file currently does not support old versions of nerfstudio (inconsistency in Colormap class)
#!/usr/bin/env python
"""
render_uncertainty.py
"""
from __future__ import annotations

import json
import os
import struct
import shutil
import sys
import types
from contextlib import ExitStack
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Literal, Optional, Union
import mediapy as media
import numpy as np
import torch
import tyro
import nerfstudio
import pkg_resources

from jaxtyping import Float
from rich import box, style
from rich.panel import Panel
from rich.progress import (
    BarColumn,
    Progress,
    TaskProgressColumn,
    TextColumn,
    TimeElapsedColumn,
    TimeRemainingColumn,
)
from rich.table import Table
from torch import Tensor
from typing_extensions import Annotated
import pickle

from nerfstudio.cameras.camera_paths import (
    get_interpolated_camera_path,
    get_path_from_json,
    get_spiral_path,
)
from nerfstudio.data.utils.dataloaders import FixedIndicesEvalDataloader
from nerfstudio.cameras.cameras import Cameras, CameraType
from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager
from nerfstudio.data.scene_box import SceneBox
from nerfstudio.model_components import renderers
from nerfstudio.pipelines.base_pipeline import Pipeline
from nerfstudio.utils import colormaps, install_checks, colors
from nerfstudio.utils.eval_utils import eval_setup
from nerfstudio.utils.rich_utils import CONSOLE, ItersPerSecColumn
from nerfstudio.utils.scripts import run_command
from bayesrays.scripts.output_uncertainty import get_output_fn, get_uncertainty

if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
    from nerfstudio.utils.colormaps import ColormapOptions

#basically the same script as render.py but first adding uncertainty output to model.get_output output dict and then rendering

def _render_trajectory_video(
    pipeline: Pipeline,
    cameras: Cameras,
    output_filename: Path,
    rendered_output_names: List[str],
    crop_data: Optional[CropData] = None,
    rendered_resolution_scaling_factor: float = 1.0,
    seconds: float = 5.0,
    output_format: Literal["images", "video"] = "video",
    image_format: Literal["jpeg", "png"] = "jpeg",
    jpeg_quality: int = 100,
    depth_near_plane: Optional[float] = None,
    depth_far_plane: Optional[float] = None,
    gt_visibility_path=None,
    colormap_options=None
) -> None:
    """Helper function to create a video of the spiral trajectory.

    Args:
        pipeline: Pipeline to evaluate with.
        cameras: Cameras to render.
        output_filename: Name of the output file.
        rendered_output_names: List of outputs to visualise.
        crop_data: Crop data to apply to the rendered images.
        rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution.
        seconds: Length of output video.
        output_format: How to save output data.
        colormap_options: Options for colormap.
    """
    CONSOLE.print("[bold green]Creating trajectory " + output_format)
    cameras.rescale_output_resolution(rendered_resolution_scaling_factor)
    cameras = cameras.to(pipeline.device)
    fps = len(cameras) / seconds

    progress = Progress(
        TextColumn(":movie_camera: Rendering :movie_camera:"),
        BarColumn(),
        TaskProgressColumn(
            text_format="[progress.percentage]{task.completed}/{task.total:>.0f}({task.percentage:>3.1f}%)",
            show_speed=True,
        ),
        ItersPerSecColumn(suffix="fps"),
        TimeRemainingColumn(elapsed_when_finished=False, compact=False),
        TimeElapsedColumn(),
    )
    output_image_dir = output_filename.parent / output_filename.stem
    if output_format == "images":
        output_image_dir.mkdir(parents=True, exist_ok=True)
    if output_format == "video":
        # make the folder if it doesn't exist
        output_filename.parent.mkdir(parents=True, exist_ok=True)
        
    with ExitStack() as stack:
        writer = None

        with progress:
            for camera_idx in progress.track(range(cameras.size), description=""):
                if gt_visibility_path:
                    pseudo_gt_visibility = media.read_image(str(gt_visibility_path)+"/{:05d}.png".format(camera_idx))
                    pseudo_gt_visibility = (pseudo_gt_visibility[..., 0] >= 1).astype('float')[...,None]
                else:
                    pseudo_gt_visibility = 1

                aabb_box = None
                if crop_data is not None:
                    bounding_box_min = crop_data.center - crop_data.scale / 2.0
                    bounding_box_max = crop_data.center + crop_data.scale / 2.0
                    aabb_box = SceneBox(torch.stack([bounding_box_min, bounding_box_max]).to(pipeline.device))
                camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx, aabb_box=aabb_box)

                if crop_data is not None:
                    with renderers.background_color_override_context(
                        crop_data.background_color.to(pipeline.device)
                    ), torch.no_grad():
                        outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
                else:
                    with torch.no_grad():
                        outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)

                render_image = []
                for rendered_output_name in rendered_output_names:
                    if rendered_output_name not in outputs:
                        CONSOLE.rule("Error", style="red")
                        CONSOLE.print(f"Could not find {rendered_output_name} in the model outputs", justify="center")
                        CONSOLE.print(
                            f"Please set --rendered_output_name to one of: {outputs.keys()}", justify="center"
                        )
                        sys.exit(1)
                    output_image = outputs[rendered_output_name]
                    if rendered_output_name == 'uncertainty':

                        if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
                            colormap_options1 = ColormapOptions(
                                colormap='inferno',
                                normalize=colormap_options.normalize,
                                colormap_min=colormap_options.colormap_min,
                                colormap_max=colormap_options.colormap_max,
                                invert=colormap_options.invert,
                            )
                            output_image = colormaps.apply_colormap(
                                    image=output_image,
                                    colormap_options=colormap_options1,
                                ).cpu().numpy()*pseudo_gt_visibility
                        else:
                            output_image = colormaps.apply_colormap(output_image, cmap="inferno").cpu().numpy()*pseudo_gt_visibility
                    elif rendered_output_name == 'rgb':
                        output_image = (output_image).cpu().numpy() * pseudo_gt_visibility
                    else:
                        output_image = colormaps.apply_depth_colormap(
                                outputs[rendered_output_name],
                                near_plane=depth_near_plane,
                                far_plane=depth_far_plane,
                            ).cpu().numpy()*pseudo_gt_visibility


                    render_image.append(output_image)
                render_image = np.concatenate(render_image, axis=1)
                if output_format == "images":
                    if image_format == "png":
                        media.write_image(output_image_dir / f"{camera_idx:05d}.png", render_image, fmt="png")
                    if image_format == "jpeg":
                        media.write_image(
                            output_image_dir / f"{camera_idx:05d}.jpg", render_image, fmt="jpeg", quality=jpeg_quality
                        )
                if output_format == "video":
                    if writer is None:
                        render_width = int(render_image.shape[1])
                        render_height = int(render_image.shape[0])
                        writer = stack.enter_context(
                            media.VideoWriter(
                                path=output_filename,
                                shape=(render_height, render_width),
                                fps=fps,
                            )
                        )
                    writer.add_image(render_image)
    
    table = Table(
        title=None,
        show_header=False,
        box=box.MINIMAL,
        title_style=style.Style(bold=True),
    )
    if output_format == "video":
        if cameras.camera_type[0] == CameraType.EQUIRECTANGULAR.value:
            CONSOLE.print("Adding spherical camera data")
            insert_spherical_metadata_into_file(output_filename)
        table.add_row("Video", str(output_filename))
    else:
        table.add_row("Images", str(output_image_dir))
    CONSOLE.print(Panel(table, title="[bold][green]:tada: Render Complete :tada:[/bold]", expand=False))


def insert_spherical_metadata_into_file(
    output_filename: Path,
) -> None:
    """Inserts spherical metadata into MP4 video file in-place.
    Args:
        output_filename: Name of the (input and) output file.
    """
    # NOTE:
    # because we didn't use faststart, the moov atom will be at the end;
    # to insert our metadata, we need to find (skip atoms until we get to) moov.
    # we should have 0x00000020 ftyp, then 0x00000008 free, then variable mdat.
    spherical_uuid = b"\xff\xcc\x82\x63\xf8\x55\x4a\x93\x88\x14\x58\x7a\x02\x52\x1f\xdd"
    spherical_metadata = bytes(
        """<rdf:SphericalVideo
xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'
xmlns:GSpherical='http://ns.google.com/videos/1.0/spherical/'>
<GSpherical:ProjectionType>equirectangular</GSpherical:ProjectionType>
<GSpherical:Spherical>True</GSpherical:Spherical>
<GSpherical:Stitched>True</GSpherical:Stitched>
<GSpherical:StitchingSoftware>nerfstudio</GSpherical:StitchingSoftware>
</rdf:SphericalVideo>""",
        "utf-8",
    )
    insert_size = len(spherical_metadata) + 8 + 16
    with open(output_filename, mode="r+b") as mp4file:
        try:
            # get file size
            mp4file_size = os.stat(output_filename).st_size

            # find moov container (probably after ftyp, free, mdat)
            while True:
                pos = mp4file.tell()
                size, tag = struct.unpack(">I4s", mp4file.read(8))
                if tag == b"moov":
                    break
                mp4file.seek(pos + size)
            # if moov isn't at end, bail
            if pos + size != mp4file_size:
                # TODO: to support faststart, rewrite all stco offsets
                raise Exception("moov container not at end of file")
            # go back and write inserted size
            mp4file.seek(pos)
            mp4file.write(struct.pack(">I", size + insert_size))
            # go inside moov
            mp4file.seek(pos + 8)
            # find trak container (probably after mvhd)
            while True:
                pos = mp4file.tell()
                size, tag = struct.unpack(">I4s", mp4file.read(8))
                if tag == b"trak":
                    break
                mp4file.seek(pos + size)
            # go back and write inserted size
            mp4file.seek(pos)
            mp4file.write(struct.pack(">I", size + insert_size))
            # we need to read everything from end of trak to end of file in order to insert
            # TODO: to support faststart, make more efficient (may load nearly all data)
            mp4file.seek(pos + size)
            rest_of_file = mp4file.read(mp4file_size - pos - size)
            # go to end of trak (again)
            mp4file.seek(pos + size)
            # insert our uuid atom with spherical metadata
            mp4file.write(struct.pack(">I4s16s", insert_size, b"uuid", spherical_uuid))
            mp4file.write(spherical_metadata)
            # write rest of file
            mp4file.write(rest_of_file)
        finally:
            mp4file.close()


@dataclass
class CropData:
    """Data for cropping an image."""

    background_color: Float[Tensor, "3"] = torch.Tensor([0.0, 0.0, 0.0])
    """background color"""
    center: Float[Tensor, "3"] = torch.Tensor([0.0, 0.0, 0.0])
    """center of the crop"""
    scale: Float[Tensor, "3"] = torch.Tensor([2.0, 2.0, 2.0])
    """scale of the crop"""


def get_crop_from_json(camera_json: Dict[str, Any]) -> Optional[CropData]:
    """Load crop data from a camera path JSON

    args:
        camera_json: camera path data
    returns:
        Crop data
    """
    if "crop" not in camera_json or camera_json["crop"] is None:
        return None

    bg_color = camera_json["crop"]["crop_bg_color"]

    return CropData(
        background_color=torch.Tensor([bg_color["r"] / 255.0, bg_color["g"] / 255.0, bg_color["b"] / 255.0]),
        center=torch.Tensor(camera_json["crop"]["crop_center"]),
        scale=torch.Tensor(camera_json["crop"]["crop_scale"]),
    )


@dataclass
class BaseRender:
    """Base class for rendering."""

    load_config: Path
    """Path to config YAML file."""
    output_path: Path = Path("renders/output.mp4")
    """Path to output video file."""
    unc_path: Path = Path("unc.npy")
    """Path to output video file."""
    visibility_path: Path = None
    """Path to nerfbuster visibility masks, if needed to be rendered."""
    image_format: Literal["jpeg", "png"] = "jpeg"
    """Image format"""
    jpeg_quality: int = 100
    """JPEG quality"""
    downscale_factor: float = 1.0
    """Scaling factor to apply to the camera image resolution."""
    eval_num_rays_per_chunk: Optional[int] = None
    """Specifies number of rays per chunk during eval. If None, use the value in the config file."""
    if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
        colormap_options: colormaps.ColormapOptions = colormaps.ColormapOptions()
    else:
        colormap_options = None
    """Colormap options."""
    depth_near_plane: Optional[float] = None
    """Closest depth to consider when using the colormap for depth. If None, use min value."""
    depth_far_plane: Optional[float] = None
    """Furthest depth to consider when using the colormap for depth. If None, use max value."""
    filter_out: bool = False
    """Render with filtering"""
    filter_thresh: float = 0.5
    """filter threshold"""
    white_bg: bool = False
    """ Render empty space as white when filtering"""
    black_bg: bool = False
    """ Render empty space as black when filtering"""


@dataclass
class RenderCameraPath(BaseRender):
    """Render a camera path generated by the viewer or blender add-on."""

    rendered_output_names: List[str] = field(default_factory=lambda: ["rgb"])
    """Name of the renderer outputs to use. rgb, depth, etc. concatenates them along y axis"""
    camera_path_filename: Path = Path("camera_path.json")
    """Filename of the camera path to render."""
    output_format: Literal["images", "video"] = "video"
    """How to save output data."""
    

    def main(self) -> None:
        """Main function."""
        if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
            _, pipeline, _, _ = eval_setup(
                self.load_config,
                eval_num_rays_per_chunk=self.eval_num_rays_per_chunk,
                test_mode="inference",
            )
        else:
            _, pipeline, _ = eval_setup(
                self.load_config,
                eval_num_rays_per_chunk=self.eval_num_rays_per_chunk,
                test_mode="inference",
            )
        #dynamically change get_outputs method to include uncertainty
        self.device = pipeline.device
        pipeline.model.filter_out = self.filter_out
        pipeline.model.filter_thresh = self.filter_thresh
        pipeline.model.hessian = torch.tensor(np.load(str(self.unc_path))).to(self.device)
        pipeline.model.lod = np.log2(round(pipeline.model.hessian.shape[0]**(1/3))-1)
        pipeline.model.get_uncertainty = types.MethodType(get_uncertainty, pipeline.model)
        pipeline.model.white_bg = self.white_bg
        pipeline.model.black_bg = self.black_bg
        pipeline.model.N =  4096*1000  #approx ray dataset size (train batch size x number of query iterations in uncertainty extraction step)
        new_method = get_output_fn(pipeline.model)
        pipeline.model.get_outputs = types.MethodType(new_method, pipeline.model)

        install_checks.check_ffmpeg_installed()

        with open(self.camera_path_filename, "r", encoding="utf-8") as f:
            camera_path = json.load(f)
        seconds = camera_path["seconds"]
        crop_data = get_crop_from_json(camera_path)
        camera_path = get_path_from_json(camera_path)

        if camera_path.camera_type[0] == CameraType.OMNIDIRECTIONALSTEREO_L.value:
            # temp folder for writing left and right view renders
            temp_folder_path = self.output_path.parent / (self.output_path.stem + "_temp")

            Path(temp_folder_path).mkdir(parents=True, exist_ok=True)
            left_eye_path = temp_folder_path / "ods_render_Left.mp4"

            self.output_path = left_eye_path

            CONSOLE.print("[bold green]:goggles: Omni-directional Stereo VR :goggles:")
            CONSOLE.print("Rendering left eye view")

        # add mp4 suffix to video output if none is specified
        if self.output_format == "video" and str(self.output_path.suffix) == "":
            self.output_path = self.output_path.with_suffix(".mp4")

        _render_trajectory_video(
            pipeline,
            camera_path,
            output_filename=self.output_path,
            rendered_output_names=self.rendered_output_names,
            rendered_resolution_scaling_factor=1.0 / self.downscale_factor,
            crop_data=crop_data,
            seconds=seconds,
            output_format=self.output_format,
            image_format=self.image_format,
            jpeg_quality=self.jpeg_quality,
            depth_near_plane=self.depth_near_plane,
            depth_far_plane=self.depth_far_plane,
            colormap_options=self.colormap_options
        )

        if camera_path.camera_type[0] == CameraType.OMNIDIRECTIONALSTEREO_L.value:
            # declare paths for left and right renders

            left_eye_path = self.output_path
            right_eye_path = left_eye_path.parent / "ods_render_Right.mp4"

            self.output_path = right_eye_path
            camera_path.camera_type[0] = CameraType.OMNIDIRECTIONALSTEREO_R.value

            CONSOLE.print("Rendering right eye view")
            _render_trajectory_video(
                pipeline,
                camera_path,
                output_filename=self.output_path,
                rendered_output_names=self.rendered_output_names,
                rendered_resolution_scaling_factor=1.0 / self.downscale_factor,
                crop_data=crop_data,
                seconds=seconds,
                output_format=self.output_format,
                image_format=self.image_format,
                jpeg_quality=self.jpeg_quality,
                depth_near_plane=self.depth_near_plane,
                depth_far_plane=self.depth_far_plane,
                gt_visibility_path=self.visibility_path,
                colormap_options=self.colormap_options
            )

            # stack the left and right eye renders for final output
            self.output_path = Path(str(left_eye_path.parent)[:-5] + ".mp4")
            ffmpeg_ods_command = ""
            if self.output_format == "video":
                ffmpeg_ods_command = f'ffmpeg -y -i "{left_eye_path}" -i "{right_eye_path}" -filter_complex "[0:v]pad=iw:2*ih[int];[int][1:v]overlay=0:h" -c:v libx264 -crf 23 -preset veryfast "{self.output_path}"'
                run_command(ffmpeg_ods_command, verbose=False)
            if self.output_format == "images":
                # create a folder for the stacked renders
                self.output_path = Path(str(left_eye_path.parent)[:-5])
                self.output_path.mkdir(parents=True, exist_ok=True)
                if self.image_format == "png":
                    ffmpeg_ods_command = f'ffmpeg -y -pattern_type glob -i "{str(left_eye_path.with_suffix("") / "*.png")}"  -pattern_type glob -i "{str(right_eye_path.with_suffix("") / "*.png")}" -filter_complex vstack -start_number 0 "{str(self.output_path)+"//%05d.png"}"'
                elif self.image_format == "jpeg":
                    ffmpeg_ods_command = f'ffmpeg -y -pattern_type glob -i "{str(left_eye_path.with_suffix("") / "*.jpg")}"  -pattern_type glob -i "{str(right_eye_path.with_suffix("") / "*.jpg")}" -filter_complex vstack -start_number 0 "{str(self.output_path)+"//%05d.jpg"}"'
                run_command(ffmpeg_ods_command, verbose=False)

            # remove the temp files directory
            if str(left_eye_path.parent)[-5:] == "_temp":
                shutil.rmtree(left_eye_path.parent, ignore_errors=True)
            CONSOLE.print("[bold green]Final ODS Render Complete")


@dataclass
class RenderInterpolated(BaseRender):
    """Render a trajectory that interpolates between training or eval dataset images."""

    rendered_output_names: List[str] = field(default_factory=lambda: ["rgb"])
    """Name of the renderer outputs to use. rgb, depth, etc. concatenates them along y axis"""
    pose_source: Literal["eval", "train"] = "eval"
    """Pose source to render."""
    interpolation_steps: int = 10
    """Number of interpolation steps between eval dataset cameras."""
    order_poses: bool = False
    """Whether to order camera poses by proximity."""
    frame_rate: int = 24
    """Frame rate of the output video."""
    output_format: Literal["images", "video"] = "video"
    """How to save output data."""
    

    def main(self) -> None:
        """Main function."""
        if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
            _, pipeline, _, _ = eval_setup(
                self.load_config,
                eval_num_rays_per_chunk=self.eval_num_rays_per_chunk,
                test_mode="test",
            )
        else:
             _, pipeline, _ = eval_setup(
                self.load_config,
                eval_num_rays_per_chunk=self.eval_num_rays_per_chunk,
                test_mode="test",
            )

        #dynamically change get_outputs method to include uncertainty
        self.device = pipeline.device
        pipeline.model.filter_out = self.filter_out
        pipeline.model.filter_thresh = self.filter_thresh
        pipeline.model.hessian = torch.tensor(np.load(str(self.unc_path))).to(self.device)
        pipeline.model.lod = np.log2(round(pipeline.model.hessian.shape[0]**(1/3))-1)
        pipeline.model.get_uncertainty = types.MethodType(get_uncertainty, pipeline.model)
        pipeline.model.white_bg = self.white_bg
        pipeline.model.black_bg = self.black_bg
        pipeline.model.N =  4096*1000  #approx ray dataset size (train batch size x number of query iterations in uncertainty extraction step)
        new_method = get_output_fn(pipeline.model)
        pipeline.model.get_outputs = types.MethodType(new_method, pipeline.model)


        install_checks.check_ffmpeg_installed()

        if self.pose_source == "eval":
            assert pipeline.datamanager.eval_dataset is not None
            cameras = pipeline.datamanager.eval_dataset.cameras
        else:
            assert pipeline.datamanager.train_dataset is not None
            cameras = pipeline.datamanager.train_dataset.cameras

        seconds = self.interpolation_steps * len(cameras) / self.frame_rate

        if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
            camera_path = get_interpolated_camera_path(
                cameras=cameras,
                steps=self.interpolation_steps,
                order_poses=self.order_poses,
            )
        else:
             camera_path = get_interpolated_camera_path(
                cameras=cameras,
                steps=self.interpolation_steps
            )
        Ks = cameras.get_intrinsics_matrices()
        if self.interpolation_steps == 1: #include the last camera
            camera_path = pipeline.datamanager.fixed_indices_eval_dataloader.cameras


        _render_trajectory_video(
            pipeline,
            camera_path,
            output_filename=self.output_path,
            rendered_output_names=self.rendered_output_names,
            rendered_resolution_scaling_factor=1.0 / self.downscale_factor,
            seconds=seconds,
            output_format=self.output_format,
            depth_near_plane=self.depth_near_plane,
            depth_far_plane=self.depth_far_plane,
            gt_visibility_path=self.visibility_path,
            colormap_options=self.colormap_options
        )



@dataclass
class SpiralRender(BaseRender):
    """Render a spiral trajectory (often not great)."""

    rendered_output_names: List[str] = field(default_factory=lambda: ["rgb"])
    """Name of the renderer outputs to use. rgb, depth, etc. concatenates them along y axis"""
    seconds: float = 3.0
    """How long the video should be."""
    output_format: Literal["images", "video"] = "video"
    """How to save output data."""
    frame_rate: int = 24
    """Frame rate of the output video (only for interpolate trajectory)."""
    radius: float = 0.1
    """Radius of the spiral."""

    def main(self) -> None:
        """Main function."""
        if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
            _, pipeline, _, _ = eval_setup(
                self.load_config,
                eval_num_rays_per_chunk=self.eval_num_rays_per_chunk,
                test_mode="test",
            )
        else:
             _, pipeline, _ = eval_setup(
                self.load_config,
                eval_num_rays_per_chunk=self.eval_num_rays_per_chunk,
                test_mode="test",
            )

        #dynamically change get_outputs method to include uncertainty
        self.device = pipeline.device
        pipeline.model.filter_out = self.filter_out
        pipeline.model.filter_thresh = self.filter_thresh
        pipeline.model.hessian = torch.tensor(np.load(str(self.unc_path))).to(self.device)
        pipeline.model.lod = np.log2(round(pipeline.model.hessian.shape[0]**(1/3))-1)
        pipeline.model.get_uncertainty = types.MethodType(get_uncertainty, pipeline.model)
        pipeline.model.white_bg = self.white_bg
        pipeline.model.black_bg = self.black_bg
        pipeline.model.N =  4096*1000  #approx ray dataset size (train batch size x number of query iterations in uncertainty extraction step)
        new_method = get_output_fn(pipeline.model)
        pipeline.model.get_outputs = types.MethodType(new_method, pipeline.model)

        install_checks.check_ffmpeg_installed()

        assert isinstance(pipeline.datamanager, VanillaDataManager)
        steps = int(self.frame_rate * self.seconds)
        camera_start = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0).flatten()
        camera_path = get_spiral_path(camera_start, steps=steps, radius=self.radius)

        _render_trajectory_video(
            pipeline,
            camera_path,
            output_filename=self.output_path,
            rendered_output_names=self.rendered_output_names,
            rendered_resolution_scaling_factor=1.0 / self.downscale_factor,
            seconds=self.seconds,
            output_format=self.output_format,
            depth_near_plane=self.depth_near_plane,
            depth_far_plane=self.depth_far_plane,
            gt_visibility_path=self.visibility_path,
            colormap_options=self.colormap_options
        )


Commands = tyro.conf.FlagConversionOff[
    Union[
        Annotated[RenderCameraPath, tyro.conf.subcommand(name="camera-path")],
        Annotated[RenderInterpolated, tyro.conf.subcommand(name="interpolate")],
        Annotated[SpiralRender, tyro.conf.subcommand(name="spiral")],
    ]
]


def entrypoint():
    """Entrypoint for use with pyproject scripts."""
    tyro.extras.set_accent_color("bright_yellow")
    tyro.cli(Commands).main()


if __name__ == "__main__":
    entrypoint()


def get_parser_fn():
    """Get the parser function for the sphinx docs."""
    return tyro.extras.get_parser(Commands)  # noqa

bayesrays/scripts/render_uncertainty.py
BayesRays/BayesRays

#!/usr/bin/env python
"""
uncertainty.py
"""
from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
import types
import time

import tyro
import torch
import numpy as np
import nerfstudio
import nerfacc
import pkg_resources
from nerfstudio.utils.eval_utils import eval_setup
from nerfstudio.utils.rich_utils import CONSOLE
from nerfstudio.models.nerfacto import NerfactoModel
from nerfstudio.models.instant_ngp import NGPModel
from nerfstudio.models.mipnerf import MipNerfModel
from nerfstudio.field_components.field_heads import FieldHeadNames
from nerfstudio.field_components.encodings import HashEncoding
from bayesrays.utils.utils import normalize_point_coords, find_grid_indices, get_gaussian_blob_new


@dataclass
class ComputeUncertainty:
    """Load a checkpoint, compute uncertainty, and save it to a npy file."""

    # Path to config YAML file.
    load_config: Path
    # Name of the output file.
    output_path: Path = Path("unc.npy")
    # Uncertainty level of detail (log2 of it)
    lod: int = 8
    # number of iterations on the trainset    
    iters: int = 1000         
    
    def find_uncertainty(self, points, deform_points, rgb, distortion):
        inds, coeffs = find_grid_indices(points, self.aabb, distortion, self.lod, self.device)
        #because deformation params are detached for each point on each ray from the grid, summation does not affect derivative
        colors = torch.sum(rgb, dim=0)
        colors[0].backward(retain_graph=True)
        r = deform_points.grad.clone().detach().view(-1,3)
        deform_points.grad.zero_()
        colors[1].backward(retain_graph=True)
        g = deform_points.grad.clone().detach().view(-1,3)
        deform_points.grad.zero_()
        colors[2].backward()
        b = deform_points.grad.clone().detach().view(-1,3)
        deform_points.grad.zero_()
        dmy = (torch.arange(points.shape[0])[...,None]).repeat((1,points.shape[1])).flatten().to(self.device)
        first = True
        for corner in range(8):
            if first:
                all_ind = torch.cat((dmy.unsqueeze(-1),inds[corner].unsqueeze(-1)), dim=-1) 
                all_r = coeffs[corner].unsqueeze(-1)*r
                all_g = coeffs[corner].unsqueeze(-1)*g
                all_b = coeffs[corner].unsqueeze(-1)*b
                first = False
            else:
                all_ind = torch.cat((all_ind, torch.cat((dmy.unsqueeze(-1),inds[corner].unsqueeze(-1)), dim=-1)), dim=0)
                all_r = torch.cat((all_r, coeffs[corner].unsqueeze(-1)*r), dim=0)
                all_g = torch.cat((all_g, coeffs[corner].unsqueeze(-1)*g), dim=0)
                all_b = torch.cat((all_b, coeffs[corner].unsqueeze(-1)*b ), dim=0)
        keys_all, inds_all = torch.unique(all_ind, dim=0, return_inverse=True)
        grad_r_1 = torch.bincount(inds_all, weights=all_r[...,0]) #for first element of deformation field
        grad_g_1 = torch.bincount(inds_all, weights=all_g[...,0])
        grad_b_1 = torch.bincount(inds_all, weights=all_b[...,0])
        grad_r_2 = torch.bincount(inds_all, weights=all_r[...,1]) #for second element of deformation field
        grad_g_2 = torch.bincount(inds_all, weights=all_g[...,1])
        grad_b_2 = torch.bincount(inds_all, weights=all_b[...,1])
        grad_r_3 = torch.bincount(inds_all, weights=all_r[...,2]) #for third element of deformation field
        grad_g_3 = torch.bincount(inds_all, weights=all_g[...,2])
        grad_b_3 = torch.bincount(inds_all, weights=all_b[...,2])
        grad_1 = grad_r_1**2+grad_g_1**2+grad_b_1**2
        grad_2 = grad_r_2**2+grad_g_2**2+grad_b_2**2
        grad_3 = grad_r_3**2+grad_g_3**2+grad_b_3**2 #will consider the trace of each submatrix for each deformation
        #vector as indicator of hessian wrt the whole vector
        
        grads_all = torch.cat((keys_all[:,1].unsqueeze(-1), (grad_1+grad_2+grad_3).unsqueeze(-1)), dim=-1)
        hessian = torch.zeros(((2**self.lod)+1)**3).to(self.device)
        hessian = hessian.put((grads_all[:,0]).long(), grads_all[:,1], True)

        return hessian
    
    
    def get_unc_nerfacto(self, ray_bundle, model):
        ''' reimplementation of get_output function from models because of lack of proper interface to ray_samples'''
        if model.collider is not None:
            ray_bundle = model.collider(ray_bundle)
        
        if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":  
            ray_samples, weights_list, ray_samples_list = model.proposal_sampler(ray_bundle, density_fns=model.density_fns)
        else:
            ray_samples,_, weights_list, ray_samples_list = model.proposal_sampler(ray_bundle, density_fns=model.density_fns)
        points = ray_samples.frustums.get_positions()
        pos, _ = normalize_point_coords(points, self.aabb, model.field.spatial_distortion)
        #find offset value in the deformation field
        offsets = self.deform_field(pos).clone().detach()
        offsets.requires_grad = True
        
        ray_samples.frustums.set_offsets(offsets)    
        
        
        field_outputs = model.field(ray_samples, compute_normals=model.config.predict_normals)
        weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])
        weights_list.append(weights)
        ray_samples_list.append(ray_samples)

        rgb = model.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
        depth = model.renderer_depth(weights=weights, ray_samples=ray_samples)
        accumulation = model.renderer_accumulation(weights=weights)

        outputs = {
            "rgb": rgb,
            "accumulation": accumulation,
            "depth": depth,
        }

        if model.config.predict_normals:
            normals = model.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
            pred_normals = model.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)
            outputs["normals"] = model.normals_shader(normals)
            outputs["pred_normals"] = model.normals_shader(pred_normals)
        # These use a lot of GPU memory, so we avoid storing them for eval.
        if model.training:
            outputs["weights_list"] = weights_list
            outputs["ray_samples_list"] = ray_samples_list

        if model.training and model.config.predict_normals:
            outputs["rendered_orientation_loss"] = orientation_loss(
                weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
            )

            outputs["rendered_pred_normal_loss"] = pred_normal_loss(
                weights.detach(),
                field_outputs[FieldHeadNames.NORMALS].detach(),
                field_outputs[FieldHeadNames.PRED_NORMALS],
            )

        for i in range(model.config.num_proposal_iterations):
            outputs[f"prop_depth_{i}"] = model.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])
        return outputs, points, offsets
    
    
    def get_unc_ngp(self, ray_bundle, model):
        assert model.field is not None
        assert pkg_resources.get_distribution("nerfstudio").version >= "0.3.1"
        
        num_rays = len(ray_bundle)
        if model.collider is not None:
            ray_bundle = model.collider(ray_bundle)
        with torch.no_grad():
            ray_samples, ray_indices = model.sampler(
                ray_bundle=ray_bundle,
                near_plane=model.config.near_plane,
                far_plane=model.config.far_plane,
                render_step_size=model.config.render_step_size*0.001,
                alpha_thre=model.config.alpha_thre,
                cone_angle=model.config.cone_angle,
            )
        points = ray_samples.frustums.get_positions()
        pos, _ = normalize_point_coords(points, self.aabb, model.field.spatial_distortion)
        #find offset value in the deformation field
        offsets = self.deform_field(pos).clone().detach()
        offsets.requires_grad = True
        
        ray_samples.frustums.set_offsets(offsets) 
        
        field_outputs = model.field(ray_samples)

        # accumulation
        packed_info = nerfacc.pack_info(ray_indices, num_rays)
        weights = nerfacc.render_weight_from_density(
            t_starts=ray_samples.frustums.starts[..., 0],
            t_ends=ray_samples.frustums.ends[..., 0],
            sigmas=field_outputs[FieldHeadNames.DENSITY][..., 0],
            packed_info=packed_info,
        )[0]
        weights = weights[..., None]

        rgb = model.renderer_rgb(
            rgb=field_outputs[FieldHeadNames.RGB],
            weights=weights,
            ray_indices=ray_indices,
            num_rays=num_rays,
        )
        depth = model.renderer_depth(
            weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=num_rays
        )
        accumulation = model.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=num_rays)

        outputs = {
            "rgb": rgb,
            "accumulation": accumulation,
            "depth": depth,
            "num_samples_per_ray": packed_info[:, 1],
        }
        return outputs, points.unsqueeze(0), offsets

    
    
    def get_unc_mipnerf(self, ray_bundle, model):
        ''' reimplementation of get_output function from models because of lack of proper interface to ray_samples'''
        if model.field is None:
            raise ValueError("populate_fields() must be called before get_outputs")
            
        
        if model.collider is not None:
            ray_bundle = model.collider(ray_bundle)
        
        
        
        # uniform sampling
        ray_samples = model.sampler_uniform(ray_bundle)

        points_coarse = ray_samples.frustums.get_positions()
        pos, _ = normalize_point_coords(points_coarse, self.aabb, model.field.spatial_distortion)
        #find offset value in the deformation field
        offsets_coarse = self.deform_field(pos).clone().detach()
        offsets_coarse.requires_grad = True
        
        ray_samples.frustums.set_offsets(offsets_coarse) 
        ray_samples.frustums.get_gaussian_blob = types.MethodType(get_gaussian_blob_new, ray_samples.frustums)


        # First pass:
        field_outputs_coarse = model.field.forward(ray_samples)
        weights_coarse = ray_samples.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY])
        rgb_coarse = model.renderer_rgb(
            rgb=field_outputs_coarse[FieldHeadNames.RGB],
            weights=weights_coarse,
        )
        accumulation_coarse = model.renderer_accumulation(weights_coarse)
        depth_coarse = model.renderer_depth(weights_coarse, ray_samples)

        # pdf sampling
        ray_samples_pdf = model.sampler_pdf(ray_bundle, ray_samples, weights_coarse)
        
        points_fine = ray_samples_pdf.frustums.get_positions()
        pos, _ = normalize_point_coords(points_fine, self.aabb, model.field.spatial_distortion)
        #find offset value in the deformation field
        offsets_fine = self.deform_field(pos).clone().detach()
        offsets_fine.requires_grad = True
        
        ray_samples_pdf.frustums.set_offsets(offsets_fine) 
        ray_samples_pdf.frustums.get_gaussian_blob = types.MethodType(get_gaussian_blob_new, ray_samples_pdf.frustums)

        # Second pass:
        field_outputs_fine = model.field.forward(ray_samples_pdf)
        weights_fine = ray_samples_pdf.get_weights(field_outputs_fine[FieldHeadNames.DENSITY])
        rgb_fine = model.renderer_rgb(
            rgb=field_outputs_fine[FieldHeadNames.RGB],
            weights=weights_fine,
        )
        accumulation_fine = model.renderer_accumulation(weights_fine)
        depth_fine = model.renderer_depth(weights_fine, ray_samples_pdf)
        
        outputs = {
            "rgb_coarse": rgb_coarse,
            "rgb_fine": rgb_fine,
            "accumulation_coarse": accumulation_coarse,
            "accumulation_fine": accumulation_fine,
            "depth_coarse": depth_coarse,
            "depth_fine": depth_fine,
        }
        return outputs, points_fine, offsets_fine, points_coarse, offsets_coarse
    
    def get_output_fn(self, model):
        
        if isinstance(model, NerfactoModel):
            return self.get_unc_nerfacto
        elif isinstance(model, NGPModel):
            return self.get_unc_ngp
        elif isinstance(model, MipNerfModel):
            return self.get_unc_mipnerf
        else:
            raise Exception("Sorry, this model is not currently supported.")
            
    def main(self) -> None:
        """Main function."""
        
        if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
            config, pipeline, checkpoint_path, _ = eval_setup(self.load_config)
        else:
            config, pipeline, checkpoint_path = eval_setup(self.load_config)
        
        self.output_path.parent.mkdir(parents=True, exist_ok=True)

        start_time = time.time()
        
        self.device = pipeline.device
        self.aabb = pipeline.model.scene_box.aabb.to(self.device)
        self.hessian = torch.zeros(((2**self.lod)+1)**3).to(self.device)
        self.deform_field = HashEncoding(num_levels = 1, 
                            min_res = 2**self.lod,
                            max_res = 2**self.lod,
                            log2_hashmap_size = self.lod*3+1, #simple regular grid (hash table size > grid size)
                            features_per_level = 3,
                            hash_init_scale = 0.,
                            implementation = "torch",
                            interpolation = "Linear")
        self.deform_field.to(self.device)
        self.deform_field.scalings = torch.tensor([2**self.lod]).to(self.device)
        
        pipeline.eval()
        len_train = max(pipeline.datamanager.train_dataset.__len__(), self.iters)
        for step in range(len_train):
            print("step",step)
            ray_bundle, batch = pipeline.datamanager.next_train(step)
            output_fn = self.get_output_fn(pipeline.model)
            if not isinstance(pipeline.model, MipNerfModel):
                outputs, points, offsets = output_fn(ray_bundle, pipeline.model)
                hessian = self.find_uncertainty(points, offsets, outputs['rgb'], pipeline.model.field.spatial_distortion)    
                self.hessian += hessian.clone().detach()
            else:
                outputs, points_fine, offsets_fine, points_coarse, offsets_coarse = output_fn(ray_bundle, pipeline.model)
                hessian = self.find_uncertainty(points_fine, offsets_fine, outputs['rgb_fine'], pipeline.model.field.spatial_distortion)    
                self.hessian += hessian.clone().detach()
                hessian = self.find_uncertainty(points_coarse, offsets_coarse, outputs['rgb_coarse'], pipeline.model.field.spatial_distortion)    
                self.hessian += hessian.clone().detach()
        end_time = time.time()    
        print("Done")
        with open(str(self.output_path), 'wb') as f:
            np.save(f, self.hessian.cpu().numpy())
        execution_time = end_time - start_time
        print(f"Execution time: {execution_time:.6f} seconds")    


def entrypoint():
    """Entrypoint for use with pyproject scripts."""
    tyro.extras.set_accent_color("bright_yellow")
    tyro.cli(ComputeUncertainty).main()


if __name__ == "__main__":
    entrypoint()

# For sphinx docs
get_parser_fn = lambda: tyro.extras.get_parser(ComputeUncertainty)  # noqa

bayesrays/scripts/uncertainty.py
BayesRays/BayesRays
# Copyright 2022 the Regents of the University of California, Nerfstudio Team and contributors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

#!/usr/bin/env python
# for older versions of nerfstudio filtering slider is currently not supported.
"""
Starts viewer in eval mode.
"""
from __future__ import annotations

import time
from dataclasses import dataclass, field, fields
from pathlib import Path

import tyro
import torch
import numpy as np
import types

import nerfstudio
import pkg_resources
import nerfacc

from nerfstudio.configs.base_config import ViewerConfig
from nerfstudio.engine.trainer import TrainerConfig
from nerfstudio.pipelines.base_pipeline import Pipeline
from nerfstudio.utils import writer, colors
from nerfstudio.utils.eval_utils import eval_setup
from nerfstudio.models.nerfacto import NerfactoModel
from nerfstudio.models.instant_ngp import NGPModel
from nerfstudio.models.mipnerf import MipNerfModel
from nerfstudio.field_components.field_heads import FieldHeadNames
from bayesrays.utils.utils import find_grid_indices
import torch.nn.functional as F
from bayesrays.scripts.output_uncertainty import get_uncertainty

if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
    from nerfstudio.viewer.server.viewer_state import ViewerState
    from nerfstudio.viewer.server.viewer_elements import  ViewerSlider
else:
    from nerfstudio.viewer.server import viewer_utils
    from nerfstudio.utils.writer import EventName, TimeWriter


def get_output_nerfacto_new(self, ray_bundle):
    ''' reimplementation of get_output function from models because of lack of proper interface to outputs dict'''
    
    N = self.N 
    reg_lambda = 1e-4 /( (2**self.lod)**3)
    H = self.hessian/N + reg_lambda
    self.un = 1/H
            
    max_uncertainty = 6 #approximate upper bound of the function log10(1/(x+lambda)) when lambda=1e-4/(256^3) and x is the hessian
    min_uncertainty = -3 #approximate lower bound of that function (cutting off at hessian = 1000)
    density_fns_new = []
    num_fns = len(self.density_fns) 
    for i in self.density_fns:
        density_fns_new.append(lambda x, i=i: i(x) * (self.get_uncertainty(x)<= self.filter_thresh*max_uncertainty))
        
    if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
        ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=density_fns_new)
    else:
        ray_samples,_, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=density_fns_new)
    field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals)
    points = ray_samples.frustums.get_positions()
    un_points = self.get_uncertainty(points)
    #get weights
    density = field_outputs[FieldHeadNames.DENSITY] * (un_points <= self.filter_thresh*max_uncertainty)
    weights = ray_samples.get_weights(density)
    
    
    uncertainty = torch.sum(weights * un_points, dim=-2) 
    uncertainty += (1-torch.sum(weights,dim=-2)) * min_uncertainty #alpha blending
    
    #normalize into acceptable range for rendering
    uncertainty = torch.clip(uncertainty, min_uncertainty, max_uncertainty)
    uncertainty = (uncertainty-min_uncertainty)/(max_uncertainty-min_uncertainty)
    
    if self.white_bg:
        self.renderer_rgb.background_color=colors.WHITE    
    elif self.black_bg:
        self.renderer_rgb.background_color=colors.BLACK        
    rgb = self.renderer_rgb(rgb=field_outputs[FieldHeadNames.RGB], weights=weights)
    depth = self.renderer_depth(weights=weights, ray_samples=ray_samples)
    accumulation = self.renderer_accumulation(weights=weights)
    
    
    # this is based on https://arxiv.org/pdf/2211.12656.pdf and the summation is not normalized. Check normalized in the viewer.
    # Uniform sampler (weights_list[0]) is used for getting uniform samples on each ray frustrum and find sum of entropy
    # of ray termination probabilities. Change sum to average for a more correct form. 
    ww = torch.clamp(weights_list[0], 1e-10 ,1.)
    entropy = -torch.sum(ww* torch.log2(ww) + (1-ww) * torch.log2(1-ww), dim=1)
   
    original_outputs = {
        "rgb": rgb,
        "accumulation": accumulation,
        "depth": depth,
    }
    original_outputs['uncertainty'] = uncertainty 
    original_outputs['entropy'] = entropy

    if self.config.predict_normals:
        normals = self.renderer_normals(normals=field_outputs[FieldHeadNames.NORMALS], weights=weights)
        pred_normals = self.renderer_normals(field_outputs[FieldHeadNames.PRED_NORMALS], weights=weights)
        original_outputs["normals"] = self.normals_shader(normals)
        original_outputs["pred_normals"] = self.normals_shader(pred_normals)    

    if self.training and self.config.predict_normals:
        original_outputs["rendered_orientation_loss"] = orientation_loss(
            weights.detach(), field_outputs[FieldHeadNames.NORMALS], ray_bundle.directions
        )

        original_outputs["rendered_pred_normal_loss"] = pred_normal_loss(
            weights.detach(),
            field_outputs[FieldHeadNames.NORMALS].detach(),
            field_outputs[FieldHeadNames.PRED_NORMALS],
        )

    for i in range(self.config.num_proposal_iterations):
        original_outputs[f"prop_depth_{i}"] = self.renderer_depth(weights=weights_list[i], ray_samples=ray_samples_list[i])


    return original_outputs

def get_output_ngp_new(self, ray_bundle):
    assert self.field is not None
    assert pkg_resources.get_distribution("nerfstudio").version >= "0.3.1"
    
    N = self.N
    reg_lambda = 1e-4 /( (2**self.lod)**3)
    H = self.hessian/N + reg_lambda
    self.un = 1/H 
            
    max_uncertainty = 6 #approximate upper bound of the function log10(1/(x+lambda)) when lambda=1e-4/(256^3) and x is the hessian
    min_uncertainty = -3 #approximate lower bound of that function (cutting off at hessian = 1000)
    density_fn_new = lambda x: self.field.density_fn(x) * (self.get_uncertainty(x.unsqueeze(0)).squeeze(0)<= self.filter_thresh*max_uncertainty)
    self.sampler.density_fn = density_fn_new    
    num_rays = len(ray_bundle)
    if self.collider is not None:
        ray_bundle = self.collider(ray_bundle)
    with torch.no_grad():
        ray_samples, ray_indices = self.sampler(
            ray_bundle=ray_bundle,
            near_plane=self.config.near_plane,
            far_plane=self.config.far_plane,
            render_step_size=self.config.render_step_size*0.001,
            alpha_thre=self.config.alpha_thre,
            cone_angle=self.config.cone_angle,
        )

    field_outputs = self.field(ray_samples)
    points = ray_samples.frustums.get_positions()
    un_points = self.get_uncertainty(points.unsqueeze(0)).squeeze(0)
    density = field_outputs[FieldHeadNames.DENSITY] * (un_points <= self.filter_thresh*max_uncertainty)
    # accumulation
    packed_info = nerfacc.pack_info(ray_indices, num_rays)
    weights = nerfacc.render_weight_from_density(
        t_starts=ray_samples.frustums.starts[..., 0],
        t_ends=ray_samples.frustums.ends[..., 0],
        sigmas=density[..., 0],
        packed_info=packed_info,
    )[0]
        
    weights = weights[..., None]
    
    comp_uncertainty = nerfacc.accumulate_along_rays(
                weights[..., 0], values=un_points, ray_indices=ray_indices, n_rays=num_rays
            )
    accumulated_weight = nerfacc.accumulate_along_rays(
        weights[..., 0], values=None, ray_indices=ray_indices, n_rays=num_rays
    )
    uncertainty = comp_uncertainty + min_uncertainty * (1.0 - accumulated_weight)#alpha blending

    #normalize into acceptable range for rendering
    uncertainty = torch.clip(uncertainty, min_uncertainty, max_uncertainty)
    uncertainty = (uncertainty-min_uncertainty)/(max_uncertainty-min_uncertainty)
    
    if self.white_bg:
        self.renderer_rgb.background_color=colors.WHITE    
    elif self.black_bg:
        self.renderer_rgb.background_color=colors.BLACK  

    rgb = self.renderer_rgb(
        rgb=field_outputs[FieldHeadNames.RGB],
        weights=weights,
        ray_indices=ray_indices,
        num_rays=num_rays,
    )
    depth = self.renderer_depth(
        weights=weights, ray_samples=ray_samples, ray_indices=ray_indices, num_rays=num_rays
    )
    accumulation = self.renderer_accumulation(weights=weights, ray_indices=ray_indices, num_rays=num_rays)
    
    outputs = {
        "rgb": rgb,
        "accumulation": accumulation,
        "depth": depth,
        "uncertainty": uncertainty
    }
    return outputs

def get_output_mipnerf_new(self, ray_bundle):
    ''' reimplementation of get_output function from models because of lack of proper interface to outputs dict'''
    
    N = self.N
    reg_lambda = 1e-4 /( (2**self.lod)**3)
    H = self.hessian/N + reg_lambda
    self.un = 1/H
            
    max_uncertainty = 6 #approximate upper bound of the function log10(1/(x+lambda)) when lambda=1e-4/(256^3) and x is the hessian
    min_uncertainty = -3 #approximate lower bound of that function (cutting off at hessian = 1000)
    
    # uniform sampling
    ray_samples = self.sampler_uniform(ray_bundle)
        
    # First pass:
    field_outputs_coarse = self.field.forward(ray_samples)
    weights_coarse = ray_samples.get_weights(field_outputs_coarse[FieldHeadNames.DENSITY])
    rgb_coarse = self.renderer_rgb(
        rgb=field_outputs_coarse[FieldHeadNames.RGB],
        weights=weights_coarse,
    )
    accumulation_coarse = self.renderer_accumulation(weights_coarse)
    depth_coarse = self.renderer_depth(weights_coarse, ray_samples)

    # pdf sampling
    ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples, weights_coarse)

    # Second pass:
    field_outputs_fine = self.field.forward(ray_samples_pdf)
        
    points = ray_samples_pdf.frustums.get_positions()
    un_points = self.get_uncertainty(points)
    
    density = field_outputs_fine[FieldHeadNames.DENSITY] * (un_points <= self.filter_thresh*max_uncertainty)
    #get weights
    weights_fine = ray_samples_pdf.get_weights(density)
    rgb_fine = self.renderer_rgb(
        rgb=field_outputs_fine[FieldHeadNames.RGB],
        weights=weights_fine,
    )
    accumulation_fine = self.renderer_accumulation(weights_fine)
    depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf)
    
    
    uncertainty = torch.sum(weights_fine * un_points, dim=-2) 
    uncertainty += (1-torch.sum(weights_fine,dim=-2)) * min_uncertainty #alpha blending
    
    #normalize into acceptable range for rendering
    uncertainty = torch.clip(uncertainty, min_uncertainty, max_uncertainty)
    uncertainty = (uncertainty-min_uncertainty)/(max_uncertainty-min_uncertainty)
    
    if self.white_bg:
        self.renderer_rgb.background_color=colors.WHITE    
    elif self.black_bg:
        self.renderer_rgb.background_color=colors.BLACK        
   
    original_outputs = {
            "rgb_coarse": rgb_coarse,
            "rgb_fine": rgb_fine,
            "accumulation_coarse": accumulation_coarse,
            "accumulation_fine": accumulation_fine,
            "depth_coarse": depth_coarse,
            "depth_fine": depth_fine,
        }
    original_outputs['uncertainty'] = uncertainty 

    return original_outputs

def get_output_fn(model):

    if isinstance(model, NerfactoModel):
        return get_output_nerfacto_new
    elif isinstance(model, NGPModel):
        return get_output_ngp_new
    elif isinstance(model, MipNerfModel):
        return get_output_mipnerf_new
    else:
        raise Exception("Sorry, this model is not currently supported.")

        
@dataclass
class ViewerConfigWithoutNumRays(ViewerConfig):
    """Configuration for viewer instantiation"""

    num_rays_per_chunk: tyro.conf.Suppress[int] = -1

    def as_viewer_config(self):
        """Converts the instance to ViewerConfig"""
        return ViewerConfig(**{x.name: getattr(self, x.name) for x in fields(self)})


@dataclass
class RunViewerU:
    """Load a checkpoint and start the viewer."""

    load_config: Path
    """Path to config YAML file."""
    viewer: ViewerConfigWithoutNumRays = field(default_factory=ViewerConfigWithoutNumRays)
    """Viewer configuration"""
    unc_path: Path = Path("unc.npy")
    """Path to output video file."""
    white_bg: bool = True
    """ Render empty space as white when filtering""" 
    black_bg: bool = False
    """ Render empty space as black when filtering""" 

    def main(self) -> None:
        """Main function."""
        if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
            config, pipeline, _, step = eval_setup(
                self.load_config,
                eval_num_rays_per_chunk=None,
                test_mode="test",
            )
        else:
            config, pipeline, _ = eval_setup(
                self.load_config,
                eval_num_rays_per_chunk=None,
                test_mode="test",
            )
            step = 0
        
        self.device = pipeline.device
        pipeline.model.filter_thresh = 1.
        pipeline.model.hessian = torch.tensor(np.load(str(self.unc_path)), device=self.device)
        pipeline.model.N =  4096*1000  #approx ray dataset size (train batch size x number of query iterations in uncertainty extraction step)
        pipeline.model.lod = np.log2(round(pipeline.model.hessian.shape[0]**(1/3))-1)
        pipeline.model.get_uncertainty = types.MethodType(get_uncertainty, pipeline.model)
        pipeline.model.white_bg = self.white_bg
        pipeline.model.black_bg = self.black_bg
        new_method = get_output_fn(pipeline.model)
        pipeline.model.get_outputs = types.MethodType(new_method, pipeline.model)
        
        
        num_rays_per_chunk = config.viewer.num_rays_per_chunk
        assert self.viewer.num_rays_per_chunk == -1
        config.vis = "viewer"
        config.viewer = self.viewer.as_viewer_config()
        config.viewer.num_rays_per_chunk = num_rays_per_chunk
    
        self._start_viewer(config, pipeline, step)

    def save_checkpoint(self, *args, **kwargs):
        """
        Mock method because we pass this instance to viewer_state.update_scene
        """
    def _update_viewer_state(self, viewer_state: viewer_utils.ViewerState, pipeline: Pipeline):
        """Updates the viewer state by rendering out scene with current pipeline
        Returns the time taken to render scene.

        """
        # NOTE: step must be > 0 otherwise the rendering would not happen
        step = 1
        num_rays_per_batch = pipeline.datamanager.get_train_rays_per_batch()
        with TimeWriter(writer, EventName.ITER_VIS_TIME) as _:
            viewer_state.update_scene = types.MethodType(update_scene_new, viewer_state)
            try:
                viewer_state.update_scene(self, step, pipeline.model, num_rays_per_batch)
            except RuntimeError:
                time.sleep(0.03)  # sleep to allow buffer to reset
                assert viewer_state.vis is not None
                viewer_state.vis["renderingState/log_errors"].write(
                    "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing."
                )

    def _start_viewer(self, config: TrainerConfig, pipeline: Pipeline, step: int):
        """Starts the viewer

        Args:
            config: Configuration of pipeline to load
            pipeline: Pipeline instance of which to load weights
            step: Step at which the pipeline was saved
        """
        base_dir = config.get_base_dir()
        viewer_log_path = base_dir / config.viewer.relative_log_filename

        if pkg_resources.get_distribution("nerfstudio").version >= "0.3.1":
            viewer_state = ViewerState(
            config.viewer,
            log_filename=viewer_log_path,
            datapath=pipeline.datamanager.get_datapath(),
            pipeline=pipeline,
            )
            viewer_state.control_panel._filter = ViewerSlider(
                    "Filter Threshold",
                    default_value=1.,
                    min_value=0.0,
                    max_value=1,
                    step=0.05,
                    hint="Filtering threshold for uncertain areas.",
                )
            viewer_state.control_panel.add_element(viewer_state.control_panel._filter)
            banner_messages = [f"Viewer at: {viewer_state.viewer_url}"]
        else:
            viewer_state, banner_messages = viewer_utils.setup_viewer(
                config.viewer, log_filename=viewer_log_path, datapath=pipeline.datamanager.get_datapath()
            )

        # We don't need logging, but writer.GLOBAL_BUFFER needs to be populated
        config.logging.local_writer.enable = False
        writer.setup_local_writer(config.logging, max_iter=config.max_num_iterations, banner_messages=banner_messages)

        assert viewer_state and pipeline.datamanager.train_dataset

        nerfstudio_version = pkg_resources.get_distribution("nerfstudio").version

        if nerfstudio_version >= "0.3.1":
            if nerfstudio_version >=  "0.3.3":
                viewer_state.init_scene(
                    train_dataset=pipeline.datamanager.train_dataset,
                    train_state="completed",
                )
            else:
                viewer_state.init_scene(
                    dataset=pipeline.datamanager.train_dataset,
                    train_state="completed",
                )
            viewer_state.viser_server.set_training_state("completed")
            viewer_state.update_scene(step=step)
            while True:
                time.sleep(0.01)
                pipeline.model.filter_thresh = viewer_state.control_panel._filter.value
        else:
            viewer_state.init_scene(
                dataset=pipeline.datamanager.train_dataset,
                start_train=False,
            )   
            while True:
                viewer_state.vis["renderingState/isTraining"].write(False)
                self._update_viewer_state(viewer_state, pipeline)

                
# this function is redefined to allow support of filter threshold slider in the viewer for old nerfstudio versions.
# The "Train Util." slider in eval time will control the filter threshold instead. 
def update_scene_new(self, trainer, step: int, graph: Model, num_rays_per_batch: int) -> None:
    """updates the scene based on the graph weights

    Args:
        step: iteration step of training
        graph: the current checkpoint of the model
    """
    has_temporal_distortion = getattr(graph, "temporal_distortion", None) is not None
    self.vis["model/has_temporal_distortion"].write(str(has_temporal_distortion).lower())

    is_training = self.vis["renderingState/isTraining"].read()
    self.step = step

    self._check_camera_path_payload(trainer, step)
    self._check_populate_paths_payload(trainer, step)

    camera_object = self._get_camera_object()
    if camera_object is None:
        return

    if is_training is None or is_training:
        # in training mode

        if self.camera_moving:
            # if the camera is moving, then we pause training and update camera continuously

            while self.camera_moving:
                self._render_image_in_viewer(camera_object, graph, is_training)
                camera_object = self._get_camera_object()
        else:
            # if the camera is not moving, then we approximate how many training steps need to be taken
            # to render at a FPS defined by self.static_fps.

            if EventName.TRAIN_RAYS_PER_SEC.value in GLOBAL_BUFFER["events"]:
                train_rays_per_sec = GLOBAL_BUFFER["events"][EventName.TRAIN_RAYS_PER_SEC.value]["avg"]
                target_train_util = self.vis["renderingState/targetTrainUtil"].read()
                if target_train_util is None:
                    target_train_util = 0.9

                batches_per_sec = train_rays_per_sec / num_rays_per_batch

                num_steps = max(int(1 / self.static_fps * batches_per_sec), 1)
            else:
                num_steps = 1

            if step % num_steps == 0:
                self._render_image_in_viewer(camera_object, graph, is_training)

    else:
        # in pause training mode, enter render loop with set graph
        local_step = step
        run_loop = not is_training
        while run_loop:
            # if self._is_render_step(local_step) and step > 0:
            if step > 0:
                self._render_image_in_viewer(camera_object, graph, is_training)
                camera_object = self._get_camera_object()
            th =  self.vis["renderingState/targetTrainUtil"].read() 
            graph.filter_thresh = th if th is not None else 1.    
            is_training = self.vis["renderingState/isTraining"].read()
            self._check_populate_paths_payload(trainer, step)
            self._check_camera_path_payload(trainer, step)
            run_loop = not is_training
            local_step += 1
                

def entrypoint():
    """Entrypoint for use with pyproject scripts."""
    tyro.extras.set_accent_color("bright_yellow")
    tyro.cli(RunViewerU).main()


if __name__ == "__main__":
    entrypoint()

# For sphinx docs
get_parser_fn = lambda: tyro.extras.get_parser(RunViewerU)  # noqa
bayesrays/scripts/viewer_uncertainty/run_viewer_u.py
BayesRays/BayesRays
import struct
import numpy as np
import torch
import argparse
from pathlib import Path
from scipy.optimize import least_squares

#parts of code from https://github.com/dunbar12138/DSNeRF/blob/main/colmapUtils


def read_next_bytes(fid, num_bytes, format_char_sequence, endian_character="<"):
    """Read and unpack the next bytes from a binary file.
    :param fid:
    :param num_bytes: Sum of combination of {2, 4, 8}, e.g. 2, 6, 16, 30, etc.
    :param format_char_sequence: List of {c, e, f, d, h, H, i, I, l, L, q, Q}.
    :param endian_character: Any of {@, =, <, >, !}
    :return: Tuple of read and unpacked values.
    """
    data = fid.read(num_bytes)
    return struct.unpack(endian_character + format_char_sequence, data)


def read_points_bin(base_dir):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::ReadPoints3DBinary(const std::string& path)
        void Reconstruction::WritePoints3DBinary(const std::string& path)
    """
    points_dir = Path('points3D.bin')
    final_dir = base_dir / points_dir
    points3D = {}
    points3D_loc = {}
    xyzs = []
    with open(final_dir, "rb") as fid:
        num_points = read_next_bytes(fid, 8, "Q")[0]
        for _ in range(num_points):
            binary_point_line_properties = read_next_bytes(
                fid, num_bytes=43, format_char_sequence="QdddBBBd")
            point3D_id = binary_point_line_properties[0]
            xyz = np.array(binary_point_line_properties[1:4])
            rgb = np.array(binary_point_line_properties[4:7])
            error = np.array(binary_point_line_properties[7])
            track_length = read_next_bytes(
                fid, num_bytes=8, format_char_sequence="Q")[0]
            track_elems = read_next_bytes(
                fid, num_bytes=8*track_length,
                format_char_sequence="ii"*track_length)
            image_ids = np.array(tuple(map(int, track_elems[0::2])))
            point2D_idxs = np.array(tuple(map(int, track_elems[1::2])))
            points3D[point3D_id] = {
                'id':point3D_id, 'xyz':xyz, 'rgb':rgb,
                'error':error, 'image_ids':image_ids,
                'point2D_idxs':point2D_idxs}
            points3D_loc[point3D_id] = xyz
    return points3D, points3D_loc      


def read_images_bin(base_dir):
    """
    see: src/base/reconstruction.cc
        void Reconstruction::ReadImagesBinary(const std::string& path)
        void Reconstruction::WriteImagesBinary(const std::string& path)
    """
    images_dir = Path('images.bin')
    final_dir = base_dir / images_dir
    images = {}
    with open(final_dir, "rb") as fid:
        num_reg_images = read_next_bytes(fid, 8, "Q")[0]
        for _ in range(num_reg_images):
            binary_image_properties = read_next_bytes(
                fid, num_bytes=64, format_char_sequence="idddddddi")
            image_id = binary_image_properties[0]
            qvec = np.array(binary_image_properties[1:5])
            tvec = np.array(binary_image_properties[5:8])
            camera_id = binary_image_properties[8]
            image_name = ""
            current_char = read_next_bytes(fid, 1, "c")[0]
            while current_char != b"\x00":   # look for the ASCII 0 entry
                image_name += current_char.decode("utf-8")
                current_char = read_next_bytes(fid, 1, "c")[0]
            num_points2D = read_next_bytes(fid, num_bytes=8,
                                           format_char_sequence="Q")[0]
            x_y_id_s = read_next_bytes(fid, num_bytes=24*num_points2D,
                                       format_char_sequence="ddq"*num_points2D)
            xys = np.column_stack([tuple(map(float, x_y_id_s[0::3])),
                                   tuple(map(float, x_y_id_s[1::3]))])
            point3D_ids = np.array(tuple(map(int, x_y_id_s[2::3])))
            images[image_id] = {
                'id':image_id, 'qvec':qvec, 'tvec':tvec,
                'camera_id':camera_id, 'name':image_name,
                'xys':xys, 'point3D_ids':point3D_ids}
    return images       

def quaternion_to_rotation_matrix(qvec):
    qw, qx, qy, qz = qvec
    
    # Compute terms for clarity
    qx2 = qx * qx
    qy2 = qy * qy
    qz2 = qz * qz
    qwqx = qw * qx
    qwqy = qw * qy
    qwqz = qw * qz
    qxqy = qx * qy
    qxqz = qx * qz
    qyqz = qy * qz
    
    # Compute rotation matrix
    R = np.array([[1 - 2*qy2 - 2*qz2, 2*qxqy - 2*qz*qw, 2*qxqz + 2*qy*qw],
                  [2*qxqy + 2*qz*qw, 1 - 2*qx2 - 2*qz2, 2*qyqz - 2*qx*qw],
                  [2*qxqz - 2*qy*qw, 2*qyqz + 2*qx*qw, 1 - 2*qx2 - 2*qy2]])
    
    return R

def find_depths(images, points,scale_factor, downscale, auto_scaled):
    camera_centers = []
    Rs = []
    depths_all = []
    xys_all = []
    scale = 1
    for i in range(len(images)):     
        R = quaternion_to_rotation_matrix(images[i+1]['qvec'])
        Rs += [R.T]
        camera_center = -R.T @ images[i+1]['tvec']
        camera_centers += [list(camera_center)]
    camera_centers = np.array(camera_centers)
    if auto_scaled:
        scale *= 1/ float(np.max(np.abs(camera_centers)))
    scale *= scale_factor
        
    for i in i_val:
        xys = np.flip(images[i+1]['xys'],-1) #flip the order of x and y (result: x indexes rows, and y columns)
        xys = xys [images[i+1]['point3D_ids'] != -1] #mask out non 3D keypoints
        
        xys = xys / downscale
        camera_center = camera_centers[i]
        points3D = np.array([list(points[j]) for j in images[i+1]['point3D_ids'] if j!=-1])
        depths =  (points3D - camera_center) @ Rs[i][:3,2]
        depths = depths * scale 
        depths_all += [depths]
        xys_all += [xys]
    
    return xys_all, depths_all #num_images ( point of image x 3|1)

def read_gt_depth(base_dir, xys):
    depths = []
    for c,i in enumerate(i_val):
        depth_gt_dir =  base_dir / Path('depth_gt_{:02d}.npy'.format(c))
        depth_gt = np.load(str(depth_gt_dir))
        xs = np.floor(xys[c][:, 0]).astype(int)
        ys = np.floor(xys[c][:, 1]).astype(int)
        depth_gt_points = depth_gt[xs, ys]
        depths+=[depth_gt_points]
    return depths   #num_images ( point of image x 1)




def main(args):
    
    print(args)
    base_dir = Path('./data') / Path(args.dataset) / Path(args.scene) 
    colmap_dir = base_dir / Path('colmap/sparse/0')
    _, points = read_points_bin(colmap_dir)
    images = read_images_bin(colmap_dir)
    xys, colmap_depths = find_depths(images, points, args.scale_factor, args.downscale, args.autoscale)
    gt_depths = read_gt_depth(base_dir, xys)
    
    colmap_depths = np.concatenate(colmap_depths)
    gt_depths = np.concatenate(gt_depths)
    
    
    def objective(x):
        a = x
        residual = a * colmap_depths - gt_depths
        return residual

    # Initial guess for parameters
    x0 = [1.0]

    # Solve the least squares problem
    result = least_squares(objective, x0)
    
    # Extract the optimal parameters
    a_opt = result.x
    a = a_opt

    print("Optimal parameters:")
    print("a =", a_opt)
    print(np.mean(colmap_depths), np.max(colmap_depths), np.min(colmap_depths), np.median(colmap_depths))
    print(np.mean(gt_depths), np.max(gt_depths), np.min(gt_depths),np.median(gt_depths))
    print(np.mean(a * colmap_depths ), np.max(a * colmap_depths ), np.min(a * colmap_depths), np.median(gt_depths)/np.median(colmap_depths))
    
    
    #save a and b
    save_dir = base_dir / Path('scale_parameters.txt')
    np.savetxt(str(save_dir), [a_opt], delimiter=',')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Your script description')
    
    # Add arguments to the parser
    parser.add_argument('--dataset', type=str, default='scannet', help='dataset name')
    parser.add_argument('--scene', type=str, default='scene_079', help='scene name')
    parser.add_argument('--scale-factor', type=float, default=1., help='scale factor')
    parser.add_argument('--downscale', type=float, default=2., help='image downscale')
    parser.add_argument("--autoscale", action="store_const", const=True, default=False, help="auto scaled")
    
    # Parse the arguments
    args = parser.parse_args()
    
    if args.dataset == 'scannet':
        i_val = [4, 12, 20, 28, 36] 
    elif args.dataset == 'LF':
        if args.scene == 'basket':
            i_val = list(np.arange(42,50,2))
        elif args.scene == 'africa':
            i_val = list(np.arange(6,14,2))
        elif args.scene == "torch":
            i_val = list(np.arange(9,17,2))
        elif args.scene == "statue":
            i_val = list(np.arange(68,76,2))
        
    
    main(args)

bayesrays/utils/scale_solver.py
BayesRays/BayesRays
import torch
from nerfstudio.data.scene_box import SceneBox
from nerfstudio.utils.math import Gaussians, conical_frustum_to_gaussian

def normalize_point_coords(points, aabb, distortion):
    ''' coordinate normalization process according to density_feild.py in nerfstudio'''
    if distortion is not None:
        pos = distortion(points)
        pos = (pos + 2.0) / 4.0
    else:        
        pos = SceneBox.get_normalized_positions(points, aabb)
    selector = ((pos > 0.0) & (pos < 1.0)).all(dim=-1) #points outside aabb are filtered out
    pos = pos * selector[..., None]
    return pos, selector

def find_grid_indices(points, aabb, distortion, lod, device, zero_out=True):
    pos, selector = normalize_point_coords(points, aabb, distortion)
    pos, selector = pos.view(-1, 3), selector[..., None].view(-1, 1)
    uncertainty_lod = 2 ** lod
    coords = (pos * uncertainty_lod).unsqueeze(0)
    inds = torch.zeros((8, pos.shape[0]), dtype=torch.int32, device=device)
    coefs = torch.zeros((8, pos.shape[0]), device=device)
    corners = torch.tensor(
        [[0, 0, 0, 0], [1, 0, 0, 1], [2, 0, 1, 0], [3, 0, 1, 1], [4, 1, 0, 0], [5, 1, 0, 1], [6, 1, 1, 0],
         [7, 1, 1, 1]], device=device)
    corners = corners.unsqueeze(1)
    inds[corners[:, :, 0].squeeze(1)] = (
            (torch.floor(coords[..., 0]) + corners[:, :, 1]) * uncertainty_lod * uncertainty_lod +
            (torch.floor(coords[..., 1]) + corners[:, :, 2]) * uncertainty_lod +
            (torch.floor(coords[..., 2]) + corners[:, :, 3])).int()
    coefs[corners[:, :, 0].squeeze(1)] = torch.abs(
        coords[..., 0] - (torch.floor(coords[..., 0]) + (1 - corners[:, :, 1]))) * torch.abs(
        coords[..., 1] - (torch.floor(coords[..., 1]) + (1 - corners[:, :, 2]))) * torch.abs(
        coords[..., 2] - (torch.floor(coords[..., 2]) + (1 - corners[:, :, 3])))
    if zero_out:
        coefs[corners[:, :, 0].squeeze(1)] *= selector[..., 0].unsqueeze(0)  # zero out the contribution of points outside aabb box

    return inds, coefs

def get_gaussian_blob_new(self) -> Gaussians: #for mipnerf
    """Calculates guassian approximation of conical frustum.

    Returns:
        Conical frustums approximated by gaussian distribution.
    """
    # Cone radius is set such that the square pixel_area matches the cone area.
    cone_radius = torch.sqrt(self.pixel_area) / 1.7724538509055159  # r = sqrt(pixel_area / pi)
    
    return conical_frustum_to_gaussian(
        origins=self.origins + self.offsets, #deforms Gaussian mean
        directions=self.directions,
        starts=self.starts,
        ends=self.ends,
        radius=cone_radius,
    )




bayesrays/utils/utils.py
BayesRays/BayesRays