Now we can save video by right-click
This commit is contained in:
parent
0026885730
commit
fe7e4d5a11
@ -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()
|
||||||
|
Loading…
Reference in New Issue
Block a user