def test_get_last_ix_for_one_record(tub_path):
inputs = ['cam/image_array', 'angle', 'throttle']
types = ['image_array', 'float', 'float']
t = Tub(tub_path, inputs=inputs, types=types)
record = create_sample_record()
t.put_record(record)
assert t.get_last_ix() == 0
def create_sample_tub(path, records=128):
inputs = [
'cam/image_array', 'user/angle', 'user/throttle',
'location/one_hot_state_array'
]
types = ['image_array', 'float', 'float', 'vector']
t = Tub(path, inputs=inputs, types=types)
cam = SquareBoxCamera()
tel = MovingSquareTelemetry()
num_loc = 10
for _ in range(records):
x, y = tel.run()
img_arr = cam.run(x, y)
loc = [0 for i in range(num_loc)]
loc[1] = 1.0
t.put_record({
'cam/image_array': img_arr,
'user/angle': x,
'user/throttle': y,
'location/one_hot_state_array': loc
})
global temp_tub_path
temp_tub_path = t
print("setting temp tub path to:", temp_tub_path)
return t
def benchmark():
# Change with a non SSD storage path
path = Path('/media/rahulrav/Cruzer/tub')
if os.path.exists(path.absolute().as_posix()):
shutil.rmtree(path)
inputs = ['input']
types = ['int']
tub = Tub(path.absolute().as_posix(), inputs, types)
write_count = 1000
for i in range(write_count):
record = {'input': i}
tub.put_record(record)
# old tub starts counting at 1
deletions = set(np.random.randint(1, write_count + 1, 100))
for index in deletions:
index = int(index)
tub.remove_record(index)
files = path.glob('*.json')
for record_file in files:
contents = record_file.read_text()
if contents:
contents = json.loads(contents)
print('Record %s' % contents)
def test_tub_put_unknown_type(tub_path):
""" Creating a record with unknown type should fail """
inputs = ['user/speed']
types = ['bob']
t=Tub(path=tub_path, inputs=inputs, types=types)
with pytest.raises(TypeError):
t.put_record({'user/speed': 0.2, })
def create_sample_tub(path, records=10):
inputs = ['cam/image_array', 'angle', 'throttle']
types = ['image_array', 'float', 'float']
t = Tub(path, inputs=inputs, types=types)
for _ in range(records):
record = create_sample_record()
t.put_record(record)
return t
def test_tub_put_image(tub_path):
""" Add an encoded image to the tub """
inputs = ['user/speed', 'cam/image']
types = ['float', 'image']
img = Image.new('RGB', (120, 160))
t=Tub(path=tub_path, inputs=inputs, types=types)
t.put_record({'cam/image': img, 'user/speed': 0.2, })
assert t.get_record(t.get_last_ix())['user/speed'] == 0.2
def create_sample_tub(path, records=10):
inputs = ['cam/image_array', 'angle', 'throttle']
types = ['image_array', 'float', 'float']
t = Tub(path, inputs=inputs, types=types)
cam = SquareBoxCamera()
tel = MovingSquareTelemetry()
for _ in range(records):
x, y = tel.run()
img_arr = cam.run(x, y)
t.put_record({'cam/image_array': img_arr, 'angle': x, 'throttle': y})
return t
def create_sample_tub(path, records=128):
inputs = ['cam/image_array', 'user/angle', 'user/throttle']
types = ['image_array', 'float', 'float']
t = Tub(path, inputs=inputs, types=types)
cam = SquareBoxCamera()
tel = MovingSquareTelemetry()
for _ in range(records):
x, y = tel.run()
img_arr = cam.run(x, y)
t.put_record({
'cam/image_array': img_arr,
'user/angle': x,
'user/throttle': y
})
global temp_tub_path
temp_tub_path = t
print("setting temp tub path to:", temp_tub_path)
return t
def augment(tub_names, new_data_dir, args):
new_data_dir = os.path.expanduser(new_data_dir)
tubgroup = TubGroup(tub_names)
# If tub directory does not exist, create directory
if not os.path.exists(new_data_dir):
os.makedirs(new_data_dir)
# If directory does not contain meta.json, copy one from the first source tub
if not os.path.exists(os.path.join(new_data_dir, 'meta.json')):
copyfile(src=tubgroup.tubs[0].meta_path,
dst=os.path.join(new_data_dir, 'meta.json'))
new_tub = Tub(new_data_dir)
for tub in tubgroup.tubs:
for ix in tub.get_index(shuffled=False):
record = tub.get_record(ix)
for augmented_record in augment_single_record(record, args):
new_tub.put_record(augmented_record)