class TestTub(unittest.TestCase):
def setUp(self):
self._path = tempfile.mkdtemp()
inputs = ['input']
types = ['int']
self.tub = Tub(self._path, inputs, types)
def test_basic_tub_operations(self):
entries = list(self.tub)
self.assertEqual(len(entries), 0)
write_count = 10
delete_indexes = [0, 8]
records = [{'input': i} for i in range(write_count)]
for record in records:
self.tub.write_record(record)
for index in delete_indexes:
self.tub.delete_records(index)
count = 0
for record in self.tub:
print('Record %s' % (record))
count += 1
self.assertEqual(count, (write_count - len(delete_indexes)))
self.assertEqual(len(self.tub), (write_count - len(delete_indexes)))
def tearDown(self):
shutil.rmtree(self._path)
def car_dir(tmpdir_factory, base_config, imu_fields) -> str:
""" Creating car dir with sub dirs and extracting tub """
car_dir = tmpdir_factory.mktemp('mycar')
os.mkdir(os.path.join(car_dir, 'models'))
# extract tub.tar.gz into car_dir/tub
this_dir = os.path.dirname(os.path.abspath(__file__))
with tarfile.open(os.path.join(this_dir, 'tub', 'tub.tar.gz')) as file:
file.extractall(car_dir)
# now create a second tub with additonal imu data
tub_dir = os.path.join(car_dir, 'tub')
tub = Tub(base_path=tub_dir)
full_dir = os.path.join(car_dir, 'tub_full')
tub_full = Tub(base_path=full_dir,
inputs=tub.manifest.inputs + imu_fields +
['behavior/one_hot_state_array', 'localizer/location'],
types=tub.manifest.types + ['float'] * 6 + ['list', 'int'])
count = 0
for record in tub:
t = TubRecord(base_config, tub.base_path, record)
img = t.image()
record['cam/image_array'] = img
for field in imu_fields:
record[field] = np.random.rand()
# add behavioural input
bhv = [1., 0.] if count < len(tub) // 2 else [0., 1.]
record["behavior/one_hot_state_array"] = bhv
record['localizer/location'] = 3 * count // len(tub)
tub_full.write_record(record)
count += 1
return car_dir
def benchmark():
# Change to a non SSD storage path
path = Path('/media/rahulrav/Cruzer/benchmark')
# Recreate paths
if os.path.exists(path.absolute().as_posix()):
shutil.rmtree(path)
inputs = ['input']
types = ['int']
tub = Tub(path.as_posix(), inputs, types, max_catalog_len=1000)
write_count = 1000
for i in range(write_count):
record = {'input': i}
tub.write_record(record)
deletions = np.random.randint(0, write_count, 100)
for index in deletions:
index = int(index)
tub.delete_record(index)
for record in tub:
print('Record %s' % record)
tub.close()
def convert_to_tub_v2(paths, output_path):
"""
Convert from old tubs to new one
:param paths: legacy tub paths
:param output_path: new tub output path
:return: None
"""
empty_record = {'__empty__': True}
if type(paths) is str:
paths = [paths]
legacy_tubs = [LegacyTub(path) for path in paths]
print(f'Total number of tubs: {len(legacy_tubs)}')
for legacy_tub in legacy_tubs:
# add input and type for empty records recording
inputs = legacy_tub.inputs + ['__empty__']
types = legacy_tub.types + ['boolean']
output_tub = Tub(output_path, inputs, types,
list(legacy_tub.meta.items()))
record_paths = legacy_tub.gather_records()
bar = IncrementalBar('Converting', max=len(record_paths))
previous_index = None
for record_path in record_paths:
try:
contents = Path(record_path).read_text()
record = json.loads(contents)
image_path = record['cam/image_array']
ms = record['milliseconds']
current_index = int(image_path.split('_')[0])
image_path = os.path.join(legacy_tub.path, image_path)
image_data = Image.open(image_path)
record['cam/image_array'] = image_data
# first record or they are continuous, just append
if not previous_index or current_index == previous_index + 1:
output_tub.write_record(record, ms)
previous_index = current_index
# otherwise fill the gap with empty records
else:
# Skipping over previous record here because it has
# already been written.
previous_index += 1
# Adding empty record nodes, and marking them deleted
# until the next valid record.
delete_list = []
while previous_index < current_index:
idx = output_tub.manifest.current_index
output_tub.write_record(empty_record, ms)
delete_list.append(idx)
previous_index += 1
output_tub.delete_records(delete_list)
bar.next()
except Exception as exception:
print(f'Ignoring record path {record_path}\n', exception)
traceback.print_exc()
# writing session id into manifest metadata
output_tub.close()