Merge pull request #93 from ManifoldFR/feature/capture_image

rdeits · web-flow · commit 20e4652a1d34 · 2021-08-02T23:12:32.000-04:00
New feature: capture image
diff --git a/.gitignore b/.gitignore
@@ -9,3 +9,4 @@ dist
 
 venv*
 .pytest_cache
+.vscode
diff --git a/setup.py b/setup.py
@@ -24,6 +24,7 @@
       "tornado >= 4.0.0",
       "pyzmq >= 17.0.0",
       "pyngrok >= 4.1.6",
+      "pillow >= 7.0.0"
     ],
     zip_safe=False,
     include_package_data=True
diff --git a/src/meshcat/commands.py b/src/meshcat/commands.py
@@ -44,6 +44,13 @@ def lower(self):
             u"matrix": list(self.matrix.T.flatten())
         }
 
+class CaptureImage:
+
+    def lower(self):
+        return {
+            u"type": u"capture_image"
+        }
+
 
 class Delete:
     __slots__ = ["path"]
diff --git a/src/meshcat/servers/zmqserver.py b/src/meshcat/servers/zmqserver.py
@@ -7,6 +7,7 @@
 import sys
 import subprocess
 import multiprocessing
+import json
 
 import tornado.web
 import tornado.ioloop
@@ -128,7 +129,13 @@ def open(self):
         self.bridge.send_scene(self)
 
     def on_message(self, message):
-        pass
+        try:
+            message = json.loads(message)
+            self.bridge.send_image(message['data'])
+            return
+        except Exception as err:
+            print(err)
+            raise
 
     def on_close(self):
         self.bridge.websocket_pool.remove(self)
@@ -240,12 +247,23 @@ def wait_for_websockets(self):
         else:
             self.ioloop.call_later(0.1, self.wait_for_websockets)
 
+    def send_image(self, data):
+        import base64
+        mime, img_code = data.split(",", 1)
+        img_bytes = base64.b64decode(img_code)
+        self.zmq_stream.send(img_bytes)
+
     def handle_zmq(self, frames):
         cmd = frames[0].decode("utf-8")
         if cmd == "url":
             self.zmq_socket.send(self.web_url.encode("utf-8"))
         elif cmd == "wait":
             self.ioloop.add_callback(self.wait_for_websockets)
+        elif cmd == "capture_image":
+            if len(self.websocket_pool) > 0:
+                self.forward_to_websockets(frames)  # on_message callback should handle the pb
+            else:
+                self.ioloop.call_later(0.3, lambda: self.handle_zmq(frames))
         elif cmd in MESHCAT_COMMANDS:
             if len(frames) != 3:
                 self.zmq_socket.send(b"error: expected 3 frames")
diff --git a/src/meshcat/viewer b/src/meshcat/viewer
@@ -1 +1 @@
-Subproject commit 85fd641e8897ffe453502e0ec152225562999b0b
+Subproject commit 2c809d463d86bafe8ef471f419be855cf3e842d6
diff --git a/src/meshcat/visualizer.py b/src/meshcat/visualizer.py
@@ -2,10 +2,13 @@
 import umsgpack
 import numpy as np
 import zmq
+import io
+from PIL import Image        
 from IPython.display import HTML
 
+
 from .path import Path
-from .commands import SetObject, SetTransform, Delete, SetProperty, SetAnimation
+from .commands import SetObject, SetTransform, Delete, SetProperty, SetAnimation, CaptureImage
 from .geometry import MeshPhongMaterial
 from .servers.zmqserver import start_zmq_server_as_subprocess
 
@@ -65,6 +68,17 @@ def get_scene(self):
         # we receive the HTML as utf-8-encoded, so decode here
         return self.zmq_socket.recv().decode('utf-8')
 
+    def get_image(self):
+        cmd_data = CaptureImage().lower()
+        self.zmq_socket.send_multipart([
+            cmd_data["type"].encode("utf-8"),
+            "".encode("utf-8"),
+            umsgpack.packb(cmd_data)
+        ])
+        img_bytes = self.zmq_socket.recv()
+        img = Image.open(io.BytesIO(img_bytes))
+        return img
+
 
 def srcdoc_escape(x):
     return x.replace("&", "&amp;").replace('"', "&quot;")
@@ -143,6 +157,10 @@ def set_property(self, key, value):
     def set_animation(self, animation, play=True, repetitions=1):
         return self.window.send(SetAnimation(animation, play=play, repetitions=repetitions))
 
+    def get_image(self):
+        """Save an image"""
+        return self.window.get_image()
+
     def delete(self):
         return self.window.send(Delete(self.path))
 

Original file line number	Diff line number	Diff line change
`@@ -9,3 +9,4 @@ dist`
`9`	`9`
`10`	`10`	`venv*`
`11`	`11`	`.pytest_cache`
	`12`	`+.vscode`