def test_basic_datastore_operations(self):
# 2 records per catalog entry in the manifest
manifest = Manifest(self._path, max_len=2)
count = 10
for i in range(count):
manifest.write_record(self._newRecord())
read_records = 0
for entry in manifest:
print('Entry %s' % (entry))
read_records += 1
self.assertEqual(count, read_records)
def test_memory_mapped_read(self):
manifest = Manifest(self._path, max_len=2)
for i in range(10):
manifest.write_record(self._newRecord())
manifest.close()
manifest_2 = Manifest(self._path, read_only=True)
read_records = 0
for _ in manifest_2:
read_records += 1
manifest_2.close()
self.assertEqual(10, read_records)
def __init__(self,
base_path,
inputs=[],
types=[],
metadata=[],
max_catalog_len=1000,
read_only=False):
self.base_path = base_path
self.images_base_path = os.path.join(self.base_path, Tub.images())
self.inputs = inputs
self.types = types
self.metadata = metadata
self.manifest = Manifest(base_path,
inputs=inputs,
types=types,
metadata=metadata,
max_len=max_catalog_len,
read_only=read_only)
self.input_types = dict(zip(self.inputs, self.types))
# Create images folder if necessary
if not os.path.exists(self.images_base_path):
os.makedirs(self.images_base_path, exist_ok=True)
def test_delete_and_restore_by_set(self):
manifest = Manifest(self._path, max_len=2)
count = 10
deleted = range(3, 7)
for i in range(count):
manifest.write_record(self._newRecord())
manifest.delete_records(deleted)
read_records = 0
for entry in manifest:
print(f'Entry {entry}')
read_records += 1
self.assertEqual(count - len(deleted), read_records)
manifest.restore_records(deleted)
read_records = 0
for entry in manifest:
print(f'Entry {entry}')
read_records += 1
self.assertEqual(count, read_records)
def test_deletion(self):
manifest = Manifest(self._path, max_len=2)
count = 10
deleted = 5
for i in range(count):
manifest.write_record(self._newRecord())
for i in range(deleted):
manifest.delete_record(i)
read_records = 0
for entry in manifest:
print('Entry %s' % (entry))
read_records += 1
self.assertEqual((count - deleted), read_records)
class Tub(object):
"""
A datastore to store sensor data in a key, value format. \n
Accepts str, int, float, image_array, image, and array data types.
"""
def __init__(self,
base_path,
inputs=[],
types=[],
metadata=[],
max_catalog_len=1000,
read_only=False):
self.base_path = base_path
self.images_base_path = os.path.join(self.base_path, Tub.images())
self.inputs = inputs
self.types = types
self.metadata = metadata
self.manifest = Manifest(base_path,
inputs=inputs,
types=types,
metadata=metadata,
max_len=max_catalog_len,
read_only=read_only)
self.input_types = dict(zip(self.inputs, self.types))
# Create images folder if necessary
if not os.path.exists(self.images_base_path):
os.makedirs(self.images_base_path, exist_ok=True)
def write_record(self, record=None):
"""
Can handle various data types including images.
"""
contents = dict()
for key, value in record.items():
if value is None:
continue
elif key not in self.input_types:
continue
else:
input_type = self.input_types[key]
if input_type == 'float':
# Handle np.float() types gracefully
contents[key] = float(value)
elif input_type == 'str':
contents[key] = value
elif input_type == 'int':
contents[key] = int(value)
elif input_type == 'boolean':
contents[key] = bool(value)
elif input_type == 'list' or input_type == 'vector':
contents[key] = list(value)
elif input_type == 'image_array':
# Handle image array
image = Image.fromarray(np.uint8(value))
name = Tub._image_file_name(self.manifest.current_index,
key)
image_path = os.path.join(self.images_base_path, name)
image.save(image_path)
contents[key] = name
# Private properties
contents['_timestamp_ms'] = int(round(time.time() * 1000))
contents['_index'] = self.manifest.current_index
self.manifest.write_record(contents)
def delete_record(self, record_index):
self.manifest.delete_record(record_index)
def delete_last_n_records(self, n):
last_index = self.manifest.current_index
first_index = last_index - n
for index in range(first_index, last_index):
if index < 0:
continue
else:
self.manifest.delete_record(index)
def close(self):
self.manifest.close()
def __iter__(self):
return ManifestIterator(self.manifest)
def __len__(self):
return self.manifest.__len__()
@classmethod
def images(cls):
return 'images'
@classmethod
def _image_file_name(cls, index, key, extension='.jpg'):
key_prefix = key.replace('/', '_')
name = '_'.join([str(index), key_prefix, extension])
# Return relative paths to maintain portability
return name