From fe7e4d5a11d47c0f9398ad6b055ce0f46ed85b4d Mon Sep 17 00:00:00 2001 From: kngwyu Date: Tue, 18 May 2021 15:44:34 +0900 Subject: [PATCH] Now we can save video by right-click --- mujoco_maze/websock_viewer.py | 35 ++++++++++++++++++++++++++++++----- 1 file changed, 30 insertions(+), 5 deletions(-) diff --git a/mujoco_maze/websock_viewer.py b/mujoco_maze/websock_viewer.py index a355b3f..647d5e1 100644 --- a/mujoco_maze/websock_viewer.py +++ b/mujoco_maze/websock_viewer.py @@ -6,6 +6,7 @@ import fastapi import uvicorn from PIL import Image + HTML = """ @@ -13,10 +14,11 @@ HTML = """ MuJoCo maze visualizer +

MuJoCo Maze Visualizer

+ Video """ @@ -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()