Merge pull request #27 from Geode-solutions/feat/screenshot

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,51 @@ 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.SetInputBufferTypeToRGBA()
+        else:
+            renderWindow.SetAlphaBitPlanes(0)
+            w2if.SetInputBufferTypeToRGB()
+
+        renderWindow.Render()
+
+        w2if.SetInput(renderWindow)
+        w2if.ReadFrontBufferOff()
+        w2if.Update()
+
+        if output_extension == "png":
+            writer = vtkPNGWriter()
+        elif output_extension in ["jpg", "jpeg"]:
+            if not include_background:
+                raise Exception("output_extension not supported with background")
+            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

@@ -30,8 +30,8 @@ def __init__(self, log):
         self.call("viewport.image.push.observer.add", [-1])
         for i in range(5):
             print(f"{i=}", flush=True)
-            reponse = self.ws.recv()
-            print(f"{reponse=}", flush=True)
+            response = self.ws.recv()
+            print(f"{response=}", flush=True)
 
     def call(self, rpc, params=[{}]):
         print(f"{rpc=} {params=}", flush=True)
@@ -57,8 +57,32 @@ 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 images_diff(self, first_image_path, second_image_path):
+        if ".png" in first_image_path:
+            first_reader = vtk.vtkPNGReader()
+        elif (".jpg" in first_image_path) or (".jpeg" in first_image_path):
+            first_reader = vtk.vtkJPEGReader()
+        first_reader.SetFileName(first_image_path)
+
+        if ".png" in second_image_path:
+            second_reader = vtk.vtkPNGReader()
+        elif (".jpg" in second_image_path) or (".jpeg" in second_image_path):
+            second_reader = vtk.vtkJPEGReader()
+        second_reader.SetFileName(second_image_path)
+
+        images_diff = vtk.vtkImageDifference()
+        images_diff.SetInputConnection(first_reader.GetOutputPort())
+        images_diff.SetImageConnection(second_reader.GetOutputPort())
+        images_diff.Update()
+
+        print(f"{images_diff.GetThresholdedError()=}")
+        return images_diff.GetThresholdedError()
 
     def compare_image(self, nb_messages, filename):
         for message in range(nb_messages):
@@ -83,20 +107,10 @@ def compare_image(self, nb_messages, filename):
                 f.write(image)
                 f.close()
 
-            test_reader = vtk.vtkJPEGReader()
-            test_reader.SetFileName(test_file_path)
-
             path_image = os.path.join(self.images_dir_path, filename)
-            answer_reader = vtk.vtkJPEGReader()
-            answer_reader.SetFileName(path_image)
 
-            images_diff = vtk.vtkImageDifference()
-            images_diff.SetInputConnection(test_reader.GetOutputPort())
-            images_diff.SetImageConnection(answer_reader.GetOutputPort())
-            images_diff.Update()
+            return self.images_diff(test_file_path, path_image)==0.0
 
-            print(f"{images_diff.GetThresholdedError()=}")
-            return images_diff.GetThresholdedError() == 0.0
 
 
 class FixtureHelper:

src/tests/test_protocol.py

@@ -103,3 +103,72 @@ 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):
+    # Create an object
+    server.call(
+        "create_object_pipeline",
+        [{"id": "123456789", "file_name": "hat.vtp"}],
+    )
+    assert server.compare_image(3, "create_object_pipeline.jpeg") == True
+
+
+    # Take a screenshot with background jpg
+    server.call(
+        "take_screenshot",
+        [{"filename": "take_screenshot_with_background", "output_extension": "jpg", "include_background": True}],
+    )
+
+    response = server.get_response()
+    blob = server.get_response()
+    assert type(blob) is bytes
+
+    with open(os.path.join(server.test_output_dir, "test.jpg"), "wb") as f:
+        f.write(blob)
+        f.close()
+    first_image_path = os.path.join(server.test_output_dir, "test.jpg")
+    second_image_path = os.path.join(server.images_dir_path, "take_screenshot_with_background.jpg")
+
+    assert server.images_diff(first_image_path, second_image_path) == 0.0
+
+    # Take a screenshot without background png
+    server.call(
+        "take_screenshot",
+        [{"filename": "take_screenshot_without_background", "output_extension": "png", "include_background": True}],
+    )
+
+    response = server.get_response()
+    response = server.get_response()
+    blob = server.get_response()
+    print(f"{blob=}", flush=True)
+    assert type(blob) is bytes
+
+    with open(os.path.join(server.test_output_dir, "test.png"), "wb") as f:
+        f.write(blob)
+        f.close()
+    first_image_path = os.path.join(server.test_output_dir, "test.png")
+    second_image_path = os.path.join(server.images_dir_path, "take_screenshot_without_background.png")
+
+    assert server.images_diff(first_image_path, second_image_path) == 0.0
+
+    # Take a screenshot with background png
+    server.call(
+        "take_screenshot",
+        [{"filename": "take_screenshot_with_background", "output_extension": "png", "include_background": True}],
+    )
+
+    response = server.get_response()
+    response = server.get_response()
+    blob = server.get_response()
+    print(f"{blob=}", flush=True)
+    assert type(blob) is bytes
+
+    with open(os.path.join(server.test_output_dir, "test.png"), "wb") as f:
+        f.write(blob)
+        f.close()
+    first_image_path = os.path.join(server.test_output_dir, "test.png")
+    second_image_path = os.path.join(server.images_dir_path, "take_screenshot_with_background.png")
+
+    assert server.images_diff(first_image_path, second_image_path) == 0.0