Issue #11: Record video using opencv - to be refined

sksurgerydavinci/ui/gui.py

@@ -4,7 +4,8 @@
 # pylint: disable=no-name-in-module, too-many-public-methods
 import os
 import datetime
-from PySide2.QtCore import Qt, Signal
+import cv2
+from PySide2.QtCore import Qt, Signal, QTimer
 from PySide2.QtWidgets import QSlider, QVBoxLayout, QHBoxLayout, \
                               QLabel, QPushButton, \
                               QMessageBox, QApplication, \
@@ -46,6 +47,9 @@ class UI(QWidget):
         self.screen = None
         self.screen_geometry = None
 
+        self.timer = QTimer()
+        self.timer.timeout.connect(self.save_frame_to_video)
+
         self.show()
 
     def add_vtk_overlay_window(self, overlay_window):
@@ -89,7 +93,7 @@ class UI(QWidget):
         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_button.clicked.connect(self.on_record_start_clicked)
 
         self.record_buttons_layout = QHBoxLayout()
         self.record_buttons_layout.addWidget(self.screenshot_button)
@@ -115,9 +119,28 @@ class UI(QWidget):
         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):
+    def on_record_start_clicked(self):
         """ Start recording a video using current time as filename """
-        pass
+        self.record_button.setText("Stop recording")
+        self.record_button.clicked.connect(self.on_record_stop_clicked)
+        fname = datetime.datetime.now().strftime("%Y-%m-%d.%H-%M-%S") + '.avi'
+        fps = 20
+        width, height = self.vtk_overlay.width(), self.vtk_overlay.height()
+        self.video_out = cv2.VideoWriter(fname, cv2.VideoWriter_fourcc('M','J','P','G'), fps, (width, height))
+
+
+        self.timer.start(1000/fps)
+
+    def on_record_stop_clicked(self):
+        """ Stop recording data to file and restore button settings. """
+        self.timer.stop()
+        self.video_out.release()
+
+        self.record_button.setText("Record video")
+        self.record_button.clicked.connect(self.on_record_start_clicked)
+
+    def save_frame_to_video(self):
+        self.video_out.write(self.vtk_overlay.convert_scene_to_numpy_array())
 
     def add_vtk_models(self, vtk_models):
         """ Add VTK models to the UI and vtk_overlay_window