nerfstudio/nerfactory/viewer/server/viewer_utils.py

# Copyright 2022 The Plenoptix 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.

"""Code to interface with the `vis/` (the JS viewer).
"""
from __future__ import annotations

import enum
import os
import sys
import threading
import time
from pathlib import Path
from typing import Any, Dict, Optional

import numpy as np
import torch
from rich.console import Console

from nerfactory.cameras.cameras import Cameras
from nerfactory.cameras.rays import RayBundle
from nerfactory.configs import base as cfg
from nerfactory.datamanagers.datasets import InputDataset
from nerfactory.models.base import Model
from nerfactory.utils import profiler, visualization, writer
from nerfactory.utils.io import load_from_json, write_to_json
from nerfactory.utils.misc import get_dict_to_torch
from nerfactory.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter
from nerfactory.viewer.server.subprocess import run_viewer_bridge_server_as_subprocess
from nerfactory.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h
from nerfactory.viewer.server.visualizer import Viewer

console = Console(width=120)


def setup_viewer(config: cfg.ViewerConfig, log_filename: Path):
    """Sets up the viewer if enabled

    Args:
        config: the configuration to instantiate viewer
    """
    if config.enable:
        viewer_state = ViewerState(config, log_filename=log_filename)
        banner_messages = [f"Viewer at: {viewer_state.viewer_url}"]
        return viewer_state, banner_messages
    return None, None


class OutputTypes(str, enum.Enum):
    """Noncomprehsnive list of output render types"""

    INIT = "init"
    RGB = "rgb"
    RGB_FINE = "rgb_fine"
    ACCUMULATION = "accumulation"
    ACCUMULATION_FINE = "accumulation_fine"


class ColormapTypes(str, enum.Enum):
    """Noncomprehsnive list of colormap render types"""

    INIT = "init"
    DEFAULT = "default"
    TURBO = "turbo"
    DEPTH = "depth"
    SEMANTIC = "semantic"
    BOOLEAN = "boolean"


class IOChangeException(Exception):
    """Basic camera exception to interrupt viewer"""


class SetTrace:
    """Basic trace function"""

    def __init__(self, func):
        self.func = func

    def __enter__(self):
        sys.settrace(self.func)
        return self

    def __exit__(self, ext_type, exc_value, traceback):
        sys.settrace(None)


class RenderThread(threading.Thread):
    """Thread that does all the rendering calls while listening for interrupts

    Args:
        state: current viewer state object
        graph: current checkpoint of model
        camera_ray_bundle: input rays to pass through the graph to render out
    """

    def __init__(self, state: "ViewerState", graph: Model, camera_ray_bundle: RayBundle):
        threading.Thread.__init__(self)
        self.state = state
        self.graph = graph
        self.camera_ray_bundle = camera_ray_bundle
        self.exc = None
        self.vis_outputs = None

    def run(self):
        """run function that renders out images given the current graph and ray bundles.
        Interlaced with a trace function that checks to see if any I/O changes were registered.
        Exits and continues program if IOChangeException thrown.
        """
        outputs = None
        try:
            with SetTrace(self.state.check_interrupt):
                with torch.no_grad():
                    outputs = self.graph.get_outputs_for_camera_ray_bundle(self.camera_ray_bundle)
        except Exception as e:  # pylint: disable=broad-except
            self.exc = e

        if outputs:
            outputs = get_dict_to_torch(outputs)
            self.vis_outputs = outputs

        self.state.check_done_render = True
        self.state.check_interrupt_vis = False

    def join(self, timeout=None):
        threading.Thread.join(self)
        if self.exc:
            raise self.exc


