Commit 8cace2f8 authored by Amir MOHAMMADI's avatar Amir MOHAMMADI

Merge branch 'handle-none-data' into 'master'

pickle VideoLikeContainers when data dtype is not supported by hdf5 format

Closes #18

See merge request !45
parents 3b4f0a95 d9e65a6f
Pipeline #53690 passed with stage
in 16 minutes and 38 seconds
import tempfile
import time import time
import bob.bio.video
import numpy as np
from bob.bio.base.test.utils import is_library_available from bob.bio.base.test.utils import is_library_available
from bob.io.base.test_utils import datafile from bob.io.base.test_utils import datafile
from bob.io.video import reader from bob.io.video import reader
import numpy as np
import bob.bio.video
regenerate_refs = False regenerate_refs = False
...@@ -24,9 +26,7 @@ def test_video_as_array(): ...@@ -24,9 +26,7 @@ def test_video_as_array():
video = video[None, ...] video = video[None, ...]
np.testing.assert_allclose(video, video_slice) np.testing.assert_allclose(video, video_slice)
video = bob.bio.video.VideoAsArray( video = bob.bio.video.VideoAsArray(path, max_number_of_frames=3)
path, max_number_of_frames=3
)
assert len(video) == 3, len(video) assert len(video) == 3, len(video)
assert video.indices == [13, 41, 69], video.indices assert video.indices == [13, 41, 69], video.indices
assert video.shape == (3, 3, 480, 640), video.shape assert video.shape == (3, 3, 480, 640), video.shape
...@@ -68,8 +68,14 @@ def test_video_like_container(): ...@@ -68,8 +68,14 @@ def test_video_like_container():
container.save(container_path) container.save(container_path)
loaded_container = bob.bio.video.VideoLikeContainer.load(container_path) loaded_container = bob.bio.video.VideoLikeContainer.load(container_path)
assert container == loaded_container
np.testing.assert_equal(np.array(container.data), np.array(loaded_container.data)) # test saving and loading None arrays
np.testing.assert_equal( with tempfile.NamedTemporaryFile(suffix=".pkl") as f:
np.array(container.indices), np.array(loaded_container.indices) data = [None] * 10 + [1]
) indices = range(11)
frame_container = bob.bio.video.VideoLikeContainer(data, indices)
frame_container.save(f.name)
loaded = bob.bio.video.VideoLikeContainer.load(f.name)
assert loaded == frame_container
import logging import logging
import pickle
from bob.bio.base import selected_indices
from bob.io.video import reader
import h5py import h5py
import numpy as np import numpy as np
from bob.bio.base import selected_indices
from bob.io.video import reader
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
...@@ -226,21 +227,34 @@ class VideoLikeContainer: ...@@ -226,21 +227,34 @@ class VideoLikeContainer:
def __array__(self, dtype=None, *args, **kwargs): def __array__(self, dtype=None, *args, **kwargs):
return np.asarray(self.data, dtype, *args, **kwargs) return np.asarray(self.data, dtype, *args, **kwargs)
def __eq__(self, o: object) -> bool:
return np.array_equal(self.data, o.data) and np.array_equal(
self.indices, o.indices
)
def save(self, file): def save(self, file):
self.save_function(self, file) self.save_function(self, file)
@staticmethod @staticmethod
def save_function(other, file): def save_function(other, file):
try:
with h5py.File(file, mode="w") as f: with h5py.File(file, mode="w") as f:
f["data"] = other.data f["data"] = other.data
f["indices"] = other.indices f["indices"] = other.indices
# revert to saving data in pickles when the dtype is not supported by hdf5
except TypeError:
with open(file, "wb") as f:
pickle.dump({"data": other.data, "indices": other.indices}, f)
@classmethod @classmethod
def load(cls, file): def load(cls, file):
try:
# weak closing of the hdf5 file so we don't load all the data into # weak closing of the hdf5 file so we don't load all the data into
# memory https://docs.h5py.org/en/stable/high/file.html#closing-files # memory https://docs.h5py.org/en/stable/high/file.html#closing-files
f = h5py.File(file, mode="r") f = h5py.File(file, mode="r")
data = f["data"] loaded = {"data": f["data"], "indices": f["indices"]}
indices = f["indices"] except OSError:
self = cls(data=data, indices=indices) with open(file, "rb") as f:
loaded = pickle.load(f)
self = cls(**loaded)
return self return self
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