def _write_tfrecord(pkl_fname, tfrecord_fname, obs_shape):
writer = tf.python_io.TFRecordWriter(tfrecord_fname)
rollouts = mypickle.load(pkl_fname)['rollouts']
for r in rollouts:
images = r['env_infos']['rgb']
semantics = r['env_infos']['semantic']
for image, semantic in zip(images, semantics):
rightlane_mask = CarlaCollSpeedRoadEnv.get_rightlane(
semantic)[0].astype(np.uint8)
image = utils.imresize(image, obs_shape)
rightlane_mask = utils.imresize(rightlane_mask[..., None],
obs_shape[:2] + [1],
PIL.Image.BILINEAR)[..., -1]
example = tf.train.Example(features=tf.train.Features(
feature={
'image':
RoadLabeller._bytes_feature(image.tostring()),
'label':
RoadLabeller._bytes_feature(rightlane_mask.tostring()),
}))
writer.write(example.SerializeToString())
writer.close()
def create_training_data(save_folder, image_shape, rescale, bordersize,
holdout_pct):
"""
:param save_folder: where images are saved
:param image_shape: shape of the image
:param rescale: make rescale times bigger, for ease of labelling
:param bordersize: how much to pad with 0s, for ease of labelling
:param holdout_pct: how much data to holdout
"""
### read image, label pairs
label_fnames = glob.glob(os.path.join(save_folder, 'label*'))
random.shuffle(label_fnames)
height, width, channels = image_shape
images_train, labels_train = [], []
images_holdout, labels_holdout = [], []
for i, label_fname in enumerate(label_fnames):
image_fname = label_fname.replace('label_', '')
image = np.asarray(Image.open(image_fname))
label = np.asarray(Image.open(label_fname))
# reduce image back down
image = image[bordersize:-bordersize, bordersize:-bordersize]
image = utils.imresize(image, (height, width, channels))
assert (tuple(image.shape) == tuple(image_shape))
# reduce label back down
label = label[bordersize:-bordersize, bordersize:-bordersize]
label = utils.imresize(label, (height, width, 1), Image.BILINEAR)
label = label[:, :, 0]
label = (label > 0.5)
assert (tuple(label.shape) == (height, width))
if i / float(len(label_fnames)) > holdout_pct:
images_train.append(image)
labels_train.append(label)
else:
images_holdout.append(image)
labels_holdout.append(label)
np.save(os.path.join(save_folder, 'data_train_images.npy'),
np.array(images_train))
np.save(os.path.join(save_folder, 'data_train_labels.npy'),
np.array(labels_train))
np.save(os.path.join(save_folder, 'data_holdout_images.npy'),
np.array(images_holdout))
np.save(os.path.join(save_folder, 'data_holdout_labels.npy'),
np.array(labels_holdout))
logger.info('Saved train and holdout')
def get_model_outputs(self, observations):
observations = np.asarray(observations)
assert observations.shape[-1] == self._image_shape[
-1], 'do not have same number of channels'
if tuple(observations.shape[1:3]) != tuple(self._image_shape[:2]):
observations = np.array([
utils.imresize(o_t, self._image_shape) for o_t in observations
])
feed_dict = {self._tf_dict['obs_ph']: observations}
probs, = self._tf_dict['sess'].run([self._tf_dict['probs']],
feed_dict=feed_dict)
return probs
def store_observation(self, step, observation, goal, use_labeller=True):
self._data_step = step
self._steps[self._index] = step
obs_im_full, obs_vec = observation
if self._labeller and use_labeller:
assert (np.any(np.isnan(goal)))
goal = self._labeller.label(([obs_im_full], [obs_vec]), [goal])[0]
#import IPython; IPython.embed()
obs_im = utils.imresize(obs_im_full, self._obs_im_shape)
assert (obs_im.shape == self._obs_im_shape)
self._observations_im[self._index, :] = obs_im.reshape(
(self._obs_im_dim, ))
self._observations_vec[self._index, :] = obs_vec.reshape(
(self._obs_vec_dim, ))
self._goals[self._index, :] = goal
return goal
def store_observation(self, step, observation, goal, use_labeller=True):
self._data_step = step
self._steps[self._index] = step
obs_im_full, obs_vec = observation
if self._labeller and use_labeller:
assert (np.any(np.isnan(goal)))
goal = self._labeller.label(([obs_im_full], [obs_vec]), [goal])[0]
if obs_im_full.shape[-1] == 3 and self._obs_im_shape[-1] == 1:
obs_im_full = utils.im2gray(obs_im_full)
assert obs_im_full.shape[-1] == 1 or obs_im_full.shape[-1] == 3
obs_im = utils.imresize(obs_im_full, self._obs_im_shape)
assert(obs_im.shape == self._obs_im_shape)
self._observations_im[self._index, :] = obs_im.reshape((self._obs_im_dim,))
self._observations_vec[self._index, :] = obs_vec.reshape((self._obs_vec_dim,))
self._goals[self._index, :] = goal
return goal
def extract_images_from_pkls(pkl_folder, save_folder, maxsaved, image_shape,
rescale, bordersize):
"""
:param pkl_folder: folder containing pkls with training images
:param save_folder: where to save the resulting images
:param maxsaved: how many images to save
:param image_shape: shape of the image
:param rescale: make rescale times bigger, for ease of labelling
:param bordersize: how much to pad with 0s, for ease of labelling
"""
random.seed(0)
fnames = glob.glob(os.path.join(pkl_folder, '*.pkl'))
random.shuffle(fnames)
logger.info('{0} files to read'.format(len(fnames)))
fnames = itertools.cycle(fnames)
im_num = 0
while im_num < maxsaved:
fname = next(fnames)
rollout = random.choice(mypickle.load(fname)['rollouts'])
obs = random.choice(rollout['observations_im'])
height, width, channels = image_shape
im = np.reshape(obs, image_shape)
im = utils.imresize(im, (rescale * height, rescale * width, channels))
im = np.pad(im, ((bordersize, bordersize), (bordersize, bordersize),
(0, 0)), 'constant')
if im.shape[-1] == 1:
im = im[:, :, 0]
Image.fromarray(im).save(
os.path.join(save_folder, 'image_{0:06d}.jpg'.format(im_num)))
im_num += 1
logger.info('Saved {0} images'.format(im_num))
def create_rollout(self, r, labeller):
# observations_im
# observations_vec - fixed
# actions - fixed
# rewards - fixed
# dones - fixed
# goals
# steps - fixed
# env_infos - use and get rid of
r_new = {
'observations_vec': r['observations_vec'],
'actions': r['actions'],
'dones': r['dones'],
'steps': r['steps'],
'env_infos': None,
}
# observations_im
obs_im = []
for t in range(len(r['dones']) - 1):
obs_im_t = []
for camera_name in self._params['cameras']:
camera_params = self._params[camera_name]
if camera_params['include_in_obs']:
camera_data = self._convert_camera_data(
r['env_infos'][camera_name][t], camera_params)
obs_im_t.append(camera_data)
obs_im_t = np.concatenate(obs_im_t, axis=2)
obs_im_t = utils.imresize(obs_im_t,
self.observation_im_space.shape)
obs_im.append(obs_im_t)
obs_im.append(obs_im[-1]) # b/c env_infos is short by 1
r_new['observations_im'] = np.array(obs_im)
# goals
_, rightlane_seen, rightlane_diff = zip(*[
CarlaCollSpeedRoadEnv.get_rightlane(s)
for s in r['env_infos']['semantic']
])
# b/c env_infos are one short
rightlane_seen = list(rightlane_seen) + [rightlane_seen[-1]]
rightlane_diff = list(rightlane_diff) + [rightlane_diff[-1]]
r_new['goals'] = np.hstack([
r['goals'],
np.array([rightlane_seen], dtype=np.float32).T,
np.array([rightlane_diff], dtype=np.float32).T
])
# rewards
rewards = []
coll_idx = list(self.observation_vec_spec.keys()).index('coll')
speed_idx = list(self.observation_vec_spec.keys()).index('speed')
goal_speed_idx = list(self.goal_spec.keys()).index('speed')
for t in range(len(r['dones'])):
if r['observations_vec'][t, coll_idx] > 0:
r_t = -self.horizon
else:
r_t = -abs((r['goals'][t, goal_speed_idx] -
r['observations_vec'][t, speed_idx]) /
r['goals'][t, goal_speed_idx])
r_t += 10 * rightlane_seen[t] * (1 - abs(rightlane_diff[t]))
rewards.append(r_t)
rewards = rewards[1:] + [0]
r_new['rewards'] = np.asarray(rewards, dtype=np.float32)
return r_new
def create_rollout(self, r, labeller, keep_env_infos=False):
# observations_im
# observations_vec - fixed
# actions - fixed
# rewards - fixed
# dones - fixed
# goals
# steps - fixed
# env_infos - use and get rid of
r_new = {
'observations_vec': r['observations_vec'],
'actions': r['actions'],
'dones': r['dones'],
'steps': r['steps'],
'env_infos': r['env_infos'] if keep_env_infos else None,
}
# observations_im
obs_im = []
for t in range(len(r['dones']) - 1):
obs_im_t = []
for camera_name in self._params['cameras']:
camera_params = self._params[camera_name]
if camera_params['include_in_obs']:
camera_data = self._convert_camera_data(
r['env_infos'][camera_name][t], camera_params)
obs_im_t.append(camera_data)
obs_im_t = np.concatenate(obs_im_t, axis=2)
obs_im_t = utils.imresize(obs_im_t,
self.observation_im_space.shape)
obs_im.append(obs_im_t)
obs_im.append(obs_im[-1]) # b/c env_infos is short by 1
obs_im = np.array(obs_im)
r_new['observations_im'] = obs_im
r_new['goals'] = np.hstack([
r['goals'],
np.nan * np.ones([len(r['goals']), 1], dtype=np.float32),
np.nan * np.ones([len(r['goals']), 1], dtype=np.float32)
])
# goals
batch_size = 350
rgb = np.append(r['env_infos']['rgb'],
r['env_infos']['rgb'][-1][None],
axis=0)
goals = r_new['goals']
i = 0
while i < len(rgb):
goals[i:i + batch_size] = labeller.label(
(rgb[i:i + batch_size], None), goals[i:i + batch_size])
i += batch_size
r_new['goals'] = goals
# rewards
rewards = []
coll_idx = list(self.observation_vec_spec.keys()).index('coll')
speed_idx = list(self.observation_vec_spec.keys()).index('speed')
goal_speed_idx = list(self.goal_spec.keys()).index('speed')
rightlane_seen_idx = list(
self.goal_spec.keys()).index('rightlane_seen')
rightlane_diff_idx = list(
self.goal_spec.keys()).index('rightlane_diff')
for t in range(len(r_new['dones'])):
if r_new['observations_vec'][t, coll_idx] > 0:
r_t = -self.horizon
else:
r_t = -abs((r_new['goals'][t, goal_speed_idx] -
r_new['observations_vec'][t, speed_idx]) /
r_new['goals'][t, goal_speed_idx])
r_t += 10 * r_new['goals'][t, rightlane_seen_idx] * (
1 - abs(r_new['goals'][t, rightlane_diff_idx]))
rewards.append(r_t)
rewards = rewards[1:] + [0]
r_new['rewards'] = np.asarray(rewards, dtype=np.float32)
return r_new