class CheckThread(threading.Thread):
    """Thread the constantly checks for io changes and sets a flag indicating interrupt

    Args:
        state: current viewer state object
    """

    def __init__(self, state):
        threading.Thread.__init__(self)
        self.state = state

    def run(self):
        """Run function that checks to see if any of the existing state has changed
        (e.g. camera pose/output type/resolutions).
        Sets the viewer state flag to true to signal
        to render thread that an interrupt was registered.
        """
        self.state.check_done_render = False
        while not self.state.check_done_render:
            # check camera
            data = self.state.vis["renderingState/camera"].read()
            if data is not None:
                camera_object = data["object"]
                if self.state.prev_camera_matrix is None or (
                    not np.allclose(camera_object["matrix"], self.state.prev_camera_matrix)
                    and not self.state.prev_moving
                ):
                    self.state.check_interrupt_vis = True
                    self.state.prev_moving = True
                    return
                self.state.prev_moving = False

            # check output type
            output_type = self.state.vis["renderingState/output_choice"].read()
            if output_type is None:
                output_type = OutputTypes.INIT
            if self.state.prev_output_type != output_type:
                self.state.check_interrupt_vis = True
                return

            # check colormap type
            colormap_type = self.state.vis["renderingState/colormap_choice"].read()
            if colormap_type is None:
                colormap_type = ColormapTypes.INIT
            if self.state.prev_colormap_type != colormap_type:
                self.state.check_interrupt_vis = True
                return

            # check max render
            max_resolution = self.state.vis["renderingState/maxResolution"].read()
            if max_resolution is not None:
                if self.state.max_resolution != max_resolution:
                    self.state.check_interrupt_vis = True
                    return


