nerfstudio/nerfactory/viewer/server/server.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.

"""Server bridge to faciliate interactions between python backend and javascript front end"""

from __future__ import absolute_import, division, print_function

import sys
from typing import Callable, List, Optional

import msgpack
import msgpack_numpy
import tornado.gen
import tornado.ioloop
import tornado.web
import tornado.websocket
import umsgpack
import zmq
import zmq.eventloop.ioloop
from aiortc import RTCPeerConnection, RTCSessionDescription
from aiortc.rtcrtpsender import RTCRtpSender
from zmq.eventloop.zmqstream import ZMQStream

from nerfactory.viewer.server.tree import SceneTree, find_node, walk
from nerfactory.viewer.server.video_stream import SingleFrameStreamTrack

MAX_ATTEMPTS = 1000
DEFAULT_ZMQ_METHOD = "tcp"
DEFAULT_ZMQ_PORT = 6000
DEFAULT_WEBSOCKET_PORT = 8051
WEBSOCKET_COMMANDS = [
    "set_transform",
    "set_object",
    "set_output_options",
    "set_output_type",
    "set_training_state",
    "get_object",
    "set_property",
    "delete",
]
WEBRTC_COMMANDS = ["set_image"]


def find_available_port(func: Callable, default_port: int, max_attempts: int = MAX_ATTEMPTS, **kwargs) -> None:
    """Finds and attempts to connect to a port

    Args:
        func: function used on connecting to port
        default_port: the default port
        max_attempts: max number of attempts to try connection. Defaults to MAX_ATTEMPTS.
    """
    for i in range(max_attempts):
        port = default_port + i
        try:
            return func(port, **kwargs), port
        except (OSError, zmq.error.ZMQError):
            print(f"Port: {port:d} in use, trying another...", file=sys.stderr)
        except Exception as e:
            print(type(e))
            raise
    raise (
        Exception(f"Could not find an available port in the range: [{default_port:d}, {max_attempts + default_port:d})")
    )


def force_codec(pc: RTCPeerConnection, sender: RTCRtpSender, forced_codec: str) -> None:
    """Sets the codec preferences on a connection between sender and reciever

    Args:
        pc: peer connection point
        sender: sender that will send to connection point
        forced_codec: codec to set
    """
    kind = forced_codec.split("/")[0]
    codecs = RTCRtpSender.getCapabilities(kind).codecs
    transceiver = next(t for t in pc.getTransceivers() if t.sender == sender)
    transceiver.setCodecPreferences([codec for codec in codecs if codec.mimeType == forced_codec])


class WebSocketHandler(tornado.websocket.WebSocketHandler):  # pylint: disable=abstract-method
    """Tornado websocket handler for receiving and sending commands from/to the viewer."""

    def __init__(self, *args, **kwargs):
        self.bridge = kwargs.pop("bridge")
        super().__init__(*args, **kwargs)

    def check_origin(self, origin):
        """This disables CORS."""
        return True

    def open(self, *args: str, **kwargs: str):
        """open websocket bridge"""
        self.bridge.websocket_pool.add(self)
        print("opened:", self, file=sys.stderr)
        self.bridge.send_scene(self)

    async def on_message(self, message: bytearray):  # pylint: disable=invalid-overridden-method
        """On reception of message, parses the message and calls the appropriate function based on the type of command

        Args:
            message: byte message to parse
        """
        data = message
        m = umsgpack.unpackb(message)
        type_ = m["type"]
        path = list(filter(lambda x: len(x) > 0, m["path"].split("/")))

        if type_ == "set_transform":
            find_node(self.bridge.tree, path).transform = data
        elif type_ == "set_object":
            find_node(self.bridge.tree, path).object = data
            find_node(self.bridge.tree, path).properties = []
        elif type_ == "set_output_options":
            find_node(self.bridge.tree, path).object = data
        elif type_ == "set_output_type":
            find_node(self.bridge.tree, path).object = data
        elif type_ == "set_max_resolution":
            find_node(self.bridge.tree, path).object = data
        elif type_ == "set_min_resolution":
            find_node(self.bridge.tree, path).object = data
        elif type_ == "set_training_state":
            find_node(self.bridge.tree, path).object = data
        elif type_ == "offer":
            offer = RTCSessionDescription(m["data"]["sdp"], m["data"]["type"])

            pc = RTCPeerConnection()
            self.bridge.pcs.add(pc)

            video = SingleFrameStreamTrack()
            self.bridge.video_tracks.add(video)
            _ = pc.addTrack(video)
            # TODO(eventually do something with the codec)
            # video_sender = pc.addTrack(video)
            # force_codec(pc, video_sender, video_codec)

            await pc.setRemoteDescription(offer)
            answer = await pc.createAnswer()
            await pc.setLocalDescription(answer)

            cmd_data = {
                "type": "answer",
                "path": "",
                "data": {"sdp": pc.localDescription.sdp, "type": pc.localDescription.type},
            }
            data = umsgpack.packb(cmd_data)
            self.write_message(data, binary=True)

    def on_close(self):
        self.bridge.websocket_pool.remove(self)
        print("closed:", self, file=sys.stderr)


