Commit 2cca566a authored by Thomas Dowrick's avatar Thomas Dowrick

Issue #13: Update OverlayOnVideoFeed to be able to record video using TimestampedVideoWriter

parent acf16d91
Pipeline #1775 failed with stages
in 1 minute and 22 seconds
......@@ -63,32 +63,53 @@ class OverlayBaseApp():
class OverlayOnVideoFeed(OverlayBaseApp):
"""
Uses the acquired video feed as the background image,
with no additional processing.
with no additional processing. Can also record a video of the scene.
"""
def __init__(self, video_source, output_filename=None):
super().__init__(video_source)
self.output_filename = output_filename
self.video_writer = None
def update(self):
""" Get the next frame of input and write to file (if enabled). """
_, self.img = self.video_source.read()
self.vtk_overlay_window.set_video_image(self.img)
self.vtk_overlay_window._RenderWindow.Render()
if self.save_frame:
output_frame = self.vtk_overlay_window.convert_scene_to_numpy_array()
output_frame = cv2.cvtColor(output_frame, cv2.COLOR_RGB2BGR)
self.writer.write_frame(output_frame, self.video_source.timestamp)
output_frame = self.get_output_frame()
self.video_writer.write_frame(output_frame,
self.video_source.timestamp)
def get_output_frame(self):
""" Get the output frame to write in numpy format."""
output_frame = \
self.vtk_overlay_window.convert_scene_to_numpy_array()
output_frame = cv2.cvtColor(output_frame, cv2.COLOR_RGB2BGR)
return output_frame
def on_record_start(self):
""" Start recording data on each frame update.
It is expected that this will be triggered using a Qt signal e.g. from
a button click. (see sksurgerydavinci.ui.Viewers for examples) """
fname = datetime.datetime.now().strftime("%Y-%m-%d.%H-%M-%S") + '.avi'
dims = (self.vtk_overlay_window.width(), self.vtk_overlay_window.height())
self.writer = TimestampedVideoWriter(fname, self.update_rate, dims)
# Set the filename to current date/time if no name specified.
if not self.output_filename:
self.output_filename = 'outputs/' + \
datetime.datetime.now().strftime("%Y-%m-%d.%H-%M-%S") + '.avi'
output_frame = self.get_output_frame()
height, width = output_frame.shape[:2]
self.video_writer = TimestampedVideoWriter(self.output_filename,
self.update_rate, width,
height)
self.save_frame = True
logging.debug("Recording started.")
def on_record_stop(self):
""" Stop recording data. """
self.save_frame = False
self.writer.close()
self.video_writer.close()
logging.debug("Recording stopped.")
# -*- coding: utf-8 -*-
import cv2
import pytest
import numpy as np
import sksurgeryutils.common_overlay_apps as coa
def test_OverlayOnVideoFeed(setup_qt):
def test_OverlayOnVideoFeed_from_file(setup_qt):
# Try to open a camera. If one isn't available, the rest of test
# will be skipped.
input_file = 'tests/data/100x50_100_frames.avi'
out_file = 'tests/output/overlay_test.avi'
overlay_app = coa.OverlayOnVideoFeed(input_file, out_file)
# Start app and get a frame from input, so that
# the window is showing something, before we start
# recording.
overlay_app.start()
overlay_app.update()
overlay_app.on_record_start()
for i in range(50):
overlay_app.update()
overlay_app.on_record_stop()
overlay_app.stop()
# Check that 50 frames were actually written to the output file
output_video = cv2.VideoCapture(out_file)
for i in range(50):
ret, _ = output_video.read()
assert ret
# Trying to read 51st frame should return False
ret, _ = output_video.read()
assert not ret
output_video.release()
def test_OverlayOnVideoFeed_from_webcam(setup_qt):
"""
Test will only run if there is a camera avilable.
"""
# Try to open a camera. If one isn't available, the rest of test
# will be skipped.
camera_source = 0
x = cv2.VideoCapture(camera_source)
if x.isOpened():
x.release()
source = 0
cam = cv2.VideoCapture(source)
if not cam.isOpened():
pytest.skip("No camera available")
cam.release()
out_file = 'tests/output/overlay_on_webcam_test.avi'
overlay_app = coa.OverlayOnVideoFeed(0, out_file)
overlay_app = coa.OverlayOnVideoFeed(camera_source)
overlay_app.start()
overlay_app.update()
overlay_app.stop()
# Start app and get a frame from input, so that
# the window is showing something, before we start
# recording.
overlay_app.start()
overlay_app.update()
overlay_app.on_record_start()
else:
pytest.skip("No camera available, skipping.")
for i in range(50):
overlay_app.update()
overlay_app.on_record_stop()
overlay_app.stop()
def test_OverlayBaseAppRaisesNotImplementedError(setup_qt):
"""
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment