Implement test for slicing of streams, saving data through streams, and the "general" filter

bob/io/stream/test/test.py

@@ -5,15 +5,16 @@
 Test Units
 """
 # ==============================================================================
-from bob.io.stream import StreamFile, Stream
-from bob.ip.stereo import StereoParameters
-from bob.ip.stereo.test.test import is_close_enough
+import os
 
-from bob.io.base import load
+import numpy as np
 
+from bob.io.base import load, HDF5File
 from pkg_resources import resource_filename
+from bob.ip.stereo import StereoParameters
+from bob.ip.stereo.test.test import is_close_enough
 
-import numpy as np
+from bob.io.stream import StreamFile, Stream
 
 # ==============================================================================
 
@@ -29,11 +30,115 @@ def resource_path(relative_path, package="bob.io.stream"):
     return resource_filename(package, relative_path)
 
 
+def test_stream_write():
+    """Test writing and reading back a file through Stream and StreamFile."""
+
+    test_data = [
+        np.arange(5 * 3 * 40 * 52, dtype=np.int16).reshape((5, 3, 40, 52)),
+        np.arange(5 * 1 * 5 * 5, dtype=np.int8).reshape((5, 1, 5, 5)),
+        np.arange(1 * 1 * 500 * 400).astype(np.float).reshape((1, 1, 500, 400)),
+        np.arange(12 * 52).astype(np.float64).reshape((12, 52)),
+    ]
+
+    for a_test_data in test_data:
+        with StreamFile(resource_path("test/data/save_test.h5"), mode="w") as output_file:
+            stream = Stream("test_data")
+            save_filter = stream.save(output_file)
+            for i in range(a_test_data.shape[0]):
+                stream.put(a_test_data[i])
+
+        with StreamFile(resource_path("test/data/save_test.h5"), mode="r") as input_file:
+            stream = Stream("test_data", input_file)
+            data = stream.load()
+
+        assert np.array_equal(data, a_test_data)
+        assert data.dtype == a_test_data.dtype
+
+    os.remove(resource_path("test/data/save_test.h5"))
+
+
 def test_stream():
-    """
-    Test that a few transforms (stereo, reproject, adjust, colormap, select, warp) applied with streams provide a 
-    similar result to saved groundtruth.
-    """
+    """Test some functionality of the stream class: shape, ndim, slicing (view), etc..."""
+
+    # create data sets
+    data_shape = (10, 3, 40, 30)  # #frames, #channels, with, height
+    data_a = np.random.random_sample(data_shape)
+    data_b = np.random.random_integers(5000, size=data_shape)
+
+    # create data file
+    f = HDF5File(resource_path("test/data/stream_test.h5"), "w")
+    f.set("data_a", data_a)
+    f.set("data_b", data_b)
+    del f
+
+    # Streams attributes when config is specified
+    f = StreamFile(
+        resource_path("test/data/input_example.h5"),
+        resource_path("config/idiap_face_streams.json"),
+        resource_path("config/idiap_face_calibration.json", "bob.ip.stereo"),
+    )
+    color = Stream("color", f)
+    assert color.shape == (1, 3, 1920, 1200)
+    assert color.timestamps[0] == 46399548
+    assert color.camera is not None
+
+    # Streams attributes when not specified.
+    f = StreamFile(resource_path("test/data/stream_test.h5"))
+    stream_a = Stream("data_a", f)
+    stream_b = Stream("data_b", f)
+    assert stream_a.shape == data_a.shape
+    assert stream_b.shape == data_b.shape
+    assert stream_a.timestamps == None
+    assert stream_b.timestamps == None
+    assert stream_a.camera == None
+    assert stream_b.camera == None
+
+    # Test loading entire datasets
+    ld_a = stream_a.load()
+    ld_b = stream_b.load()
+    assert ld_a.shape == data_a.shape
+    assert ld_b.shape == data_b.shape
+    assert np.array_equal(ld_a, data_a)
+    assert np.array_equal(ld_b, data_b)
+
+    # Test slicing over the first dimension
+    test_slices = [
+        slice(None, None, None),
+        slice(5, None, None),
+        slice(None, 3, None),
+        slice(None, None, 3),
+        slice(1, 10, 3),
+        slice(9, 0, -3),
+        slice(-5, -1, None),
+        slice(10, 0, -3),
+    ]
+
+    for a_slice in test_slices:
+        gt_slice = data_a[a_slice]
+        s_slice = stream_a[a_slice].load()
+        s_l_slice = stream_a.load(a_slice)
+        assert np.array_equal(s_slice, gt_slice), "Slice " + str(a_slice) + " assertation failed."
+        assert np.array_equal(s_l_slice, gt_slice), "Slice " + str(a_slice) + " assertation failed."
+
+    # test slicing over other dimensions:
+    test_indices = [
+        (slice(None, None, None), 1),
+        (slice(None, None, None), slice(None, None, None), slice(1, 2, 3), slice(4, 5, 6)),
+        (slice(None, None, None), slice(None, None, None), slice(-1, -6, -3), slice(-4, 5, -6)),
+    ]
+
+    for index in test_indices:
+        gt_slice = data_a[index]
+        s_slice = stream_a[index]
+        assert s_slice.shape == gt_slice.shape, "index " + str(index) + " assertation failed."
+        assert s_slice.ndim == gt_slice.ndim, "index " + str(index) + " assertation failed."
+        assert np.array_equal(s_slice.load(), gt_slice), "index " + str(index) + " assertation failed."
+
+    os.remove(resource_path("test/data/stream_test.h5"))
+
+
+def test_filters():
+    """Test that a few filters provide a similar result to saved groundtruth."""
 
     gt_color = load(resource_path("test/data/reprojection_color.png"))
     gt_depth = load(resource_path("test/data/reprojection_depth.png"))
@@ -80,6 +185,10 @@ def test_stream():
     warp_swir_norm = swir_norm.warp(color)
     warp_thermal = thermal.normalize().warp(color)
 
+    # User defined operations through "filter": eg clipping values bellow average
+    test_func = lambda data: np.where(data > data.mean(), data - data.mean(), data.mean())
+    user_filter = color.filter(process_frame=test_func)
+
     # landmarks and bounding box
     color.bounding_box[0] = bounding_box
     color.image_points[0] = landmarks
@@ -104,3 +213,6 @@ def test_stream():
 
     assert np.array_equal(rep_color.bounding_box[0], gt_bounding_box)
     assert np.allclose(rep_color.image_points[0], gt_landmarks)
+
+    # user filter results
+    assert np.array_equal(user_filter[0], test_func(color[0]))