Issue #11: Add button to take screenshot

sksurgerydavinci/ui/gui.py

@@ -3,6 +3,7 @@
 # pylint: disable=fixme, attribute-defined-outside-init
 # pylint: disable=no-name-in-module, too-many-public-methods
 import os
+import datetime
 from PySide2.QtCore import Qt, Signal
 from PySide2.QtWidgets import QSlider, QVBoxLayout, QHBoxLayout, \
                               QLabel, QPushButton, \
@@ -33,13 +34,14 @@ class UI(QWidget):
         self.controls_vbox = QVBoxLayout()
         self.opacity_box = QHBoxLayout()
         self.visiblity_buttons = QHBoxLayout()
-        self.other_buttons = QHBoxLayout()
+        self.other_buttons = QVBoxLayout()
 
         self.add_vtk_overlay_window(overlay_window)
         self.add_opacity_slider()
         self.layout.addLayout(self.visiblity_buttons)
         self.layout.addLayout(self.other_buttons)
         self.add_exit_button()
+        self.add_record_buttons()
 
         self.screen = None
         self.screen_geometry = None
@@ -81,6 +83,19 @@ class UI(QWidget):
         self.exit_button.setStyleSheet(style_sheet)
         self.other_buttons.addWidget(self.exit_button)
 
+    def add_record_buttons(self):
+        """ Add buttons to take screenshot and record video. """
+        self.screenshot_button = QPushButton('Take screenshot')
+        self.screenshot_button.clicked.connect(self.on_screenshot_clicked)
+
+        self.record_button = QPushButton('Record video')
+        self.record_button.clicked.connect(self.on_record_clicked)
+
+        self.record_buttons_layout = QHBoxLayout()
+        self.record_buttons_layout.addWidget(self.screenshot_button)
+        self.record_buttons_layout.addWidget(self.record_button)
+        self.other_buttons.addLayout(self.record_buttons_layout)
+
     def on_exit_clicked(self):
         """ Close the application when the exit button has been clicked"""
 
@@ -95,6 +110,15 @@ class UI(QWidget):
             self.exit_signal.emit()
             QApplication.exit()
 
+    def on_screenshot_clicked(self):
+        """ Save a screenshot to disk, using date and time as filename """
+        fname = datetime.datetime.now().strftime("%Y-%m-%d.%H-%M-%S") + '.png'
+        self.vtk_overlay.save_scene_to_file(fname)
+
+    def on_record_clicked(self):
+        """ Start recording a video using current time as filename """
+        pass
+
     def add_vtk_models(self, vtk_models):
         """ Add VTK models to the UI and vtk_overlay_window