mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-12 19:04:03 +08:00
evonneng
a78394c49d
* first round adding typing and linting to server files * continued... * added new formatting conventions * committing current progress... * finished with geometry * homerun :')
314 lines
11 KiB
Python
314 lines
11 KiB
Python
# 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 tornado.gen
|
|
import tornado.ioloop
|
|
import tornado.web
|
|
import tornado.websocket
|
|
import msgpack
|
|
import msgpack_numpy
|
|
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 pyrad.viewer.server.tree import SceneTree, find_node, walk
|
|
from pyrad.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()
|