class ViewerState:
    """Class to hold state for viewer variables

    Args:
        config: viewer setup configuration
    """

    def __init__(self, config: cfg.ViewerConfig, log_filename: Path):
        self.config = config
        self.vis = None
        self.viewer_url = None
        self.log_filename = log_filename
        if self.config.launch_bridge_server:
            # start the viewer bridge server
            zmq_port = int(self.config.zmq_url.split(":")[-1])
            websocket_port = self.config.websocket_port
            self.log_filename.parent.mkdir(exist_ok=True)
            run_viewer_bridge_server_as_subprocess(zmq_port, websocket_port, log_filename=str(self.log_filename))
            # TODO(ethan): move this into the writer such that it's at the bottom
            # of the logging stack and easy to see and click
            # TODO(ethan): log the output of the viewer bridge server in a file where the training logs go
            console.line()
            json_filename = os.path.join(os.path.dirname(__file__), "../app/package.json")
            version = load_from_json(Path(json_filename))["version"]
            self.viewer_url = f"https://viewer.nerfactory.com/{version}/?websocket_url=localhost:{websocket_port}"
            viewer_url_local = f"http://localhost:4000/?websocket_url=localhost:{websocket_port}"
            pub_open_viewer_instructions_string = f"[Public] Open the viewer at {self.viewer_url}"
            dev_open_viewer_instructions_string = f"[Local] Open the viewer at {viewer_url_local}"
            console.rule(characters="=")
            console.print(pub_open_viewer_instructions_string)
            console.print(dev_open_viewer_instructions_string)
            console.rule(characters="=")
            console.line()
        self.vis = Viewer(zmq_url=self.config.zmq_url)

        # viewer specific variables
        self.prev_camera_matrix = None
        self.prev_output_type = OutputTypes.INIT
        self.prev_colormap_type = ColormapTypes.INIT
        self.prev_moving = False
        self.output_type_changed = True
        self.max_resolution = 1000
        self.check_interrupt_vis = False
        self.check_done_render = True
        self.step = 0
        self.static_fps = 1
        self.moving_fps = 24
        self.camera_moving = False
        self.prev_camera_timestamp = 0

        self.output_list = None

    def init_scene(self, dataset: InputDataset, start_train=True) -> None:
        """Draw some images and the scene aabb in the viewer.

        Args:
            dataset: dataset to render in the scene
            start_train: whether to start train when viewer init;
                if False, only displays dataset until resume train is toggled
        """
        # set the config base dir
        self.vis["renderingState/config_base_dir"].write(str(self.log_filename.parents[0]))

        # clear the current scene
        self.vis["sceneState/sceneBounds"].delete()
        self.vis["sceneState/cameras"].delete()

        # draw the training cameras and images
        image_indices = range(len(dataset))
        for idx in image_indices:
            image = dataset[idx]["image"]
            bgr = image[..., [2, 1, 0]]
            camera_json = dataset.dataset_inputs.cameras.to_json(camera_idx=idx, image=bgr, max_size=100)
            self.vis[f"sceneState/cameras/{idx:06d}"].write(camera_json)

        # draw the scene bounds (i.e., the bounding box)
        json_ = dataset.dataset_inputs.scene_bounds.to_json()
        self.vis["sceneState/sceneBounds"].write(json_)

        # set the initial state whether to train or not
        self.vis["renderingState/isTraining"].write(start_train)

        # set the properties of the camera
        # self.vis["renderingState/camera"].write(json_)

        # set the main camera intrinsics to one from the dataset
        # K = camera.get_intrinsics_matrix()
        # set_persp_intrinsics_matrix(self.vis, K.double().numpy())

    def update_scene(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
        """

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

        # check if we should interrupt from a button press?
        camera_path_payload = self.vis["camera_path_payload"].read()
        if camera_path_payload:
            # save a model checkpoint
            trainer.save_checkpoint(step)
            # write to json file
            camera_path_filename = camera_path_payload["camera_path_filename"]
            camera_path = camera_path_payload["camera_path"]
            write_to_json(Path(camera_path_filename), camera_path)
            self.vis["camera_path_payload"].delete()

        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()
                is_training = self.vis["renderingState/isTraining"].read()
                run_loop = not is_training
                local_step += 1

    def check_interrupt(self, frame, event, arg):  # pylint: disable=unused-argument
        """Raises interrupt when flag has been set and not already on lowest resolution.
        Used in conjunction with SetTrace.
        """
        if event == "line":
            if self.check_interrupt_vis and not self.camera_moving:
                raise IOChangeException
        return self.check_interrupt

    def _get_camera_object(self):
        """Gets the camera object from the viewer and updates the movement state if it has changed."""

        data = self.vis["renderingState/camera"].read()
        if data is None:
            return None

        camera_object = data["object"]

        if self.prev_camera_matrix is not None and np.allclose(camera_object["matrix"], self.prev_camera_matrix):
            self.camera_moving = False
        else:
            self.prev_camera_matrix = camera_object["matrix"]
            self.camera_moving = True
        return camera_object

    def _apply_colormap(self, outputs: Dict[str, Any], stuff_colors: torch.Tensor = None, eps=1e-6):
        """Determines which colormap to use based on set colormap type

        Args:
            outputs: the output tensors for which to apply colormaps on
            stuff_colors: is only set if colormap is for semantics. Defaults to None.
            eps: epsilon to handle floating point comparisons
        """
        if self.output_list:
            reformatted_output = self._process_invalid_output(self.prev_output_type)

        # default for rgb images
        if self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[reformatted_output].shape[-1] == 3:
            return outputs[reformatted_output]

        # rendering depth outputs
        if self.prev_colormap_type == ColormapTypes.DEPTH or (
            self.prev_colormap_type == ColormapTypes.DEFAULT
            and outputs[reformatted_output].dtype == torch.float
            and (torch.max(outputs[reformatted_output]) - 1.0) > eps  # handle floating point arithmetic
        ):
            accumulation_str = (
                OutputTypes.ACCUMULATION
                if OutputTypes.ACCUMULATION in self.output_list
                else OutputTypes.ACCUMULATION_FINE
            )
            return visualization.apply_depth_colormap(
                outputs[reformatted_output], accumulation=outputs[accumulation_str]
            )

        # rendering accumulation outputs
        if self.prev_colormap_type == ColormapTypes.TURBO or (
            self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[reformatted_output].dtype == torch.float
        ):
            return visualization.apply_colormap(outputs[reformatted_output])

        # rendering semantic outputs
        if self.prev_colormap_type == ColormapTypes.SEMANTIC or (
            self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[reformatted_output].dtype == torch.int
        ):
            logits = outputs[reformatted_output]
            labels = torch.argmax(torch.nn.functional.softmax(logits, dim=-1), dim=-1)  # type: ignore
            assert stuff_colors is not None
            return stuff_colors[labels]

        # rendering boolean outputs
        if self.prev_colormap_type == ColormapTypes.BOOLEAN or (
            self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[reformatted_output].dtype == torch.bool
        ):
            return visualization.apply_boolean_colormap(outputs[reformatted_output])

        raise NotImplementedError

    def _send_output_to_viewer(self, outputs: Dict[str, Any], stuff_colors: torch.Tensor = None, eps=1e-6):
        """Chooses the correct output and sends it to the viewer

        Args:
            outputs: the dictionary of outputs to choose from, from the graph
            stuff_colors: is only set if colormap is for semantics. Defaults to None.
            eps: epsilon to handle floating point comparisons
        """
        if self.output_list is None:
            self.output_list = list(outputs.keys())
            viewer_output_list = list(np.copy(self.output_list))
            # remapping rgb_fine -> rgb for all cases just so that we dont have 2 of them in the options
            if OutputTypes.RGB_FINE in self.output_list:
                viewer_output_list.remove(OutputTypes.RGB_FINE)
            viewer_output_list.insert(0, OutputTypes.RGB)
            self.vis["renderingState/output_options"].write(viewer_output_list)

        reformatted_output = self._process_invalid_output(self.prev_output_type)
        # re-register colormaps and send to viewer
        if self.output_type_changed or self.prev_colormap_type == ColormapTypes.INIT:
            self.prev_colormap_type = ColormapTypes.DEFAULT
            colormap_options = [ColormapTypes.DEFAULT]
            if (
                outputs[reformatted_output].shape[-1] != 3
                and outputs[reformatted_output].dtype == torch.float
                and (torch.max(outputs[reformatted_output]) - 1.0) <= eps  # handle floating point arithmetic
            ):
                # accumulation can also include depth
                colormap_options.extend(["depth"])
            self.output_type_changed = False
            self.vis["renderingState/colormap_choice"].write(self.prev_colormap_type)
            self.vis["renderingState/colormap_options"].write(colormap_options)
        selected_output = (self._apply_colormap(outputs, stuff_colors) * 255).type(torch.uint8)
        image = selected_output.cpu().numpy()
        self.vis.set_image(image)

    def _update_viewer_stats(self, render_time: float, num_rays: int, image_height: int) -> None:
        """Function that calculates and populates all the rendering statistics accordingly

        Args:
            render_time: total time spent rendering current view
            num_rays: number of rays rendered
            image_height: resolution of the current view
        """
        writer.put_time(
            name=EventName.VIS_RAYS_PER_SEC, duration=num_rays / render_time, step=self.step, avg_over_steps=True
        )
        is_training = self.vis["renderingState/isTraining"].read()
        if is_training is None or is_training:
            # process the  current rendering fps
            eval_fps = f"{1 / render_time:.2f} fps at {image_height} res"
            self.vis["renderingState/eval_fps"].write(eval_fps)
            # process remaining training ETA
            self.vis["renderingState/train_eta"].write(GLOBAL_BUFFER["events"].get(EventName.ETA.value, "Starting"))
            # process ratio time spent on vis vs train
            if EventName.ITER_VIS_TIME.value in GLOBAL_BUFFER["events"]:
                vis_time = GLOBAL_BUFFER["events"][EventName.ITER_VIS_TIME.value]["avg"]
                train_time = GLOBAL_BUFFER["events"][EventName.ITER_TRAIN_TIME.value]["avg"]
                vis_train_ratio = f"{int(vis_time / train_time * 100)}% spent on viewer"
                self.vis["renderingState/vis_train_ratio"].write(vis_train_ratio)
            else:
                self.vis["renderingState/vis_train_ratio"].write("Starting")
        else:
            self.vis["renderingState/eval_fps"].write("Paused")
            self.vis["renderingState/train_eta"].write("Paused")
            self.vis["renderingState/vis_train_ratio"].write("100% spent on viewer")

    def _calculate_image_height(self, camera_object, is_training: bool) -> Optional[int]:
        """Calculate the maximum image height that can be rendered in the time budget

        Args:
            camera_object: the camera object to use for rendering
            is_training: whether or not we are training
        Returns:
            image_height: the maximum image height that can be rendered in the time budget
        """
        max_resolution = self.vis["renderingState/maxResolution"].read()
        if max_resolution:
            self.max_resolution = max_resolution

        if self.camera_moving or not is_training:
            target_train_util = 0
        else:
            target_train_util = self.vis["renderingState/targetTrainUtil"].read()
            if target_train_util is None:
                target_train_util = 0.9

        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"]
        else:
            return None
        if EventName.VIS_RAYS_PER_SEC.value in GLOBAL_BUFFER["events"]:
            vis_rays_per_sec = GLOBAL_BUFFER["events"][EventName.VIS_RAYS_PER_SEC.value]["avg"]
        else:
            vis_rays_per_sec = train_rays_per_sec

        current_fps = self.moving_fps if self.camera_moving else self.static_fps

        # calculate number of rays that can be rendered given the target fps
        num_vis_rays = vis_rays_per_sec / current_fps * (1 - target_train_util)

        aspect_ratio = camera_object["aspect"]

        image_height = (num_vis_rays / aspect_ratio) ** 0.5
        image_height = int(round(image_height, -1))
        image_height = min(self.max_resolution, image_height)

        return image_height

    def _process_invalid_output(self, output_type: str) -> str:
        """Check to see whether we are in the corner case of RGB; if still invalid, throw error
        Returns correct string mapping given improperly formatted output_type.

        Args:
            output_type: reformatted output type
        """
        if output_type == OutputTypes.INIT:
            output_type = OutputTypes.RGB

        # check if rgb or rgb_fine should be the case TODO: add other checks here
        attempted_output_type = output_type
        if output_type not in self.output_list and output_type == OutputTypes.RGB:
            output_type = OutputTypes.RGB_FINE

        # check if output_type is not in list
        if output_type not in self.output_list:
            assert (
                NotImplementedError
            ), f"Output {attempted_output_type} not in list. Tried to reformat as {output_type} but still not found."
        return output_type

    @profiler.time_function
    def _render_image_in_viewer(self, camera_object, graph: Model, is_training: bool) -> None:
        """
        Draw an image using the current camera pose from the viewer.
        The image is sent of a TCP connection and then uses WebRTC to send it to the viewer.

        Args:
            graph: current checkpoint of model
        """
        # Check that timestamp is newer than the last one
        if int(camera_object["timestamp"]) < self.prev_camera_timestamp:
            return

        self.prev_camera_timestamp = int(camera_object["timestamp"])

        # check and perform output type updates
        output_type = self.vis["renderingState/output_choice"].read()
        output_type = OutputTypes.INIT if output_type is None else output_type
        self.output_type_changed = self.prev_output_type != output_type
        self.prev_output_type = output_type

        # check and perform colormap type updates
        colormap_type = self.vis["renderingState/colormap_choice"].read()
        colormap_type = ColormapTypes.INIT if colormap_type is None else colormap_type
        self.prev_colormap_type = colormap_type

        # Calculate camera pose and intrinsics
        try:
            image_height = self._calculate_image_height(camera_object, is_training)
        except ZeroDivisionError as e:
            self.vis["renderingState/log_errors"].write("Error: Screen too small; no rays intersecting scene.")
            time.sleep(0.03)  # sleep to allow buffer to reset
            print(f"Error: {e}")
            return

        if image_height is None:
            return

        intrinsics_matrix, camera_to_world_h = get_intrinsics_matrix_and_camera_to_world_h(
            camera_object, image_height=image_height
        )

        camera_to_world = camera_to_world_h[:3, :]
        camera_to_world = torch.stack(
            [
                camera_to_world[0, :],
                camera_to_world[2, :],
                camera_to_world[1, :],
            ],
            dim=0,
        )

        camera = Cameras(
            fx=intrinsics_matrix[0, 0],
            fy=intrinsics_matrix[1, 1],
            cx=intrinsics_matrix[0, 2],
            cy=intrinsics_matrix[1, 2],
            camera_to_worlds=camera_to_world[None, ...],
        )
        camera = camera.to(graph.device)

        camera_ray_bundle = camera.generate_rays(camera_indices=0)

        graph.eval()

        check_thread = CheckThread(state=self)
        render_thread = RenderThread(state=self, graph=graph, camera_ray_bundle=camera_ray_bundle)

        with TimeWriter(None, None, write=False) as vis_t:
            check_thread.start()
            render_thread.start()
            try:
                render_thread.join()
                check_thread.join()
            except IOChangeException:
                pass
            except RuntimeError as e:
                del camera_ray_bundle
                torch.cuda.empty_cache()
                time.sleep(0.5)  # sleep to allow buffer to reset
                self.vis["renderingState/log_errors"].write(
                    "Error: GPU out of memory. Reduce resolution to prevent viewer from crashing."
                )
                print(f"Error: {e}")

        graph.train()
        outputs = render_thread.vis_outputs
        if outputs is not None:
            stuff_colors = graph.stuff_colors if hasattr(graph, "stuff_colors") else None
            self._send_output_to_viewer(outputs, stuff_colors=stuff_colors)
            self._update_viewer_stats(vis_t.duration, num_rays=len(camera_ray_bundle), image_height=image_height)