filled out FileCapture and added tests

src/waggle/data/vision.py

@@ -225,17 +225,39 @@ def __init__(self, device, format):
         self.device = device
         self.format = format
 
+        self.capture = cv2.VideoCapture(self.device)
+        if not self.capture.isOpened():
+            raise RuntimeError(
+                f"unable to open video capture for device {self.device!r}"
+            )
+
     def close(self):
-        print("bye file!")
+        self.capture.release()
 
     def snapshot(self):
-        pass
+        return self._grab_frame()
 
     def stream(self):
-        pass
+        try:
+            while True:
+                yield self._grab_frame()
+        except RuntimeError:
+            return
 
     def record(self):
-        raise RuntimeError("TODO write this error")
+        raise RuntimeError(
+            "Camera already opened. Camera.record must be called outside of a with block."
+        )
+
+    def _grab_frame(self):
+        ok = self.capture.grab()
+        if not ok:
+            raise RuntimeError("failed to grab frame")
+        timestamp = get_timestamp()
+        ok, data = self.capture.retrieve()
+        if not ok:
+            raise RuntimeError("failed to decode frame")
+        return ImageSample(data=data, timestamp=timestamp, format=self.format)
 
 
 class StreamCapture:
@@ -261,14 +283,32 @@ def close(self):
         self.capture.release()
 
     def snapshot(self):
-        return self.grab_frame()
+        return self._grab_frame()
 
     def stream(self):
-        while True:
-            yield self.grab_frame()
+        try:
+            while True:
+                yield self._grab_frame()
+        except RuntimeError:
+            return
 
     def record(self):
-        raise RuntimeError("TODO write this error")
+        raise RuntimeError(
+            "Camera already opened. Camera.record must be called outside of a with block."
+        )
+
+    def _grab_frame(self):
+        if not self._ready_for_next_frame.wait(timeout=10.0):
+            raise RuntimeError(
+                "failed to grab a frame from the background thread: timed out"
+            )
+        self._ready_for_next_frame.clear()
+        with acquire_with_timeout(self.lock, timeout=1.0):
+            timestamp = self.timestamp
+            ok, data = self.capture.retrieve()
+            if not ok:
+                raise RuntimeError("failed to retrieve the taken snapshot")
+        return ImageSample(data=data, timestamp=timestamp, format=self.format)
 
     def _run(self):
         # we sleep slighly shorter than FPS to drain the buffer efficiently
@@ -292,19 +332,6 @@ def _run(self):
         finally:
             self.stopped.set()
 
-    def grab_frame(self):
-        if not self._ready_for_next_frame.wait(timeout=10.0):
-            raise RuntimeError(
-                "failed to grab a frame from the background thread: timed out"
-            )
-        self._ready_for_next_frame.clear()
-        with acquire_with_timeout(self.lock, timeout=1.0):
-            timestamp = self.timestamp
-            ok, data = self.capture.retrieve()
-            if not ok:
-                raise RuntimeError("failed to retrieve the taken snapshot")
-        return ImageSample(data=data, timestamp=timestamp, format=self.format)
-
 
 class ImageFolder:
     available_formats = {".jpg", ".jpeg", ".png"}

tests/test_data.py

@@ -1,6 +1,13 @@
 import unittest
 from waggle.data.audio import AudioFolder, AudioSample
-from waggle.data.vision import RGB, BGR, ImageFolder, ImageSample, resolve_device
+from waggle.data.vision import (
+    RGB,
+    BGR,
+    ImageFolder,
+    ImageSample,
+    resolve_device,
+    Camera,
+)
 from waggle.data.timestamp import get_timestamp
 import numpy as np
 from tempfile import TemporaryDirectory
@@ -118,6 +125,20 @@ def test_get_timestamp(self):
         ts = get_timestamp()
         self.assertIsInstance(ts, int)
 
+    def test_camera_file_snapshot(self):
+        cam = Camera("file://tests/test.mp4")
+        sample = cam.snapshot()
+        assert sample.data.shape == (640, 480, 3)
+
+    def test_camera_file_stream(self):
+        cam = Camera("file://tests/test.mp4")
+        numframes = 0
+        for sample in cam.stream():
+            assert sample.data.shape == (640, 480, 3)
+            numframes += 1
+        # NOTE test.mp4 was created with exactly 90 frames, so we check it matches
+        assert numframes == 90
+
 
 if __name__ == "__main__":
     unittest.main()