Implement screenshots in show (#1759)

* Implement screenshots

* Use previous defaults

* Add some docs

* Render off-screen when not interacting

* Add a test

* Dummify the test"

* Coverage tweak

* Add a note about possible issues

* Fix example

* Add xpos and ypos

---------

Co-authored-by: adam-urbanczyk <[email protected]>

Author: adam-urbanczyk

cadquery/vis.py

@@ -1,7 +1,7 @@
 from . import Shape, Workplane, Assembly, Sketch, Compound, Color, Vector, Location
 from .occ_impl.assembly import _loc2vtk, toVTK
 
-from typing import Union, Any, List, Tuple, Iterable, cast
+from typing import Union, Any, List, Tuple, Iterable, cast, Optional
 
 from typish import instance_of
 
@@ -18,10 +18,12 @@
     vtkPolyDataMapper,
     vtkAssembly,
     vtkRenderWindow,
+    vtkWindowToImageFilter,
 )
 from vtkmodules.vtkCommonCore import vtkPoints
 from vtkmodules.vtkCommonDataModel import vtkCellArray, vtkPolyData
 from vtkmodules.vtkCommonColor import vtkNamedColors
+from vtkmodules.vtkIOImage import vtkPNGWriter
 
 
 DEFAULT_COLOR = [1, 0.8, 0, 1]
@@ -150,10 +152,22 @@ def show(
     edges: bool = False,
     specular: bool = True,
     title: str = "CQ viewer",
+    screenshot: Optional[str] = None,
+    interact: bool = True,
+    zoom: float = 1.0,
+    roll: float = -35,
+    elevation: float = -45,
+    width: Union[int, float] = 0.5,
+    height: Union[int, float] = 0.5,
+    trihedron: bool = True,
+    bgcolor: tuple[float, float, float] = (1, 1, 1),
+    gradient: bool = True,
+    xpos: Union[int, float] = 0,
+    ypos: Union[int, float] = 0,
     **kwrags: Any,
 ):
     """
-    Show CQ objects using VTK.
+    Show CQ objects using VTK. This functions optionally allows to make screenshots.
     """
 
     # split objects
