feat(screenshot): new rpc

src/opengeodeweb_viewer/rpc/schemas/take_screenshot.json

@@ -0,0 +1,25 @@
+{
+  "rpc": "take_screenshot",
+  "type": "object",
+  "properties": {
+    "filename": {
+      "type": "string"
+    },
+    "output_extension": {
+      "type": "string",
+      "enum": [
+        "png",
+        "jpg"
+      ]
+    },
+    "include_background": {
+      "type": "boolean"
+    }
+  },
+  "required": [
+    "filename",
+    "output_extension",
+    "include_background"
+  ],
+  "additionalProperties": false
+}

src/opengeodeweb_viewer/vtk_protocol.py

@@ -1,10 +1,19 @@
+# Standard library imports
 import json
 import os
+
+# Third party imports
+import vtk
 from vtk.web import protocols as vtk_protocols
+from vtkmodules.vtkIOImage import vtkPNGWriter, vtkJPEGWriter
+from vtkmodules.vtkRenderingCore import (vtkWindowToImageFilter)
 from wslink import register as exportRpc
-import vtk
+
+# Local application imports
 from .function import validate_schemas
 
+
+
 schemas = os.path.join(os.path.dirname(__file__), "rpc/schemas")
 
 with open(os.path.join(schemas, "create_visualization.json"), "r") as file:
@@ -35,6 +44,8 @@
     set_color_json = json.load(file)
 with open(os.path.join(schemas, "set_vertex_attribute.json"), "r") as file:
     set_vertex_attribute_json = json.load(file)
+with open(os.path.join(schemas, "take_screenshot.json"), "r") as file:
+    take_screenshot_json = json.load(file)
 
 
 class VtkView(vtk_protocols.vtkWebProtocol):
@@ -271,6 +282,45 @@ def setVertexAttribute(self, params):
         mapper.SetScalarModeToUsePointFieldData()
         self.render()
 
+    @exportRpc(take_screenshot_json["rpc"])
+    def takeScreenshot(self, params):
+        validate_schemas(params, take_screenshot_json)
+        print(f"{params=}", flush=True)
+        filename = params["filename"]
+        output_extension = params["output_extension"]
+        include_background = params["include_background"]
+        renderWindow = self.getView("-1")
+        renderer = self.get_renderer()
+
+        w2if = vtkWindowToImageFilter()
+
+        if not include_background:
+            renderWindow.SetAlphaBitPlanes(1)
+
+        w2if.SetInput(renderWindow)
+        w2if.SetInputBufferTypeToRGBA() 
+        w2if.ReadFrontBufferOff()
+        w2if.Update()
+        renderWindow.SetAlphaBitPlanes(0)
+
+        if output_extension == "png":
+            writer = vtkPNGWriter()
+        elif output_extension == "jpg":
+            writer = vtkJPEGWriter()
+        else:
+            raise Exception("output_extension not supported")
+
+        new_filename = filename + '.' + output_extension
+        file_path = os.path.join(self.DATA_FOLDER_PATH, new_filename)
+        writer.SetFileName(file_path)
+        writer.SetInputConnection(w2if.GetOutputPort())
+        writer.Write()
+
+        with open(file_path, "rb") as file:
+            file_content = file.read()
+
+        return {"blob": self.addAttachment(file_content)}
+
     def get_data_base(self):
         return self.getSharedObject("db")
 

src/tests/conftest.py

@@ -57,8 +57,11 @@ def print_log(self):
         print(output)
 
     def get_response(self):
-        response = eval(self.ws.recv())
-        return response
+        response = self.ws.recv()
+        if isinstance(response, bytes):
+            return response
+        else:
+            return eval(response)
 
     def compare_image(self, nb_messages, filename):
         for message in range(nb_messages):

src/tests/test_protocol.py

@@ -103,3 +103,26 @@ def test_set_color(server):
 
     server.call("set_color", [{"id": "123456789", "red": 50, "green": 2, "blue": 250}])
     assert server.compare_image(3, "set_color.jpeg") == True
+
+
+
+def test_take_screenshot(server):
+
+    server.call(
+        "create_object_pipeline",
+        [{"id": "123456789", "file_name": "hat.vtp"}],
+    )
+    assert server.compare_image(3, "create_object_pipeline.jpeg") == True
+
+    server.call(
+        "take_screenshot",
+        [{"filename": "screenshot", "output_extension": "png", "include_background": False}],
+    )
+
+    response1 = server.get_response()
+    print(f"{response1=}", flush=True)
+    response2 = server.get_response()
+    print(f"{response2=}", flush=True)
+
+    blob = response2
+    assert type(blob) is bytes