Commit f44f12fd authored by Thomas Dowrick's avatar Thomas Dowrick

Merge branch 'crop-roi' into 'master'

Crop roi

See merge request WEISS/SoftwareRepositories/SNAPPY/scikit-surgeryutils!8
parents aac68fc6 caa42be7
Pipeline #1875 passed with stages
in 20 minutes and 50 seconds
......@@ -9,6 +9,8 @@ import cv2
from PySide2.QtCore import QTimer
from sksurgeryimage.acquire.video_source import TimestampedVideoSource
from sksurgeryimage.acquire.video_writer import TimestampedVideoWriter
from sksurgeryimage.ui.ImageCropper import ImageCropper
from sksurgeryvtk.widgets.vtk_overlay_window import VTKOverlayWindow
from sksurgeryvtk.models.vtk_surface_model_directory_loader \
import VTKSurfaceModelDirectoryLoader
......@@ -21,9 +23,9 @@ class OverlayBaseApp():
:param video_source: OpenCV compatible video source (int or filename)
"""
def __init__(self, video_source):
def __init__(self, video_source, dims=None):
self.vtk_overlay_window = VTKOverlayWindow()
self.video_source = TimestampedVideoSource(video_source)
self.video_source = TimestampedVideoSource(video_source, dims)
self.update_rate = 30
self.img = None
self.timer = None
......@@ -63,18 +65,46 @@ class OverlayBaseApp():
class OverlayOnVideoFeed(OverlayBaseApp):
"""
Uses the acquired video feed as the background image,
with no additional processing. Can also record a video of the scene.
with no additional processing.
"""
def update(self):
""" Get the next frame of input and display it. """
_, self.img = self.video_source.read()
self.vtk_overlay_window.set_video_image(self.img)
self.vtk_overlay_window._RenderWindow.Render()
class OverlayOnVideoFeedCropRecord(OverlayBaseApp):
""" Add cropping of the incoming video feed, and the ability to
record the vtk_overlay_window.
:param video_source: OpenCV compatible video source (int or filename)
:param output_filename: Location of output video file when recording.
If none specified, the current date/time is
used as the filename.
"""
def __init__(self, video_source, output_filename=None):
super().__init__(video_source)
def __init__(self, video_source, output_filename=None, dims=None):
super().__init__(video_source, dims)
self.output_filename = output_filename
self.video_writer = None
self.roi = None
def update(self):
""" Get the next frame of input and write to file (if enabled). """
""" Get the next frame of input, crop and/or
write to file (if either enabled). """
_, self.img = self.video_source.read()
self.vtk_overlay_window.set_video_image(self.img)
if self.roi:
start_x, start_y = self.roi[0]
end_x, end_y = self.roi[1]
self.vtk_overlay_window.set_video_image(self.img[start_y:end_y,
start_x:end_x,
:])
else:
self.vtk_overlay_window.set_video_image(self.img)
self.vtk_overlay_window._RenderWindow.Render()
if self.save_frame:
......@@ -82,6 +112,12 @@ class OverlayOnVideoFeed(OverlayBaseApp):
self.video_writer.write_frame(output_frame,
self.video_source.timestamp)
def set_roi(self):
"""Crop the incoming video stream using ImageCropper."""
#pylint:disable=attribute-defined-outside-init
self.roi = ImageCropper().crop(self.img)
logging.debug("Setting ROI: %i", self.roi)
def get_output_frame(self):
""" Get the output frame to write in numpy format."""
output_frame = \
......
# -*- coding: utf-8 -*-
import cv2
import pytest
import mock
import numpy as np
import sksurgeryutils.common_overlay_apps as coa
def test_OverlayOnVideoFeed_from_file(setup_qt):
def test_OverlayOnVideoFeedCropRecord_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)
overlay_app = coa.OverlayOnVideoFeedCropRecord(input_file, out_file)
# Start app and get a frame from input, so that
# the window is showing something, before we start
......@@ -39,7 +40,7 @@ def test_OverlayOnVideoFeed_from_file(setup_qt):
output_video.release()
def test_OverlayOnVideoFeed_from_webcam(setup_qt):
def test_OverlayOnVideoFeedCropRecord_from_webcam(setup_qt):
"""
Test will only run if there is a camera avilable.
"""
......@@ -53,8 +54,9 @@ def test_OverlayOnVideoFeed_from_webcam(setup_qt):
cam.release()
out_file = 'tests/output/overlay_on_webcam_test.avi'
overlay_app = coa.OverlayOnVideoFeed(0, out_file)
# Don't pass an output filename as a parameter, so that
# the code to generate a filename from current date/time is executed.
overlay_app = coa.OverlayOnVideoFeedCropRecord(0)
# Start app and get a frame from input, so that
# the window is showing something, before we start
......@@ -82,3 +84,19 @@ def test_OverlayBaseAppRaisesNotImplementedError(setup_qt):
overlay_app = ErrorApp(input_file)
overlay_app.update()
roi = [(25, 25), (50, 50)]
@mock.patch('sksurgeryutils.common_overlay_apps.ImageCropper.crop')
def test_OverlayOnVideoFeedCropRecord_set_roi(mock_crop, setup_qt):
mock_crop.return_value = roi
input_file = 'tests/data/100x50_100_frames.avi'
overlay_app = coa.OverlayOnVideoFeedCropRecord(input_file)
overlay_app.update() # Get a frame so that we can crop it
overlay_app.set_roi()
overlay_app.update() # This should apply the roi to the next frame
expected_shape = (roi[1][0] - roi[0][0], roi[1][1] - roi[0][1])
assert overlay_app.vtk_overlay_window.input.shape[:2] == expected_shape
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