def get_dataset_paths(dataset_name):
if dataset_name == "push_box_hardware":
episodes_root = os.path.join(get_data_ssd_root(),
"dataset/push_box_hardware")
episodes_config = load_yaml(
os.path.join(
get_project_root(),
'experiments/exp_22_push_box_hardware/push_box_hardware_episodes_config.yaml'
))
transporter_model_chkpt = None
dense_descriptor_model_chkpt = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_string_pull/trained_models/perception/dense_descriptors/data_aug_2020-07-02-02-39-27-400442/net_best_model.pth"
return {
'dataset_name': dataset_name,
'dataset_root': episodes_root,
'episodes_config': episodes_config,
'main_camera_name': 'd415_01',
'dense_descriptor_camera_list': ['d415_01', 'd415_02'],
'transporter_model_chkpt': transporter_model_chkpt,
'dense_descriptor_model_chkpt': dense_descriptor_model_chkpt,
}
elif dataset_name == "push_box_string_pull":
episodes_root = os.path.join(get_data_ssd_root(),
"dataset/push_box_string_pull")
episodes_config = load_yaml(
os.path.join(
get_project_root(),
'experiments/exp_22_push_box_hardware/push_box_string_pull_episodes_config.yaml'
))
transporter_model_chkpt = None
dense_descriptor_model_chkpt = "/home/manuelli/data/key_dynam/dev/experiments/22/dataset_push_box_string_pull/trained_models/perception/dense_descriptors/data_aug_2020-07-02-02-39-27-400442/net_best_model.pth"
return {
'dataset_name': dataset_name,
'dataset_root': episodes_root,
'episodes_config': episodes_config,
'main_camera_name': 'd415_01',
'dense_descriptor_camera_list': ['d415_01', 'd415_02'],
'transporter_model_chkpt': transporter_model_chkpt,
'dense_descriptor_model_chkpt': dense_descriptor_model_chkpt,
}
else:
raise ValueError("unknown dataset:", dataset_name)
def top_down_dataset_root():
dataset_name = "2020-04-20-14-58-21-418302_T_aug_random_velocity_1000"
# dataset_root = os.path.join(get_data_root(), "dev/experiments/09/data", dataset_name)
dataset_root = os.path.join(get_data_ssd_root(), 'dataset', dataset_name)
config = load_yaml(os.path.join(dataset_root, 'config.yaml'))
return {'dataset_name': dataset_name,
'dataset_root': dataset_root,
'config': config,
'main_camera_name': 'camera_1_top_down'}
def single_corelle_mug_600():
dataset_name = "single_corelle_mug_600"
dataset_root = os.path.join(get_data_ssd_root(), 'dataset', dataset_name)
config = load_yaml(os.path.join(dataset_root, 'config.yaml'))
return {
'dataset_name': dataset_name,
'dataset_root': dataset_root,
'config': config,
'main_camera_name': 'camera_1_top_down',
}
def mugs_random_colors_1000():
dataset_name = "mugs_random_colors_1000"
dataset_root = os.path.join(get_data_ssd_root(), 'dataset', dataset_name)
config = load_yaml(os.path.join(dataset_root, 'config.yaml'))
return {
'dataset_name': dataset_name,
'dataset_root': dataset_root,
'config': config,
'main_camera_name': 'camera_1_top_down',
}
def angled_cam_dataset_root():
dataset_name = "2020-04-23-20-45-12-697915_T_aug_random_velocity_1000_angled_cam"
# prepare folders
# dataset_root = os.path.join(get_data_root(), 'dev/experiments/10/data', dataset_name)
dataset_root = os.path.join(get_data_ssd_root(), 'dataset', dataset_name)
config = load_yaml(os.path.join(dataset_root, 'config.yaml'))
return {'dataset_name': dataset_name,
'dataset_root': dataset_root,
'config': config,
'main_camera_name': 'camera_angled',
}
def get_dataset_paths(dataset_name):
if dataset_name == "2020-04-20-14-58-21-418302_T_aug_random_velocity_1000":
return top_down_dataset_root()
elif dataset_name == "2020-04-23-20-45-12-697915_T_aug_random_velocity_1000_angled_cam":
return angled_cam_dataset_root()
elif dataset_name == "box_push_1000_top_down":
dataset_root = os.path.join(get_data_ssd_root(), 'dataset', "box_push_1000")
config = load_yaml(os.path.join(dataset_root, 'config.yaml'))
transporter_model_chkpt = "/home/manuelli/data/key_dynam/dev/experiments/drake_pusher_slider_v2/dataset_box_push_1000_top_down/trained_models/perception/transporter/transporter_standard_2020-06-14-22-29-31-256422/train_nKp6_invStd10.0/net_best.pth"
dense_descriptor_model_chkpt = "/home/manuelli/data/key_dynam/dev/experiments/drake_pusher_slider_v2/dataset_box_push_1000_top_down/trained_models/perception/dense_descriptors/data_aug_2020-06-14-21-47-52-389769/net_best_model.pth"
return {'dataset_name': dataset_name,
'dataset_root': dataset_root,
'config': config,
'main_camera_name': 'camera_1_top_down',
'dense_descriptor_camera_list': ['camera_1_top_down', 'camera_2_top_down_rotated'],
'transporter_model_chkpt': transporter_model_chkpt,
'dense_descriptor_model_chkpt': dense_descriptor_model_chkpt,
}
elif dataset_name == "box_push_1000_angled":
dataset_root = os.path.join(get_data_ssd_root(), 'dataset', "box_push_1000")
config = load_yaml(os.path.join(dataset_root, 'config.yaml'))
transporter_model_chkpt = "/home/manuelli/data/key_dynam/dev/experiments/drake_pusher_slider_v2/dataset_box_push_1000_angled/trained_models/perception/transporter/transporter_standard_2020-06-15-18-35-52-478769/train_nKp6_invStd10.0/net_best.pth"
dense_descriptor_model_chkpt = "/home/manuelli/data/key_dynam/dev/experiments/drake_pusher_slider_v2/dataset_box_push_1000_angled/trained_models/perception/dense_descriptors/data_aug_2020-06-15-15-39-24-127276/net_best_model.pth"
return {'dataset_name': dataset_name,
'dataset_root': dataset_root,
'config': config,
'main_camera_name': 'camera_angled',
'dense_descriptor_camera_list': ['camera_angled', 'camera_angled_rotated'],
'transporter_model_chkpt': transporter_model_chkpt,
'dense_descriptor_model_chkpt': dense_descriptor_model_chkpt,
}
else:
raise ValueError("unknown dataset:", dataset_name)
def box_on_side_dataset_root():
dataset_name = "dps_box_on_side_600"
# dataset_root = os.path.join(get_data_root(), "dev/experiments/18/data", dataset_name)
dataset_root = os.path.join(get_data_ssd_root(), 'dataset', dataset_name)
config = load_yaml(os.path.join(dataset_root, 'config.yaml'))
dense_descriptor_model_chkpt = "/home/manuelli/data/key_dynam/dev/experiments/drake_pusher_slider_box_on_side/dataset_dps_box_on_side_600/trained_models/perception/dense_descriptor/3D_loss_camera_angled_2020-05-13-23-39-35-818188/net_best_dy_model.pth"
return {'dataset_name': dataset_name,
'dataset_root': dataset_root,
'config': config,
'main_camera_name': 'camera_angled',
'dense_descriptor_model_chkpt': dense_descriptor_model_chkpt,
}
def correlle_mug_small_many_colors_600():
dataset_name = "correlle_mug-small_many_colors_600"
dataset_root = os.path.join(get_data_ssd_root(), 'dataset', dataset_name)
config = load_yaml(os.path.join(dataset_root, 'config.yaml'))
dense_descriptor_model_chkpt = "/home/manuelli/data/key_dynam/dev/experiments/20/dataset_correlle_mug-small_many_colors_600/trained_models/perception/dense_descriptors/data_aug_2020-06-03-16-41-29-740641/net_best_model.pth"
return {
'dataset_name': dataset_name,
'dataset_root': dataset_root,
'config': config,
'main_camera_name': 'camera_1_top_down',
"dense_descriptor_model_chkpt": dense_descriptor_model_chkpt,
}
def main():
start_time = time.time()
config = load_yaml(os.path.join(get_project_root(), 'experiments/drake_pusher_slider/env_config.yaml'))
config['dataset']['num_episodes'] = 1000 # half for train, half for valid
set_seed(500) # just randomly chosen
num_episodes = config['dataset']['num_episodes']
DATASET_NAME = "box_push_%d" %(num_episodes)
OUTPUT_DIR = os.path.join(get_data_ssd_root(), 'dataset', DATASET_NAME)
if not os.path.exists(OUTPUT_DIR):
os.makedirs(OUTPUT_DIR)
collect_episodes(
config,
output_dir=OUTPUT_DIR,
visualize=False,
debug=False)
elapsed = time.time() - start_time
print("Generating and saving dataset to disk took %d seconds" % (int(elapsed)))
def multiprocess_main(num_episodes=1000, num_threads=4):
set_seed(500) # just randomly chosen
start_time = time.time()
config = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_20_mugs/config.yaml'))
num_episodes_per_thread = math.ceil(num_episodes / num_threads)
num_episodes = num_threads * num_episodes_per_thread
# DATASET_NAME = "mugs_random_colors_%d" % (num_episodes)
# DATASET_NAME = "single_mug_%d"
# DATASET_NAME = "correlle_mug-small_single_color_%d" %(num_episodes)
# DATASET_NAME = "single_corelle_mug_%d" %(num_episodes)
# DATASET_NAME = "correlle_mug-small_many_colors_%d" %(num_episodes)
DATASET_NAME = "correlle_mug-small_many_colors_random_%d" % (num_episodes)
# OUTPUT_DIR = os.path.join(get_data_root(), 'sandbox', DATASET_NAME)
OUTPUT_DIR = os.path.join(get_data_ssd_root(), 'dataset', DATASET_NAME)
print("OUTPUT_DIR:", OUTPUT_DIR)
output_dir = OUTPUT_DIR
if not os.path.exists(output_dir):
os.makedirs(output_dir)
def f(q_tmp):
config = load_yaml(
os.path.join(get_project_root(),
'experiments/exp_20_mugs/config.yaml'))
config['dataset']['num_episodes'] = num_episodes_per_thread
out = collect_episodes(config,
output_dir=OUTPUT_DIR,
visualize=False,
debug=False,
run_from_thread=True)
q_tmp.put(out)
q = Queue()
process_list = []
for i in range(num_threads):
p = Process(target=f, args=(q, ))
p.start()
process_list.append(p)
metadata = {'episodes': {}}
for p in process_list:
while p.is_alive():
p.join(timeout=1)
# empty out the queue
while not q.empty():
out = q.get()
metadata['episodes'].update(out['metadata']['episodes'])
# double check
for p in process_list:
p.join()
time.sleep(1.0)
print("All threads joined")
elapsed = time.time() - start_time
# collect the metadata.yaml files
while not q.empty():
out = q.get()
metadata['episodes'].update(out['metadata']['episodes'])
save_yaml(metadata, os.path.join(OUTPUT_DIR, 'metadata.yaml'))
print("Generating and saving dataset to disk took %d seconds" %
(int(elapsed)))
