Commit aac68fc6 authored by Thomas Dowrick's avatar Thomas Dowrick

Merge branch '13-modify-overlayonvideofeed-get-timestamp-for-each-frame' into 'master'

Resolve "modify OverlayOnVideoFeed get timestamp for each frame"

Closes #13

See merge request WEISS/SoftwareRepositories/SNAPPY/scikit-surgeryutils!7
parents 2501e17a 9f8f9668
Pipeline #1780 passed with stages
in 12 minutes and 51 seconds
......@@ -2,9 +2,13 @@
#pylint: disable=no-member, no-name-in-module, protected-access
# coding=utf-8
import datetime
import logging
import cv2
from PySide2.QtCore import QTimer
from sksurgeryimage.acquire.video_source import TimestampedVideoSource
from sksurgeryimage.acquire.video_writer import TimestampedVideoWriter
from sksurgeryvtk.widgets.vtk_overlay_window import VTKOverlayWindow
from sksurgeryvtk.models.vtk_surface_model_directory_loader \
import VTKSurfaceModelDirectoryLoader
......@@ -19,10 +23,11 @@ class OverlayBaseApp():
"""
def __init__(self, video_source):
self.vtk_overlay_window = VTKOverlayWindow()
self.video_source = cv2.VideoCapture(video_source)
self.video_source = TimestampedVideoSource(video_source)
self.update_rate = 30
self.img = None
self.timer = None
self.save_frame = None
def start(self):
"""Show the overlay widget and
......@@ -58,9 +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.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) """
# 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.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 = 5
x = cv2.VideoCapture(camera_source)
if x.isOpened():
x.release()
overlay_app = coa.OverlayOnVideoFeed(camera_source)
overlay_app.add_vtk_models_from_dir('tests/data/models/Liver')
overlay_app.start()
overlay_app.update()
overlay_app.stop()
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)
# 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()
def test_OverlayBaseAppRaisesNotImplementedError(setup_qt):
"""
Test will only run if there is a camera avilable
"""
class ErrorApp(coa.OverlayBaseApp):
def something(self):
pass
with pytest.raises(NotImplementedError):
overlay_app = ErrorApp(0)
input_file = 'tests/data/100x50_100_frames.avi'
overlay_app = ErrorApp(input_file)
overlay_app.update()
# -*- coding: utf-8 -*-
import pytest
from PySide2.QtWidgets import QApplication
@pytest.fixture(scope="module")
def setup_qt():
""" Create the QT application. """
app = QApplication([])
return app
......@@ -20,7 +20,7 @@ def test_image_to_pixel_invalid_because_input_is_greyscale():
iu.image_to_pixmap(np.zeros((100, 50, 1), dtype=np.uint8))
def test_image_to_pixel_valid_rgb_example(setup_qt):
def test_image_to_pixel_valid_rgb_example():
blank_image = np.zeros((50, 100, 3), dtype=np.uint8)
pixmap = iu.image_to_pixmap(blank_image)
assert pixmap.width() == blank_image.shape[1]
......
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