|
| 1 | +# Standard library imports |
1 | 2 | import json |
2 | 3 | import os |
| 4 | + |
| 5 | +# Third party imports |
| 6 | +import vtk |
3 | 7 | from vtk.web import protocols as vtk_protocols |
| 8 | +from vtkmodules.vtkIOImage import vtkPNGWriter, vtkJPEGWriter |
| 9 | +from vtkmodules.vtkRenderingCore import (vtkWindowToImageFilter) |
4 | 10 | from wslink import register as exportRpc |
5 | | -import vtk |
| 11 | + |
| 12 | +# Local application imports |
6 | 13 | from .function import validate_schemas |
7 | 14 |
|
| 15 | + |
| 16 | + |
8 | 17 | schemas = os.path.join(os.path.dirname(__file__), "rpc/schemas") |
9 | 18 |
|
10 | 19 | with open(os.path.join(schemas, "create_visualization.json"), "r") as file: |
|
35 | 44 | set_color_json = json.load(file) |
36 | 45 | with open(os.path.join(schemas, "set_vertex_attribute.json"), "r") as file: |
37 | 46 | set_vertex_attribute_json = json.load(file) |
| 47 | +with open(os.path.join(schemas, "take_screenshot.json"), "r") as file: |
| 48 | + take_screenshot_json = json.load(file) |
38 | 49 |
|
39 | 50 |
|
40 | 51 | class VtkView(vtk_protocols.vtkWebProtocol): |
@@ -271,6 +282,51 @@ def setVertexAttribute(self, params): |
271 | 282 | mapper.SetScalarModeToUsePointFieldData() |
272 | 283 | self.render() |
273 | 284 |
|
| 285 | + @exportRpc(take_screenshot_json["rpc"]) |
| 286 | + def takeScreenshot(self, params): |
| 287 | + validate_schemas(params, take_screenshot_json) |
| 288 | + print(f"{params=}", flush=True) |
| 289 | + filename = params["filename"] |
| 290 | + output_extension = params["output_extension"] |
| 291 | + include_background = params["include_background"] |
| 292 | + renderWindow = self.getView("-1") |
| 293 | + renderer = self.get_renderer() |
| 294 | + |
| 295 | + w2if = vtkWindowToImageFilter() |
| 296 | + |
| 297 | + if not include_background: |
| 298 | + renderWindow.SetAlphaBitPlanes(1) |
| 299 | + w2if.SetInputBufferTypeToRGBA() |
| 300 | + else: |
| 301 | + renderWindow.SetAlphaBitPlanes(0) |
| 302 | + w2if.SetInputBufferTypeToRGB() |
| 303 | + |
| 304 | + renderWindow.Render() |
| 305 | + |
| 306 | + w2if.SetInput(renderWindow) |
| 307 | + w2if.ReadFrontBufferOff() |
| 308 | + w2if.Update() |
| 309 | + |
| 310 | + if output_extension == "png": |
| 311 | + writer = vtkPNGWriter() |
| 312 | + elif output_extension in ["jpg", "jpeg"]: |
| 313 | + if not include_background: |
| 314 | + raise Exception("output_extension not supported with background") |
| 315 | + writer = vtkJPEGWriter() |
| 316 | + else: |
| 317 | + raise Exception("output_extension not supported") |
| 318 | + |
| 319 | + new_filename = filename + '.' + output_extension |
| 320 | + file_path = os.path.join(self.DATA_FOLDER_PATH, new_filename) |
| 321 | + writer.SetFileName(file_path) |
| 322 | + writer.SetInputConnection(w2if.GetOutputPort()) |
| 323 | + writer.Write() |
| 324 | + |
| 325 | + with open(file_path, "rb") as file: |
| 326 | + file_content = file.read() |
| 327 | + |
| 328 | + return {"blob": self.addAttachment(file_content)} |
| 329 | + |
274 | 330 | def get_data_base(self): |
275 | 331 | return self.getSharedObject("db") |
276 | 332 |
|
|
0 commit comments