dodox/mujoco_maze
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Now we can save video by right-click

Browse Source
This commit is contained in:
kngwyu 2021-05-18 15:44:34 +09:00
parent 0026885730
commit fe7e4d5a11
1 changed files with 30 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
mujoco_maze/websock_viewer.py
View File

@ -6,6 +6,7 @@ import fastapi
import uvicorn
from PIL import Image
HTML = """
<!DOCTYPE html>
<html>
@ -13,10 +14,11 @@ HTML = """
<title>MuJoCo maze visualizer</title>
</head>
<body>
<h2>MuJoCo Maze Visualizer</h2>
<script>
var web_socket = new WebSocket('ws://127.0.0.1:{{port}}/ws');
web_socket.binaryType = "arraybuffer";
web_socket.onmessage = function(event) {
var ws_image = new WebSocket('ws://127.0.0.1:{{port}}/ws');
ws_image.binaryType = "arraybuffer";
ws_image.onmessage = function(event) {
var canvas = document.getElementById('canvas');
var ctx = canvas.getContext('2d');
var blob = new Blob([event.data], {type:'image/png'});
@ -28,10 +30,20 @@ HTML = """
console.log(url);
image.src = url;
}
function saveVideo() {
var xhr = new XMLHttpRequest();
xhr.open('GET', 'https://via.placeholder.com/150', true);
xhr.responseType = 'blob';
xhr.onerror = err => {
alert('Video is not ready');
};
xhr.send();
}
</script>
<div>
<canvas id="canvas" width="600" height="480"></canvas>
</div>
<a href="video">Video</a>
</body>
</html>
"""
@ -42,6 +54,7 @@ class _ServerWorker(mp.Process):
super().__init__()
self.pipe = pipe
self.port = port
self.video_frames = []
def _run_server(self) -> None:
@ -49,19 +62,20 @@ class _ServerWorker(mp.Process):
html = HTML.replace("{{port}}", str(self.port))
@app.get("/")
async def get():
async def root():
return fastapi.responses.HTMLResponse(html)
server = None
@app.websocket("/ws")
async def ws_send_image(websocket: fastapi.WebSocket):
async def ws(websocket: fastapi.WebSocket):
await websocket.accept()
loop = asyncio.get_running_loop()
while True:
image_array = await loop.run_in_executor(None, self.pipe.recv)
if image_array is None:
break
self.video_frames.append(image_array)
image = Image.fromarray(image_array)
with io.BytesIO() as stream:
image.save(stream, format="png")
@ -70,6 +84,17 @@ class _ServerWorker(mp.Process):
await websocket.close()
server.should_exit = True
@app.get("/video")
async def video():
import imageio
writer = imageio.get_writer("/tmp/mujoco-maze-video.mp4")
for frame in self.video_frames:
writer.append_data(frame)
writer.close()
video = open("/tmp/mujoco-maze-video.mp4", mode="rb")
return fastapi.responses.StreamingResponse(video, media_type="video/mp4")
config = uvicorn.Config(app, port=self.port)
server = uvicorn.Server(config)
server.run()