|
| 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 |
| 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,22 @@ 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 | + renderWindow = self.getView("-1") |
| 290 | + w2if = vtkWindowToImageFilter() |
| 291 | + w2if.SetInput(renderWindow) |
| 292 | + w2if.SetInputBufferTypeToRGB() |
| 293 | + w2if.ReadFrontBufferOff() |
| 294 | + w2if.Update() |
| 295 | + |
| 296 | + writer = vtkPNGWriter() |
| 297 | + writer.SetFileName(os.path.join(self.DATA_FOLDER_PATH, 'screenshot.png')) |
| 298 | + writer.SetInputConnection(w2if.GetOutputPort()) |
| 299 | + writer.Write() |
| 300 | + |
274 | 301 | def get_data_base(self):
|
275 | 302 | return self.getSharedObject("db")
|
276 | 303 |
|
|
0 commit comments