Commit 7b1318ac authored by Matt Clarkson's avatar Matt Clarkson

Merge branch '44-remove-extension'

parents 47a8629e f31434ba
Pipeline #3487 passed with stages
in 7 minutes and 5 seconds
......@@ -12,7 +12,10 @@ import cv2
import numpy as np
LOGGER = logging.getLogger(__name__)
#pylint:disable=useless-object-inheritance
# pylint:disable=useless-object-inheritance
class VideoWriter(object):
"""
Class to write images to disk using cv2.VideoWriter.
......@@ -24,6 +27,9 @@ class VideoWriter(object):
"""
def __init__(self, filename, fps=25, width=640, height=480, codec='MJPG'):
self.filename = None
# This also validates.
self.set_filename(filename)
fourcc = cv2.VideoWriter_fourcc(*codec)
......@@ -88,7 +94,6 @@ class VideoWriter(object):
self.video_writer.write(frame)
class TimestampedVideoWriter(VideoWriter):
"""
Class to write images and timestamps to disk, inherits from VideoWriter.
......@@ -103,11 +108,13 @@ class TimestampedVideoWriter(VideoWriter):
super(TimestampedVideoWriter, self).__init__(filename, fps,
width, height, codec)
timestamp_filename = filename + '.timestamps'
suffix = '.timestamps.txt'
basename, _ = os.path.splitext(filename)
timestamp_filename = basename + suffix
self.timestamp_file = open(timestamp_filename, 'w')
self.default_timestamp_message = "NO_TIMESTAMP"
def close(self):
""" Close/release the output files for video and timestamps. """
logging.debug("Closing video writer")
......@@ -137,6 +144,7 @@ class TimestampedVideoWriter(VideoWriter):
# Convert datetime object to string
self.timestamp_file.write(timestamp.isoformat() + '\n')
class ThreadedTimestampedVideoWriter(TimestampedVideoWriter):
""" TimestampedVideoWriter that can be run in a thread.
Uses Queue.Queue() to store data, which is thread safe.
......
......@@ -17,21 +17,25 @@ def test_videowriter_invalid_filename_raises_error():
with pytest.raises(ValueError):
vw.VideoWriter(invalid_filename, fps, width, height)
@mock.patch('os.makedirs')
def test_videowriter_create_dir(mocked_makedirs):
filename = 'new_dir/file.avi'
vw.VideoWriter(filename, fps, width, height)
mocked_makedirs.assert_called_with('new_dir')
def test_timestamp_videowriter_creates_timestamp_file(tmpdir):
filename = os.path.join(tmpdir.dirname, 'test.avi')
vw.TimestampedVideoWriter(filename, fps, width, height)
basename, _ = os.path.splitext(filename)
assert os.path.isfile(filename)
assert os.path.isfile(filename + '.timestamps')
assert os.path.isfile(os.path.join(basename + '.timestamps.txt'))
# Video writer is tested further in test_integration
def test_default_timestamp_written(tmpdir):
# Write a frame without passing in a timestamp.
# If no timestamp is passed, "NO_TIMESTAMP" should be written to
......@@ -45,7 +49,8 @@ def test_default_timestamp_written(tmpdir):
video_writer.write_frame(frame)
video_writer.close()
timestamp_file = filename + '.timestamps'
basename, _ = os.path.splitext(filename)
timestamp_file = os.path.join(basename + '.timestamps.txt')
expected_timestamp = video_writer.default_timestamp_message + '\n'
with open(timestamp_file) as f:
......
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