-
Notifications
You must be signed in to change notification settings - Fork 26
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[jiminy_py] Fix support of capture_frame on Windows and Jupyter. (#170)
* Enable capture frame in Jupyter notebook using background process and IPC. * Merge video recording options in a single one. * Do not throw error but warning if 'wait' is not available. Add option to disable viewer auto open gui. Bug fixes. * Do not kill meshcat server and recorder and SIGINT is raised. Fix capture frame not waiting for mesh loading. * Move meshcat server and recorder process methods in dedicated file to fix 'spawn' support (Win10). * Patch pyppeteer to properly ignore SIGINT. Make sure multiproc manager are properly closed. * Improve default meshcat camera pose. Fix viewer open_gui if already existing meshcat server. Improve running meshcat server detection. Co-authored-by: Alexis Duburcq <[email protected]>
- Loading branch information
Showing
7 changed files
with
392 additions
and
244 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Empty file.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,140 @@ | ||
| import os | ||
| import sys | ||
| import signal | ||
| import atexit | ||
| import asyncio | ||
| import subprocess | ||
| import multiprocessing | ||
| from ctypes import c_char_p, c_bool, c_int | ||
| from contextlib import redirect_stderr | ||
|
|
||
| from pyppeteer.connection import Connection | ||
| from pyppeteer.browser import Browser | ||
| from pyppeteer.launcher import Launcher, get_ws_endpoint | ||
| from requests_html import HTMLSession | ||
|
|
||
| # ================ Monkey-patch ======================= | ||
|
|
||
| # Make sure raise SIGINT does not kill chrome | ||
| # pyppeteer browser backend automatically, so that | ||
| # it allows a closing handle to be manually registered. | ||
| async def launch(self) -> Browser: | ||
| """Start chrome process and return `Browser` object.""" | ||
| self.chromeClosed = False | ||
| self.connection = None | ||
|
|
||
| options = dict() | ||
| options['env'] = self.env | ||
| if not self.dumpio: | ||
| options['stdout'] = subprocess.PIPE | ||
| options['stderr'] = subprocess.STDOUT | ||
| if sys.platform.startswith('win'): | ||
| startupflags = subprocess.DETACHED_PROCESS | \ | ||
| subprocess.CREATE_NEW_PROCESS_GROUP | ||
| self.proc = subprocess.Popen( | ||
| self.cmd, **options, creationflags=startupflags, shell=False) | ||
| else: | ||
| self.proc = subprocess.Popen( | ||
| self.cmd, **options, preexec_fn=os.setpgrp, shell=False) | ||
|
|
||
| # don't forget to close browser process | ||
| def _close_process(*args, **kwargs) -> None: | ||
| if not self.chromeClosed: | ||
| self._loop.run_until_complete(self.killChrome()) | ||
| atexit.register(_close_process) | ||
| if self.handleSIGTERM: | ||
| signal.signal(signal.SIGTERM, _close_process) | ||
| if not sys.platform.startswith('win'): | ||
| if self.handleSIGHUP: | ||
| signal.signal(signal.SIGHUP, _close_process) | ||
|
|
||
| self.browserWSEndpoint = get_ws_endpoint(self.url) | ||
| self.connection = Connection(self.browserWSEndpoint, self._loop) | ||
| browser = await Browser.create( | ||
| self.connection, | ||
| [], | ||
| self.ignoreHTTPSErrors, | ||
| self.defaultViewport, | ||
| self.proc, | ||
| self.killChrome) | ||
| await self.ensureInitialPage(browser) | ||
| return browser | ||
| Launcher.launch = launch | ||
|
|
||
| # ====================================================== | ||
|
|
||
| def capture_frame(client, width, height): | ||
| """ | ||
| @brief Send a javascript command to the hidden browser to | ||
| capture frame, then wait for it (since it is async). | ||
| """ | ||
| async def capture_frame_async(client): | ||
| nonlocal width, height | ||
| _width = client.html.page.viewport['width'] | ||
| _height = client.html.page.viewport['height'] | ||
| if not width > 0: | ||
| width = _width | ||
| if not height > 0: | ||
| height = _height | ||
| if _width != width or _height != height: | ||
| await client.html.page.setViewport( | ||
| {'width': width, 'height': height}) | ||
| return await client.html.page.evaluate(""" | ||
| () => { | ||
| return viewer.capture_image(); | ||
| } | ||
| """) | ||
| loop = asyncio.get_event_loop() | ||
| img_data = loop.run_until_complete(capture_frame_async(client)) | ||
| return img_data | ||
|
|
||
| def meshcat_recorder(meshcat_url, | ||
| take_snapshot_shm, | ||
| img_data_html_shm, | ||
| width_shm, | ||
| height_shm): | ||
| # Do not catch signal interrupt automatically, to avoid | ||
| # killing meshcat server and stopping Jupyter notebook cell. | ||
| signal.signal(signal.SIGINT, signal.SIG_IGN) | ||
|
|
||
| session = HTMLSession() | ||
| client = session.get(meshcat_url) | ||
| client.html.render(keep_page=True) | ||
|
|
||
| with open(os.devnull, 'w') as f: | ||
| with redirect_stderr(f): | ||
| while True: | ||
| if take_snapshot_shm.value: | ||
| img_data_html_shm.value = capture_frame( | ||
| client, width_shm.value, height_shm.value) | ||
| take_snapshot_shm.value = False | ||
|
|
||
| def mgr_init(): | ||
| signal.signal(signal.SIGINT, signal.SIG_IGN) | ||
|
|
||
| def start_meshcat_recorder(meshcat_url): | ||
| """ | ||
| @brief Run meshcat server in background using multiprocessing | ||
| Process to enable parallel asyncio loop execution, which | ||
| is necessary to support recording in Jupyter notebook. | ||
| """ | ||
| manager = multiprocessing.managers.SyncManager() | ||
| manager.start(mgr_init) | ||
| recorder_shm = { | ||
| 'take_snapshot': manager.Value(c_bool, False), | ||
| 'img_data_html': manager.Value(c_char_p, ""), | ||
| 'width': manager.Value(c_int, -1), | ||
| 'height': manager.Value(c_int, -1) | ||
| } | ||
|
|
||
| recorder = multiprocessing.Process( | ||
| target=meshcat_recorder, | ||
| args=(meshcat_url, | ||
| recorder_shm['take_snapshot'], | ||
| recorder_shm['img_data_html'], | ||
| recorder_shm['width'], | ||
| recorder_shm['height']), | ||
| daemon=True) | ||
| recorder.start() | ||
|
|
||
| return recorder, manager, recorder_shm |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,155 @@ | ||
| import os | ||
| import psutil | ||
| import signal | ||
| import asyncio | ||
| import umsgpack | ||
| import tornado.web | ||
| import multiprocessing | ||
| from contextlib import redirect_stderr | ||
|
|
||
| from meshcat.servers.zmqserver import ( | ||
| VIEWER_ROOT, ZMQWebSocketBridge, WebSocketHandler) | ||
|
|
||
| # ================ Monkey-patch ======================= | ||
|
|
||
| # Add support of cross-origin connection. | ||
| # It is useful to execute custom javascript commands within | ||
| # a Jupyter Notebook, and it is not an actual security flaw | ||
| # for local servers since they are not accessible from the | ||
| # outside anyway. | ||
| WebSocketHandler.check_origin = lambda self, origin: True | ||
|
|
||
| # Override the default html page to disable auto-update of | ||
| # three js "controls" of the camera, so that it can be moved | ||
| # programmatically in any position, without any constraint, as | ||
| # long as the user is not moving it manually using the mouse. | ||
| class MyFileHandler(tornado.web.StaticFileHandler): | ||
| def initialize(self, default_path, default_filename, fallback_path): | ||
| self.default_path = os.path.abspath(default_path) | ||
| self.default_filename = default_filename | ||
| self.fallback_path = os.path.abspath(fallback_path) | ||
| super().initialize(self.default_path, self.default_filename) | ||
|
|
||
| def set_extra_headers(self, path): | ||
| self.set_header('Cache-Control', | ||
| 'no-store, no-cache, must-revalidate, max-age=0') | ||
|
|
||
| def validate_absolute_path(self, root, absolute_path): | ||
| if os.path.isdir(absolute_path): | ||
| absolute_path = os.path.join(absolute_path, self.default_filename) | ||
| return self.validate_absolute_path(root, absolute_path) | ||
| if os.path.exists(absolute_path) and \ | ||
| os.path.basename(absolute_path) != self.default_filename: | ||
| return super().validate_absolute_path(root, absolute_path) | ||
| return os.path.join(self.fallback_path, absolute_path[(len(root)+1):]) | ||
|
|
||
| def make_app(self): | ||
| return tornado.web.Application([ | ||
| (r"/static/?(.*)", MyFileHandler, { | ||
| "default_path": VIEWER_ROOT, | ||
| "fallback_path": os.path.dirname(__file__), | ||
| "default_filename": "index.html"}), | ||
| (r"/", WebSocketHandler, {"bridge": self}) | ||
| ]) | ||
| ZMQWebSocketBridge.make_app = make_app | ||
|
|
||
| # Implement bidirectional communication because zmq and the | ||
| # websockets by gathering and forward messages received from | ||
| # the websockets to zmq. Note that there is currently no way | ||
| # to identify the client associated to each reply, but it is | ||
| # usually not a big deal, since the same answers is usual | ||
| # expected from each of them. Comma is used as a delimiter. | ||
| # | ||
| # It also fixes flushing issue when 'handle_zmq' is not directly | ||
| # responsible for sending a message through the zmq socket. | ||
| def handle_web(self, message): | ||
| self.bridge.websocket_messages.append(message) | ||
| if len(self.bridge.websocket_messages) == len(self.bridge.websocket_pool): | ||
| gathered_msg = ",".join(self.bridge.websocket_messages) | ||
| self.bridge.zmq_socket.send(gathered_msg.encode("utf-8")) | ||
| self.bridge.zmq_stream.flush() | ||
| WebSocketHandler.on_message = handle_web | ||
|
|
||
| def wait_for_websockets(self): | ||
| if len(self.websocket_pool) > 0: | ||
| self.zmq_socket.send(b"ok") | ||
| self.zmq_stream.flush() | ||
| else: | ||
| self.ioloop.call_later(0.1, self.wait_for_websockets) | ||
| ZMQWebSocketBridge.wait_for_websockets = wait_for_websockets | ||
|
|
||
| handle_zmq_orig = ZMQWebSocketBridge.handle_zmq | ||
| def handle_zmq(self, frames): | ||
| self.websocket_messages = [] # Used to gather websocket messages | ||
| cmd = frames[0].decode("utf-8") | ||
| if cmd == "meshes_loaded": | ||
| if not self.websocket_pool: | ||
| self.zmq_socket.send("".encode("utf-8")) | ||
| for websocket in self.websocket_pool: | ||
| websocket.write_message(umsgpack.packb({ | ||
| u"type": u"meshes_loaded" | ||
| }), binary=True) | ||
| else: | ||
| handle_zmq_orig(self, frames) | ||
| ZMQWebSocketBridge.handle_zmq = handle_zmq | ||
|
|
||
| # ====================================================== | ||
|
|
||
| def meshcat_server(info): | ||
| """ | ||
| @brief Meshcat server deamon, using in/out argument to get the | ||
| zmq url instead of reading stdout as it was. | ||
| """ | ||
| # Do not catch signal interrupt automatically, to avoid | ||
| # killing meshcat server and stopping Jupyter notebook cell. | ||
| signal.signal(signal.SIGINT, signal.SIG_IGN) | ||
|
|
||
| # Create asyncio event loop if not already existing | ||
| asyncio.get_event_loop() | ||
|
|
||
| with open(os.devnull, 'w') as f: | ||
| with redirect_stderr(f): | ||
| bridge = ZMQWebSocketBridge() | ||
| info['zmq_url'] = bridge.zmq_url | ||
| info['web_url'] = bridge.web_url | ||
| bridge.run() | ||
|
|
||
| def start_meshcat_server(): | ||
| """ | ||
| @brief Run meshcat server in background using multiprocessing | ||
| Process to enable monkey patching and proper interprocess | ||
| communication through a manager. | ||
| """ | ||
| manager = multiprocessing.Manager() | ||
| info = manager.dict() | ||
| server = multiprocessing.Process( | ||
| target=meshcat_server, args=(info,), daemon=True) | ||
| server.start() | ||
|
|
||
| # Wait for the process to finish initialization | ||
| while not info: | ||
| pass | ||
| zmq_url, web_url = info['zmq_url'], info['web_url'] | ||
| manager.shutdown() | ||
|
|
||
| return server, zmq_url, web_url | ||
|
|
||
| def start_meshcat_server_standalone(): | ||
| import argparse | ||
| argparse.ArgumentParser( | ||
| description="Serve the Jiminy MeshCat HTML files and listen for ZeroMQ commands") | ||
|
|
||
| server, zmq_url, web_url = start_meshcat_server() | ||
| print(zmq_url) | ||
| print(web_url) | ||
|
|
||
| try: | ||
| server.join() | ||
| except KeyboardInterrupt: | ||
| server.terminate() | ||
| server.join(timeout=0.5) | ||
| try: | ||
| proc = psutil.Process(server.pid) | ||
| proc.send_signal(signal.SIGKILL) | ||
| except psutil.NoSuchProcess: | ||
| pass |
Oops, something went wrong.