|
|
import tensorflow as tf |
|
|
import h5py |
|
|
import os |
|
|
import fnmatch |
|
|
import cv2 |
|
|
import numpy as np |
|
|
from tqdm import tqdm |
|
|
|
|
|
def decode_img(img): |
|
|
return cv2.cvtColor(cv2.imdecode(np.frombuffer(img, np.uint8), cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB) |
|
|
|
|
|
def decode_all_imgs(imgs): |
|
|
return [decode_img(img) for img in imgs] |
|
|
|
|
|
def _bytes_feature(value): |
|
|
"""Returns a bytes_list from a string / byte.""" |
|
|
if isinstance(value, type(tf.constant(0))): |
|
|
value = value.numpy() |
|
|
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) |
|
|
def _bool_feature(value): |
|
|
"""Returns a bool_list from a boolean.""" |
|
|
return tf.train.Feature(int64_list=tf.train.Int64List(value=[int(value)])) |
|
|
|
|
|
def serialize_example(action, base_action, qpos, qvel, cam_high, cam_left_wrist, cam_right_wrist, cam_low, instruction, terminate_episode): |
|
|
if base_action is not None: |
|
|
feature = { |
|
|
'action': _bytes_feature(tf.io.serialize_tensor(action)), |
|
|
'base_action': _bytes_feature(tf.io.serialize_tensor(base_action)), |
|
|
'qpos': _bytes_feature(tf.io.serialize_tensor(qpos)), |
|
|
'qvel': _bytes_feature(tf.io.serialize_tensor(qvel)), |
|
|
'cam_high': _bytes_feature(tf.io.serialize_tensor(cam_high)), |
|
|
'cam_left_wrist': _bytes_feature(tf.io.serialize_tensor(cam_left_wrist)), |
|
|
'cam_right_wrist': _bytes_feature(tf.io.serialize_tensor(cam_right_wrist)), |
|
|
'instruction': _bytes_feature(instruction), |
|
|
'terminate_episode': _bool_feature(terminate_episode) |
|
|
} |
|
|
else: |
|
|
feature = { |
|
|
'action': _bytes_feature(tf.io.serialize_tensor(action)), |
|
|
'qpos': _bytes_feature(tf.io.serialize_tensor(qpos)), |
|
|
'qvel': _bytes_feature(tf.io.serialize_tensor(qvel)), |
|
|
'cam_high': _bytes_feature(tf.io.serialize_tensor(cam_high)), |
|
|
'cam_left_wrist': _bytes_feature(tf.io.serialize_tensor(cam_left_wrist)), |
|
|
'cam_right_wrist': _bytes_feature(tf.io.serialize_tensor(cam_right_wrist)), |
|
|
'cam_low': _bytes_feature(tf.io.serialize_tensor(cam_low)), |
|
|
'instruction': _bytes_feature(instruction), |
|
|
'terminate_episode': _bool_feature(terminate_episode) |
|
|
} |
|
|
example_proto = tf.train.Example(features=tf.train.Features(feature=feature)) |
|
|
return example_proto.SerializeToString() |
|
|
|
|
|
def write_tfrecords(root_dir, out_dir): |
|
|
if not os.path.exists(out_dir): |
|
|
os.makedirs(out_dir) |
|
|
num_files = 0 |
|
|
for root, dirs, files in os.walk(root_dir): |
|
|
num_files += len(fnmatch.filter(files, '*.hdf5')) |
|
|
with tqdm(total=num_files) as pbar: |
|
|
for root, dirs, files in os.walk(root_dir): |
|
|
for filename in fnmatch.filter(files, '*.hdf5'): |
|
|
filepath = os.path.join(root, filename) |
|
|
with h5py.File(filepath, 'r') as f: |
|
|
if not 'instruction' in f: |
|
|
continue |
|
|
pbar.update(1) |
|
|
output_dir = os.path.join(out_dir, os.path.relpath(root, root_dir)) |
|
|
if not os.path.exists(output_dir): |
|
|
os.makedirs(output_dir) |
|
|
print(f"Writing TFRecords to {output_dir}") |
|
|
tfrecord_path = os.path.join(output_dir, filename.replace('.hdf5', '.tfrecord')) |
|
|
with tf.io.TFRecordWriter(tfrecord_path) as writer: |
|
|
num_episodes = f['action'].shape[0] |
|
|
for i in range(num_episodes): |
|
|
action = f['action'][i] |
|
|
if 'base_action' in f: |
|
|
base_action = f['base_action'][i] |
|
|
else: |
|
|
base_action = None |
|
|
qpos = f['observations']['qpos'][i] |
|
|
qvel = f['observations']['qvel'][i] |
|
|
cam_high = decode_img(f['observations']['images']['cam_high'][i]) |
|
|
cam_left_wrist = decode_img(f['observations']['images']['cam_left_wrist'][i]) |
|
|
cam_right_wrist = decode_img(f['observations']['images']['cam_right_wrist'][i]) |
|
|
if 'cam_low' in f['observations']['images']: |
|
|
cam_low = decode_img(f['observations']['images']['cam_low'][i]) |
|
|
else: |
|
|
cam_low = None |
|
|
instruction = f['instruction'][()] |
|
|
terminate_episode = i == num_episodes - 1 |
|
|
serialized_example = serialize_example(action, base_action, qpos, qvel, cam_high, cam_left_wrist, cam_right_wrist, cam_low, instruction, terminate_episode) |
|
|
writer.write(serialized_example) |
|
|
print(f"TFRecords written to {tfrecord_path}") |
|
|
print(f"TFRecords written to {out_dir}") |
|
|
|
|
|
root_dir = '../datasets/aloha/' |
|
|
output_dir = '../datasets/aloha/tfrecords/' |
|
|
|
|
|
write_tfrecords(root_dir, output_dir) |
|
|
|