@@ -171,6 +185,11 @@ def show(
 
     # VTK window boilerplate
     win = vtkRenderWindow()
+
+    # Render off-screen when not interacting
+    if not interact:
+        win.SetOffScreenRendering(1)
+
     win.SetWindowName(title)
     win.AddRenderer(renderer)
 
@@ -208,26 +227,30 @@ def show(
     axes.GetZAxisCaptionActor2D().GetCaptionTextProperty().ShallowCopy(tp)
 
     # add to an orientation widget
-    orient_widget = vtkOrientationMarkerWidget()
-    orient_widget.SetOrientationMarker(axes)
-    orient_widget.SetViewport(0.9, 0.0, 1.0, 0.2)
-    orient_widget.SetZoom(1.1)
-    orient_widget.SetInteractor(inter)
-    orient_widget.EnabledOn()
-    orient_widget.InteractiveOff()
+    if trihedron:
+        orient_widget = vtkOrientationMarkerWidget()
+        orient_widget.SetOrientationMarker(axes)
+        orient_widget.SetViewport(0.9, 0.0, 1.0, 0.2)
+        orient_widget.SetZoom(1.1)
+        orient_widget.SetInteractor(inter)
+        orient_widget.EnabledOn()
+        orient_widget.InteractiveOff()
 
     # use gradient background
-    renderer.SetBackground(1, 1, 1)
-    renderer.GradientBackgroundOn()
+    renderer.SetBackground(*bgcolor)
+
+    if gradient:
+        renderer.GradientBackgroundOn()
 
     # use FXXAA
     renderer.UseFXAAOn()
 
     # set camera
     camera = renderer.GetActiveCamera()
-    camera.Roll(-35)
-    camera.Elevation(-45)
+    camera.Roll(roll)
+    camera.Elevation(elevation)
     renderer.ResetCamera()
+    camera.Zoom(zoom)
 
     # add pts and locs
     renderer.AddActor(pts)
@@ -241,12 +264,36 @@ def show(
     inter.Initialize()
 
     w, h = win.GetScreenSize()
-    win.SetSize(w // 2, h // 2)
-    win.SetPosition(-10, 0)
+    win.SetSize(
+        int(w * width) if isinstance(width, float) else width,
+        int(h * height) if isinstance(height, float) else height,
+    )  # is height, width specified as float assume it is relative
+
+    # set position
+    win.SetPosition(
+        int(w * xpos) if isinstance(xpos, float) else xpos,
+        int(h * ypos) if isinstance(ypos, float) else ypos,
+    )
 
     # show and return
     win.Render()
-    inter.Start()
+
+    # make a screenshot
+    if screenshot:
+        win2image = vtkWindowToImageFilter()
+        win2image.SetInput(win)
+        win2image.SetInputBufferTypeToRGB()
+        win2image.ReadFrontBufferOff()
+        win2image.Update()
+
+        writer = vtkPNGWriter()
+        writer.SetFileName(screenshot)
+        writer.SetInputConnection(win2image.GetOutputPort())
+        writer.Write()
+
+    # start interaction
+    if interact:
+        inter.Start()
 
 
 # alias

doc/vis.rst

@@ -33,7 +33,7 @@ One can visualize objects of type :class:`~cadquery.Workplane`, :class:`~cadquer
 .. code-block:: python
 
    from cadquery import *
-   from cadquery.occ_impl.shapes import *
+   from cadquery.func import *
    from cadquery.vis import show
 
    w = Workplane().sphere(0.5).split(keepTop=True)
@@ -62,7 +62,7 @@ Additionally it is possible to integrate with other libraries using VTK and disp
 .. code-block:: python
 
     from cadquery.vis import show
-    from cadquery.occ_impl.shapes import torus
+    from cadquery.func import torus
 
     from vtkmodules.vtkRenderingAnnotation import vtkAnnotatedCubeActor
 
@@ -77,6 +77,23 @@ Additionally it is possible to integrate with other libraries using VTK and disp
 
 Note that currently the show function is blocking.
 
+Screenshots
+===========
+
+`:meth:~cadquery.vis.show` allows additionally to take screenshots in `png` format. One can specify zoom,
+camera position and windows size.
+
+.. code-block:: python
+
+    from cadquery.vis import show
+    from cadquery.func import box
+
+    b = box(1,1,1)
+
+    show(b, width=800, height=800, screenshot='img.png', zoom=2, roll=-20, elevation=-30, interact=False)
+
+NB: intermittent issues were observed with this functionality, please submit detailed bug reports in case
+of problems.
 
 Jupyter/JupterLab
 =================

tests/test_vis.py

@@ -3,10 +3,21 @@
 
 import cadquery.vis as vis
 
-from vtkmodules.vtkRenderingCore import vtkRenderWindow, vtkRenderWindowInteractor
+from vtkmodules.vtkRenderingCore import (
+    vtkRenderWindow,
+    vtkRenderWindowInteractor,
+    vtkWindowToImageFilter,
+)
 from vtkmodules.vtkRenderingAnnotation import vtkAnnotatedCubeActor
+from vtkmodules.vtkIOImage import vtkPNGWriter
 
 from pytest import fixture
+from path import Path
+
+
+@fixture(scope="module")
+def tmpdir(tmp_path_factory):
+    return Path(tmp_path_factory.mktemp("screenshots"))
 
 
 @fixture
@@ -54,12 +65,29 @@ def SetPosition(*args):
 
         pass
 
+    def SetOffScreenRendering(*args):
+
+        pass
+
+
+class FakeWin2Img(vtkWindowToImageFilter):
+    def Update(*args):
+
+        pass
+
+
+class FakePNGWriter(vtkPNGWriter):
+    def Write(*args):
+
+        pass
+
 
 def test_show(wp, assy, sk, monkeypatch):
 
     # use some dummy vtk objects
     monkeypatch.setattr(vis, "vtkRenderWindowInteractor", FakeInteractor)
     monkeypatch.setattr(vis, "vtkRenderWindow", FakeWindow)
+    monkeypatch.setattr(vis, "vtkWindowToImageFilter", FakeWin2Img)
 
     # simple smoke test
     show(wp)
@@ -92,3 +120,17 @@ def test_show(wp, assy, sk, monkeypatch):
 
     # show a raw vtkProp
     show(vtkAxesActor(), [vtkAnnotatedCubeActor()])
+
+
+def test_screenshot(wp, tmpdir, monkeypatch):
+
+    # smoke test for now
+
+    # use some dummy vtk objects
+    monkeypatch.setattr(vis, "vtkRenderWindowInteractor", FakeInteractor)
+    monkeypatch.setattr(vis, "vtkRenderWindow", FakeWindow)
+    monkeypatch.setattr(vis, "vtkWindowToImageFilter", FakeWin2Img)
+    monkeypatch.setattr(vis, "vtkPNGWriter", FakePNGWriter)
+
+    with tmpdir:
+        show(wp, interact=False, screenshot="img.png", trihedron=False, gradient=False)