|
5 | 5 | # Third party imports
|
6 | 6 | import vtk
|
7 | 7 | from vtk.web import protocols as vtk_protocols
|
8 |
| -from vtkmodules.vtkIOImage import vtkPNGWriter |
| 8 | +from vtkmodules.vtkIOImage import vtkPNGWriter, vtkJPEGWriter |
9 | 9 | from vtkmodules.vtkRenderingCore import (vtkWindowToImageFilter)
|
10 | 10 | from wslink import register as exportRpc
|
11 | 11 |
|
@@ -286,17 +286,34 @@ def setVertexAttribute(self, params):
|
286 | 286 | def takeScreenshot(self, params):
|
287 | 287 | validate_schemas(params, take_screenshot_json)
|
288 | 288 | print(f"{params=}", flush=True)
|
| 289 | + filename = params["filename"] |
| 290 | + output_extension = params["output_extension"] |
| 291 | + include_background = params["include_background"] |
289 | 292 | renderWindow = self.getView("-1")
|
| 293 | + renderer = self.get_renderer() |
| 294 | + |
290 | 295 | w2if = vtkWindowToImageFilter()
|
| 296 | + |
| 297 | + if not include_background: |
| 298 | + # renderer.SetBackground([255,255,255]) |
| 299 | + # renderer.SetLayer(1) |
| 300 | + renderWindow.SetAlphaBitPlanes(1) |
291 | 301 | w2if.SetInput(renderWindow)
|
292 |
| - w2if.SetInputBufferTypeToRGB() |
| 302 | + w2if.SetInputBufferTypeToRGBA() |
293 | 303 | w2if.ReadFrontBufferOff()
|
294 | 304 | w2if.Update()
|
| 305 | + renderWindow.SetAlphaBitPlanes(0) |
| 306 | + |
| 307 | + if output_extension == "png": |
| 308 | + writer = vtkPNGWriter() |
| 309 | + elif output_extension == "jpg": |
| 310 | + writer = vtkJPEGWriter() |
295 | 311 |
|
296 |
| - writer = vtkPNGWriter() |
297 |
| - writer.SetFileName(os.path.join(self.DATA_FOLDER_PATH, 'screenshot.png')) |
| 312 | + writer.SetFileName(os.path.join(self.DATA_FOLDER_PATH, filename + '.' + output_extension)) |
298 | 313 | writer.SetInputConnection(w2if.GetOutputPort())
|
299 | 314 | writer.Write()
|
| 315 | + # renderer.SetLayer(0) |
| 316 | + return |
300 | 317 |
|
301 | 318 | def get_data_base(self):
|
302 | 319 | return self.getSharedObject("db")
|
|
0 commit comments