Remove webrtc (#1579)

* Remove webrtc

* cleanup

* lint

* cleanup

* Lower quailty

* Update viewer version

* Update docs

* remove cryptography req
This commit is contained in:
Matthew Tancik authored and GitHub committed 2023-03-10 12:00:35 -08:00
1 parent a277cd4886
commit 63af268495
16 files changed
+58 -435

No files matched your search

Binary file not shown.

After

Width:  |  Height:  |  Size: 582 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 593 KiB

@@ -4,7 +4,15 @@
The viewer is built using [ThreeJS](https://threejs.org/) and packaged into a [ReactJS](https://reactjs.org/) application. This client viewer application will connect via a websocket to a server running on your machine. The following figure helps to illustrate how our viewer framework works:
![visualize_dataset](imgs/viewer_figure.png)
```{image} imgs/viewer_light.png
:align: center
:class: only-light
```
```{image} imgs/viewer_dark.png
:align: center
:class: only-dark
```
## System design
@@ -14,12 +22,10 @@ In the center, we have the Bridge Server, which facilitates the connection betwe
#### Connection between the Bridge Server and the Client App
The connection between the Bridge Server and the Client App works with WebSockets and WebRTC.
The connection between the Bridge Server and the Client App works with WebSockets.
- **WebSocket connection** - The WebSocket is used by the Bridge Server to dispatch commands coming from the nerfstudio TCP connection. Commands can be used for drawing primitives, for setting the transform of objects, for the setting various properties, and more.
- **WebRTC connection** - We use WebRTC to stream images being rendered from nerfstudio. The websocket connection if first used to establish the WebRTC connection. Then, the Client App constantly publishes camera pose information to the Bridge Server and stores the camera information (intrinsics and extrinsics). This information is then queried from the nerfstudio code, used to render an image with some Graph, and then the image is send over the TCP connection and dispatched via WebRTC to render the stream of images.
## Installing and running locally
#### Running the Client App
@@ -50,7 +56,7 @@ yarn install
yarn start
```
The local webserver runs on port 4000 by default,
The local webserver runs on port 4000 by default,
so when `ns-train` is running, you can connect to the viewer locally at
[http://localhost:4000/?websocket_url=ws://localhost:7007](http://localhost:4000/?websocket_url=ws://localhost:7007)
@@ -70,7 +76,6 @@ We thank the authors and contributors to the following repos, which we've starte
- [meshcat](https://github.com/rdeits/meshcat) - made by [Robin Deits](https://github.com/rdeits)
- [ThreeJS](https://threejs.org/)
- [ReactJS](https://reactjs.org/)
- [WebRTC](https://webrtc.org/) - WebRTC is a framework for real-time communication that allows two peers to send video, audio, or general data to and from each other with low latency. We've adopted WebRTC to stream rendered images to our viewer.
## FAQ
-9
View File
@@ -23,8 +23,6 @@ from dataclasses import dataclass
from pathlib import Path
from typing import Any, List, Optional, Tuple, Type
from typing_extensions import Literal
# model instances
from nerfstudio.utils import writer
@@ -148,10 +146,3 @@ class ViewerConfig(PrintableConfig):
actually used in training/evaluation. If -1, display all."""
quit_on_train_completion: bool = False
"""Whether to kill the training job when it has completed. Note this will stop rendering in the viewer."""
skip_openrelay: bool = False
"""Avoid using openrelay to communicate with the viewer. Try disabling if you have trouble
connecting to the viewer"""
codec: Literal["H264", "VP8"] = "VP8"
"""Video codec that viewer will use."""
local: bool = False
"""If running local server instance, avoid using relays to communicate with the viewer."""
+1 -1
View File
@@ -1,7 +1,7 @@
{
"name": "viewer",
"homepage": ".",
"version": "23-02-3-0",
"version": "23-03-9-0",
"private": true,
"dependencies": {
"@emotion/react": "^11.10.4",
+4 -10
View File
@@ -132,8 +132,8 @@ $sidebar-width: 100% !default;
padding: 10px;
}
.StatusPanel .StatusPanel-button-group {
margin: 10px;
.StatusPanel .StatusPanel-button {
margin-bottom: 4px;
}
.StatusPanel .StatusPanel-hide-scene-button {
@@ -348,7 +348,7 @@ $sidebar-width: 100% !default;
// Viewer Window
.canvas-container-main,
.WebRTCVideo,
.RenderWindow,
.ViewerWindow-render-crop-container {
position: absolute;
top: 0px;
@@ -358,7 +358,7 @@ $sidebar-width: 100% !default;
border: $custom-border;
}
.WebRTCVideo {
.RenderWindow {
z-index: 1;
}
@@ -373,12 +373,6 @@ $sidebar-width: 100% !default;
z-index: 3;
}
.WebRTCVideo-video {
width: 100%;
height: 100%;
object-fit: fill;
}
.ViewerWindow-render-crop {
display: flex;
box-shadow: rgb(0 0 0 / 70%) 0px 0px 0px 1600px;
@@ -0,0 +1,12 @@
import * as React from 'react';
import { useSelector } from 'react-redux';
export default function RenderWindow() {
const render_img = useSelector((state) => state.render_img);
return (
<div className="RenderWindow">
<img src={render_img} width="100%" height="100%" alt="Render window" />
</div>
);
}
@@ -1,7 +1,6 @@
import * as React from 'react';
import Button from '@mui/material/Button';
import { ButtonGroup } from '@mui/material';
import { useDispatch, useSelector } from 'react-redux';
import PauseIcon from '@mui/icons-material/Pause';
import PlayArrowIcon from '@mui/icons-material/PlayArrow';
@@ -25,9 +24,6 @@ export default function StatusPanel(props: StatusPanelProps) {
const isWebsocketConnected = useSelector(
(state) => state.websocketState.isConnected,
);
const isWebrtcConnected = useSelector(
(state) => state.webrtcState.isConnected,
);
const eval_res = useSelector((state) => state.renderingState.eval_res);
const vis_train_ratio = useSelector(
(state) => state.renderingState.vis_train_ratio,
@@ -78,13 +74,9 @@ export default function StatusPanel(props: StatusPanelProps) {
const training_icon = isTraining ? <PauseIcon /> : <PlayArrowIcon />;
const websocket_connected_text = isWebsocketConnected
? 'Server Connected'
: 'Server Disconnected';
const webrtc_connected_text = isWebrtcConnected
? 'Render Connected'
: 'Render Disconnected';
? 'Renderer Connected'
: 'Renderer Disconnected';
const websocket_connected_color = isWebsocketConnected ? 'success' : 'error';
const webrtc_connected_color = isWebrtcConnected ? 'success' : 'error';
return (
<div className="StatusPanel">
@@ -143,26 +135,13 @@ export default function StatusPanel(props: StatusPanelProps) {
<b>Time Allocation:</b> {vis_train_ratio}
</div>
</div>
<ButtonGroup
className="StatusPanel-button-group"
variant="text"
aria-label="text button group"
<Button
className="StatusPanel-button"
color={websocket_connected_color}
style={{ textTransform: 'none' }}
>
<Button
className="StatusPanel-button"
color={websocket_connected_color}
style={{ textTransform: 'none' }}
>
{websocket_connected_text}
</Button>
<Button
className="StatusPanel-button"
color={webrtc_connected_color}
style={{ textTransform: 'none' }}
>
{webrtc_connected_text}
</Button>
</ButtonGroup>
{websocket_connected_text}
</Button>
</div>
);
}
@@ -11,8 +11,8 @@ import PublicSharpIcon from '@mui/icons-material/PublicSharp';
import SyncOutlinedIcon from '@mui/icons-material/SyncOutlined';
import ThreeDRotationIcon from '@mui/icons-material/ThreeDRotation';
import VideoCameraBackIcon from '@mui/icons-material/VideoCameraBackOutlined';
import WebRtcWindow from '../WebRtcWindow/WebRtcWindow';
import { WebSocketContext } from '../WebSocket/WebSocket';
import RenderWindow from '../RenderWindow/RenderWindow';
const msgpack = require('msgpack-lite');
@@ -100,7 +100,7 @@ function TransformIcons(props) {
);
}
// manages a camera and the web rtc stream...
// manages a camera
export default function ViewerWindow(props) {
const sceneTree = props.sceneTree;
const scene = sceneTree.object;
@@ -274,8 +274,7 @@ export default function ViewerWindow(props) {
return (
<>
{/* the webrtc viewer needs to know the camera pose */}
<WebRtcWindow />
<RenderWindow />
<div className="canvas-container-main" ref={myRef}>
<div className="ViewerWindow-camera-toggle">
<CameraToggle />
@@ -1,186 +0,0 @@
import React, { createContext, useContext, useEffect, useRef } from 'react';
import { useDispatch } from 'react-redux';
import { WebSocketContext } from '../WebSocket/WebSocket';
const WebRtcContext = createContext(null);
const msgpack = require('msgpack-lite');
export { WebRtcContext };
export default function WebRtcWindow() {
const websocket = useContext(WebSocketContext).socket;
const pcRef = useRef(null);
const localVideoRef = useRef(null);
const dispatch = useDispatch();
const getRTCPeerConnection = () => {
const pc = new RTCPeerConnection({
iceServers: [
{ urls: 'stun:stun.l.google.com:19302' },
{
urls: 'stun:openrelay.metered.ca:80',
},
{
urls: 'turn:openrelay.metered.ca:80',
username: 'openrelayproject',
credential: 'openrelayproject',
},
{
urls: 'turn:openrelay.metered.ca:443',
username: 'openrelayproject',
credential: 'openrelayproject',
},
{
urls: 'turn:openrelay.metered.ca:443?transport=tcp',
username: 'openrelayproject',
credential: 'openrelayproject',
},
],
});
// connect video
pc.addEventListener('track', (evt) => {
if (evt.track.kind === 'video') {
[localVideoRef.current.srcObject] = evt.streams; // uses array destructuring
}
});
pc.addTransceiver('video', { direction: 'recvonly' });
// for updating the status of the peer connection
pc.oniceconnectionstatechange = () => {
// https://developer.mozilla.org/en-US/docs/Web/API/RTCPeerConnection/connectionState
console.log(`[webrtc] connectionState: ${pc.connectionState}`);
if (
pc.connectionState === 'connecting' ||
pc.connectionState === 'connected'
) {
console.log('[webrtc] connected');
dispatch({
type: 'write',
path: 'webrtcState/isConnected',
data: true,
});
} else {
dispatch({
type: 'write',
path: 'webrtcState/isConnected',
data: false,
});
}
};
pc.onclose = () => {
dispatch({
type: 'write',
path: 'webrtcState/isConnected',
data: false,
});
};
return pc;
};
const sendOffer = () => {
pcRef.current
.createOffer()
.then((offer) => {
console.log('[webrtc] created offer');
console.log(offer);
return pcRef.current.setLocalDescription(offer);
})
.then(() => {
// wait for ICE gathering to complete
console.log('[webrtc] set local description');
return new Promise((resolve) => {
if (pcRef.current.iceGatheringState === 'complete') {
console.log('[webrtc] ICE gathering complete');
resolve();
} else {
const checkState = () => {
console.log(
`[webrtc] iceGatheringState: ${pcRef.current.iceGatheringState}`,
);
if (pcRef.current.iceGatheringState === 'complete') {
pcRef.current.removeEventListener(
'icegatheringstatechange',
checkState,
);
resolve();
}
};
console.log(
'[webrtc] adding listener for `icegatheringstatechange`',
);
pcRef.current.addEventListener(
'icegatheringstatechange',
checkState,
);
}
});
})
.then(() => {
// send the offer
console.log('[webrtc] sending offer');
const offer = pcRef.current.localDescription;
const cmd = 'write';
const path = 'webrtc/offer';
const data = {
type: cmd,
path,
data: {
sdp: offer.sdp,
type: offer.type,
},
};
const message = msgpack.encode(data);
websocket.send(message);
});
};
useEffect(() => {
websocket.addEventListener('message', (originalCmd) => {
// set the remote description when the offer is received
const cmd = msgpack.decode(new Uint8Array(originalCmd.data));
if (cmd.path === '/webrtc/answer') {
console.log('[webrtc] received answer');
const answer = cmd.data;
console.log(answer);
if (answer !== null) {
pcRef.current.setRemoteDescription(answer);
}
}
});
websocket.addEventListener('open', () => {
// connect webrtc when the websocket is connected
pcRef.current = getRTCPeerConnection();
console.log('[webrtc] starting process');
sendOffer();
});
// kill the webrtc connection on dismount
return () => {
if (pcRef.current !== null) {
dispatch({
type: 'write',
path: 'webrtcState/isConnected',
data: false,
});
pcRef.current.close();
}
};
}, [websocket]); // dependency to call this whenever the websocket changes
return (
<div className="WebRTCVideo">
<video
className="WebRTCVideo-video"
autoPlay
playsInline
muted
ref={localVideoRef}
/>
</div>
);
}
+4 -7
View File
@@ -6,13 +6,13 @@ const initialState = {
isConnected: false,
websocket_url: 'ws://localhost:7007',
},
// the webrtc connection state
webrtcState: {
isConnected: false,
},
// for sending actual commands to the client
camera_path_payload: null,
populate_paths_payload: false,
render_img: null, // The rendered images
// the rendering state
renderingState: {
// cameras
@@ -65,9 +65,6 @@ const initialState = {
sceneBox: null,
cameras: null,
},
// ------------------------------------------------------------------------
// which of the state keys to synchronize with the bridge server
sync_keys: ['renderingState', 'sceneState'],
};
function setData(newState, state, path, data) {
+4 -8
View File
@@ -64,14 +64,10 @@ class WebSocketHandler(tornado.websocket.WebSocketHandler): # pylint: disable=a
if type_ == "write":
# writes the data coming from the websocket
find_node(self.bridge.state_tree, path).data = m["data"]
if m["path"] != "renderingState/camera" and m["path"] != "webrtc/offer":
# dispatch to the websockets if not a camera update!
# TODO(ethan): handle the camera properly!
# TODO: but don't update the current websocket
command = {"type": "write", "path": m["path"], "data": m["data"]}
packed_data = umsgpack.packb(command)
frames = ["write".encode("utf-8"), m["path"].encode("utf-8"), packed_data]
self.bridge.forward_to_websockets(frames, websocket_to_skip=self)
command = {"type": "write", "path": m["path"], "data": m["data"]}
packed_data = umsgpack.packb(command)
frames = ["write".encode("utf-8"), m["path"].encode("utf-8"), packed_data]
self.bridge.forward_to_websockets(frames, websocket_to_skip=self)
elif type_ == "read":
# reads and returns the data
data = find_node(self.bridge.state_tree, path).data
-16
View File
@@ -21,8 +21,6 @@ from typing import Any, Callable, Dict, List, Optional, Tuple
import numpy as np
import torch
import zmq
from aiortc import RTCPeerConnection
from aiortc.rtcrtpsender import RTCRtpSender
def get_chunks(
@@ -127,17 +125,3 @@ def find_available_port(func: Callable, default_port: int, max_attempts: int = 1
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 receiver
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])
-46
View File
@@ -1,46 +0,0 @@
# 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.
"""Video Stream objects for WebRTC"""
import numpy as np
from aiortc import VideoStreamTrack
from av import VideoFrame
class SingleFrameStreamTrack(VideoStreamTrack):
"""Single Frame stream class: pushes single frames to a stream"""
def __init__(self):
super().__init__()
self.background_frame = np.ones((480, 640, 3), dtype="uint8") * 100 # gray background
self.frame = None
self.put_frame(self.background_frame)
def put_frame(self, frame: np.ndarray) -> None:
"""Sets the current viewing frame
Args:
frame: image to be viewed
"""
self.frame = VideoFrame.from_ndarray(frame)
async def recv(self):
"""Async method to grab and wait on frame"""
pts, time_base = await self.next_timestamp()
frame = self.frame
frame.pts = pts
frame.time_base = time_base
return frame
+10 -110
View File
@@ -18,7 +18,7 @@
"""
from __future__ import annotations
import asyncio
import base64
import enum
import os
import sys
@@ -28,14 +28,9 @@ import warnings
from pathlib import Path
from typing import Any, Dict, Optional, Tuple
import cv2
import numpy as np
import torch
from aiortc import (
RTCConfiguration,
RTCIceServer,
RTCPeerConnection,
RTCSessionDescription,
)
from cryptography.utils import CryptographyDeprecationWarning
from rich.console import Console
@@ -51,11 +46,7 @@ from nerfstudio.utils.decorators import check_main_thread, decorate_all
from nerfstudio.utils.io import load_from_json, write_to_json
from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter
from nerfstudio.viewer.server.subprocess import run_viewer_bridge_server_as_subprocess
from nerfstudio.viewer.server.utils import (
force_codec,
get_intrinsics_matrix_and_camera_to_world_h,
)
from nerfstudio.viewer.server.video_stream import SingleFrameStreamTrack
from nerfstudio.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h
from nerfstudio.viewer.server.visualizer import Viewer
warnings.filterwarnings("ignore", category=CryptographyDeprecationWarning)
@@ -283,11 +274,9 @@ class ViewerState:
CONSOLE.rule(characters="=")
CONSOLE.line()
self.vis = Viewer(zmq_port=zmq_port, ip_address=self.config.ip_address)
self.vis_webrtc_thread = Viewer(zmq_port=zmq_port, ip_address=self.config.ip_address)
else:
assert self.config.zmq_port is not None
self.vis = Viewer(zmq_port=self.config.zmq_port, ip_address=self.config.ip_address)
self.vis_webrtc_thread = Viewer(zmq_port=self.config.zmq_port, ip_address=self.config.ip_address)
# viewer specific variables
self.prev_camera_matrix = None
@@ -314,12 +303,6 @@ class ViewerState:
self.output_list = None
# webrtc
self.pcs = set()
self.video_tracks = set()
self.webrtc_thread = None
self.kill_webrtc_signal = False
def _pick_drawn_image_idxs(self, total_num: int) -> list[int]:
"""Determine indicies of images to display in viewer.
@@ -417,26 +400,6 @@ class ViewerState:
self.vis["renderingState/all_camera_paths"].write(all_path_dict)
self.vis["populate_paths_payload"].delete()
def _check_webrtc_offer(self):
"""Check if there is a webrtc offer to respond to."""
data = self.vis["webrtc/offer"].read()
if data:
if self.webrtc_thread and self.webrtc_thread.is_alive():
# kill the previous thread if the webpage refreshes
self.kill_webrtc_signal = True
return
def loop_in_thread(loop):
asyncio.set_event_loop(loop)
loop.run_until_complete(self.send_webrtc_answer(data))
loop = asyncio.get_event_loop()
self.webrtc_thread = threading.Thread(target=loop_in_thread, args=(loop,))
self.webrtc_thread.daemon = True
self.webrtc_thread.start()
# remove the offer from the state tree
self.vis["webrtc/offer"].delete()
def _update_render_aabb(self, graph):
"""
update the render aabb box for the viewer:
@@ -500,7 +463,6 @@ class ViewerState:
self._check_camera_path_payload(trainer, step)
self._check_populate_paths_payload(trainer, step)
self._check_webrtc_offer()
camera_object = self._get_camera_object()
if camera_object is None:
@@ -546,7 +508,6 @@ class ViewerState:
is_training = self.vis["renderingState/isTraining"].read()
self._check_populate_paths_payload(trainer, step)
self._check_camera_path_payload(trainer, step)
self._check_webrtc_offer()
run_loop = not is_training
local_step += 1
@@ -666,66 +627,6 @@ class ViewerState:
raise NotImplementedError
async def send_webrtc_answer(self, data):
"""Setup the webrtc connection."""
# returns the description to for WebRTC to the specific websocket connection
offer = RTCSessionDescription(data["sdp"], data["type"])
if self.config.local:
pc = RTCPeerConnection()
else:
if self.config.skip_openrelay:
ice_servers = [
RTCIceServer(urls="stun:stun.l.google.com:19302"),
]
else:
ice_servers = [
RTCIceServer(urls="stun:stun.l.google.com:19302"),
RTCIceServer(urls="stun:openrelay.metered.ca:80"),
RTCIceServer(
urls="turn:openrelay.metered.ca:80", username="openrelayproject", credential="openrelayproject"
),
RTCIceServer(
urls="turn:openrelay.metered.ca:443", username="openrelayproject", credential="openrelayproject"
),
RTCIceServer(
urls="turn:openrelay.metered.ca:443?transport=tcp",
username="openrelayproject",
credential="openrelayproject",
),
]
pc = RTCPeerConnection(configuration=RTCConfiguration(iceServers=ice_servers))
self.pcs.add(pc)
video = SingleFrameStreamTrack()
self.video_tracks.add(video)
video_sender = pc.addTrack(video)
print(f"viewer using video/{self.config.codec}")
force_codec(pc, video_sender, f"video/{self.config.codec}")
await pc.setRemoteDescription(offer)
answer = await pc.createAnswer()
await pc.setLocalDescription(answer)
self.vis_webrtc_thread["webrtc/answer"].write(
{"sdp": pc.localDescription.sdp, "type": pc.localDescription.type}
)
self.vis_webrtc_thread["webrtc/answer"].delete()
# continually exchange media
while True:
await asyncio.sleep(1)
if self.kill_webrtc_signal:
self.kill_webrtc_signal = False
return
def set_image(self, image):
"""Write the image over webrtc."""
for video_track in self.video_tracks:
video_track.put_frame(image)
def _send_output_to_viewer(self, outputs: Dict[str, Any], colors: torch.Tensor = None):
"""Chooses the correct output and sends it to the viewer
@@ -760,8 +661,12 @@ class ViewerState:
self.vis["renderingState/colormap_choice"].write(self.prev_colormap_type)
self.vis["renderingState/colormap_options"].write(colormap_options)
selected_output = (self._apply_colormap(outputs, colors) * 255).type(torch.uint8)
image = selected_output.cpu().numpy()
self.set_image(image)
image = selected_output[..., [2, 1, 0]].cpu().numpy()
data = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, 75])[1].tobytes()
data = str("data:image/jpeg;base64," + base64.b64encode(data).decode("ascii"))
self.vis["render_img"].write(data)
def _update_viewer_stats(self, render_time: float, num_rays: int, image_height: int, image_width: int) -> None:
"""Function that calculates and populates all the rendering statistics accordingly
@@ -846,11 +751,6 @@ class ViewerState:
if image_width > self.max_resolution:
image_width = self.max_resolution
image_height = int(image_width / aspect_ratio)
if self.config.codec != "VP8":
# force even values to allow hardware encoder usage
quantize = 2
image_width = int(image_width / quantize) * quantize
image_height = int(image_height / quantize) * quantize
return image_height, image_width
def _process_invalid_output(self, output_type: str) -> str:
@@ -880,7 +780,7 @@ class ViewerState:
# pylint: disable=too-many-statements
"""
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.
The image is sent over a TCP connection.
Args:
graph: current checkpoint of model
+1 -3
View File
@@ -14,11 +14,9 @@ classifiers = [
"Programming Language :: Python",
]
dependencies = [
"aiohttp>=3.8.1",
"aiortc>=1.3.2",
"appdirs>=1.4",
"av>=9.2.0",
"cryptography>=38,<39",
"cryptography>=38",
"tyro>=0.3.31",
"gdown>=4.6.0",
"ninja>=1.10",