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