if k in train_image_indices:
train = 1
split = TRAIN_ID
if train:
num_objects = min(max(num_objs_rv.rvs(size=1)[0] + 1, min_num_objects), num_train)
obj_names = [objects[i] for i in np.random.choice(train_indices, size=num_objects, replace=False)]
else:
num_objects = min(max(num_objs_rv.rvs(size=1)[0] + 1, min_num_objects), num_test)
obj_names = [objects[i] for i in np.random.choice(test_indices, size=num_objects, replace=False)]
# get human consent
message = 'Please place %d objects:\n' %(num_objects)
for name in obj_names:
message += '\t{}\n'.format(name)
message += 'Hit ENTER when finished.'
utils.keyboard_input(message=message)
# capture
for sensor_name, sensor in sensors.iteritems():
logging.info('Capturing images from sensor %s' %(sensor_name))
# read pose and intrinsics
sensor_pose = sensor_poses[sensor_name]
camera_intr = camera_intrs[sensor_name]
workspace_im = workspace_ims[sensor_name]
dataset = datasets[sensor_name]
T_camera_world = sensor_pose
datapoint = dataset.datapoint_template
# read raw images
raw_color_im, raw_depth_im, _ = sensor.frames()
def __init__(self, filename, config, access_mode=WRITE_ACCESS):
# read params
self._filename = filename
self._config = config
self._datapoints_per_file = config['datapoints_per_file']
self._access_mode = access_mode
# check valid access mode
if access_mode == READ_WRITE_ACCESS:
raise ValueError('Read and write not supported simultaneously.')
# open dataset folder
# create dataset if necessary
if not os.path.exists(
self._filename) and access_mode != READ_ONLY_ACCESS:
os.mkdir(self._filename)
# throw error if dataset doesn't exist
elif not os.path.exists(
self._filename) and access_mode == READ_ONLY_ACCESS:
raise ValueError('Dataset %s does not exist!' % (self._filename))
# check dataset empty
elif os.path.exists(self._filename) and len(os.listdir(
self._filename)) > 0 and access_mode == WRITE_ACCESS:
human_input = utils.keyboard_input(
'Dataset %s exists. Overwrite?' % (self.filename), yesno=True)
if human_input.lower() == 'n':
raise ValueError('User opted not to overwrite dataset')
# save config to location
if access_mode != READ_ONLY_ACCESS:
config_filename = os.path.join(self._filename, 'config.json')
json.dump(self._config, open(config_filename, 'w'))
# init data storage
self._allocate_tensors()
# init state variables
if access_mode == WRITE_ACCESS:
# init no files
self._num_tensors = 0
self._num_datapoints = 0
if not os.path.exists(self.tensor_dir):
os.mkdir(self.tensor_dir)
elif access_mode == READ_ONLY_ACCESS:
# read the number of tensor files
tensor_dir = self.tensor_dir
tensor_filenames = utils.filenames(tensor_dir,
tag=COMPRESSED_TENSOR_EXT,
sorted=True)
file_nums = np.array(
[int(filename[-9:-4]) for filename in tensor_filenames])
self._num_tensors = np.max(file_nums) + 1
# compute the number of datapoints
last_tensor_ind = np.where(file_nums == self._num_tensors -
1)[0][0]
last_tensor_data = np.load(
tensor_filenames[last_tensor_ind])['arr_0']
self._num_datapoints_last_file = last_tensor_data.shape[0]
self._num_datapoints = self._datapoints_per_file * (
self._num_tensors - 1) + self._num_datapoints_last_file
# form index maps for each file
self._index_to_file_num = {}
self._file_num_to_indices = {}
# set file index
cur_file_num = 0
start_datapoint_index = 0
# set mapping from file num to datapoint indices
self._file_num_to_indices[cur_file_num] = np.arange(
self._datapoints_per_file) + start_datapoint_index
for ind in range(self._num_datapoints):
# set mapping from index to file num
self._index_to_file_num[ind] = cur_file_num
# update to the next file
if ind > 0 and ind % self._datapoints_per_file == 0:
cur_file_num += 1
start_datapoint_index += self._datapoints_per_file
# set mapping from file num to datapoint indices
if cur_file_num < self._num_tensors - 1:
self._file_num_to_indices[cur_file_num] = np.arange(
self._datapoints_per_file) + start_datapoint_index
else:
self._file_num_to_indices[cur_file_num] = np.arange(
self._num_datapoints_last_file
) + start_datapoint_index
else:
raise ValueError('Access mode %s not supported' % (access_mode))
n_samples=obj_config['stp_num_samples'],
threshold=obj_config['stp_min_prob']
)
for push_num in range(10):
push_idx = None
# push_idx = 33
sample_id = 0
num_toppled, pose_angles, actual_poses = [], [], []
print '\n\n\n\n\n\nSTARTING DIFFERENT PUSH', push_num
while sample_id < 10:
sample_id += 1
# sim_point_cloud_masked = get_sim_point_cloud(depth_scene, env.state.obj)
usr_input = 'n'
while usr_input != 'y' and usr_input != 'a':
usr_input = utils.keyboard_input('\n\nPut the object in position y: continue, v: vis pose, a: abort push[y/v/a]:')
if usr_input == 'v':
vis3d.figure()
env.state.obj.T_obj_world = orig_pose
env.render_3d_scene()
vis3d.show(starting_camera_pose=CAMERA_POSE)
if usr_input == 'a':
print 'trying different push'
push_idx = None
sample_id = 0
num_toppled, pose_angles, actual_poses = [], [], []
continue
usr_input = 'n'
while usr_input != 'y':
s2r = sim_to_real_tf(sim_point_cloud_masked, vis=True)