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()