From 22d7f419fc74e849d5195258a1a64467d0f27403 Mon Sep 17 00:00:00 2001 From: kngwyu Date: Mon, 19 Apr 2021 20:11:35 +0900 Subject: [PATCH] Spearate websock_viewer file --- mujoco_maze/maze_env.py | 70 +++------------------------------ mujoco_maze/static/favicon.ico | Bin 2040 -> 0 bytes mujoco_maze/websock_viewer.py | 65 ++++++++++++++++++++++++++++++ 3 files changed, 70 insertions(+), 65 deletions(-) delete mode 100644 mujoco_maze/static/favicon.ico create mode 100644 mujoco_maze/websock_viewer.py diff --git a/mujoco_maze/maze_env.py b/mujoco_maze/maze_env.py index 7e7a68f..3a455d0 100644 --- a/mujoco_maze/maze_env.py +++ b/mujoco_maze/maze_env.py @@ -362,66 +362,6 @@ class MazeEnv(gym.Env): else: return self.wrapped_env.viewer - def _setup_websock_server(self) -> None: - import multiprocessing as mp - - class _ProcWorker(mp.Process): - def __init__(self, pipe: mp.connection.Pipe, port: int) -> None: - super().__init__() - self.pipe = pipe - self.port = port - - def _run_server(self) -> None: - import asyncio - import io - import fastapi - import pathlib - import uvicorn - - from PIL import Image - - app = fastapi.FastAPI() - html_path = pathlib.Path(__file__).parent.joinpath("static/index.html") - html = html_path.read_text().replace("{{port}}", str(self.port)) - server = None - - @app.get("/") - async def get(): - return fastapi.responses.HTMLResponse(html) - - @app.websocket("/ws") - async def ws_send_image(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 - image = Image.fromarray(image_array) - with io.BytesIO() as stream: - image.save(stream, format="png") - res = stream.getvalue() - await websocket.send_bytes(res) - await websocket.close() - server.should_exit = True - - config = uvicorn.Config(app, port=self.port) - server = uvicorn.Server(config) - server.run() - - def run(self) -> None: - try: - self._run_server() - except KeyboardInterrupt: - pass - except Exception as e: - print("Exception in websocket server") - raise e - - self._websock_server_pipe, pipe = mp.Pipe() - worker = _ProcWorker(pipe, self._websock_port) - worker.start() - def _render_image(self) -> np.ndarray: self._mj_offscreen_viewer._set_mujoco_buffers() self._mj_offscreen_viewer.render(640, 480) @@ -433,12 +373,12 @@ class MazeEnv(gym.Env): def render(self, mode="human", **kwargs) -> Any: if self._websock_port is not None: if self._mj_offscreen_viewer is None: - import mujoco_py + from mujoco_py import MjRenderContextOffscreen as MjRenderOffscreen - self._mj_offscreen_viewer = mujoco_py.MjRenderContextOffscreen( - self.wrapped_env.sim - ) - self._setup_websock_server() + from mujoco_maze.websock_viewer import start_server + + self._mj_offscreen_viewer = MjRenderOffscreen(self.wrapped_env.sim) + self._websock_server_pipe = start_server(self._websock_port) self._websock_server_pipe.send(self._render_image()) return True else: diff --git a/mujoco_maze/static/favicon.ico b/mujoco_maze/static/favicon.ico deleted file mode 100644 index 45acf83b6535edeae3c9da3588404583625b926b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 2040 zcmV}tqtHPiUI&73Q=Fq(EwO$c~md>@V$3<)p|Z#@z;Mp z#W{x{D6PspopmH}g0+?~4DmdlBu+@;1mE`v!w_pNNu1ztgkgv=rp!rFf*>FaL-H&q zPP11Y4$y7X$kL1?PEcAAh9M#f^WBGiloHnKHEEKfd`+!(LQxc44XywWwDb&6eTy5D`H>`sgEfKk?$? zl0W_H0xt|{Hk;t!W^n@!>h(H);FF~po6QCU+U+(X!e+f889X7_PG3!d(wN5|u-ons z5nAmwB4ECpHy)3$X^hlr*uvs)D5b!G!<8}TD$Q7H5d;yz_yKHReZF`>iKBY1K#-RK zj7B4(#sP&=OePbw(e(R$oO6tZBNpj`dZW(K@zDz<=yW?geCJ&v57x^S&N-S<6H$tI zw_`ZH;cB&}-#!0Gaq=>N$z;Mhs?)5sIXXTf&vS;uAplVnp|m1RQij7Jg3#+7 zf(Y?0#sv+|zx@`CW&`IOtL2I$O;Jiwn@&hS|D4Hqf_0X_4`{Vom`^|b^aTJ{*F$QD zkI`D8wI{P^(;B3vao z4}bD=oxfzr2L}f)l%U;i(>?4WB310*;DD$Vk>@$jo<7Ig zf`d+nC~A^rIoDTLF9S%u2F*9$CC~19TIo`Rx8{RT3A5P0p@{5Fib`Oj6e48{+M+;2 z%IEII%r}1cW1?1s5@EgGFkSDMZZ;erA0Z-4CsS6d6>2rbzhU&GS-7IE>!6QJSU1ySUt- zS#D4i1=(JVzVGAvCB@cS@;paG@H`Kt6nUPL?y4o zqF``2V14n7*=)w?>1o+vMNv=`1)k?&j6v)2ZHi*=)y9;@;>vqLQ4}bxt9Vh^GHz`Z z6Ok$xz!-z)d01<)dB%7=u3FrjoSl@lp3Ew$j!%yPNaCa#$&F@%+uX%`!F)DHL^wJr z+cnSfDmMs2PEO0*^?J>GR*JWKc!;+-dl7)|`{a3E)mj^k=XuE9-*4-xwW{it7kBl( z)c|l8v$o0=QB_<78l4{F@v}Ps_Qs`f zlzIVBA{_S0+@i=Ck0*QjHHW>!lF=k3%PgzWr~+l;T}-RhX{`lh(K(pDPT@U~7;Q%(z zXfzvaH{1OLbTpccJ-xzuwJ!6lrO|Ab%qU^C_Xq1N&8S7*?%b82bJ!)~t1pxHmn$+l z-5!_EFLAD%Vdrn0BTBG^k^E@suE-}|v*S!D!cU}v? zY`!E(2K3J#RNoMTiwm3sglaN6=a>yH%Imz-oIN;0D~-c3ydKg!?BSf_;1RWYo!(pDXE2aW&4;%z2#m)nu^@LZSWQd;&-XB=hppEJaI;tv z;8#8mFuEE6C8+j68*R#){ None: + super().__init__() + self.pipe = pipe + self.port = port + + def _run_server(self) -> None: + + app = fastapi.FastAPI() + static = pathlib.Path(__file__).parent.joinpath("static") + html_path = static.joinpath("index.html") + html = html_path.read_text().replace("{{port}}", str(self.port)) + + @app.get("/") + async def get(): + return fastapi.responses.HTMLResponse(html) + + server = None + + @app.websocket("/ws") + async def ws_send_image(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 + image = Image.fromarray(image_array) + with io.BytesIO() as stream: + image.save(stream, format="png") + res = stream.getvalue() + await websocket.send_bytes(res) + await websocket.close() + server.should_exit = True + + config = uvicorn.Config(app, port=self.port) + server = uvicorn.Server(config) + server.run() + + def run(self) -> None: + try: + self._run_server() + except KeyboardInterrupt: + pass + except Exception as e: + print("Exception in websocket server") + raise e + + +def start_server(port: int) -> mp.connection.Connection: + mainproc_pipe, server_pipe = mp.Pipe() + worker = _ServerWorker(server_pipe, port) + worker.start() + return mainproc_pipe