class ZMQWebSocketBridge:
    """ZMQ web socket bridge class

    Args:
        zmq_url: zmq url to connect to. Defaults to None.
        host: host of server. Defaults to "127.0.0.1".
        websocket_port: websocket port to connect to. Defaults to None.
    """

    context = zmq.Context()

    def __init__(self, zmq_url: Optional[str] = None, host: str = "127.0.0.1", websocket_port: Optional[int] = None):
        self.host = host
        self.websocket_pool = set()
        self.app = self.make_app()
        self.ioloop = tornado.ioloop.IOLoop.current()
        self.pcs = set()
        self.video_tracks = set()

        if zmq_url is None:

            def f(port):
                return self.setup_zmq(f"{DEFAULT_ZMQ_METHOD:s}://{self.host:s}:{port:d}")

            (self.zmq_socket, self.zmq_stream, self.zmq_url), _ = find_available_port(f, DEFAULT_ZMQ_PORT)
        else:
            self.zmq_socket, self.zmq_stream, self.zmq_url = self.setup_zmq(zmq_url)

        listen_kwargs = {}

        if websocket_port is None:
            _, self.websocket_port = find_available_port(self.app.listen, DEFAULT_WEBSOCKET_PORT, **listen_kwargs)
        else:
            self.app.listen(websocket_port, **listen_kwargs)
            self.websocket_port = websocket_port

        self.tree = SceneTree()

    def __str__(self) -> str:
        class_name = self.__class__.__name__
        return f"{class_name} using zmq_url={self.zmq_url} and websocket_port={self.websocket_port}"

    def make_app(self):
        """Create a tornado application for the websocket server."""
        return tornado.web.Application([(r"/", WebSocketHandler, {"bridge": self})])

    def handle_zmq(self, frames: List[bytes]):
        """Switch function that places commands in tree based on websocket command

        Args:
            frames: the list containing command + object to be placed in tree
        """
        cmd = frames[0].decode("utf-8")
        print(cmd)
        if len(frames) != 3:
            self.zmq_socket.send(b"error: expected 3 frames")
            return
        path = list(filter(lambda x: len(x) > 0, frames[1].decode("utf-8").split("/")))
        data = frames[2]
        if cmd in WEBSOCKET_COMMANDS:
            if cmd != "get_object":
                self.forward_to_websockets(frames)
            if cmd == "set_transform":
                find_node(self.tree, path).transform = data
            elif cmd == "set_object":
                find_node(self.tree, path).object = data
                find_node(self.tree, path).properties = []
            elif cmd == "set_output_options":
                find_node(self.tree, path).object = data
            elif cmd == "set_output_type":
                find_node(self.tree, path).object = data
            elif cmd == "set_max_resolution":
                find_node(self.tree, path).object = data
            elif cmd == "set_min_resolution":
                find_node(self.tree, path).object = data
            elif cmd == "set_training_state":
                find_node(self.tree, path).object = data
            elif cmd == "get_object":
                data = find_node(self.tree, path).object
                if isinstance(data, type(None)):
                    data = umsgpack.packb("error: object not found")
                self.zmq_socket.send(data)
                return
            elif cmd == "set_property":
                find_node(self.tree, path).properties.append(data)
            elif cmd == "delete":
                if len(path) > 0:
                    parent = find_node(self.tree, path[:-1])
                    child = path[-1]
                    if child in parent:
                        del parent[child]
                else:
                    self.tree = SceneTree()
        elif cmd in WEBRTC_COMMANDS:
            if cmd == "set_image":
                image = msgpack.unpackb(
                    data, object_hook=msgpack_numpy.decode, use_list=False, max_bin_len=50000000, raw=False
                )
                for video_track in self.video_tracks:
                    video_track.put_frame(image)
        else:
            self.zmq_socket.send(b"error: unknown command")
            return
        self.zmq_socket.send(b"ok")
        return

    def forward_to_websockets(self, frames: List[bytes]):
        """Forward a zmq message to all websockets.

        Args:
            frames: byte messages to be sent over
        """
        _, _, data = frames  # cmd, path, data
        for websocket in self.websocket_pool:
            websocket.write_message(data, binary=True)

    def setup_zmq(self, url: str):
        """Setup a zmq socket and connect it to the given url.

        Args:
            url: point of connection
        """
        zmq_socket = self.context.socket(zmq.REP)  # pylint: disable=no-member
        zmq_socket.bind(url)
        zmq_stream = ZMQStream(zmq_socket)
        zmq_stream.on_recv(self.handle_zmq)
        return zmq_socket, zmq_stream, url

    def send_scene(self, websocket: WebSocketHandler):
        """Sends entire tree of information over the specified websocket

        Args:
            websocket: websocket to send information over
        """
        for node in walk(self.tree):
            if node.object is not None:
                websocket.write_message(node.object, binary=True)
            for p in node.properties:
                websocket.write_message(p, binary=True)
            if node.transform is not None:
                websocket.write_message(node.transform, binary=True)

    def run(self):
        """starts and runs the websocket bridge"""
        self.ioloop.start()


def start_server_as_subprocess(zmq_url=None):
    """Starts the ZMQWebSocketBridge server as a subprocess."""
    raise NotImplementedError()